setup_postgres.py

#!/usr/bin/env python3
"""
PostgreSQL Setup Script for Modal Vector Service

This script helps set up a PostgreSQL database with pgvector extension
for the Modal vector service.
"""

import os
import sys
import subprocess
import psycopg2
from urllib.parse import urlparse


def test_postgres_connection(postgres_url: str) -> bool:
    """Test PostgreSQL connection and pgvector availability"""
    try:
        print(f"🔗 Testing connection to PostgreSQL...")
        conn = psycopg2.connect(postgres_url)
        cursor = conn.cursor()

        # Test basic connection
        cursor.execute("SELECT version();")
        version = cursor.fetchone()[0]
        print(f"✅ Connected to PostgreSQL: {version}")

        # Test pgvector extension
        try:
            cursor.execute("CREATE EXTENSION IF NOT EXISTS vector;")
            cursor.execute(
                "SELECT extversion FROM pg_extension WHERE extname = 'vector';"
            )
            vector_version = cursor.fetchone()
            if vector_version:
                print(f"✅ pgvector extension available: v{vector_version[0]}")
            else:
                print("⚠️ pgvector extension not found")
                return False
        except Exception as e:
            print(f"❌ pgvector extension error: {e}")
            return False

        # Create test table to verify vector operations
        cursor.execute(
            """
            CREATE TABLE IF NOT EXISTS vector_test (
                id SERIAL PRIMARY KEY,
                embedding vector(384)
            );
        """
        )

        # Test vector operations
        test_vector = [0.1] * 384  # 384-dimensional test vector
        cursor.execute(
            "INSERT INTO vector_test (embedding) VALUES (%s) RETURNING id;",
            (test_vector,),
        )
        test_id = cursor.fetchone()[0]
        print(f"✅ Vector operations working (test ID: {test_id})")

        # Clean up test
        cursor.execute("DELETE FROM vector_test WHERE id = %s;", (test_id,))

        conn.commit()
        cursor.close()
        conn.close()

        return True

    except Exception as e:
        print(f"❌ PostgreSQL connection failed: {e}")
        return False


def setup_modal_secret(postgres_url: str):
    """Set up Modal secret for PostgreSQL"""
    try:
        print("🔐 Setting up Modal secret for PostgreSQL...")

        # Create or update the Modal secret
        result = subprocess.run(
            [
                "modal",
                "secret",
                "create",
                "postgres-secret",
                f"MODAL_POSTGRES_URL={postgres_url}",
            ],
            capture_output=True,
            text=True,
        )

        if result.returncode == 0:
            print("✅ Modal secret created successfully")
            print("\nTo use in your Modal functions, add:")
            print("@app.function(secrets=[modal.Secret.from_name('postgres-secret')])")
        else:
            # Try updating if creation failed
            result = subprocess.run(
                [
                    "modal",
                    "secret",
                    "update",
                    "postgres-secret",
                    f"MODAL_POSTGRES_URL={postgres_url}",
                ],
                capture_output=True,
                text=True,
            )

            if result.returncode == 0:
                print("✅ Modal secret updated successfully")
            else:
                print(f"❌ Failed to create/update Modal secret: {result.stderr}")
                return False

        return True

    except Exception as e:
        print(f"❌ Error setting up Modal secret: {e}")
        return False


def create_vector_tables(postgres_url: str):
    """Create the vector memory tables"""
    try:
        print("📊 Creating vector memory tables...")
        conn = psycopg2.connect(postgres_url)
        cursor = conn.cursor()

        # Create the main vector memories table
        cursor.execute(
            """
            CREATE TABLE IF NOT EXISTS vector_memories (
                id SERIAL PRIMARY KEY,
                client_id VARCHAR(255) NOT NULL,
                text TEXT NOT NULL,
                embedding vector(384),  -- all-MiniLM-L6-v2 produces 384-dim vectors
                metadata JSONB,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            );
        """
        )

        # Create indexes for performance
        cursor.execute(
            """
            CREATE INDEX IF NOT EXISTS idx_vector_memories_client_id 
            ON vector_memories(client_id);
        """
        )

        cursor.execute(
            """
            CREATE INDEX IF NOT EXISTS idx_vector_memories_created_at 
            ON vector_memories(created_at);
        """
        )

        # Create vector similarity index (HNSW for fast approximate search)
        cursor.execute(
            """
            CREATE INDEX IF NOT EXISTS idx_vector_memories_embedding 
            ON vector_memories USING hnsw (embedding vector_cosine_ops);
        """
        )

        conn.commit()
        cursor.close()
        conn.close()

        print("✅ Vector memory tables created successfully")
        return True

    except Exception as e:
        print(f"❌ Error creating vector tables: {e}")
        return False


def main():
    print("🚀 PostgreSQL Setup for Modal Vector Service")
    print("=" * 50)

    # Check if PostgreSQL URL is provided
    postgres_url = os.getenv("POSTGRES_URL")
    if not postgres_url:
        print("\n📝 PostgreSQL URL not found in environment.")
        print("\nOptions for PostgreSQL with pgvector:")
        print("1. Neon (https://neon.tech) - Free tier with pgvector")
        print("2. Supabase (https://supabase.com) - Free tier with pgvector")
        print("3. Railway (https://railway.app) - PostgreSQL with pgvector")
        print("4. Your own PostgreSQL instance")

        print("\nTo use this script:")
        print("export POSTGRES_URL='postgresql://user:password@host:port/database'")
        print("python setup_postgres.py")

        # Try to get URL from user input
        postgres_url = input(
            "\nEnter PostgreSQL URL (or press Enter to skip): "
        ).strip()
        if not postgres_url:
            print("⏭️ Skipping PostgreSQL setup")
            return

    # Test the connection
    if not test_postgres_connection(postgres_url):
        print("❌ PostgreSQL setup failed - connection test failed")
        return

    # Create vector tables
    if not create_vector_tables(postgres_url):
        print("❌ PostgreSQL setup failed - table creation failed")
        return

    # Set up Modal secret
    if not setup_modal_secret(postgres_url):
        print("❌ PostgreSQL setup failed - Modal secret setup failed")
        return

    print("\n🎉 PostgreSQL setup completed successfully!")
    print("\nNext steps:")
    print("1. Redeploy your Modal vector service")
    print("2. Test vector storage and search")
    print("3. Monitor performance in Modal dashboard")

    # Parse URL to show connection info (without password)
    parsed = urlparse(postgres_url)
    print(f"\n📊 Database Info:")
    print(f"   Host: {parsed.hostname}")
    print(f"   Port: {parsed.port or 5432}")
    print(f"   Database: {parsed.path[1:] if parsed.path else 'postgres'}")
    print(f"   User: {parsed.username}")


if __name__ == "__main__":
    main()