diff --git a/migrations/env.py b/migrations/env.py index 3455243c..c018ca19 100644 --- a/migrations/env.py +++ b/migrations/env.py @@ -15,7 +15,8 @@ Note: Uses pg8000 (sync driver) while the application uses asyncpg (async driver). This is intentional - Alembic runs synchronously, and both drivers produce -identical DDL/schema operations. +identical DDL/schema operations. An AUTOMATION_DB_URL naming an async driver is +rewritten to its sync equivalent by normalize_url_for_alembic(). """ import os @@ -23,7 +24,10 @@ from alembic import context from sqlalchemy import create_engine, text -from openhands.automation.db import _build_pg8000_connect_args +from openhands.automation.db import ( + _build_pg8000_connect_args, + normalize_url_for_alembic, +) from openhands.automation.models import Base @@ -66,11 +70,13 @@ def get_engine(database_name=DB_NAME): """ # SQLite or explicit PostgreSQL URL if DB_URL: - url = DB_URL - # For SQLite, remove async driver prefix if present (Alembic is sync) - if url.startswith("sqlite+aiosqlite"): - url = url.replace("sqlite+aiosqlite", "sqlite", 1) - return create_engine(url, pool_pre_ping=True) + url = normalize_url_for_alembic(DB_URL) + connect_args = ( + _build_pg8000_connect_args(DB_SSL_MODE) + if url.startswith("postgresql+pg8000") + else {} + ) + return create_engine(url, connect_args=connect_args, pool_pre_ping=True) # GCP Cloud SQL if GCP_DB_INSTANCE: @@ -104,9 +110,7 @@ def get_db_connection(): def run_migrations_offline(): if DB_URL: - url = DB_URL - if url.startswith("sqlite+aiosqlite"): - url = url.replace("sqlite+aiosqlite", "sqlite", 1) + url = normalize_url_for_alembic(DB_URL) else: url = f"postgresql+pg8000://{DB_USER}:{DB_PASS}@{DB_HOST}:{DB_PORT}/{DB_NAME}" diff --git a/openhands/automation/app.py b/openhands/automation/app.py index 9f8e333e..44792ee6 100644 --- a/openhands/automation/app.py +++ b/openhands/automation/app.py @@ -84,7 +84,7 @@ async def lifespan(app: FastAPI): from alembic import command from alembic.config import Config - from openhands.automation.db import normalize_sqlite_url_for_alembic + from openhands.automation.db import normalize_url_for_alembic # Find migrations folder relative to this package. # When installed via pip/uvx, migrations are bundled inside @@ -108,7 +108,7 @@ async def lifespan(app: FastAPI): alembic_cfg = Config() alembic_cfg.set_main_option("script_location", str(migrations_path)) # Set the database URL for Alembic to use (sync version) - db_url = normalize_sqlite_url_for_alembic(settings.db_url) + db_url = normalize_url_for_alembic(settings.db_url) alembic_cfg.set_main_option("sqlalchemy.url", db_url) # Run migrations synchronously (Alembic doesn't support async) diff --git a/openhands/automation/db.py b/openhands/automation/db.py index ff4ec2ff..0221b8dd 100644 --- a/openhands/automation/db.py +++ b/openhands/automation/db.py @@ -66,14 +66,22 @@ def is_sqlite_url(url: str) -> bool: return url.startswith("sqlite") -def normalize_sqlite_url_for_alembic(url: str) -> str: - """Convert async SQLite URL to sync version for Alembic. +ALEMBIC_SYNC_DRIVERS = { + "sqlite+aiosqlite": "sqlite", + "postgresql+asyncpg": "postgresql+pg8000", +} - Alembic doesn't support async drivers, so we need to convert - sqlite+aiosqlite:// URLs to plain sqlite:// URLs. + +def normalize_url_for_alembic(url: str) -> str: + """Convert an async database URL to the sync driver Alembic uses. + + Alembic runs synchronously and cannot drive aiosqlite or asyncpg, so + sqlite+aiosqlite:// becomes sqlite:// and postgresql+asyncpg:// becomes + postgresql+pg8000://. Any other URL is returned unchanged. """ - if url.startswith("sqlite+aiosqlite"): - return url.replace("sqlite+aiosqlite", "sqlite", 1) + for async_driver, sync_driver in ALEMBIC_SYNC_DRIVERS.items(): + if url.startswith(async_driver): + return url.replace(async_driver, sync_driver, 1) return url diff --git a/tests/test_db.py b/tests/test_db.py index 8df9f9cc..c1846cb2 100644 --- a/tests/test_db.py +++ b/tests/test_db.py @@ -14,7 +14,7 @@ _build_pg8000_connect_args, _create_sqlite_engine, is_sqlite_url, - normalize_sqlite_url_for_alembic, + normalize_url_for_alembic, set_sqlite_mode, using_sqlite, ) @@ -173,38 +173,54 @@ def fake_create_async_engine(*args: Any, **kwargs: Any) -> Any: assert captured["kwargs"]["pool_pre_ping"] is True -class TestNormalizeSqliteUrlForAlembic: - """Tests for normalize_sqlite_url_for_alembic helper function.""" +class TestNormalizeUrlForAlembic: + """Tests for normalize_url_for_alembic helper function.""" def test_converts_aiosqlite_to_sqlite(self): """Converts sqlite+aiosqlite:// to sqlite://.""" url = "sqlite+aiosqlite:///test.db" - assert normalize_sqlite_url_for_alembic(url) == "sqlite:///test.db" + assert normalize_url_for_alembic(url) == "sqlite:///test.db" def test_converts_aiosqlite_with_absolute_path(self): """Converts sqlite+aiosqlite with absolute path.""" url = "sqlite+aiosqlite:////data/automations.db" - assert normalize_sqlite_url_for_alembic(url) == "sqlite:////data/automations.db" + assert normalize_url_for_alembic(url) == "sqlite:////data/automations.db" def test_preserves_plain_sqlite_url(self): """Plain sqlite:// URL is unchanged.""" url = "sqlite:///test.db" - assert normalize_sqlite_url_for_alembic(url) == "sqlite:///test.db" + assert normalize_url_for_alembic(url) == "sqlite:///test.db" def test_preserves_postgresql_url(self): """PostgreSQL URLs are unchanged.""" url = "postgresql://user:pass@host/db" - assert normalize_sqlite_url_for_alembic(url) == url + assert normalize_url_for_alembic(url) == url - def test_preserves_postgresql_asyncpg_url(self): - """PostgreSQL+asyncpg URLs are unchanged.""" + def test_converts_asyncpg_to_pg8000(self): + """Converts postgresql+asyncpg:// to postgresql+pg8000://.""" url = "postgresql+asyncpg://user:pass@host/db" - assert normalize_sqlite_url_for_alembic(url) == url + assert normalize_url_for_alembic(url) == "postgresql+pg8000://user:pass@host/db" + + def test_converts_only_the_asyncpg_driver_prefix(self): + """A password containing the driver name is left alone.""" + url = "postgresql+asyncpg://user:postgresql+asyncpg@host/db" + assert normalize_url_for_alembic(url) == ( + "postgresql+pg8000://user:postgresql+asyncpg@host/db" + ) + + def test_preserves_postgresql_psycopg_url(self): + """Other PostgreSQL drivers are unchanged.""" + url = "postgresql+psycopg2://user:pass@host/db" + assert normalize_url_for_alembic(url) == url def test_handles_memory_database(self): """Memory database URL is converted correctly.""" url = "sqlite+aiosqlite:///:memory:" - assert normalize_sqlite_url_for_alembic(url) == "sqlite:///:memory:" + assert normalize_url_for_alembic(url) == "sqlite:///:memory:" + + def test_preserves_empty_url(self): + """Empty URL is unchanged.""" + assert normalize_url_for_alembic("") == "" class TestSqliteMigrations: @@ -270,7 +286,7 @@ def test_auto_migration_applies_schema(self): """Auto-migration on startup creates all required tables. This tests the auto-migration behavior added in app.py for SQLite, - using the normalize_sqlite_url_for_alembic helper function. + using the normalize_url_for_alembic helper function. """ import subprocess @@ -281,7 +297,7 @@ def test_auto_migration_applies_schema(self): try: # Test the URL normalization helper async_url = f"sqlite+aiosqlite:///{db_path}" - sync_url = normalize_sqlite_url_for_alembic(async_url) + sync_url = normalize_url_for_alembic(async_url) assert sync_url == f"sqlite:///{db_path}" # Run migrations using subprocess to avoid env.py PostgreSQL defaults