Merge pull request #973 from julep-ai/x/doc-file-owner-migration

fix(memory-store): Change association structure of files and docs

memory-store/migrations/000005_files.down.sql

@@ -1,14 +1,12 @@
 BEGIN;
 
--- Drop agent_files table and its dependencies
-DROP TABLE IF EXISTS agent_files;
-
--- Drop user_files table and its dependencies
-DROP TABLE IF EXISTS user_files;
+-- Drop file_owners table and its dependencies
+DROP TRIGGER IF EXISTS trg_validate_file_owner ON file_owners;
+DROP FUNCTION IF EXISTS validate_file_owner();
+DROP TABLE IF EXISTS file_owners;
 
 -- Drop files table and its dependencies
 DROP TRIGGER IF EXISTS trg_files_updated_at ON files;
-
 DROP TABLE IF EXISTS files;
 
 COMMIT;

memory-store/migrations/000005_files.up.sql

@@ -56,30 +56,48 @@ DO $$ BEGIN
     END IF;
 END $$;
 
--- Create the user_files table
-CREATE TABLE IF NOT EXISTS user_files (
+-- Create the file_owners table
+CREATE TABLE IF NOT EXISTS file_owners (
     developer_id UUID NOT NULL,
-    user_id UUID NOT NULL,
     file_id UUID NOT NULL,
-    CONSTRAINT pk_user_files PRIMARY KEY (developer_id, user_id, file_id),
-    CONSTRAINT fk_user_files_user FOREIGN KEY (developer_id, user_id) REFERENCES users (developer_id, user_id),
-    CONSTRAINT fk_user_files_file FOREIGN KEY (developer_id, file_id) REFERENCES files (developer_id, file_id)
+    owner_type TEXT NOT NULL,  -- 'user' or 'agent'
+    owner_id UUID NOT NULL,
+    CONSTRAINT pk_file_owners PRIMARY KEY (developer_id, file_id),
+    CONSTRAINT fk_file_owners_file FOREIGN KEY (developer_id, file_id) REFERENCES files (developer_id, file_id),
+    CONSTRAINT ct_file_owners_owner_type CHECK (owner_type IN ('user', 'agent'))
 );
 
--- Create index if it doesn't exist
-CREATE INDEX IF NOT EXISTS idx_user_files_user ON user_files (developer_id, user_id);
+-- Create indexes
+CREATE INDEX IF NOT EXISTS idx_file_owners_owner 
+    ON file_owners (developer_id, owner_type, owner_id);
 
--- Create the agent_files table
-CREATE TABLE IF NOT EXISTS agent_files (
-    developer_id UUID NOT NULL,
-    agent_id UUID NOT NULL,
-    file_id UUID NOT NULL,
-    CONSTRAINT pk_agent_files PRIMARY KEY (developer_id, agent_id, file_id),
-    CONSTRAINT fk_agent_files_agent FOREIGN KEY (developer_id, agent_id) REFERENCES agents (developer_id, agent_id),
-    CONSTRAINT fk_agent_files_file FOREIGN KEY (developer_id, file_id) REFERENCES files (developer_id, file_id)
-);
+-- Create function to validate owner reference
+CREATE OR REPLACE FUNCTION validate_file_owner()
+RETURNS TRIGGER AS $$
+BEGIN
+    IF NEW.owner_type = 'user' THEN
+        IF NOT EXISTS (
+            SELECT 1 FROM users 
+            WHERE developer_id = NEW.developer_id AND user_id = NEW.owner_id
+        ) THEN
+            RAISE EXCEPTION 'Invalid user reference';
+        END IF;
+    ELSIF NEW.owner_type = 'agent' THEN
+        IF NOT EXISTS (
+            SELECT 1 FROM agents 
+            WHERE developer_id = NEW.developer_id AND agent_id = NEW.owner_id
+        ) THEN
+            RAISE EXCEPTION 'Invalid agent reference';
+        END IF;
+    END IF;
+    RETURN NEW;
+END;
+$$ LANGUAGE plpgsql;
 
--- Create index if it doesn't exist
-CREATE INDEX IF NOT EXISTS idx_agent_files_agent ON agent_files (developer_id, agent_id);
+-- Create trigger for validation
+CREATE TRIGGER trg_validate_file_owner
+BEFORE INSERT OR UPDATE ON file_owners
+FOR EACH ROW
+EXECUTE FUNCTION validate_file_owner();
 
 COMMIT;

memory-store/migrations/000006_docs.down.sql

@@ -1,41 +1,27 @@
 BEGIN;
 
+-- Drop doc_owners table and its dependencies
+DROP TRIGGER IF EXISTS trg_validate_doc_owner ON doc_owners;
+DROP FUNCTION IF EXISTS validate_doc_owner();
+DROP TABLE IF EXISTS doc_owners;
+
+-- Drop docs table and its dependencies
+DROP TRIGGER IF EXISTS trg_docs_search_tsv ON docs;
+DROP TRIGGER IF EXISTS trg_docs_updated_at ON docs;
+DROP FUNCTION IF EXISTS docs_update_search_tsv();
+
 -- Drop indexes
 DROP INDEX IF EXISTS idx_docs_content_trgm;
-
 DROP INDEX IF EXISTS idx_docs_title_trgm;
-
 DROP INDEX IF EXISTS idx_docs_search_tsv;
-
 DROP INDEX IF EXISTS idx_docs_metadata;
-
-DROP INDEX IF EXISTS idx_agent_docs_agent;
-
-DROP INDEX IF EXISTS idx_user_docs_user;
-
 DROP INDEX IF EXISTS idx_docs_developer;
-
 DROP INDEX IF EXISTS idx_docs_id_sorted;
 
--- Drop triggers
-DROP TRIGGER IF EXISTS trg_docs_search_tsv ON docs;
-
-DROP TRIGGER IF EXISTS trg_docs_updated_at ON docs;
-
--- Drop the constraint that depends on is_valid_language function
-ALTER TABLE IF EXISTS docs
-DROP CONSTRAINT IF EXISTS ct_docs_valid_language;
-
--- Drop functions
-DROP FUNCTION IF EXISTS docs_update_search_tsv ();
-
-DROP FUNCTION IF EXISTS is_valid_language (text);
-
--- Drop tables (in correct order due to foreign key constraints)
-DROP TABLE IF EXISTS agent_docs;
-
-DROP TABLE IF EXISTS user_docs;
-
+-- Drop docs table
 DROP TABLE IF EXISTS docs;
 
+-- Drop language validation function
+DROP FUNCTION IF EXISTS is_valid_language(text);
+
 COMMIT;

memory-store/migrations/000006_docs.up.sql

@@ -63,31 +63,51 @@ BEGIN
     END IF;
 END $$;
 
--- Create the user_docs table
-CREATE TABLE IF NOT EXISTS user_docs (
+-- Create the doc_owners table
+CREATE TABLE IF NOT EXISTS doc_owners (
     developer_id UUID NOT NULL,
-    user_id UUID NOT NULL,
     doc_id UUID NOT NULL,
-    CONSTRAINT pk_user_docs PRIMARY KEY (developer_id, user_id, doc_id),
-    CONSTRAINT fk_user_docs_user FOREIGN KEY (developer_id, user_id) REFERENCES users (developer_id, user_id),
-    CONSTRAINT fk_user_docs_doc FOREIGN KEY (developer_id, doc_id) REFERENCES docs (developer_id, doc_id)
+    owner_type TEXT NOT NULL,  -- 'user' or 'agent'
+    owner_id UUID NOT NULL,
+    CONSTRAINT pk_doc_owners PRIMARY KEY (developer_id, doc_id),
+    CONSTRAINT fk_doc_owners_doc FOREIGN KEY (developer_id, doc_id) REFERENCES docs (developer_id, doc_id),
+    CONSTRAINT ct_doc_owners_owner_type CHECK (owner_type IN ('user', 'agent'))
 );
 
--- Create the agent_docs table
-CREATE TABLE IF NOT EXISTS agent_docs (
-    developer_id UUID NOT NULL,
-    agent_id UUID NOT NULL,
-    doc_id UUID NOT NULL,
-    CONSTRAINT pk_agent_docs PRIMARY KEY (developer_id, agent_id, doc_id),
-    CONSTRAINT fk_agent_docs_agent FOREIGN KEY (developer_id, agent_id) REFERENCES agents (developer_id, agent_id),
-    CONSTRAINT fk_agent_docs_doc FOREIGN KEY (developer_id, doc_id) REFERENCES docs (developer_id, doc_id)
-);
+-- Create indexes
+CREATE INDEX IF NOT EXISTS idx_doc_owners_owner 
+    ON doc_owners (developer_id, owner_type, owner_id);
 
--- Create indexes if not exists
-CREATE INDEX IF NOT EXISTS idx_user_docs_user ON user_docs (developer_id, user_id);
+-- Create function to validate owner reference
+CREATE OR REPLACE FUNCTION validate_doc_owner()
+RETURNS TRIGGER AS $$
+BEGIN
+    IF NEW.owner_type = 'user' THEN
+        IF NOT EXISTS (
+            SELECT 1 FROM users 
+            WHERE developer_id = NEW.developer_id AND user_id = NEW.owner_id
+        ) THEN
+            RAISE EXCEPTION 'Invalid user reference';
+        END IF;
+    ELSIF NEW.owner_type = 'agent' THEN
+        IF NOT EXISTS (
+            SELECT 1 FROM agents 
+            WHERE developer_id = NEW.developer_id AND agent_id = NEW.owner_id
+        ) THEN
+            RAISE EXCEPTION 'Invalid agent reference';
+        END IF;
+    END IF;
+    RETURN NEW;
+END;
+$$ LANGUAGE plpgsql;
 
-CREATE INDEX IF NOT EXISTS idx_agent_docs_agent ON agent_docs (developer_id, agent_id);
+-- Create trigger for validation
+CREATE TRIGGER trg_validate_doc_owner
+BEFORE INSERT OR UPDATE ON doc_owners
+FOR EACH ROW
+EXECUTE FUNCTION validate_doc_owner();
 
+-- Create indexes if not exists
 CREATE INDEX IF NOT EXISTS idx_docs_metadata ON docs USING GIN (metadata);
 
 -- Enable necessary PostgreSQL extensions