phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 10:44:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

test(test_db_schema_migration.py): ci/cd test to enforce schema migrations are documented in .sql files

Browse source
This commit is contained in:
Krrish Dholakia 2025-03-26 16:59:50 -07:00
parent 37265c8b1c
commit d4adc9764b
1 changed files with 45 additions and 24 deletions
Download patch file Download diff file Expand all files Collapse all files

69
tests/proxy_unit_tests/test_db_schema_migration.py
View file

@ -3,6 +3,8 @@ import os
import subprocess
from pathlib import Path
from pytest_postgresql import factories
import shutil
import tempfile
# Create postgresql fixture
postgresql_my_proc = factories.postgresql_proc(port=None)
@ -19,44 +21,63 @@ def test_schema_migration_check(schema_setup):
"""Test to check if schema requires migration"""
# Set test database URL
test_db_url = f"postgresql://{schema_setup.info.user}:@{schema_setup.info.host}:{schema_setup.info.port}/{schema_setup.info.dbname}"
# test_db_url = "postgresql://neondb_owner:npg_JiZPS0DAhRn4@ep-delicate-wave-a55cvbuc.us-east-2.aws.neon.tech/neondb?sslmode=require"
os.environ["DATABASE_URL"] = test_db_url
deploy_dir = Path("./deploy")
migrations_dir = deploy_dir / "migrations"
deploy_dir = Path("../../deploy")
source_migrations_dir = deploy_dir / "migrations"
schema_path = Path("../../schema.prisma")
print(migrations_dir)
if not migrations_dir.exists() or not any(migrations_dir.iterdir()):
print("No existing migrations found - first migration needed")
pytest.fail("No existing migrations found - first migration needed")
# Create temporary migrations directory next to schema.prisma
temp_migrations_dir = schema_path.parent / "migrations"
# If migrations exist, check for changes
result = subprocess.run(
["prisma", "migrate", "status"], capture_output=True, text=True
)
try:
# Copy migrations to correct location
if temp_migrations_dir.exists():
shutil.rmtree(temp_migrations_dir)
shutil.copytree(source_migrations_dir, temp_migrations_dir)
status_output = result.stdout.lower()
needs_migration = any(
state in status_output for state in ["drift detected", "pending"]
)
if not temp_migrations_dir.exists() or not any(temp_migrations_dir.iterdir()):
print("No existing migrations found - first migration needed")
pytest.fail("No existing migrations found - first migration needed")
if needs_migration:
print("Schema changes detected. New migration needed.")
# Show the differences
# Apply all existing migrations
subprocess.run(
["prisma", "migrate", "deploy", "--schema", str(schema_path)], check=True
)
# Compare current database state against schema
diff_result = subprocess.run(
[
"prisma",
"migrate",
"diff",
"--from-migrations",
"--from-url",
test_db_url,
"--to-schema-datamodel",
"schema.prisma",
str(schema_path),
"--script", # Show the SQL diff
"--exit-code", # Return exit code 2 if there are differences
],
capture_output=True,
text=True,
)
print("Schema differences:")
print(diff_result.stdout)
else:
print("No schema changes detected. Migration not needed.")
assert not needs_migration, "Schema changes detected - new migration required"
print("Exit code:", diff_result.returncode)
print("Stdout:", diff_result.stdout)
print("Stderr:", diff_result.stderr)
if diff_result.returncode == 2:
print("Schema changes detected. New migration needed.")
print("Schema differences:")
print(diff_result.stdout)
pytest.fail(
"Schema changes detected - new migration required. Run 'prisma migrate dev' to create new migration."
)
else:
print("No schema changes detected. Migration not needed.")
finally:
# Clean up: remove temporary migrations directory
if temp_migrations_dir.exists():
shutil.rmtree(temp_migrations_dir)