From 6bf33b6673b59768d8f2295ba3d7e7e8bdec7d3e Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Tue, 3 Feb 2026 10:32:09 -0800 Subject: [PATCH] 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 | 7 +- ci_cd/migration_utils.py | 75 ++++++ ci_cd/run_migration.py | 7 +- .../migration.sql | 11 +- .../litellm_proxy_extras/utils.py | 30 +++ tests/test_ci_cd/test_migration_utils.py | 216 ++++++++++++++++++ 6 files changed, 340 insertions(+), 6 deletions(-) create mode 100644 ci_cd/migration_utils.py create mode 100644 tests/test_ci_cd/test_migration_utils.py diff --git a/ci_cd/baseline_db.py b/ci_cd/baseline_db.py index ecc080abedd..259928e9021 100644 --- a/ci_cd/baseline_db.py +++ b/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 diff --git a/ci_cd/migration_utils.py b/ci_cd/migration_utils.py new file mode 100644 index 00000000000..f945c557961 --- /dev/null +++ b/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 diff --git a/ci_cd/run_migration.py b/ci_cd/run_migration.py index b11a38395c1..b7b08845509 100644 --- a/ci_cd/run_migration.py +++ b/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 diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260131150814_add_team_user_to_vector_stores/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260131150814_add_team_user_to_vector_stores/migration.sql index 2032f76a5de..b69ecad74aa 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260131150814_add_team_user_to_vector_stores/migration.sql +++ b/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"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index f3155722187..7bdd85593e0 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/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 diff --git a/tests/test_ci_cd/test_migration_utils.py b/tests/test_ci_cd/test_migration_utils.py new file mode 100644 index 00000000000..1851969e184 --- /dev/null +++ b/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