chore: bug fixes for file queries + added tests

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

@@ -19,19 +19,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
 )
@@ -40,7 +60,6 @@
 RETURNING developer_id, agent_id;
 """).sql(pretty=True)
 
-
 # @rewrap_exceptions(
 # @rewrap_exceptions(
 #     {

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

@@ -27,7 +27,7 @@
     description,
     mime_type,
     size,
-    hash,
+    hash
 )
 VALUES (
     $1, -- developer_id
@@ -36,34 +36,28 @@
     $4, -- description
     $5, -- mime_type
     $6, -- size
-    $7, -- hash
+    $7  -- hash
 )
 RETURNING *;
 """).sql(pretty=True)
 
-# Create user file association
-user_file_query = parse_one("""
-INSERT INTO user_files (
-    developer_id,
-    user_id,
-    file_id
-)
-VALUES ($1, $2, $3)
-ON CONFLICT (developer_id, user_id, file_id) DO NOTHING;  -- Uses primary key index
-""").sql(pretty=True)
-
-# Create agent file association
-agent_file_query = parse_one("""
-INSERT INTO agent_files (
-    developer_id,
-    agent_id,
-    file_id
+# 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
 )
-VALUES ($1, $2, $3)
-ON CONFLICT (developer_id, agent_id, file_id) DO NOTHING;  -- Uses primary key index
+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(
 #     {
@@ -90,6 +84,7 @@
     transform=lambda d: {
         **d,
         "id": d["file_id"],
+        "hash": d["hash"].hex(),
         "content": "DUMMY: NEED TO FETCH CONTENT FROM BLOB STORAGE",
     },
 )
@@ -121,8 +116,8 @@ async def create_file(
 
     # Calculate size and hash
     content_bytes = base64.b64decode(data.content)
-    data.size = len(content_bytes)
-    data.hash = hashlib.sha256(content_bytes).digest()
+    size = len(content_bytes)
+    hash_bytes = hashlib.sha256(content_bytes).digest()
 
     # Base file parameters
     file_params = [
@@ -131,21 +126,18 @@ async def create_file(
         data.name,
         data.description,
         data.mime_type,
-        data.size,
-        data.hash,
+        size,
+        hash_bytes,
     ]
 
     queries = []
 
-    # Create the file
+    # Create the file first
     queries.append((file_query, file_params))
 
-    # Create the association only if both owner_type and owner_id are provided
+    # Then create the association if owner info provided
     if owner_type and owner_id:
-        assoc_params = [developer_id, owner_id, file_id]
-        if owner_type == "user":
-            queries.append((user_file_query, assoc_params))
-        else:  # agent
-            queries.append((agent_file_query, assoc_params))
+        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

@@ -16,53 +16,40 @@
 from ...metrics.counters import increase_counter
 from ..utils import partialclass, pg_query, rewrap_exceptions, wrap_in_class
 
-# Simple query to delete file (when no associations exist)
+# 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 NOT EXISTS (
-    SELECT 1 
-    FROM user_files uf 
-    WHERE uf.file_id = $2
-    LIMIT 1
-)
-AND NOT EXISTS (
-    SELECT 1 
-    FROM agent_files af 
-    WHERE af.file_id = $2
-    LIMIT 1
-)
-RETURNING file_id;
-""").sql(pretty=True)
-
-# Query to delete owner's association
-delete_user_assoc_query = parse_one("""
-DELETE FROM user_files 
-WHERE developer_id = $1 
-AND file_id = $2 
-AND user_id = $3
-RETURNING file_id;
-""").sql(pretty=True)
-
-delete_agent_assoc_query = parse_one("""
-DELETE FROM agent_files 
-WHERE developer_id = $1 
-AND file_id = $2 
-AND agent_id = $3
+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",
-#         ),
-#     }
-# )
+@rewrap_exceptions(
+    {
+        asyncpg.NoDataFoundError: partialclass(
+            HTTPException,
+            status_code=404,
+            detail="File not found",
+        ),
+    }
+)
 @wrap_in_class(
     ResourceDeletedResponse,
     one=True,
@@ -77,46 +64,24 @@
 @beartype
 async def delete_file(
     *,
-    file_id: UUID,
     developer_id: UUID,
-    owner_id: UUID | None = None,
+    file_id: UUID,
     owner_type: Literal["user", "agent"] | None = None,
-) -> list[tuple[str, list] | tuple[str, list, str]]:
+    owner_id: UUID | None = None,
+) -> tuple[str, list]:
     """
-    Deletes a file and/or its association using simple, efficient queries.
-
-    If owner details provided:
-        1. Deletes the owner's association
-        2. Checks for remaining associations
-        3. Deletes file if no associations remain
-    If no owner details:
-        - Deletes file only if it has no associations
+    Deletes a file and its ownership records.
 
     Args:
-        file_id (UUID): The UUID of the file to be deleted.
-        developer_id (UUID): The UUID of the developer owning the file.
-        owner_id (UUID | None): Optional owner ID to verify ownership
-        owner_type (str | None): Optional owner type to verify ownership
+        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:
-        list[tuple[str, list] | tuple[str, list, str]]: List of SQL queries, their parameters, and fetch type
+        tuple[str, list]: SQL query and parameters
     """
-    queries = []
-
-    if owner_id and owner_type:
-        # Delete specific association
-        assoc_params = [developer_id, file_id, owner_id]
-        assoc_query = (
-            delete_user_assoc_query
-            if owner_type == "user"
-            else delete_agent_assoc_query
-        )
-        queries.append((assoc_query, assoc_params))
-
-        # If no associations, delete file
-        queries.append((delete_file_query, [developer_id, file_id]))
-    else:
-        # Try to delete file if it has no associations
-        queries.append((delete_file_query, [developer_id, file_id]))
-
-    return queries
+    return (
+        delete_file_query,
+        [developer_id, file_id, owner_type, owner_id],
+    )