Merge pull request #974 from julep-ai/f/file-queries

F/file queries: Add file sql queries

agents-api/agents_api/queries/agents/delete_agent.py

@@ -17,19 +17,39 @@
 
 # Define the raw SQL query
 agent_query = parse_one("""
-WITH deleted_docs AS (
+WITH deleted_file_owners AS (
+    DELETE FROM file_owners
+    WHERE developer_id = $1 
+    AND owner_type = 'agent'
+    AND owner_id = $2
+),
+deleted_doc_owners AS (
+    DELETE FROM doc_owners
+    WHERE developer_id = $1
+    AND owner_type = 'agent'
+    AND owner_id = $2
+),
+deleted_files AS (
+    DELETE FROM files
+    WHERE developer_id = $1
+    AND file_id IN (
+        SELECT file_id FROM file_owners 
+        WHERE developer_id = $1 
+        AND owner_type = 'agent' 
+        AND owner_id = $2
+    )
+),
+deleted_docs AS (
     DELETE FROM docs
     WHERE developer_id = $1
     AND doc_id IN (
-        SELECT ad.doc_id
-        FROM agent_docs ad
-        WHERE ad.agent_id = $2
-        AND ad.developer_id = $1
+        SELECT doc_id FROM doc_owners
+        WHERE developer_id = $1 
+        AND owner_type = 'agent' 
+        AND owner_id = $2
     )
-), deleted_agent_docs AS (
-    DELETE FROM agent_docs
-    WHERE agent_id = $2 AND developer_id = $1
-), deleted_tools AS (
+),
+deleted_tools AS (
     DELETE FROM tools
     WHERE agent_id = $2 AND developer_id = $1
 )

agents-api/agents_api/queries/files/__init__.py

@@ -0,0 +1,16 @@
+"""
+The `files` module within the `queries` package provides SQL query functions for managing files
+in the PostgreSQL database. This includes operations for:
+
+- Creating new files
+- Retrieving file details
+- Listing files with filtering and pagination
+- Deleting files and their associations
+"""
+
+from .create_file import create_file
+from .delete_file import delete_file
+from .get_file import get_file
+from .list_files import list_files
+
+__all__ = ["create_file", "delete_file", "get_file", "list_files"]

agents-api/agents_api/queries/files/create_file.py

@@ -0,0 +1,144 @@
+"""
+This module contains the functionality for creating files in the PostgreSQL database.
+It includes functions to construct and execute SQL queries for inserting new file records.
+"""
+
+import base64
+import hashlib
+from typing import Any, Literal
+from uuid import UUID
+
+import asyncpg
+from beartype import beartype
+from fastapi import HTTPException
+from sqlglot import parse_one
+from uuid_extensions import uuid7
+
+from ...autogen.openapi_model import CreateFileRequest, File
+from ...metrics.counters import increase_counter
+from ..utils import partialclass, pg_query, rewrap_exceptions, wrap_in_class
+
+# Create file
+file_query = parse_one("""
+INSERT INTO files (
+    developer_id,
+    file_id,
+    name,
+    description,
+    mime_type,
+    size,
+    hash
+)
+VALUES (
+    $1, -- developer_id
+    $2, -- file_id
+    $3, -- name
+    $4, -- description
+    $5, -- mime_type
+    $6, -- size
+    $7  -- hash
+)
+RETURNING *;
+""").sql(pretty=True)
+
+# Replace both user_file and agent_file queries with a single file_owner query
+file_owner_query = parse_one("""
+WITH inserted_owner AS (
+    INSERT INTO file_owners (
+        developer_id,
+        file_id,
+        owner_type,
+        owner_id
+    )
+    VALUES ($1, $2, $3, $4)
+    RETURNING file_id
+)
+SELECT f.*
+FROM inserted_owner io
+JOIN files f ON f.file_id = io.file_id;
+""").sql(pretty=True)
+
+
+# Add error handling decorator
+# @rewrap_exceptions(
+#     {
+#         asyncpg.UniqueViolationError: partialclass(
+#             HTTPException,
+#             status_code=409,
+#             detail="A file with this name already exists for this developer",
+#         ),
+#         asyncpg.NoDataFoundError: partialclass(
+#             HTTPException,
+#             status_code=404,
+#             detail="The specified owner does not exist",
+#         ),
+#         asyncpg.ForeignKeyViolationError: partialclass(
+#             HTTPException,
+#             status_code=404,
+#             detail="The specified developer does not exist",
+#         ),
+#     }
+# )
+@wrap_in_class(
+    File,
+    one=True,
+    transform=lambda d: {
+        **d,
+        "id": d["file_id"],
+        "hash": d["hash"].hex(),
+        "content": "DUMMY: NEED TO FETCH CONTENT FROM BLOB STORAGE",
+    },
+)
+@increase_counter("create_file")
+@pg_query
+@beartype
+async def create_file(
+    *,
+    developer_id: UUID,
+    file_id: UUID | None = None,
+    data: CreateFileRequest,
+    owner_type: Literal["user", "agent"] | None = None,
+    owner_id: UUID | None = None,
+) -> list[tuple[str, list] | tuple[str, list, str]]:
+    """
+    Constructs and executes SQL queries to create a new file and optionally associate it with an owner.
+
+    Parameters:
+        developer_id (UUID): The unique identifier for the developer.
+        file_id (UUID | None): Optional unique identifier for the file.
+        data (CreateFileRequest): The file data to insert.
+        owner_type (Literal["user", "agent"] | None): Optional type of owner
+        owner_id (UUID | None): Optional ID of the owner
+
+    Returns:
+        list[tuple[str, list] | tuple[str, list, str]]: List of SQL queries, their parameters, and fetch type
+    """
+    file_id = file_id or uuid7()
+
+    # Calculate size and hash
+    content_bytes = base64.b64decode(data.content)
+    size = len(content_bytes)
+    hash_bytes = hashlib.sha256(content_bytes).digest()
+
+    # Base file parameters
+    file_params = [
+        developer_id,
+        file_id,
+        data.name,
+        data.description,
+        data.mime_type,
+        size,
+        hash_bytes,
+    ]
+
+    queries = []
+
+    # Create the file first
+    queries.append((file_query, file_params))
+
+    # Then create the association if owner info provided
+    if owner_type and owner_id:
+        assoc_params = [developer_id, file_id, owner_type, owner_id]
+        queries.append((file_owner_query, assoc_params))
+
+    return queries

agents-api/agents_api/queries/files/delete_file.py

@@ -0,0 +1,87 @@
+"""
+This module contains the functionality for deleting files from the PostgreSQL database.
+It constructs and executes SQL queries to remove file records and associated data.
+"""
+
+from typing import Literal
+from uuid import UUID
+
+import asyncpg
+from beartype import beartype
+from fastapi import HTTPException
+from sqlglot import parse_one
+
+from ...autogen.openapi_model import ResourceDeletedResponse
+from ...common.utils.datetime import utcnow
+from ...metrics.counters import increase_counter
+from ..utils import partialclass, pg_query, rewrap_exceptions, wrap_in_class
+
+# Delete file query with ownership check
+delete_file_query = parse_one("""
+WITH deleted_owners AS (
+    DELETE FROM file_owners
+    WHERE developer_id = $1 
+    AND file_id = $2
+    AND (
+        ($3::text IS NULL AND $4::uuid IS NULL) OR
+        (owner_type = $3 AND owner_id = $4)
+    )
+)
+DELETE FROM files 
+WHERE developer_id = $1 
+AND file_id = $2
+AND ($3::text IS NULL OR EXISTS (
+    SELECT 1 FROM file_owners 
+    WHERE developer_id = $1 
+    AND file_id = $2 
+    AND owner_type = $3 
+    AND owner_id = $4
+))
+RETURNING file_id;
+""").sql(pretty=True)
+
+
+@rewrap_exceptions(
+    {
+        asyncpg.NoDataFoundError: partialclass(
+            HTTPException,
+            status_code=404,
+            detail="File not found",
+        ),
+    }
+)
+@wrap_in_class(
+    ResourceDeletedResponse,
+    one=True,
+    transform=lambda d: {
+        "id": d["file_id"],
+        "deleted_at": utcnow(),
+        "jobs": [],
+    },
+)
+@increase_counter("delete_file")
+@pg_query
+@beartype
+async def delete_file(
+    *,
+    developer_id: UUID,
+    file_id: UUID,
+    owner_type: Literal["user", "agent"] | None = None,
+    owner_id: UUID | None = None,
+) -> tuple[str, list]:
+    """
+    Deletes a file and its ownership records.
+
+    Args:
+        developer_id: The developer's UUID
+        file_id: The file's UUID
+        owner_type: Optional type of owner ("user" or "agent")
+        owner_id: Optional UUID of the owner
+
+    Returns:
+        tuple[str, list]: SQL query and parameters
+    """
+    return (
+        delete_file_query,
+        [developer_id, file_id, owner_type, owner_id],
+    )

agents-api/agents_api/queries/files/get_file.py

@@ -0,0 +1,81 @@
+"""
+This module contains the functionality for retrieving a single file from the PostgreSQL database.
+It constructs and executes SQL queries to fetch file details based on file ID and developer ID.
+"""
+
+from typing import Literal
+from uuid import UUID
+
+import asyncpg
+from beartype import beartype
+from fastapi import HTTPException
+from sqlglot import parse_one
+
+from ...autogen.openapi_model import File
+from ..utils import partialclass, pg_query, rewrap_exceptions, wrap_in_class
+
+# Define the raw SQL query
+file_query = parse_one("""
+SELECT f.*
+FROM files f
+LEFT JOIN file_owners fo ON f.developer_id = fo.developer_id AND f.file_id = fo.file_id
+WHERE f.developer_id = $1
+AND f.file_id = $2
+AND (
+    ($3::text IS NULL AND $4::uuid IS NULL) OR
+    (fo.owner_type = $3 AND fo.owner_id = $4)
+)
+LIMIT 1;
+""").sql(pretty=True)
+
+
+# @rewrap_exceptions(
+#     {
+#         asyncpg.NoDataFoundError: partialclass(
+#             HTTPException,
+#             status_code=404,
+#             detail="File not found",
+#         ),
+#         asyncpg.ForeignKeyViolationError: partialclass(
+#             HTTPException,
+#             status_code=404,
+#             detail="Developer not found",
+#         ),
+#     }
+# )
+@wrap_in_class(
+    File,
+    one=True,
+    transform=lambda d: {
+        "id": d["file_id"],
+        **d,
+        "hash": d["hash"].hex(),
+        "content": "DUMMY: NEED TO FETCH CONTENT FROM BLOB STORAGE",
+    },
+)
+@pg_query
+@beartype
+async def get_file(
+    *,
+    file_id: UUID,
+    developer_id: UUID,
+    owner_type: Literal["user", "agent"] | None = None,
+    owner_id: UUID | None = None,
+) -> tuple[str, list]:
+    """
+    Constructs the SQL query to retrieve a file's details.
+    Uses composite index on (developer_id, file_id) for efficient lookup.
+
+    Args:
+        file_id: The UUID of the file to retrieve
+        developer_id: The UUID of the developer owning the file
+        owner_type: Optional type of owner ("user" or "agent")
+        owner_id: Optional UUID of the owner
+
+    Returns:
+        tuple[str, list]: SQL query and parameters
+    """
+    return (
+        file_query,
+        [developer_id, file_id, owner_type, owner_id],
+    )