wip: initial set of exceptions added

julep-ai · Dec 20, 2024 · dc0ec36 · dc0ec36
1 parent 7b0be5c
commit dc0ec36
Show file tree

Hide file tree

Showing 28 changed files with 283 additions and 167 deletions.
diff --git a/agents-api/agents_api/queries/agents/create_agent.py b/agents-api/agents_api/queries/agents/create_agent.py
@@ -8,13 +8,16 @@
 from beartype import beartype
 from sqlglot import parse_one
 from uuid_extensions import uuid7
-
+import asyncpg
+from fastapi import HTTPException
 from ...autogen.openapi_model import Agent, CreateAgentRequest
 from ...metrics.counters import increase_counter
 from ..utils import (
     generate_canonical_name,
     pg_query,
     wrap_in_class,
+    rewrap_exceptions,
+    partialclass,
 )
 
 # Define the raw SQL query
@@ -45,35 +48,30 @@
 """).sql(pretty=True)
 
 
-# @rewrap_exceptions(
-#     {
-#         psycopg_errors.ForeignKeyViolation: partialclass(
-#             HTTPException,
-#             status_code=404,
-#             detail="The specified developer does not exist.",
-#         ),
-#         psycopg_errors.UniqueViolation: partialclass(
-#             HTTPException,
-#             status_code=409,
-#             detail="An agent with this canonical name already exists for this developer.",
-#         ),
-#         psycopg_errors.CheckViolation: partialclass(
-#             HTTPException,
-#             status_code=400,
-#             detail="The provided data violates one or more constraints. Please check the input values.",
-#         ),
-#         ValidationError: partialclass(
-#             HTTPException,
-#             status_code=400,
-#             detail="Input validation failed. Please check the provided data.",
-#         ),
-#         TypeError: partialclass(
-#             HTTPException,
-#             status_code=400,
-#             detail="A type mismatch occurred. Please review the input.",
-#         ),
-#     }
-# )
+@rewrap_exceptions(
+    {
+        asyncpg.exceptions.ForeignKeyViolationError: partialclass(
+            HTTPException,
+            status_code=404,
+            detail="The specified developer does not exist.",
+        ),
+        asyncpg.exceptions.UniqueViolationError: partialclass(
+            HTTPException,
+            status_code=409,
+            detail="An agent with this canonical name already exists for this developer.",
+        ),
+        asyncpg.exceptions.CheckViolationError: partialclass(
+            HTTPException,
+            status_code=400,
+            detail="The provided data violates one or more constraints. Please check the input values.",
+        ),
+        asyncpg.exceptions.DataError: partialclass(
+            HTTPException,
+            status_code=400,
+            detail="Invalid data provided. Please check the input values.",
+        ),
+    }
+)
 @wrap_in_class(
     Agent,
     one=True,

diff --git a/agents-api/agents_api/queries/agents/create_or_update_agent.py b/agents-api/agents_api/queries/agents/create_or_update_agent.py
@@ -7,13 +7,17 @@
 
 from beartype import beartype
 from sqlglot import parse_one
+from fastapi import HTTPException
+import asyncpg
 
 from ...autogen.openapi_model import Agent, CreateOrUpdateAgentRequest
 from ...metrics.counters import increase_counter
 from ..utils import (
     generate_canonical_name,
     pg_query,
     wrap_in_class,
+    rewrap_exceptions,
+    partialclass,
 )
 
 # Define the raw SQL query
@@ -44,15 +48,30 @@
 """).sql(pretty=True)
 
 
-# @rewrap_exceptions(
-#     {
-#         psycopg_errors.ForeignKeyViolation: partialclass(
-#             HTTPException,
-#             status_code=404,
-#             detail="The specified developer does not exist.",
-#         )
-#     }
-# )
+@rewrap_exceptions(
+    {
+        asyncpg.exceptions.ForeignKeyViolationError: partialclass(
+            HTTPException,
+            status_code=404,
+            detail="The specified developer does not exist.",
+        ),
+        asyncpg.exceptions.UniqueViolationError: partialclass(
+            HTTPException,
+            status_code=409,
+            detail="An agent with this canonical name already exists for this developer.",
+        ),
+        asyncpg.exceptions.CheckViolationError: partialclass(
+            HTTPException,
+            status_code=400,
+            detail="The provided data violates one or more constraints. Please check the input values.",
+        ),
+        asyncpg.exceptions.DataError: partialclass(
+            HTTPException,
+            status_code=400,
+            detail="Invalid data provided. Please check the input values.",
+        ),
+    }
+)
 @wrap_in_class(
     Agent,
     one=True,

diff --git a/agents-api/agents_api/queries/agents/delete_agent.py b/agents-api/agents_api/queries/agents/delete_agent.py
@@ -7,12 +7,16 @@
 
 from beartype import beartype
 from sqlglot import parse_one
+from fastapi import HTTPException
+import asyncpg
 
 from ...autogen.openapi_model import ResourceDeletedResponse
 from ...common.utils.datetime import utcnow
 from ..utils import (
     pg_query,
     wrap_in_class,
+    rewrap_exceptions,
+    partialclass,
 )
 
 # Define the raw SQL query
@@ -59,17 +63,30 @@
 """).sql(pretty=True)
 
 
-# @rewrap_exceptions(
-# @rewrap_exceptions(
-#     {
-#         psycopg_errors.ForeignKeyViolation: partialclass(
-#             HTTPException,
-#             status_code=404,
-#             detail="The specified developer does not exist.",
-#         )
-#     }
-#     # TODO: Add more exceptions
-# )
+@rewrap_exceptions(
+    {
+        asyncpg.exceptions.ForeignKeyViolationError: partialclass(
+            HTTPException,
+            status_code=404,
+            detail="The specified developer does not exist.",
+        ),
+        asyncpg.exceptions.UniqueViolationError: partialclass(
+            HTTPException,
+            status_code=409,
+            detail="An agent with this canonical name already exists for this developer.",
+        ),
+        asyncpg.exceptions.CheckViolationError: partialclass(
+            HTTPException,
+            status_code=400,
+            detail="The provided data violates one or more constraints. Please check the input values.",
+        ),
+        asyncpg.exceptions.DataError: partialclass(
+            HTTPException,
+            status_code=400,
+            detail="Invalid data provided. Please check the input values.",
+        ),
+    }
+)
 @wrap_in_class(
     ResourceDeletedResponse,
     one=True,

diff --git a/agents-api/agents_api/queries/agents/get_agent.py b/agents-api/agents_api/queries/agents/get_agent.py
@@ -7,11 +7,15 @@
 
 from beartype import beartype
 from sqlglot import parse_one
+from fastapi import HTTPException
+import asyncpg
 
 from ...autogen.openapi_model import Agent
 from ..utils import (
     pg_query,
     wrap_in_class,
+    rewrap_exceptions,
+    partialclass,
 )
 
 # Define the raw SQL query
@@ -35,16 +39,20 @@
 """).sql(pretty=True)
 
 
