fix(queries-user): major bug fixes, refactor and added init user sql …

…test
julep-ai · Dec 16, 2024 · afc51ab · afc51ab
1 parent ceca67a
commit afc51ab
Show file tree

Hide file tree

Showing 10 changed files with 467 additions and 175 deletions.
diff --git a/agents-api/agents_api/queries/users/__init__.py b/agents-api/agents_api/queries/users/__init__.py
@@ -9,20 +9,20 @@
 - Deleting users
 """
 
-from .create_or_update_user import create_or_update_user_query
+from .create_or_update_user import create_or_update_user
 from .create_user import create_user
-from .delete_user import delete_user_query
-from .get_user import get_user_query
-from .list_users import list_users_query
-from .patch_user import patch_user_query
-from .update_user import update_user_query
+from .get_user import get_user
+from .list_users import list_users
+from .patch_user import patch_user
+from .update_user import update_user
+from .delete_user import delete_user
 
 __all__ = [
     "create_user",
-    "create_or_update_user_query",
-    "delete_user_query",
-    "get_user_query",
-    "list_users_query",
-    "patch_user_query",
-    "update_user_query",
+    "create_or_update_user",
+    "delete_user",
+    "get_user",
+    "list_users",
+    "patch_user",
+    "update_user",
 ]
diff --git a/agents-api/agents_api/queries/users/create_or_update_user.py b/agents-api/agents_api/queries/users/create_or_update_user.py
@@ -1,30 +1,72 @@
-from typing import Any
 from uuid import UUID
 
-from asyncpg import exceptions as asyncpg_exceptions
+import asyncpg
 from beartype import beartype
 from fastapi import HTTPException
 from sqlglot import parse_one
+from sqlglot.optimizer import optimize
 
 from ...autogen.openapi_model import CreateUserRequest, User
 from ...metrics.counters import increase_counter
 from ..utils import partialclass, pg_query, rewrap_exceptions, wrap_in_class
 
+# Optimize the raw query by using COALESCE for metadata to avoid explicit check
+raw_query = """
+INSERT INTO users (
+    developer_id,
+    user_id,
+    name,
+    about,
+    metadata
+)
+VALUES (
+    %(developer_id)s,
+    %(user_id)s,
+    %(name)s,
+    %(about)s,
+    COALESCE(%(metadata)s, '{}'::jsonb)
+)
+ON CONFLICT (developer_id, user_id) DO UPDATE SET
+    name = EXCLUDED.name,
+    about = EXCLUDED.about,
+    metadata = EXCLUDED.metadata
+RETURNING *;
+"""
+
+# Add index hint for better performance
+query = optimize(
+    parse_one(raw_query),
+    schema={
+        "users": {
+            "developer_id": "UUID",
+            "user_id": "UUID",
+            "name": "STRING",
+            "about": "STRING",
+            "metadata": "JSONB",
+        }
+    },
+).sql(pretty=True)
+
 
 @rewrap_exceptions(
     {
-        asyncpg_exceptions.ForeignKeyViolationError: partialclass(
+        asyncpg.ForeignKeyViolationError: partialclass(
             HTTPException,
             status_code=404,
             detail="The specified developer does not exist.",
-        )
+        ),
+        asyncpg.UniqueViolationError: partialclass(  # Add handling for potential race conditions
+            HTTPException,
+            status_code=409,
+            detail="A user with this ID already exists.",
+        ),
     }
 )
 @wrap_in_class(User)
 @increase_counter("create_or_update_user")
 @pg_query
 @beartype
-def create_or_update_user_query(
+def create_or_update_user(
     *, developer_id: UUID, user_id: UUID, data: CreateUserRequest
 ) -> tuple[str, dict]:
     """
