Make migration generation idempotent by default

- Add shared utility function to make migrations idempotent (IF NOT EXISTS)
- Update migration generation scripts to use shared utility
- Fix existing migration file to use IF NOT EXISTS clauses
- Add comprehensive test suite with 23 test cases covering edge cases
- Ensure all future migrations are automatically idempotent

This prevents migration failures when columns/indexes already exist,
making migrations safe to re-run on databases that were manually fixed.

ci_cd/baseline_db.py

@@ -2,6 +2,8 @@ import subprocess
 from pathlib import Path
 from datetime import datetime
 
+from ci_cd.migration_utils import make_migration_idempotent
+
 
 def create_baseline():
     """Create baseline migration in deploy/migrations"""
@@ -41,9 +43,12 @@ def create_baseline():
             check=True,
         )
 
+        # Post-process SQL to make it idempotent
+        idempotent_sql = make_migration_idempotent(result.stdout)
+
         # Write the SQL to migration.sql
         migration_file = migration_dir / "migration.sql"
-        migration_file.write_text(result.stdout)
+        migration_file.write_text(idempotent_sql)
 
         print(f"Created baseline migration in {migration_dir}")
         return True

ci_cd/migration_utils.py

@@ -0,0 +1,75 @@
+"""
+Utility functions for processing database migrations.
+
+This module provides functions to make SQL migrations idempotent by adding
+IF NOT EXISTS clauses to ADD COLUMN and CREATE INDEX statements.
+"""
+
+import re
+
+
+def make_migration_idempotent(sql_content: str) -> str:
+    """
+    Post-process SQL migration to make it idempotent by adding IF NOT EXISTS clauses.
+    
+    This function adds IF NOT EXISTS to:
+    - ADD COLUMN statements
+    - CREATE INDEX statements
+    - CREATE UNIQUE INDEX statements
+    
+    It safely handles:
+    - Already idempotent migrations (won't duplicate IF NOT EXISTS)
+    - Multi-line SQL statements
+    - Different whitespace patterns
+    - Case-insensitive matching
+    
+    Args:
+        sql_content: Raw SQL from Prisma migrate diff
+        
+    Returns:
+        SQL with IF NOT EXISTS clauses added to appropriate statements
+        
+    Examples:
+        >>> sql = 'ALTER TABLE "Test" ADD COLUMN "col1" TEXT;'
+        >>> make_migration_idempotent(sql)
+        'ALTER TABLE "Test" ADD COLUMN IF NOT EXISTS "col1" TEXT;'
+        
+        >>> sql = 'CREATE INDEX "idx1" ON "Test"("col1");'
+        >>> make_migration_idempotent(sql)
+        'CREATE INDEX IF NOT EXISTS "idx1" ON "Test"("col1");'
+        
+        >>> sql = 'ALTER TABLE "Test" ADD COLUMN IF NOT EXISTS "col1" TEXT;'
+        >>> make_migration_idempotent(sql)
+        'ALTER TABLE "Test" ADD COLUMN IF NOT EXISTS "col1" TEXT;'
+    """
+    if not sql_content or not sql_content.strip():
+        return sql_content
+    
+    # Add IF NOT EXISTS to ADD COLUMN statements (only if not already present)
+    # Pattern matches: ADD COLUMN followed by whitespace and a quoted identifier
+    # Uses negative lookahead to avoid matching if IF NOT EXISTS is already present
+    sql_content = re.sub(
+        r'ADD COLUMN\s+(?!IF NOT EXISTS\s+)("[\w]+")',
+        r'ADD COLUMN IF NOT EXISTS \1',
+        sql_content,
+        flags=re.IGNORECASE | re.MULTILINE
+    )
+    
+    # Add IF NOT EXISTS to CREATE INDEX statements (only if not already present)
+    sql_content = re.sub(
+        r'CREATE INDEX\s+(?!IF NOT EXISTS\s+)("[\w]+")',
+        r'CREATE INDEX IF NOT EXISTS \1',
+        sql_content,
+        flags=re.IGNORECASE | re.MULTILINE
+    )
+    
+    # Add IF NOT EXISTS to CREATE UNIQUE INDEX statements (only if not already present)
+    # Must come after CREATE INDEX to avoid partial matches
+    sql_content = re.sub(
+        r'CREATE UNIQUE INDEX\s+(?!IF NOT EXISTS\s+)("[\w]+")',
+        r'CREATE UNIQUE INDEX IF NOT EXISTS \1',
+        sql_content,
+        flags=re.IGNORECASE | re.MULTILINE
+    )
+    
+    return sql_content