-# @rewrap_exceptions(
-# {
-#     psycopg_errors.ForeignKeyViolation: partialclass(
-#         HTTPException,
-#         status_code=404,
-#         detail="The specified developer does not exist.",
-#     )
-# }
-# # TODO: Add more exceptions
-# )
+@rewrap_exceptions(
+    {
+        asyncpg.exceptions.ForeignKeyViolationError: partialclass(
+            HTTPException,
+            status_code=404,
+            detail="The specified developer does not exist.",
+        ),
+        asyncpg.exceptions.DataError: partialclass(
+            HTTPException,
+            status_code=400,
+            detail="Invalid data provided. Please check the input values.",
+        ),
+    }
+)
 @wrap_in_class(Agent, one=True, transform=lambda d: {"id": d["agent_id"], **d})
 @pg_query
 @beartype

diff --git a/agents-api/agents_api/queries/agents/list_agents.py b/agents-api/agents_api/queries/agents/list_agents.py
@@ -8,11 +8,13 @@
 
 from beartype import beartype
 from fastapi import HTTPException
-
+import asyncpg
 from ...autogen.openapi_model import Agent
 from ..utils import (
     pg_query,
     wrap_in_class,
+    rewrap_exceptions,
+    partialclass,
 )
 
 # Define the raw SQL query
@@ -39,17 +41,20 @@
 LIMIT $2 OFFSET $3;
 """
 
-
-#  @rewrap_exceptions(
-#     {
-#         psycopg_errors.ForeignKeyViolation: partialclass(
-#             HTTPException,
-#             status_code=404,
-#             detail="The specified developer does not exist.",
-#         )
-#     }
-#     # TODO: Add more exceptions
-# )
+@rewrap_exceptions(
+    {
+        asyncpg.exceptions.ForeignKeyViolationError: partialclass(
+            HTTPException,
+            status_code=404,
+            detail="The specified developer does not exist.",
+        ),
+        asyncpg.exceptions.DataError: partialclass(
+            HTTPException,
+            status_code=400,
+            detail="Invalid data provided. Please check the input values.",
+        ),
+    }
+)
 @wrap_in_class(Agent, transform=lambda d: {"id": d["agent_id"], **d})
 @pg_query
 @beartype

diff --git a/agents-api/agents_api/queries/agents/patch_agent.py b/agents-api/agents_api/queries/agents/patch_agent.py
@@ -7,12 +7,16 @@
 
 from beartype import beartype
 from sqlglot import parse_one
+from fastapi import HTTPException
+import asyncpg
 
 from ...autogen.openapi_model import PatchAgentRequest, ResourceUpdatedResponse
 from ...metrics.counters import increase_counter
 from ..utils import (
     pg_query,
     wrap_in_class,
+    rewrap_exceptions,
+    partialclass,
 )
 
 # Define the raw SQL query
@@ -44,16 +48,30 @@
 """).sql(pretty=True)
 
 
-# @rewrap_exceptions(
-#     {
-#         psycopg_errors.ForeignKeyViolation: partialclass(
-#             HTTPException,
-#             status_code=404,
-#             detail="The specified developer does not exist.",
-#         )
-#     }
-#     # TODO: Add more exceptions
-# )
+@rewrap_exceptions(
+    {
+        asyncpg.exceptions.ForeignKeyViolationError: partialclass(
+            HTTPException,
+            status_code=404,
+            detail="The specified developer does not exist.",
+        ),
+        asyncpg.exceptions.UniqueViolationError: partialclass(
+            HTTPException,
+            status_code=409,
+            detail="An agent with this canonical name already exists for this developer.",
+        ),
+        asyncpg.exceptions.CheckViolationError: partialclass(
+            HTTPException,
+            status_code=400,
+            detail="The provided data violates one or more constraints. Please check the input values.",
+        ),
+        asyncpg.exceptions.DataError: partialclass(
+            HTTPException,
+            status_code=400,
+            detail="Invalid data provided. Please check the input values.",
+        ),
+    }
+)
 @wrap_in_class(
     ResourceUpdatedResponse,
     one=True,