@@ -37,35 +79,16 @@ def create_or_update_user_query(
 
     Returns:
         tuple[str, dict]: SQL query and parameters.
-    """
-    query = parse_one("""
-    INSERT INTO users (
-        developer_id,
-        user_id,
-        name,
-        about,
-        metadata
-    )
-    VALUES (
-        %(developer_id)s,
-        %(user_id)s,
-        %(name)s,
-        %(about)s,
-        %(metadata)s
-    )
-    ON CONFLICT (developer_id, user_id) DO UPDATE SET
-        name = EXCLUDED.name,
-        about = EXCLUDED.about,
-        metadata = EXCLUDED.metadata
-    RETURNING *;
-    """).sql()
 
+    Raises:
+        HTTPException: If developer doesn't exist (404) or on unique constraint violation (409)
+    """
     params = {
         "developer_id": developer_id,
         "user_id": user_id,
         "name": data.name,
         "about": data.about,
-        "metadata": data.metadata or {},
+        "metadata": data.metadata,  # Let COALESCE handle None case in SQL
     }
 
     return query, params
diff --git a/agents-api/agents_api/queries/users/create_user.py b/agents-api/agents_api/queries/users/create_user.py
@@ -1,29 +1,60 @@
-from typing import Any
 from uuid import UUID
 
+import asyncpg
 from beartype import beartype
 from fastapi import HTTPException
-from psycopg import errors as psycopg_errors
-from pydantic import ValidationError
-from sqlglot import parse_one
+from sqlglot import optimize, parse_one
 from uuid_extensions import uuid7
 
 from ...autogen.openapi_model import CreateUserRequest, User
 from ...metrics.counters import increase_counter
 from ..utils import partialclass, pg_query, rewrap_exceptions, wrap_in_class
 
+# Define the raw SQL query outside the function
+raw_query = """
+INSERT INTO users (
+    developer_id,
+    user_id,
+    name,
+    about,
+    metadata
+)
+VALUES (
+    %(developer_id)s,
+    %(user_id)s,
+    %(name)s,
+    %(about)s,
+    %(metadata)s
+)
+RETURNING *;
+"""
+
+# Parse and optimize the query
+query = optimize(
+    parse_one(raw_query),
+    schema={
+        "users": {
+            "developer_id": "UUID",
+            "user_id": "UUID",
+            "name": "STRING",
+            "about": "STRING",
+            "metadata": "JSONB",
+        }
+    },
+).sql(pretty=True)
+
 
 @rewrap_exceptions(
     {
-        psycopg_errors.ForeignKeyViolation: partialclass(
+        asyncpg.ForeignKeyViolationError: partialclass(
             HTTPException,
             status_code=404,
             detail="The specified developer does not exist.",
         ),
-        ValidationError: partialclass(
+        asyncpg.NullValueNoIndicatorParameterError: partialclass(
             HTTPException,
-            status_code=400,
-            detail="Input validation failed. Please check the provided data.",
+            status_code=404,
+            detail="The specified developer does not exist.",
         ),
     }
 )
@@ -50,24 +81,6 @@ def create_user(
     """
     user_id = user_id or uuid7()
 
-    query = parse_one("""
-    INSERT INTO users (
-        developer_id,
-        user_id,
-        name,
-        about,
-        metadata
-    )
-    VALUES (
-        %(developer_id)s,
-        %(user_id)s,
-        %(name)s,
-        %(about)s,
-        %(metadata)s
-    )
-    RETURNING *;
-    """).sql()
-
     params = {
         "developer_id": developer_id,
         "user_id": user_id,

diff --git a/agents-api/agents_api/queries/users/delete_user.py b/agents-api/agents_api/queries/users/delete_user.py
@@ -1,19 +1,44 @@
-from typing import Any
 from uuid import UUID
 
+import asyncpg
 from beartype import beartype
 from fastapi import HTTPException
-from psycopg import errors as psycopg_errors
 from sqlglot import parse_one
+from sqlglot.optimizer import optimize
 
 from ...autogen.openapi_model import ResourceDeletedResponse
 from ...metrics.counters import increase_counter
 from ..utils import partialclass, pg_query, rewrap_exceptions, wrap_in_class
 
+# Define the raw SQL query outside the function
+raw_query = """
+WITH deleted_data AS (
+    DELETE FROM user_files 
+    WHERE developer_id = %(developer_id)s AND user_id = %(user_id)s
+),
+deleted_docs AS (
+    DELETE FROM user_docs 
+    WHERE developer_id = %(developer_id)s AND user_id = %(user_id)s
+)
+DELETE FROM users 
+WHERE developer_id = %(developer_id)s AND user_id = %(user_id)s
+RETURNING user_id as id, developer_id;
+"""
+
+# Parse and optimize the query
+query = optimize(
+    parse_one(raw_query),
+    schema={
+        "user_files": {"developer_id": "UUID", "user_id": "UUID"},
+        "user_docs": {"developer_id": "UUID", "user_id": "UUID"},
+        "users": {"developer_id": "UUID", "user_id": "UUID"},
+    },
+).sql(pretty=True)
+
 
 @rewrap_exceptions(
     {
-        psycopg_errors.ForeignKeyViolation: partialclass(
+        asyncpg.ForeignKeyViolationError: partialclass(
             HTTPException,
             status_code=404,
             detail="The specified developer does not exist.",
@@ -24,25 +49,17 @@
 @increase_counter("delete_user")
 @pg_query
 @beartype
-def delete_user_query(*, developer_id: UUID, user_id: UUID) -> tuple[list[str], dict]:
+def delete_user(*, developer_id: UUID, user_id: UUID) -> tuple[str, dict]:
     """
-    Constructs optimized SQL queries to delete a user and related data.
+    Constructs optimized SQL query to delete a user and related data.
     Uses primary key for efficient deletion.
 
     Args:
         developer_id (UUID): The developer's UUID
         user_id (UUID): The user's UUID
 
     Returns:
-        tuple[list[str], dict]: List of SQL queries and parameters
+        tuple[str, dict]: SQL query and parameters
     """
-    query = parse_one("""
-    BEGIN;
-    DELETE FROM user_files WHERE developer_id = %(developer_id)s AND user_id = %(user_id)s;
-    DELETE FROM user_docs WHERE developer_id = %(developer_id)s AND user_id = %(user_id)s;
-    DELETE FROM users WHERE developer_id = %(developer_id)s AND user_id = %(user_id)s
-    RETURNING user_id as id, developer_id;
-    COMMIT;
-    """).sql()
-
-    return [query], {"developer_id": developer_id, "user_id": user_id}
+
+    return query, {"developer_id": developer_id, "user_id": user_id}