fix(memory-store): Misc fixes; Switch to golang-migrate

Signed-off-by: Diwank Singh Tomer <[email protected]>

memory-store/README.md

@@ -0,0 +1,7 @@
+### prototyping flow:
+
+1. Install `pgmigrate` (until I move to golang-migrate)
+2. In a separate window, `docker compose up db vectorizer-worker` to start db instances
+3. `cd memory-store` and `pgmigrate migrate --database "postgres://postgres:[email protected]:5432/postgres" --migrations ./migrations` to apply the migrations
+4. `pip install --user -U pgcli`
+5. `pgcli "postgres://postgres:postgres@localhost:5432/postgres"`

memory-store/migrations/000001_initial.down.sql

@@ -0,0 +1,17 @@
+-- Drop the update_updated_at_column function
+DROP FUNCTION IF EXISTS update_updated_at_column();
+
+-- Drop misc extensions
+DROP EXTENSION IF EXISTS "uuid-ossp" CASCADE;
+DROP EXTENSION IF EXISTS citext CASCADE;
+DROP EXTENSION IF EXISTS btree_gist CASCADE;
+DROP EXTENSION IF EXISTS btree_gin CASCADE;
+
+-- Drop timescale's pgai extensions
+DROP EXTENSION IF EXISTS ai CASCADE;
+DROP EXTENSION IF EXISTS vectorscale CASCADE;
+DROP EXTENSION IF EXISTS vector CASCADE;
+
+-- Drop timescaledb extensions
+DROP EXTENSION IF EXISTS timescaledb_toolkit CASCADE;
+DROP EXTENSION IF EXISTS timescaledb CASCADE;

memory-store/migrations/00001_initial.sql → memory-store/migrations/000001_initial.up.sql

@@ -1,3 +1,5 @@
+BEGIN;
+
 -- init timescaledb
 CREATE EXTENSION IF NOT EXISTS timescaledb CASCADE;
 CREATE EXTENSION IF NOT EXISTS timescaledb_toolkit CASCADE;
@@ -23,3 +25,5 @@ END;
 $$ language 'plpgsql';
 
 COMMENT ON FUNCTION update_updated_at_column() IS 'Trigger function to automatically update updated_at timestamp';
+
+COMMIT;

memory-store/migrations/000002_developers.down.sql

@@ -0,0 +1,4 @@
+-- Drop the table (this will automatically drop associated indexes and triggers)
+DROP TABLE IF EXISTS developers CASCADE;
+
+-- Note: The update_updated_at_column() function is not dropped as it might be used by other tables

memory-store/migrations/00002_developers.sql → memory-store/migrations/000002_developers.up.sql

@@ -1,5 +1,7 @@
+BEGIN;
+
 -- Create developers table
