diff --git a/hub_core/migrations/env.py b/hub_core/migrations/env.py index 90e95c3..d322e7b 100644 --- a/hub_core/migrations/env.py +++ b/hub_core/migrations/env.py @@ -5,6 +5,7 @@ from alembic import context from sqlalchemy import engine_from_config, pool from hub_core.models import Base +from hub_core.migrations.roles import migration_role_statement config = context.config @@ -38,6 +39,8 @@ def run_migrations_online() -> None: poolclass=pool.NullPool, ) with connectable.connect() as connection: + if statement := migration_role_statement(os.environ.get("HUB_CORE_MIGRATION_ROLE")): + connection.exec_driver_sql(statement) context.configure(connection=connection, target_metadata=target_metadata) with context.begin_transaction(): context.run_migrations() diff --git a/hub_core/migrations/roles.py b/hub_core/migrations/roles.py new file mode 100644 index 0000000..7f93ae3 --- /dev/null +++ b/hub_core/migrations/roles.py @@ -0,0 +1,12 @@ +from __future__ import annotations + +import re + + +def migration_role_statement(role: str | None) -> str | None: + """Return a safely quoted SET ROLE statement for an admitted owner role.""" + if not role: + return None + if not re.fullmatch(r"[a-z_][a-z0-9_]{0,62}", role): + raise ValueError("HUB_CORE_MIGRATION_ROLE is not a safe PostgreSQL role name") + return f'SET ROLE "{role}"' diff --git a/hub_core/runtime/postgres_store.py b/hub_core/runtime/postgres_store.py index 3d0d373..d58d2bd 100644 --- a/hub_core/runtime/postgres_store.py +++ b/hub_core/runtime/postgres_store.py @@ -22,6 +22,8 @@ from hub_core.runtime.models import ( RegistryRegistration, ) from hub_core.runtime.tables import ( + compat_api_keys, + compat_hubs, runtime_audit_ledger, runtime_interaction_events, runtime_messages, @@ -47,6 +49,13 @@ class PostgresPortStore: try: async with self.engine.connect() as connection: await connection.execute(sa.text("SELECT 1")) + # A live connection alone can hide a mis-granted runtime lease. + # Exercise the durable port, compatibility, and authorization + # tables that protected traffic actually depends on. + for table in (runtime_audit_ledger, compat_hubs, compat_api_keys): + await connection.execute( + sa.select(sa.literal(1)).select_from(table).limit(1) + ) except Exception: return {"database": "unavailable"} return {"database": "ok"} diff --git a/tests/test_migration_environment.py b/tests/test_migration_environment.py new file mode 100644 index 0000000..8eec3b4 --- /dev/null +++ b/tests/test_migration_environment.py @@ -0,0 +1,12 @@ +from __future__ import annotations + +import pytest + +from hub_core.migrations.roles import migration_role_statement + + +def test_migration_role_is_quoted_and_validated() -> None: + assert migration_role_statement("hub_runtime_owner") == 'SET ROLE "hub_runtime_owner"' + assert migration_role_statement(None) is None + with pytest.raises(ValueError, match="safe PostgreSQL role"): + migration_role_statement('owner"; DROP SCHEMA public; --') diff --git a/tests/test_postgres_store.py b/tests/test_postgres_store.py index 62e8d81..ba2854c 100644 --- a/tests/test_postgres_store.py +++ b/tests/test_postgres_store.py @@ -89,6 +89,22 @@ def test_postgresql_readiness_fails_when_database_is_unavailable() -> None: asyncio.run(store.aclose()) +def test_postgresql_readiness_fails_when_runtime_tables_are_unavailable(tmp_path) -> None: + database_url = f"sqlite+aiosqlite:///{tmp_path / 'empty.db'}" + store = PostgresPortStore.from_url(database_url) + settings = RuntimeSettings( + environment="production", + backend="postgresql", + allow_ephemeral=False, + database_url=database_url, + ) + response = TestClient(create_app(settings=settings, port_store=store)).get("/readyz") + + assert response.status_code == 503 + assert response.json()["checks"]["database"] == "unavailable" + asyncio.run(store.aclose()) + + def _event(event_type: str) -> dict: return { "schema_version": "0.1.0",