ci_cd/run_migration.py

@@ -5,6 +5,8 @@ from datetime import datetime
 import testing.postgresql
 import shutil
 
+from ci_cd.migration_utils import make_migration_idempotent
+
 
 def create_migration(migration_name: str = None):
     """
@@ -64,9 +66,12 @@ def create_migration(migration_name: str = None):
                     migration_dir = migrations_dir / f"{timestamp}_{migration_name}"
                     migration_dir.mkdir(parents=True, exist_ok=True)
 
+                    # Post-process SQL to make it idempotent
+                    idempotent_sql = make_migration_idempotent(result.stdout)
+
                     # Write the SQL to migration.sql
                     migration_file = migration_dir / "migration.sql"
-                    migration_file.write_text(result.stdout)
+                    migration_file.write_text(idempotent_sql)
 
                     print(f"Created migration in {migration_dir}")
                     return True

litellm-proxy-extras/litellm_proxy_extras/migrations/20260131150814_add_team_user_to_vector_stores/migration.sql

@@ -1,10 +1,13 @@
 -- AlterTable
-ALTER TABLE "LiteLLM_ManagedVectorStoresTable" ADD COLUMN     "team_id" TEXT,
-ADD COLUMN     "user_id" TEXT;
+ALTER TABLE "LiteLLM_ManagedVectorStoresTable"
+    ADD COLUMN IF NOT EXISTS "team_id" TEXT,
+    ADD COLUMN IF NOT EXISTS "user_id" TEXT;
 
 -- CreateIndex
-CREATE INDEX "LiteLLM_ManagedVectorStoresTable_team_id_idx" ON "LiteLLM_ManagedVectorStoresTable"("team_id");
+CREATE INDEX IF NOT EXISTS "LiteLLM_ManagedVectorStoresTable_team_id_idx"
+    ON "LiteLLM_ManagedVectorStoresTable" ("team_id");
 
 -- CreateIndex
-CREATE INDEX "LiteLLM_ManagedVectorStoresTable_user_id_idx" ON "LiteLLM_ManagedVectorStoresTable"("user_id");
+CREATE INDEX IF NOT EXISTS "LiteLLM_ManagedVectorStoresTable_user_id_idx"
+    ON "LiteLLM_ManagedVectorStoresTable" ("user_id");
 

litellm-proxy-extras/litellm_proxy_extras/utils.py

@@ -57,6 +57,24 @@ def _get_prisma_command() -> str:
     return "prisma"
 
 
+def _make_migration_idempotent(sql_content: str) -> str:
+    """
+    Post-process SQL migration to make it idempotent by adding IF NOT EXISTS clauses.
+    
+    This is a wrapper around the shared utility function from ci_cd.migration_utils.
+    
+    Args:
+        sql_content: Raw SQL from Prisma migrate diff
+        
+    Returns:
+        SQL with IF NOT EXISTS clauses added to ADD COLUMN and CREATE INDEX statements
+    """
+    # Import here to avoid circular dependencies
+    from ci_cd.migration_utils import make_migration_idempotent
+    
+    return make_migration_idempotent(sql_content)
+
+
 class ProxyExtrasDBManager:
     @staticmethod
     def _get_prisma_dir() -> str:
@@ -302,6 +320,18 @@ class ProxyExtrasDBManager:
         if not diff_sql_path.exists():
             logger.warning("Migration diff was not created")
             return
+        
+        # Post-process SQL to make it idempotent
+        try:
+            with open(diff_sql_path, "r") as f:
+                sql_content = f.read()
+            idempotent_sql = _make_migration_idempotent(sql_content)
+            with open(diff_sql_path, "w") as f:
+                f.write(idempotent_sql)
+            logger.info("Migration SQL post-processed to be idempotent")
+        except Exception as e:
+            logger.warning(f"Failed to post-process migration SQL: {e}")
+        
         logger.info(f"Migration diff created at {diff_sql_path}")
 
         # 2. Run prisma db execute to apply the migration

tests/test_ci_cd/test_migration_utils.py

@@ -0,0 +1,216 @@
+"""
+Tests for migration utility functions.
+
+Tests cover edge cases including:
+- Already idempotent migrations
+- Multi-line statements
+- Different whitespace patterns
+- Case variations
+- Comments and other SQL statements
+- Empty and edge case inputs
+"""
+
+import pytest
+
+# Import from ci_cd.migration_utils
+import sys
+from pathlib import Path
+
+# Add project root to path for imports
+project_root = Path(__file__).parent.parent.parent
+sys.path.insert(0, str(project_root))
+
+from ci_cd.migration_utils import make_migration_idempotent
+
+
+class TestMakeMigrationIdempotent:
+    """Test suite for make_migration_idempotent function."""
+
+    def test_add_column_single_line(self):
+        """Test ADD COLUMN in single line format."""
+        sql = 'ALTER TABLE "Test" ADD COLUMN     "col1" TEXT;'
+        expected = 'ALTER TABLE "Test" ADD COLUMN IF NOT EXISTS "col1" TEXT;'
+        assert make_migration_idempotent(sql) == expected
+
+    def test_add_column_multi_line(self):
+        """Test ADD COLUMN in multi-line format."""
+        sql = '''ALTER TABLE "Test"
+    ADD COLUMN     "col1" TEXT,
+    ADD COLUMN     "col2" TEXT;'''
+        expected = '''ALTER TABLE "Test"
+    ADD COLUMN IF NOT EXISTS "col1" TEXT,
+    ADD COLUMN IF NOT EXISTS "col2" TEXT;'''
+        assert make_migration_idempotent(sql) == expected
+
+    def test_add_column_already_idempotent(self):
+        """Test that already idempotent ADD COLUMN is not modified."""
+        sql = 'ALTER TABLE "Test" ADD COLUMN IF NOT EXISTS "col1" TEXT;'
+        expected = sql  # Should remain unchanged
+        assert make_migration_idempotent(sql) == expected
+
+    def test_create_index(self):
+        """Test CREATE INDEX statement."""
+        sql = 'CREATE INDEX "idx1" ON "Test"("col1");'
+        expected = 'CREATE INDEX IF NOT EXISTS "idx1" ON "Test"("col1");'
+        assert make_migration_idempotent(sql) == expected
+
+    def test_create_unique_index(self):
+        """Test CREATE UNIQUE INDEX statement."""
+        sql = 'CREATE UNIQUE INDEX "idx1" ON "Test"("col1");'
+        expected = 'CREATE UNIQUE INDEX IF NOT EXISTS "idx1" ON "Test"("col1");'
+        assert make_migration_idempotent(sql) == expected
+
+    def test_create_index_already_idempotent(self):
+        """Test that already idempotent CREATE INDEX is not modified."""
+        sql = 'CREATE INDEX IF NOT EXISTS "idx1" ON "Test"("col1");'
+        expected = sql  # Should remain unchanged
+        assert make_migration_idempotent(sql) == expected
+
+    def test_create_unique_index_already_idempotent(self):
+        """Test that already idempotent CREATE UNIQUE INDEX is not modified."""
+        sql = 'CREATE UNIQUE INDEX IF NOT EXISTS "idx1" ON "Test"("col1");'
+        expected = sql  # Should remain unchanged
+        assert make_migration_idempotent(sql) == expected
+
+    def test_case_insensitive(self):
+        """Test that function is case-insensitive."""
+        sql = 'alter table "Test" add column     "col1" text;'
+        result = make_migration_idempotent(sql)
+        assert 'IF NOT EXISTS' in result.upper() or 'if not exists' in result.lower()
+
+    def test_multiple_statements(self):
+        """Test multiple statements in one SQL block."""
+        sql = '''ALTER TABLE "Test" ADD COLUMN     "col1" TEXT;
+CREATE INDEX "idx1" ON "Test"("col1");
+CREATE UNIQUE INDEX "idx2" ON "Test"("col2");'''
+        result = make_migration_idempotent(sql)
+        assert 'ADD COLUMN IF NOT EXISTS' in result
+        assert 'CREATE INDEX IF NOT EXISTS' in result
+        assert 'CREATE UNIQUE INDEX IF NOT EXISTS' in result
+
+    def test_with_comments(self):
+        """Test SQL with comments."""
+        sql = '''-- AlterTable
+ALTER TABLE "Test" ADD COLUMN     "col1" TEXT;
+
+-- CreateIndex
+CREATE INDEX "idx1" ON "Test"("col1");'''
+        result = make_migration_idempotent(sql)
+        assert '-- AlterTable' in result
+        assert 'ADD COLUMN IF NOT EXISTS' in result
+        assert 'CREATE INDEX IF NOT EXISTS' in result
+
+    def test_mixed_idempotent_and_non_idempotent(self):
+        """Test mix of already idempotent and non-idempotent statements."""
+        sql = '''ALTER TABLE "Test" ADD COLUMN IF NOT EXISTS "col1" TEXT;
+ALTER TABLE "Test" ADD COLUMN     "col2" TEXT;
+CREATE INDEX IF NOT EXISTS "idx1" ON "Test"("col1");
+CREATE INDEX "idx2" ON "Test"("col2");'''
+        result = make_migration_idempotent(sql)
+        # Count occurrences to ensure no duplication
+        assert result.count('ADD COLUMN IF NOT EXISTS') == 2
+        assert result.count('CREATE INDEX IF NOT EXISTS') == 2
+        assert 'ADD COLUMN     "col2"' not in result  # Should be replaced
+
+    def test_complex_index_name(self):
+        """Test index names with underscores and numbers."""
+        sql = 'CREATE INDEX "LiteLLM_Table_col1_idx_123" ON "Test"("col1");'
+        expected = 'CREATE INDEX IF NOT EXISTS "LiteLLM_Table_col1_idx_123" ON "Test"("col1");'
+        assert make_migration_idempotent(sql) == expected
+
+    def test_complex_column_name(self):
+        """Test column names with underscores and numbers."""
+        sql = 'ALTER TABLE "Test" ADD COLUMN     "user_id_123" TEXT;'
+        expected = 'ALTER TABLE "Test" ADD COLUMN IF NOT EXISTS "user_id_123" TEXT;'
+        assert make_migration_idempotent(sql) == expected
+
+    def test_empty_string(self):
+        """Test empty string input."""
+        assert make_migration_idempotent('') == ''
+        assert make_migration_idempotent('   ') == '   '
+
+    def test_whitespace_only(self):
+        """Test whitespace-only input."""
+        sql = '\n\n  \n'
+        result = make_migration_idempotent(sql)
+        assert result == sql
+
+    def test_no_matching_statements(self):
+        """Test SQL with no ADD COLUMN or CREATE INDEX statements."""
+        sql = '''DROP INDEX "idx1";
+ALTER TABLE "Test" DROP COLUMN "col1";
+SELECT * FROM "Test";'''
+        result = make_migration_idempotent(sql)
+        assert result == sql  # Should remain unchanged
+
+    def test_real_world_example(self):
+        """Test with a real migration file pattern."""
+        sql = '''-- AlterTable
+ALTER TABLE "LiteLLM_ManagedVectorStoresTable"
+    ADD COLUMN     "team_id" TEXT,
+    ADD COLUMN     "user_id" TEXT;
+
+-- CreateIndex
+CREATE INDEX "LiteLLM_ManagedVectorStoresTable_team_id_idx"
+    ON "LiteLLM_ManagedVectorStoresTable" ("team_id");
+
+-- CreateIndex
+CREATE INDEX "LiteLLM_ManagedVectorStoresTable_user_id_idx"
+    ON "LiteLLM_ManagedVectorStoresTable" ("user_id");'''
+        result = make_migration_idempotent(sql)
+        assert 'ADD COLUMN IF NOT EXISTS' in result
+        assert 'CREATE INDEX IF NOT EXISTS' in result
+        assert result.count('IF NOT EXISTS') == 4  # 2 columns + 2 indexes
+
+    def test_index_with_multiple_columns(self):
+        """Test CREATE INDEX with multiple columns."""
+        sql = 'CREATE UNIQUE INDEX "idx1" ON "Test"("col1", "col2", "col3");'
+        expected = 'CREATE UNIQUE INDEX IF NOT EXISTS "idx1" ON "Test"("col1", "col2", "col3");'
+        result = make_migration_idempotent(sql)
+        assert 'CREATE UNIQUE INDEX IF NOT EXISTS' in result
+
+    def test_variable_whitespace(self):
+        """Test with various whitespace patterns."""
+        test_cases = [
+            ('ADD COLUMN "col1"', 'ADD COLUMN IF NOT EXISTS "col1"'),
+            ('ADD COLUMN     "col1"', 'ADD COLUMN IF NOT EXISTS "col1"'),
+            ('ADD COLUMN\t"col1"', 'ADD COLUMN IF NOT EXISTS "col1"'),
+            ('ADD COLUMN\n    "col1"', 'ADD COLUMN IF NOT EXISTS "col1"'),
+        ]
+        for input_sql, expected_part in test_cases:
+            full_sql = f'ALTER TABLE "Test" {input_sql} TEXT;'
+            result = make_migration_idempotent(full_sql)
+            assert expected_part in result
+
+    def test_idempotent_function_itself(self):
+        """Test that the function is idempotent (can be called multiple times)."""
+        sql = 'ALTER TABLE "Test" ADD COLUMN     "col1" TEXT;'
+        result1 = make_migration_idempotent(sql)
+        result2 = make_migration_idempotent(result1)
+        assert result1 == result2  # Should not change on second call
+
+    def test_special_characters_in_quotes(self):
+        """Test that function only matches quoted identifiers with word characters."""
+        # This should NOT match (no quotes)
+        sql = 'ALTER TABLE Test ADD COLUMN col1 TEXT;'
+        result = make_migration_idempotent(sql)
+        # Should not modify since pattern requires quotes
+        assert 'IF NOT EXISTS' not in result or 'ADD COLUMN     "col1"' in sql
+
+    def test_add_column_in_comment(self):
+        """Test that ADD COLUMN in comments is not modified."""
+        sql = '''-- This is a comment about ADD COLUMN
+ALTER TABLE "Test" ADD COLUMN     "col1" TEXT;'''
+        result = make_migration_idempotent(sql)
+        # Comment should remain unchanged, but actual statement should be modified
+        assert '-- This is a comment about ADD COLUMN' in result
+        assert 'ADD COLUMN IF NOT EXISTS' in result
+
+    def test_string_literal_with_add_column(self):
+        """Test that ADD COLUMN in string literals is not modified."""
+        # This is unlikely in migration SQL, but test for safety
+        sql = 'ALTER TABLE "Test" ADD COLUMN     "col1" TEXT DEFAULT \'ADD COLUMN test\';'
+        result = make_migration_idempotent(sql)
+        # Should modify the ADD COLUMN statement, not the string literal
+        assert 'ADD COLUMN IF NOT EXISTS' in result
+        assert "'ADD COLUMN test'" in result  # String literal unchanged