-CREATE TABLE developers (
+CREATE TABLE IF NOT EXISTS developers (
     developer_id UUID NOT NULL,
     email TEXT NOT NULL CONSTRAINT ct_developers_email_format CHECK (email ~* '^[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}$'),
     active BOOLEAN NOT NULL DEFAULT true,
@@ -12,22 +14,29 @@ CREATE TABLE developers (
 );
 
 -- Create sorted index on developer_id (optimized for UUID v7)
-CREATE INDEX idx_developers_id_sorted ON developers (developer_id DESC);
+CREATE INDEX IF NOT EXISTS idx_developers_id_sorted ON developers (developer_id DESC);
 
 -- Create index on email
-CREATE INDEX idx_developers_email ON developers (email);
+CREATE INDEX IF NOT EXISTS idx_developers_email ON developers (email);
 
 -- Create GIN index for tags array
-CREATE INDEX idx_developers_tags ON developers USING GIN (tags);
+CREATE INDEX IF NOT EXISTS idx_developers_tags ON developers USING GIN (tags);
 
 -- Create partial index for active developers
-CREATE INDEX idx_developers_active ON developers (developer_id) WHERE active = true;
+CREATE INDEX IF NOT EXISTS idx_developers_active ON developers (developer_id) WHERE active = true;
 
 -- Create trigger to automatically update updated_at
-CREATE TRIGGER trg_developers_updated_at
-    BEFORE UPDATE ON developers
-    FOR EACH ROW
-    EXECUTE FUNCTION update_updated_at_column();
+DO $$ 
+BEGIN
+    IF NOT EXISTS (SELECT 1 FROM pg_trigger WHERE tgname = 'trg_developers_updated_at') THEN
+        CREATE TRIGGER trg_developers_updated_at
+            BEFORE UPDATE ON developers
+            FOR EACH ROW
+            EXECUTE FUNCTION update_updated_at_column();
+    END IF;
+END
+$$;
 
 -- Add comment to table
-COMMENT ON TABLE developers IS 'Stores developer information including their settings and tags';
+COMMENT ON TABLE developers IS 'Stores developer information including their settings and tags';
+COMMIT;

memory-store/migrations/000003_users.down.sql

@@ -0,0 +1,18 @@
+BEGIN;
+
+-- Drop trigger first
+DROP TRIGGER IF EXISTS update_users_updated_at ON users;
+
+-- Drop indexes
+DROP INDEX IF EXISTS users_metadata_gin_idx;
+DROP INDEX IF EXISTS users_developer_id_idx;
+DROP INDEX IF EXISTS users_id_sorted_idx;
+
+-- Drop foreign key constraint
+ALTER TABLE IF EXISTS users
+    DROP CONSTRAINT IF EXISTS users_developer_id_fkey;
+
+-- Finally drop the table
+DROP TABLE IF EXISTS users;
+
+COMMIT;

memory-store/migrations/000003_users.up.sql

@@ -0,0 +1,49 @@
+BEGIN;
+
+-- Create users table if it doesn't exist
+CREATE TABLE IF NOT EXISTS users (
+    developer_id UUID NOT NULL,
+    user_id UUID NOT NULL,
+    name TEXT NOT NULL,
+    about TEXT,
+    created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    updated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    metadata JSONB NOT NULL DEFAULT '{}'::JSONB,
+    CONSTRAINT pk_users PRIMARY KEY (developer_id, user_id)
+);
+
+-- Create sorted index on user_id if it doesn't exist
+CREATE INDEX IF NOT EXISTS users_id_sorted_idx ON users (user_id DESC);
+
+-- Create foreign key constraint and index if they don't exist
+DO $$ BEGIN
+    IF NOT EXISTS (
+        SELECT 1 FROM pg_constraint WHERE conname = 'users_developer_id_fkey'
+    ) THEN
+        ALTER TABLE users 
+            ADD CONSTRAINT users_developer_id_fkey 
+            FOREIGN KEY (developer_id) 
+            REFERENCES developers(developer_id);
+    END IF;
+END $$;
+
+CREATE INDEX IF NOT EXISTS users_developer_id_idx ON users (developer_id);
+
+-- Create a GIN index on the entire metadata column if it doesn't exist
+CREATE INDEX IF NOT EXISTS users_metadata_gin_idx ON users USING GIN (metadata);
+
+-- Create trigger if it doesn't exist
+DO $$ BEGIN
+    IF NOT EXISTS (
+        SELECT 1 FROM pg_trigger WHERE tgname = 'update_users_updated_at'
+    ) THEN
+        CREATE TRIGGER update_users_updated_at
+            BEFORE UPDATE ON users
+            FOR EACH ROW
+            EXECUTE FUNCTION update_updated_at_column();
+    END IF;
+END $$;
+
+-- Add comment to table (comments are idempotent by default)
+COMMENT ON TABLE users IS 'Stores user information linked to developers';
+COMMIT;

memory-store/migrations/000004_agents.down.sql

@@ -0,0 +1,14 @@
+BEGIN;
+
+-- Drop trigger first
+DROP TRIGGER IF EXISTS trg_agents_updated_at ON agents;
+
+-- Drop indexes
+DROP INDEX IF EXISTS idx_agents_metadata;
+DROP INDEX IF EXISTS idx_agents_developer;
+DROP INDEX IF EXISTS idx_agents_id_sorted;
+
+-- Drop table (this will automatically drop associated constraints)
+DROP TABLE IF EXISTS agents;
+
+COMMIT;

memory-store/migrations/00004_agents.sql → memory-store/migrations/000004_agents.up.sql

@@ -1,5 +1,14 @@
+BEGIN;
+
+-- Drop existing objects if they exist
+DROP TRIGGER IF EXISTS trg_agents_updated_at ON agents;
+DROP INDEX IF EXISTS idx_agents_metadata;
+DROP INDEX IF EXISTS idx_agents_developer;
+DROP INDEX IF EXISTS idx_agents_id_sorted;
+DROP TABLE IF EXISTS agents;
+
 -- Create agents table
-CREATE TABLE agents (
+CREATE TABLE IF NOT EXISTS agents (
     developer_id UUID NOT NULL,
     agent_id UUID NOT NULL,
     canonical_name citext NOT NULL CONSTRAINT ct_agents_canonical_name_length CHECK (length(canonical_name) >= 1 AND length(canonical_name) <= 255),
@@ -17,24 +26,26 @@ CREATE TABLE agents (
 );
 
 -- Create sorted index on agent_id (optimized for UUID v7)
-CREATE INDEX idx_agents_id_sorted ON agents (agent_id DESC);
+CREATE INDEX IF NOT EXISTS idx_agents_id_sorted ON agents (agent_id DESC);
 
 -- Create foreign key constraint and index on developer_id
 ALTER TABLE agents 
+    DROP CONSTRAINT IF EXISTS fk_agents_developer,
     ADD CONSTRAINT fk_agents_developer
     FOREIGN KEY (developer_id) 
     REFERENCES developers(developer_id);
 
-CREATE INDEX idx_agents_developer ON agents (developer_id);
+CREATE INDEX IF NOT EXISTS idx_agents_developer ON agents (developer_id);
 
 -- Create a GIN index on the entire metadata column
-CREATE INDEX idx_agents_metadata ON agents USING GIN (metadata);
+CREATE INDEX IF NOT EXISTS idx_agents_metadata ON agents USING GIN (metadata);
 
 -- Create trigger to automatically update updated_at
-CREATE TRIGGER trg_agents_updated_at
+CREATE OR REPLACE TRIGGER trg_agents_updated_at
     BEFORE UPDATE ON agents
     FOR EACH ROW
     EXECUTE FUNCTION update_updated_at_column();
 
 -- Add comment to table
-COMMENT ON TABLE agents IS 'Stores AI agent configurations and metadata for developers';
+COMMENT ON TABLE agents IS 'Stores AI agent configurations and metadata for developers';
+COMMIT;

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

@@ -0,0 +1,13 @@
+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 files table and its dependencies
+DROP TRIGGER IF EXISTS trg_files_updated_at ON files;
+DROP TABLE IF EXISTS files;
+
+COMMIT;

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

@@ -1,5 +1,7 @@
+BEGIN;
+
 -- Create files table
-CREATE TABLE files (
+CREATE TABLE IF NOT EXISTS files (
     developer_id UUID NOT NULL,
     file_id UUID NOT NULL,
     name TEXT NOT NULL CONSTRAINT ct_files_name_length CHECK (length(name) >= 1 AND length(name) <= 255),
@@ -12,32 +14,41 @@ CREATE TABLE files (
     CONSTRAINT pk_files PRIMARY KEY (developer_id, file_id)
 );
 
--- Create sorted index on file_id (optimized for UUID v7)
-CREATE INDEX idx_files_id_sorted ON files (file_id DESC);
-
--- Create foreign key constraint and index on developer_id
-ALTER TABLE files 
-    ADD CONSTRAINT fk_files_developer 
-    FOREIGN KEY (developer_id) 
-    REFERENCES developers(developer_id);
+-- Create sorted index on file_id if it doesn't exist
+CREATE INDEX IF NOT EXISTS idx_files_id_sorted ON files (file_id DESC);
 
-CREATE INDEX idx_files_developer ON files (developer_id);
+-- Create foreign key constraint and index if they don't exist
+DO $$ BEGIN
+    IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'fk_files_developer') THEN
+        ALTER TABLE files 
+            ADD CONSTRAINT fk_files_developer 
+            FOREIGN KEY (developer_id) 
+            REFERENCES developers(developer_id);
+    END IF;
+END $$;
 
--- Before creating the user_files and agent_files tables, we need to ensure that the file_id is unique for each developer
-ALTER TABLE files
-    ADD CONSTRAINT uq_files_developer_id_file_id UNIQUE (developer_id, file_id);
+CREATE INDEX IF NOT EXISTS idx_files_developer ON files (developer_id);
 
--- Create trigger to automatically update updated_at
-CREATE TRIGGER trg_files_updated_at
-    BEFORE UPDATE ON files
-    FOR EACH ROW
-    EXECUTE FUNCTION update_updated_at_column();
+-- Add unique constraint if it doesn't exist
+DO $$ BEGIN
+    IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'uq_files_developer_id_file_id') THEN
+        ALTER TABLE files
+            ADD CONSTRAINT uq_files_developer_id_file_id UNIQUE (developer_id, file_id);
+    END IF;
+END $$;
 
--- Add comment to table
-COMMENT ON TABLE files IS 'Stores file metadata and references for developers';
+-- Create trigger if it doesn't exist
+DO $$ BEGIN
+    IF NOT EXISTS (SELECT 1 FROM pg_trigger WHERE tgname = 'trg_files_updated_at') THEN
+        CREATE TRIGGER trg_files_updated_at
+            BEFORE UPDATE ON files
+            FOR EACH ROW
+            EXECUTE FUNCTION update_updated_at_column();
+    END IF;
+END $$;
 
 -- Create the user_files table
-CREATE TABLE user_files (
+CREATE TABLE IF NOT EXISTS user_files (
     developer_id UUID NOT NULL,
     user_id UUID NOT NULL,
     file_id UUID NOT NULL,
@@ -46,11 +57,11 @@ CREATE TABLE user_files (
     CONSTRAINT fk_user_files_file FOREIGN KEY (developer_id, file_id) REFERENCES files(developer_id, file_id)
 );
 
--- Indexes for efficient querying
-CREATE INDEX idx_user_files_user ON user_files (developer_id, user_id);
+-- Create index if it doesn't exist
+CREATE INDEX IF NOT EXISTS idx_user_files_user ON user_files (developer_id, user_id);
 
 -- Create the agent_files table
-CREATE TABLE agent_files (
+CREATE TABLE IF NOT EXISTS agent_files (
     developer_id UUID NOT NULL,
     agent_id UUID NOT NULL,
     file_id UUID NOT NULL,
@@ -59,5 +70,7 @@ CREATE TABLE agent_files (
     CONSTRAINT fk_agent_files_file FOREIGN KEY (developer_id, file_id) REFERENCES files(developer_id, file_id)
 );
 
--- Indexes for efficient querying
-CREATE INDEX idx_agent_files_agent ON agent_files (developer_id, agent_id);
+-- Create index if it doesn't exist
+CREATE INDEX IF NOT EXISTS idx_agent_files_agent ON agent_files (developer_id, agent_id);
+
+COMMIT;

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

@@ -0,0 +1,29 @@
+BEGIN;
+
+-- 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 TABLE IF EXISTS docs;
+
+COMMIT;