fix: improve UUID handling in flow helpers (#5346)

* Update flow.py

solving error when user id is string and not uuid

* Update flow.py

* solve lint errors
This commit is contained in:
Edwin Jose 2024-12-18 20:14:41 -05:00 committed by GitHub
commit d5498cf3cb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -70,7 +70,8 @@ async def load_flow(
async def find_flow(flow_name: str, user_id: str) -> str | None:
async with session_scope() as session:
stmt = select(Flow).where(Flow.name == flow_name).where(Flow.user_id == user_id)
uuid_user_id = UUID(user_id) if isinstance(user_id, str) else user_id
stmt = select(Flow).where(Flow.name == flow_name).where(Flow.user_id == uuid_user_id)
flow = (await session.exec(stmt)).first()
return flow.id if flow else None
@ -274,7 +275,7 @@ def get_arg_names(inputs: list[Vertex]) -> list[dict[str, str]]:
]
async def get_flow_by_id_or_endpoint_name(flow_id_or_name: str, user_id: UUID | None = None) -> FlowRead | None:
async def get_flow_by_id_or_endpoint_name(flow_id_or_name: str, user_id: str | UUID | None = None) -> FlowRead | None:
async with session_scope() as session:
endpoint_name = None
try:
@ -284,7 +285,8 @@ async def get_flow_by_id_or_endpoint_name(flow_id_or_name: str, user_id: UUID |
endpoint_name = flow_id_or_name
stmt = select(Flow).where(Flow.endpoint_name == endpoint_name)
if user_id:
stmt = stmt.where(Flow.user_id == user_id)
uuid_user_id = UUID(user_id) if isinstance(user_id, str) else user_id
stmt = stmt.where(Flow.user_id == uuid_user_id)
flow = (await session.exec(stmt)).first()
if flow is None:
raise HTTPException(status_code=404, detail=f"Flow identifier {flow_id_or_name} not found")