add organization_id filter for get_workflow and get_workflow_run (#1422)

This commit is contained in:
Shuchang Zheng
2024-12-22 17:49:33 -08:00
committed by GitHub
parent b256bace6a
commit 2e37542218
7 changed files with 44 additions and 18 deletions

View File

@@ -1283,12 +1283,13 @@ class AgentDB:
)
return None
async def get_workflow_run(self, workflow_run_id: str) -> WorkflowRun | None:
async def get_workflow_run(self, workflow_run_id: str, organization_id: str | None = None) -> WorkflowRun | None:
try:
async with self.Session() as session:
if workflow_run := (
await session.scalars(select(WorkflowRunModel).filter_by(workflow_run_id=workflow_run_id))
).first():
get_workflow_run_query = select(WorkflowRunModel).filter_by(workflow_run_id=workflow_run_id)
if organization_id:
get_workflow_run_query = get_workflow_run_query.filter_by(organization_id=organization_id)
if workflow_run := (await session.scalars(get_workflow_run_query)).first():
return convert_to_workflow_run(workflow_run)
return None
except SQLAlchemyError: