tcsenpai/pensieve
1
0
Fork 0
mirror of https://github.com/tcsenpai/pensieve.git synced 2025-06-10 13:07:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: support scan and trigger specific plugins

Browse Source
This commit is contained in:
arkohut 2024-07-10 18:35:44 +08:00
parent d1b83571e6
commit b5d524f7b0
3 changed files with 25 additions and 26 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
memos/commands.py
View File

@ -98,7 +98,7 @@ def show(library_id: int):
@lib_app.command("scan") @lib_app.command("scan")
def scan(library_id: int, force: bool = False): def scan(library_id: int, force: bool = False, plugins: List[int] = typer.Option(None, "--plugin", "-p")):
response = httpx.get(f"{BASE_URL}/libraries/{library_id}") response = httpx.get(f"{BASE_URL}/libraries/{library_id}")
if response.status_code != 200: if response.status_code != 200:
@ -187,7 +187,7 @@ def scan(library_id: int, force: bool = False):
update_response = httpx.put( update_response = httpx.put(
f"{BASE_URL}/entities/{existing_entity['id']}", f"{BASE_URL}/entities/{existing_entity['id']}",
json=new_entity, json=new_entity,
params={"trigger_webhooks_flag": "true"}, params={"trigger_webhooks_flag": "true", **({"plugins": plugins} if plugins else {})},
timeout=30 timeout=30
) )
if 200 <= update_response.status_code < 300: if 200 <= update_response.status_code < 300:
@ -207,6 +207,7 @@ def scan(library_id: int, force: bool = False):
post_response = httpx.post( post_response = httpx.post(
f"{BASE_URL}/libraries/{library_id}/entities", f"{BASE_URL}/libraries/{library_id}/entities",
json=new_entity, json=new_entity,
params={"plugins": plugins} if plugins else {},
timeout=30 timeout=30
) )
if 200 <= post_response.status_code < 300: if 200 <= post_response.status_code < 300:

40
memos/server.py
View File

@ -152,29 +152,31 @@ def new_folder(
return crud.add_folder(library_id=library.id, folder=folder, db=db) return crud.add_folder(library_id=library.id, folder=folder, db=db)
async def trigger_webhooks(library: Library, entity: Entity, request: Request): async def trigger_webhooks(library: Library, entity: Entity, request: Request, plugins: List[int] = None):
async with httpx.AsyncClient() as client: async with httpx.AsyncClient() as client:
tasks = [] tasks = []
for plugin in library.plugins: for plugin in library.plugins:
if plugin.webhook_url: if plugins is None or plugin.id in plugins:
location = str(request.url_for("get_entity_by_id", entity_id=entity.id)) if plugin.webhook_url:
task = client.post( location = str(request.url_for("get_entity_by_id", entity_id=entity.id))
plugin.webhook_url, task = client.post(
json=entity.model_dump(mode="json"), plugin.webhook_url,
headers={"Location": location}, json=entity.model_dump(mode="json"),
timeout=10.0, # Adding a timeout of 10 seconds headers={"Location": location},
) timeout=10.0, # Adding a timeout of 10 seconds
tasks.append(task) )
tasks.append(task)
responses = await asyncio.gather(*tasks, return_exceptions=True) responses = await asyncio.gather(*tasks, return_exceptions=True)
for plugin, response in zip(library.plugins, responses): for plugin, response in zip(library.plugins, responses):
if isinstance(response, Exception): if plugins is None or plugin.id in plugins:
print(f"Error triggering webhook for plugin {plugin.id}: {response}") if isinstance(response, Exception):
elif response.status_code >= 400: print(f"Error triggering webhook for plugin {plugin.id}: {response}")
print( elif response.status_code >= 400:
f"Error triggering webhook for plugin {plugin.id}: {response.status_code} - {response.text}" print(
) f"Error triggering webhook for plugin {plugin.id}: {response.status_code} - {response.text}"
)
@app.post("/libraries/{library_id}/entities", response_model=Entity, tags=["entity"]) @app.post("/libraries/{library_id}/entities", response_model=Entity, tags=["entity"])
@ -183,6 +185,7 @@ async def new_entity(
library_id: int, library_id: int,
request: Request, request: Request,
db: Session = Depends(get_db), db: Session = Depends(get_db),
plugins: Annotated[List[int] | None, Query()] = None,
): ):
library = crud.get_library_by_id(library_id, db) library = crud.get_library_by_id(library_id, db)
if library is None: if library is None:
@ -191,7 +194,7 @@ async def new_entity(
) )
entity = crud.create_entity(library_id, new_entity, db) entity = crud.create_entity(library_id, new_entity, db)
await trigger_webhooks(library, entity, request) await trigger_webhooks(library, entity, request, plugins)
return entity return entity
@ -271,6 +274,7 @@ async def update_entity(
request: Request, request: Request,
db: Session = Depends(get_db), db: Session = Depends(get_db),
trigger_webhooks_flag: bool = False, trigger_webhooks_flag: bool = False,
plugins: Annotated[List[int] | None, Query()] = None,
): ):
entity = crud.get_entity_by_id(entity_id, db) entity = crud.get_entity_by_id(entity_id, db)
if entity is None: if entity is None:
@ -286,7 +290,7 @@ async def update_entity(
raise HTTPException( raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND, detail="Library not found" status_code=status.HTTP_404_NOT_FOUND, detail="Library not found"
) )
await trigger_webhooks(library, entity, request) await trigger_webhooks(library, entity, request, plugins)
return entity return entity

6
memos/test_server.py
View File

@ -93,12 +93,6 @@ def client():
Base.metadata.drop_all(bind=engine) Base.metadata.drop_all(bind=engine)
def test_read_main(client):
response = client.get("/")
assert response.status_code == 200
assert response.json() == {"healthy": True}
# Test the new_library endpoint # Test the new_library endpoint
def test_new_library(client): def test_new_library(client):
library_param = NewLibraryParam(name="Test Library") library_param = NewLibraryParam(name="Test Library")