From f4e6b4861857514c60c406b1414334d925ef8dcb Mon Sep 17 00:00:00 2001 From: vedantsahai18 Date: Sun, 15 Dec 2024 00:52:53 -0500 Subject: [PATCH] feat(agents-api): add user queries --- .../agents_api/queries/users/__init__.py | 28 +++++++ .../queries/users/create_or_update_user.py | 72 +++++++++++++++++ .../agents_api/queries/users/create_user.py | 76 +++++++++++++++++ .../agents_api/queries/users/delete_user.py | 45 +++++++++++ .../agents_api/queries/users/get_user.py | 50 ++++++++++++ .../agents_api/queries/users/list_users.py | 81 +++++++++++++++++++ .../agents_api/queries/users/patch_user.py | 73 +++++++++++++++++ .../agents_api/queries/users/update_user.py | 68 ++++++++++++++++ 8 files changed, 493 insertions(+) create mode 100644 agents-api/agents_api/queries/users/__init__.py create mode 100644 agents-api/agents_api/queries/users/create_or_update_user.py create mode 100644 agents-api/agents_api/queries/users/create_user.py create mode 100644 agents-api/agents_api/queries/users/delete_user.py create mode 100644 agents-api/agents_api/queries/users/get_user.py create mode 100644 agents-api/agents_api/queries/users/list_users.py create mode 100644 agents-api/agents_api/queries/users/patch_user.py create mode 100644 agents-api/agents_api/queries/users/update_user.py diff --git a/agents-api/agents_api/queries/users/__init__.py b/agents-api/agents_api/queries/users/__init__.py new file mode 100644 index 000000000..4e810a4fe --- /dev/null +++ b/agents-api/agents_api/queries/users/__init__.py @@ -0,0 +1,28 @@ +""" +The `user` module within the `queries` package provides SQL query functions for managing users +in the TimescaleDB database. This includes operations for: + +- Creating new users +- Updating existing users +- Retrieving user details +- Listing users with filtering and pagination +- Deleting users +""" + +from .create_user import create_user +from .create_or_update_user import create_or_update_user_query +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 + +__all__ = [ + "create_user", + "create_or_update_user_query", + "delete_user_query", + "get_user_query", + "list_users_query", + "patch_user_query", + "update_user_query", +] \ No newline at end of file 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 new file mode 100644 index 000000000..a6312b243 --- /dev/null +++ b/agents-api/agents_api/queries/users/create_or_update_user.py @@ -0,0 +1,72 @@ +from typing import Any +from uuid import UUID + +from beartype import beartype +from fastapi import HTTPException +from asyncpg import exceptions as asyncpg_exceptions +from sqlglot import parse_one + +from ...autogen.openapi_model import CreateUserRequest, User +from ...metrics.counters import increase_counter +from ..utils import partialclass, pg_query, rewrap_exceptions, wrap_in_class + + +@rewrap_exceptions({ + asyncpg_exceptions.ForeignKeyViolationError: partialclass( + HTTPException, + status_code=404, + detail="The specified developer does not exist.", + ) +}) +@wrap_in_class(User) +@increase_counter("create_or_update_user") +@pg_query +@beartype +def create_or_update_user_query( + *, + developer_id: UUID, + user_id: UUID, + data: CreateUserRequest +) -> tuple[str, dict]: + """ + Constructs an SQL query to create or update a user. + + Args: + developer_id (UUID): The UUID of the developer. + user_id (UUID): The UUID of the user. + data (CreateUserRequest): The user data to insert or update. + + 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() + + params = { + "developer_id": developer_id, + "user_id": user_id, + "name": data.name, + "about": data.about, + "metadata": data.metadata or {}, + } + + 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 new file mode 100644 index 000000000..194d9bf03 --- /dev/null +++ b/agents-api/agents_api/queries/users/create_user.py @@ -0,0 +1,76 @@ +from typing import Any +from uuid import UUID + +from beartype import beartype +from fastapi import HTTPException +from psycopg import errors as psycopg_errors +from sqlglot import parse_one +from pydantic import ValidationError +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 + +@rewrap_exceptions({ + psycopg_errors.ForeignKeyViolation: partialclass( + HTTPException, + status_code=404, + detail="The specified developer does not exist.", + ), + ValidationError: partialclass( + HTTPException, + status_code=400, + detail="Input validation failed. Please check the provided data.", + ), +}) +@wrap_in_class(User) +@increase_counter("create_user") +@pg_query +@beartype +def create_user( + *, + developer_id: UUID, + user_id: UUID | None = None, + data: CreateUserRequest, +) -> tuple[str, dict]: + """ + Constructs the SQL query to create a new user. + + Args: + developer_id (UUID): The UUID of the developer creating the user. + user_id (UUID, optional): The UUID for the new user. If None, one will be generated. + data (CreateUserRequest): The user data to insert. + + Returns: + tuple[str, dict]: A tuple containing the SQL query and its parameters. + """ + 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, + "name": data.name, + "about": data.about, + "metadata": data.metadata or {}, + } + + return query, params diff --git a/agents-api/agents_api/queries/users/delete_user.py b/agents-api/agents_api/queries/users/delete_user.py new file mode 100644 index 000000000..551129f00 --- /dev/null +++ b/agents-api/agents_api/queries/users/delete_user.py @@ -0,0 +1,45 @@ +from typing import Any +from uuid import UUID + +from beartype import beartype +from fastapi import HTTPException +from psycopg import errors as psycopg_errors +from sqlglot import parse_one + +from ...autogen.openapi_model import ResourceDeletedResponse +from ...metrics.counters import increase_counter +from ..utils import partialclass, pg_query, rewrap_exceptions, wrap_in_class + +@rewrap_exceptions({ + psycopg_errors.ForeignKeyViolation: partialclass( + HTTPException, + status_code=404, + detail="The specified developer does not exist.", + ) +}) +@wrap_in_class(ResourceDeletedResponse, one=True) +@increase_counter("delete_user") +@pg_query +@beartype +def delete_user_query(*, developer_id: UUID, user_id: UUID) -> tuple[list[str], dict]: + """ + Constructs optimized SQL queries 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 + """ + 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} \ No newline at end of file diff --git a/agents-api/agents_api/queries/users/get_user.py b/agents-api/agents_api/queries/users/get_user.py new file mode 100644 index 000000000..3982ea46e --- /dev/null +++ b/agents-api/agents_api/queries/users/get_user.py @@ -0,0 +1,50 @@ +from typing import Any +from uuid import UUID + +from beartype import beartype +from fastapi import HTTPException +from psycopg import errors as psycopg_errors +from sqlglot import parse_one + +from ...autogen.openapi_model import User +from ...metrics.counters import increase_counter +from ..utils import partialclass, pg_query, rewrap_exceptions, wrap_in_class + +@rewrap_exceptions({ + psycopg_errors.ForeignKeyViolation: partialclass( + HTTPException, + status_code=404, + detail="The specified developer does not exist.", + ) +}) +@wrap_in_class(User, one=True) +@increase_counter("get_user") +@pg_query +@beartype +def get_user_query(*, developer_id: UUID, user_id: UUID) -> tuple[str, dict]: + """ + Constructs an optimized SQL query to retrieve a user's details. + Uses the primary key index (developer_id, user_id) for efficient lookup. + + Args: + developer_id (UUID): The UUID of the developer. + user_id (UUID): The UUID of the user to retrieve. + + Returns: + tuple[str, dict]: SQL query and parameters. + """ + query = parse_one(""" + SELECT + user_id as id, + developer_id, + name, + about, + metadata, + created_at, + updated_at + FROM users + WHERE developer_id = %(developer_id)s + AND user_id = %(user_id)s; + """).sql() + + return query, {"developer_id": developer_id, "user_id": user_id} \ No newline at end of file diff --git a/agents-api/agents_api/queries/users/list_users.py b/agents-api/agents_api/queries/users/list_users.py new file mode 100644 index 000000000..312299082 --- /dev/null +++ b/agents-api/agents_api/queries/users/list_users.py @@ -0,0 +1,81 @@ +from typing import Any, Literal +from uuid import UUID + +from beartype import beartype +from fastapi import HTTPException +from psycopg import errors as psycopg_errors +from sqlglot import parse_one + +from ...autogen.openapi_model import User +from ...metrics.counters import increase_counter +from ..utils import partialclass, pg_query, rewrap_exceptions, wrap_in_class + +@rewrap_exceptions({ + psycopg_errors.ForeignKeyViolation: partialclass( + HTTPException, + status_code=404, + detail="The specified developer does not exist.", + ) +}) +@wrap_in_class(User) +@increase_counter("list_users") +@pg_query +@beartype +def list_users_query( + *, + developer_id: UUID, + limit: int = 100, + offset: int = 0, + sort_by: Literal["created_at", "updated_at"] = "created_at", + direction: Literal["asc", "desc"] = "desc", + metadata_filter: dict | None = None, +) -> tuple[str, dict]: + """ + Constructs an optimized SQL query for listing users with pagination and filtering. + Uses indexes on developer_id and metadata for efficient querying. + + Args: + developer_id (UUID): The developer's UUID + limit (int): Maximum number of records to return + offset (int): Number of records to skip + sort_by (str): Field to sort by + direction (str): Sort direction + metadata_filter (dict, optional): Metadata-based filters + + Returns: + tuple[str, dict]: SQL query and parameters + """ + if limit < 1 or limit > 1000: + raise HTTPException(status_code=400, detail="Limit must be between 1 and 1000") + if offset < 0: + raise HTTPException(status_code=400, detail="Offset must be non-negative") + + metadata_clause = "" + params = { + "developer_id": developer_id, + "limit": limit, + "offset": offset + } + + if metadata_filter: + metadata_clause = "AND metadata @> %(metadata_filter)s" + params["metadata_filter"] = metadata_filter + + query = parse_one(f""" + SELECT + user_id as id, + developer_id, + name, + about, + metadata, + created_at, + updated_at + FROM users + WHERE developer_id = %(developer_id)s + {metadata_clause} + ORDER BY {sort_by} {direction} + LIMIT %(limit)s + OFFSET %(offset)s; + """).sql() + + return query, params \ No newline at end of file diff --git a/agents-api/agents_api/queries/users/patch_user.py b/agents-api/agents_api/queries/users/patch_user.py new file mode 100644 index 000000000..468b38b00 --- /dev/null +++ b/agents-api/agents_api/queries/users/patch_user.py @@ -0,0 +1,73 @@ +from typing import Any +from uuid import UUID + +from beartype import beartype +from fastapi import HTTPException +from psycopg import errors as psycopg_errors +from sqlglot import parse_one + +from ...autogen.openapi_model import PatchUserRequest, ResourceUpdatedResponse +from ...metrics.counters import increase_counter +from ..utils import partialclass, pg_query, rewrap_exceptions, wrap_in_class + +@rewrap_exceptions({ + psycopg_errors.ForeignKeyViolation: partialclass( + HTTPException, + status_code=404, + detail="The specified developer does not exist.", + ) +}) +@wrap_in_class(ResourceUpdatedResponse, one=True) +@increase_counter("patch_user") +@pg_query +@beartype +def patch_user_query( + *, + developer_id: UUID, + user_id: UUID, + data: PatchUserRequest +) -> tuple[str, dict]: + """ + Constructs an optimized SQL query for partial user updates. + Uses primary key for efficient update and jsonb_merge for metadata. + + Args: + developer_id (UUID): The developer's UUID + user_id (UUID): The user's UUID + data (PatchUserRequest): Partial update data + + Returns: + tuple[str, dict]: SQL query and parameters + """ + update_parts = [] + params = { + "developer_id": developer_id, + "user_id": user_id, + } + + if data.name is not None: + update_parts.append("name = %(name)s") + params["name"] = data.name + if data.about is not None: + update_parts.append("about = %(about)s") + params["about"] = data.about + if data.metadata is not None: + update_parts.append("metadata = metadata || %(metadata)s") + params["metadata"] = data.metadata + + query = parse_one(f""" + UPDATE users + SET {", ".join(update_parts)} + WHERE developer_id = %(developer_id)s + AND user_id = %(user_id)s + RETURNING + user_id as id, + developer_id, + name, + about, + metadata, + created_at, + updated_at; + """).sql() + + return query, params \ No newline at end of file diff --git a/agents-api/agents_api/queries/users/update_user.py b/agents-api/agents_api/queries/users/update_user.py new file mode 100644 index 000000000..ed33e3e42 --- /dev/null +++ b/agents-api/agents_api/queries/users/update_user.py @@ -0,0 +1,68 @@ +from typing import Any +from uuid import UUID + +from beartype import beartype +from fastapi import HTTPException +from psycopg import errors as psycopg_errors +from sqlglot import parse_one + +from ...autogen.openapi_model import UpdateUserRequest, ResourceUpdatedResponse +from ...metrics.counters import increase_counter +from ..utils import partialclass, pg_query, rewrap_exceptions, wrap_in_class + +@rewrap_exceptions({ + psycopg_errors.ForeignKeyViolation: partialclass( + HTTPException, + status_code=404, + detail="The specified developer does not exist.", + ) +}) +@wrap_in_class(ResourceUpdatedResponse, one=True) +@increase_counter("update_user") +@pg_query +@beartype +def update_user_query( + *, + developer_id: UUID, + user_id: UUID, + data: UpdateUserRequest +) -> tuple[str, dict]: + """ + Constructs an optimized SQL query to update a user's details. + Uses primary key for efficient update. + + Args: + developer_id (UUID): The developer's UUID + user_id (UUID): The user's UUID + data (UpdateUserRequest): Updated user data + + Returns: + tuple[str, dict]: SQL query and parameters + """ + query = parse_one(""" + UPDATE users + SET + name = %(name)s, + about = %(about)s, + metadata = %(metadata)s + WHERE developer_id = %(developer_id)s + AND user_id = %(user_id)s + RETURNING + user_id as id, + developer_id, + name, + about, + metadata, + created_at, + updated_at; + """).sql() + + params = { + "developer_id": developer_id, + "user_id": user_id, + "name": data.name, + "about": data.about, + "metadata": data.metadata or {}, + } + + return query, params \ No newline at end of file