diff --git a/AGENTS.md b/AGENTS.md index 4e4bdf62..19ec6e14 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -205,7 +205,7 @@ Supports **PostgreSQL** (cloud) and **SQLite** (local/self-hosted). | Config | `AUTOMATION_DB_HOST`, `AUTOMATION_DB_PORT`, etc. | `AUTOMATION_DB_URL=sqlite+aiosqlite:///path.db` | | Driver | asyncpg | aiosqlite | | Row locking | `FOR UPDATE SKIP LOCKED` | Skipped (single-process) | -| Migrations | `alembic upgrade head` (manual) | Auto-run on startup | +| Migrations | `alembic upgrade head` (manual), or auto-run on startup with `AUTOMATION_AUTO_MIGRATE=true` | Auto-run on startup (set `AUTOMATION_AUTO_MIGRATE=false` to opt out) | ### Writing Migrations 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..2a574afd 100644 --- a/openhands/automation/app.py +++ b/openhands/automation/app.py @@ -77,14 +77,16 @@ async def lifespan(app: FastAPI): # Set SQLite mode flag for scheduler/dispatcher to use set_sqlite_mode(engine_result.is_sqlite) - # Auto-run migrations for SQLite on startup - # This ensures the schema is always up-to-date for local deployments - # For PostgreSQL, migrations are typically run separately via `alembic upgrade head` - if engine_result.is_sqlite: + should_migrate = ( + settings.auto_migrate + if settings.auto_migrate is not None + else engine_result.is_sqlite + ) + if should_migrate: 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,16 +110,16 @@ 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) try: command.upgrade(alembic_cfg, "head") - logger.info("SQLite database migrations applied successfully") + logger.info("Database migrations applied successfully") except Exception as e: - logger.error(f"Failed to apply SQLite migrations: {e}") - msg = f"SQLite migration failed. Database may be inconsistent: {e}" + logger.error(f"Failed to apply migrations: {e}") + msg = f"Migration failed. Database may be inconsistent: {e}" raise RuntimeError(msg) from e # Start the background scheduler and dispatcher diff --git a/openhands/automation/config.py b/openhands/automation/config.py index e86aa8f3..15633585 100644 --- a/openhands/automation/config.py +++ b/openhands/automation/config.py @@ -471,6 +471,11 @@ class ServiceSettings(BaseSettings): # Database URL (alternative to host/port config, supports SQLite for local mode) AUTOMATION_DB_URL: Full database URL (e.g., sqlite+aiosqlite:////data/automations.db) + AUTOMATION_AUTO_MIGRATE: Run `alembic upgrade head` during startup. + Unset (the default) migrates SQLite only, which is what local + deployments have always done; true also migrates PostgreSQL; + false never migrates on startup, for deployments that run + migrations from their own pipeline. # GCP Cloud SQL AUTOMATION_GCP_DB_INSTANCE: Cloud SQL instance (optional) @@ -546,6 +551,10 @@ class ServiceSettings(BaseSettings): # - postgresql+asyncpg://user:pass@host/db (PostgreSQL) db_url: str = "" + # None keeps the historical behaviour: SQLite migrates itself, Postgres + # does not. + auto_migrate: bool | None = None + # GCP Cloud SQL (if set, takes precedence over host/port) gcp_db_instance: str | None = None gcp_project: str | None = None 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_config.py b/tests/test_config.py index 657cc3d0..22a30502 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -71,6 +71,20 @@ def test_failure_disable_threshold_uses_documented_env_var(self, monkeypatch): assert settings.failure_disable_threshold == 0 + def test_auto_migrate_defaults_to_unset(self, monkeypatch): + monkeypatch.delenv("AUTOMATION_AUTO_MIGRATE", raising=False) + + assert ServiceSettings().auto_migrate is None + + @pytest.mark.parametrize( + ("value", "expected"), + [("true", True), ("1", True), ("false", False), ("0", False)], + ) + def test_auto_migrate_parses_env_var(self, monkeypatch, value, expected): + monkeypatch.setenv("AUTOMATION_AUTO_MIGRATE", value) + + assert ServiceSettings().auto_migrate is expected + class TestBasePath: """Verify base_path is derived from base_url path + /api/automation.""" 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 diff --git a/tests/test_startup_migrations.py b/tests/test_startup_migrations.py new file mode 100644 index 00000000..ccc4988c --- /dev/null +++ b/tests/test_startup_migrations.py @@ -0,0 +1,101 @@ +"""Tests for the startup migration gate in the application lifespan.""" + +import asyncio +from unittest.mock import MagicMock, patch + +import pytest +from fastapi import FastAPI +from sqlalchemy.ext.asyncio import create_async_engine + +from openhands.automation import app as app_module +from openhands.automation.config import ServiceSettings +from openhands.automation.db import EngineResult, set_sqlite_mode + + +POSTGRES_URL = "postgresql+asyncpg://user:pass@db.example.com/automations" +SQLITE_URL = "sqlite+aiosqlite:///:memory:" + +BACKGROUND_LOOPS = ( + "scheduler_loop", + "dispatcher_loop", + "watchdog_loop", + "git_sync_loop", + "stream_supervisor_loop", +) + + +async def _idle_loop(*args, shutdown_event: asyncio.Event, **kwargs): + await shutdown_event.wait() + + +@pytest.fixture(autouse=True) +def reset_sqlite_mode(): + yield + set_sqlite_mode(False) + + +async def _run_lifespan(monkeypatch, *, db_url, is_sqlite, auto_migrate): + """Start and stop the lifespan, returning the mocked alembic upgrade.""" + settings = ServiceSettings(db_url=db_url, auto_migrate=auto_migrate) + engine_result = EngineResult( + engine=create_async_engine(SQLITE_URL), is_sqlite=is_sqlite + ) + + async def fake_create_engine(_settings=None): + return engine_result + + monkeypatch.setattr(app_module, "get_settings", lambda: settings) + monkeypatch.setattr(app_module, "create_engine", fake_create_engine) + for loop_name in BACKGROUND_LOOPS: + monkeypatch.setattr(app_module, loop_name, _idle_loop) + + upgrade = MagicMock() + with patch("alembic.command.upgrade", upgrade): + async with app_module.lifespan(FastAPI()): + pass + return upgrade + + +class TestStartupMigrationGate: + """AUTOMATION_AUTO_MIGRATE decides whether startup runs migrations.""" + + async def test_postgres_does_not_migrate_by_default(self, monkeypatch): + upgrade = await _run_lifespan( + monkeypatch, db_url=POSTGRES_URL, is_sqlite=False, auto_migrate=None + ) + + upgrade.assert_not_called() + + async def test_postgres_migrates_when_opted_in(self, monkeypatch): + upgrade = await _run_lifespan( + monkeypatch, db_url=POSTGRES_URL, is_sqlite=False, auto_migrate=True + ) + + upgrade.assert_called_once() + config = upgrade.call_args.args[0] + assert config.get_main_option("sqlalchemy.url") == ( + "postgresql+pg8000://user:pass@db.example.com/automations" + ) + + async def test_postgres_does_not_migrate_when_opted_out(self, monkeypatch): + upgrade = await _run_lifespan( + monkeypatch, db_url=POSTGRES_URL, is_sqlite=False, auto_migrate=False + ) + + upgrade.assert_not_called() + + async def test_sqlite_migrates_by_default(self, monkeypatch): + upgrade = await _run_lifespan( + monkeypatch, db_url=SQLITE_URL, is_sqlite=True, auto_migrate=None + ) + + upgrade.assert_called_once() + config = upgrade.call_args.args[0] + assert config.get_main_option("sqlalchemy.url") == "sqlite:///:memory:" + + async def test_sqlite_does_not_migrate_when_opted_out(self, monkeypatch): + upgrade = await _run_lifespan( + monkeypatch, db_url=SQLITE_URL, is_sqlite=True, auto_migrate=False + ) + + upgrade.assert_not_called()