mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-07-21 20:06:06 +00:00
Junyan Qin
724 lines
34 KiB
Python
724 lines
34 KiB
Python
"""Real PostgreSQL coverage for the one-shot Cloud release migration job."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
import uuid
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import sqlalchemy as sa
|
|
from sqlalchemy import text
|
|
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
|
|
|
|
from langbot.pkg.persistence import release_migration
|
|
from langbot.pkg.persistence.alembic_runner import (
|
|
get_alembic_current,
|
|
get_alembic_head,
|
|
run_alembic_stamp,
|
|
run_alembic_upgrade,
|
|
)
|
|
from langbot.pkg.persistence.mgr import (
|
|
PersistenceManager,
|
|
PersistenceMode,
|
|
_RELEASE_MIGRATION_ADVISORY_LOCK_ID,
|
|
)
|
|
from langbot.pkg.utils import constants
|
|
|
|
|
|
pytestmark = [pytest.mark.integration, pytest.mark.slow, pytest.mark.asyncio]
|
|
_RUNTIME_PASSWORD = 'runtime-secret-not-used-by-migration'
|
|
|
|
|
|
@pytest.fixture
|
|
def postgres_url() -> str:
|
|
url = os.environ.get('TEST_POSTGRES_URL')
|
|
if not url:
|
|
pytest.skip('TEST_POSTGRES_URL not set')
|
|
return url
|
|
|
|
|
|
@pytest.fixture
|
|
async def postgres_engine(postgres_url: str):
|
|
engine = create_async_engine(postgres_url, isolation_level='AUTOCOMMIT')
|
|
yield engine
|
|
await engine.dispose()
|
|
|
|
|
|
@pytest.fixture
|
|
async def clean_database(postgres_engine: AsyncEngine):
|
|
async def clean() -> None:
|
|
async with postgres_engine.begin() as conn:
|
|
table_names = await conn.run_sync(lambda sync_conn: sa.inspect(sync_conn).get_table_names())
|
|
quote = postgres_engine.dialect.identifier_preparer.quote
|
|
for table_name in table_names:
|
|
await conn.execute(text(f'DROP TABLE {quote(table_name)} CASCADE'))
|
|
|
|
await clean()
|
|
yield
|
|
await clean()
|
|
|
|
|
|
def _restore_postgres_registry(monkeypatch) -> None:
|
|
from langbot.pkg.persistence import mgr as persistence_mgr_module
|
|
from langbot.pkg.persistence.databases.postgresql import PostgreSQLDatabaseManager
|
|
|
|
monkeypatch.setattr(
|
|
persistence_mgr_module.database,
|
|
'preregistered_managers',
|
|
[PostgreSQLDatabaseManager],
|
|
)
|
|
|
|
|
|
def _application(postgres_url: str, *, runtime_role: str = 'langbot_runtime_not_used_by_migration') -> SimpleNamespace:
|
|
url = sa.engine.make_url(postgres_url)
|
|
return SimpleNamespace(
|
|
instance_config=SimpleNamespace(
|
|
data={
|
|
'database': {
|
|
'use': 'postgresql',
|
|
'postgresql': {
|
|
'host': url.host,
|
|
'port': url.port,
|
|
# This is deliberately not the operator role in the DSN.
|
|
'user': runtime_role,
|
|
'password': _RUNTIME_PASSWORD,
|
|
'database': url.database,
|
|
},
|
|
'cloud_migration': {'operator_dsn_env': 'TEST_RELEASE_OPERATOR_DSN'},
|
|
},
|
|
'vdb': {
|
|
'use': 'pgvector',
|
|
'pgvector': {
|
|
'use_business_database': True,
|
|
'allowed_dimensions': [384, 512, 768, 1024, 1536],
|
|
},
|
|
},
|
|
}
|
|
),
|
|
logger=logging.getLogger('cloud-release-migration-entrypoint-test'),
|
|
persistence_mgr=None,
|
|
)
|
|
|
|
|
|
async def test_release_entrypoint_holds_lock_migrates_validates_and_disposes(
|
|
postgres_url: str,
|
|
postgres_engine: AsyncEngine,
|
|
clean_database,
|
|
monkeypatch,
|
|
) -> None:
|
|
_restore_postgres_registry(monkeypatch)
|
|
monkeypatch.setattr(constants, 'instance_id', 'release-migration-entrypoint-test')
|
|
original_validate = PersistenceManager._validate_release_schema
|
|
validation_observed_lock = False
|
|
runtime_role = f'lb_release_runtime_{uuid.uuid4().hex[:12]}'
|
|
quote = postgres_engine.dialect.identifier_preparer.quote
|
|
|
|
async def validate_while_asserting_lock(self: PersistenceManager) -> None:
|
|
nonlocal validation_observed_lock
|
|
async with postgres_engine.connect() as conn:
|
|
acquired = await conn.scalar(
|
|
text('SELECT pg_try_advisory_lock(:lock_id)'),
|
|
{'lock_id': _RELEASE_MIGRATION_ADVISORY_LOCK_ID},
|
|
)
|
|
if acquired:
|
|
await conn.scalar(
|
|
text('SELECT pg_advisory_unlock(:lock_id)'),
|
|
{'lock_id': _RELEASE_MIGRATION_ADVISORY_LOCK_ID},
|
|
)
|
|
assert acquired is False
|
|
validation_observed_lock = True
|
|
await original_validate(self)
|
|
|
|
monkeypatch.setattr(PersistenceManager, '_validate_release_schema', validate_while_asserting_lock)
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(text(f"CREATE ROLE {quote(runtime_role)} LOGIN PASSWORD '{_RUNTIME_PASSWORD}'"))
|
|
# The release job must make this otherwise bare LOGIN usable without
|
|
# relying on pre-provisioned object ACLs.
|
|
assert (
|
|
await conn.scalar(
|
|
text(
|
|
"""
|
|
SELECT NOT EXISTS (
|
|
SELECT 1
|
|
FROM pg_class c
|
|
JOIN pg_namespace n ON n.oid = c.relnamespace
|
|
CROSS JOIN LATERAL aclexplode(c.relacl) acl
|
|
JOIN pg_roles grantee ON grantee.oid = acl.grantee
|
|
WHERE n.nspname = current_schema()
|
|
AND grantee.rolname = :runtime_role
|
|
)
|
|
"""
|
|
),
|
|
{'runtime_role': runtime_role},
|
|
)
|
|
is True
|
|
)
|
|
ap = _application(postgres_url, runtime_role=runtime_role)
|
|
try:
|
|
await release_migration.run_cloud_release_migration(
|
|
ap,
|
|
environ={'TEST_RELEASE_OPERATOR_DSN': postgres_url},
|
|
)
|
|
|
|
assert validation_observed_lock is True
|
|
assert await get_alembic_current(postgres_engine) == get_alembic_head()
|
|
manager = ap.persistence_mgr
|
|
assert isinstance(manager, PersistenceManager)
|
|
business_tables = set(manager._runtime_business_table_names())
|
|
async with postgres_engine.connect() as conn:
|
|
assert await conn.scalar(text("SELECT to_regclass('langbot_vectors') IS NOT NULL")) is True
|
|
assert (
|
|
await conn.scalar(text("SELECT EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'vector')")) is True
|
|
)
|
|
runtime_role_state = (
|
|
(
|
|
await conn.execute(
|
|
text(
|
|
"""
|
|
SELECT
|
|
rolcanlogin,
|
|
rolsuper,
|
|
rolbypassrls,
|
|
rolcreatedb,
|
|
rolcreaterole,
|
|
rolreplication
|
|
FROM pg_roles
|
|
WHERE rolname = :runtime_role
|
|
"""
|
|
),
|
|
{'runtime_role': runtime_role},
|
|
)
|
|
)
|
|
.mappings()
|
|
.one()
|
|
)
|
|
assert dict(runtime_role_state) == {
|
|
'rolcanlogin': True,
|
|
'rolsuper': False,
|
|
'rolbypassrls': False,
|
|
'rolcreatedb': False,
|
|
'rolcreaterole': False,
|
|
'rolreplication': False,
|
|
}
|
|
|
|
database_privileges = set(
|
|
(
|
|
await conn.execute(
|
|
text(
|
|
"""
|
|
SELECT acl.privilege_type
|
|
FROM pg_database database
|
|
CROSS JOIN LATERAL aclexplode(database.datacl) acl
|
|
JOIN pg_roles grantee ON grantee.oid = acl.grantee
|
|
WHERE database.datname = current_database()
|
|
AND grantee.rolname = :runtime_role
|
|
AND acl.is_grantable IS FALSE
|
|
"""
|
|
),
|
|
{'runtime_role': runtime_role},
|
|
)
|
|
)
|
|
.scalars()
|
|
.all()
|
|
)
|
|
schema_privileges = set(
|
|
(
|
|
await conn.execute(
|
|
text(
|
|
"""
|
|
SELECT acl.privilege_type
|
|
FROM pg_namespace namespace
|
|
CROSS JOIN LATERAL aclexplode(namespace.nspacl) acl
|
|
JOIN pg_roles grantee ON grantee.oid = acl.grantee
|
|
WHERE namespace.nspname = current_schema()
|
|
AND grantee.rolname = :runtime_role
|
|
AND acl.is_grantable IS FALSE
|
|
"""
|
|
),
|
|
{'runtime_role': runtime_role},
|
|
)
|
|
)
|
|
.scalars()
|
|
.all()
|
|
)
|
|
assert database_privileges == {'CONNECT'}
|
|
assert schema_privileges == {'USAGE'}
|
|
assert (
|
|
await conn.scalar(
|
|
text("SELECT has_database_privilege(:runtime_role, current_database(), 'CREATE')"),
|
|
{'runtime_role': runtime_role},
|
|
)
|
|
is False
|
|
)
|
|
assert (
|
|
await conn.scalar(
|
|
text("SELECT has_schema_privilege(:runtime_role, current_schema(), 'CREATE')"),
|
|
{'runtime_role': runtime_role},
|
|
)
|
|
is False
|
|
)
|
|
# PostgreSQL grants TEMP to PUBLIC by default. The first Cloud
|
|
# release deliberately tolerates that inherited compatibility
|
|
# privilege while granting no direct TEMP ACL to the runtime role.
|
|
assert (
|
|
await conn.scalar(
|
|
text("SELECT has_database_privilege(:runtime_role, current_database(), 'TEMP')"),
|
|
{'runtime_role': runtime_role},
|
|
)
|
|
is True
|
|
)
|
|
|
|
object_grants = (
|
|
(
|
|
await conn.execute(
|
|
text(
|
|
"""
|
|
SELECT
|
|
c.relname,
|
|
c.relkind::text AS relkind,
|
|
acl.privilege_type,
|
|
acl.is_grantable
|
|
FROM pg_class c
|
|
JOIN pg_namespace n ON n.oid = c.relnamespace
|
|
CROSS JOIN LATERAL aclexplode(c.relacl) acl
|
|
JOIN pg_roles grantee ON grantee.oid = acl.grantee
|
|
WHERE n.nspname = current_schema()
|
|
AND c.relkind IN ('r', 'p', 'S')
|
|
AND grantee.rolname = :runtime_role
|
|
ORDER BY c.relname, acl.privilege_type
|
|
"""
|
|
),
|
|
{'runtime_role': runtime_role},
|
|
)
|
|
)
|
|
.mappings()
|
|
.all()
|
|
)
|
|
direct_table_grants: dict[str, set[str]] = {}
|
|
direct_sequence_grants: dict[str, set[str]] = {}
|
|
for grant in object_grants:
|
|
assert grant['is_grantable'] is False
|
|
target = direct_sequence_grants if grant['relkind'] == 'S' else direct_table_grants
|
|
target.setdefault(grant['relname'], set()).add(grant['privilege_type'])
|
|
|
|
assert set(direct_table_grants) == business_tables | {'alembic_version'}
|
|
assert all(
|
|
privileges == {'SELECT', 'INSERT', 'UPDATE', 'DELETE'}
|
|
for table_name, privileges in direct_table_grants.items()
|
|
if table_name != 'alembic_version'
|
|
)
|
|
assert direct_table_grants['alembic_version'] == {'SELECT'}
|
|
assert direct_sequence_grants
|
|
assert all(privileges == {'USAGE', 'SELECT'} for privileges in direct_sequence_grants.values())
|
|
# The session-level lock must be released before the one-shot job exits.
|
|
assert (
|
|
await conn.scalar(
|
|
text('SELECT pg_try_advisory_lock(:lock_id)'),
|
|
{'lock_id': _RELEASE_MIGRATION_ADVISORY_LOCK_ID},
|
|
)
|
|
is True
|
|
)
|
|
assert (
|
|
await conn.scalar(
|
|
text('SELECT pg_advisory_unlock(:lock_id)'),
|
|
{'lock_id': _RELEASE_MIGRATION_ADVISORY_LOCK_ID},
|
|
)
|
|
is True
|
|
)
|
|
|
|
runtime_url = (
|
|
sa.engine.make_url(postgres_url)
|
|
.set(username=runtime_role, password=_RUNTIME_PASSWORD)
|
|
.render_as_string(hide_password=False)
|
|
)
|
|
runtime_engine = create_async_engine(runtime_url)
|
|
try:
|
|
async with runtime_engine.begin() as conn:
|
|
assert await conn.scalar(text('SELECT current_user')) == runtime_role
|
|
await conn.execute(text("INSERT INTO metadata (key, value) VALUES ('runtime-grant-smoke', 'created')"))
|
|
await conn.execute(text("UPDATE metadata SET value = 'updated' WHERE key = 'runtime-grant-smoke'"))
|
|
assert (
|
|
await conn.scalar(text("SELECT value FROM metadata WHERE key = 'runtime-grant-smoke'")) == 'updated'
|
|
)
|
|
await conn.execute(text("DELETE FROM metadata WHERE key = 'runtime-grant-smoke'"))
|
|
await conn.scalar(
|
|
text('SELECT nextval(CAST(:sequence_name AS regclass))'),
|
|
{'sequence_name': f'public.{sorted(direct_sequence_grants)[0]}'},
|
|
)
|
|
finally:
|
|
await runtime_engine.dispose()
|
|
|
|
runtime_application = _application(postgres_url, runtime_role=runtime_role)
|
|
runtime_manager = PersistenceManager(runtime_application, mode=PersistenceMode.CLOUD_RUNTIME)
|
|
runtime_application.persistence_mgr = runtime_manager
|
|
try:
|
|
# This reads alembic_version as the actual runtime role and then
|
|
# reruns the complete grant/catalog validator before startup.
|
|
await runtime_manager.initialize()
|
|
finally:
|
|
await runtime_manager.get_db_engine().dispose()
|
|
|
|
# The catalog validator must fail closed if the role is no longer
|
|
# deployable, even though the remaining grants still look plausible.
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(text(f'REVOKE DELETE ON TABLE public.metadata FROM {quote(runtime_role)}'))
|
|
with pytest.raises(RuntimeError, match="table 'metadata' grants are incomplete"):
|
|
await manager._validate_configured_runtime_postgres_role(require_grants=True)
|
|
finally:
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(text(f'DROP OWNED BY {quote(runtime_role)}'))
|
|
await conn.execute(text(f'DROP ROLE IF EXISTS {quote(runtime_role)}'))
|
|
|
|
|
|
async def test_release_entrypoint_rejects_privileged_or_table_owning_runtime_role(
|
|
postgres_url: str,
|
|
postgres_engine: AsyncEngine,
|
|
clean_database,
|
|
monkeypatch,
|
|
) -> None:
|
|
_restore_postgres_registry(monkeypatch)
|
|
monkeypatch.setattr(constants, 'instance_id', 'release-runtime-role-validation-test')
|
|
runtime_role = f'lb_release_runtime_{uuid.uuid4().hex[:12]}'
|
|
quote = postgres_engine.dialect.identifier_preparer.quote
|
|
async with postgres_engine.connect() as conn:
|
|
operator_role = await conn.scalar(text('SELECT current_user'))
|
|
await conn.execute(text(f'CREATE ROLE {quote(runtime_role)} LOGIN'))
|
|
|
|
ap = _application(postgres_url, runtime_role=runtime_role)
|
|
try:
|
|
await release_migration.run_cloud_release_migration(
|
|
ap,
|
|
environ={'TEST_RELEASE_OPERATOR_DSN': postgres_url},
|
|
)
|
|
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} SUPERUSER'))
|
|
with pytest.raises(RuntimeError, match='must not be superuser or BYPASSRLS'):
|
|
await release_migration.run_cloud_release_migration(
|
|
ap,
|
|
environ={'TEST_RELEASE_OPERATOR_DSN': postgres_url},
|
|
)
|
|
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} NOSUPERUSER BYPASSRLS'))
|
|
with pytest.raises(RuntimeError, match='must not be superuser or BYPASSRLS'):
|
|
await release_migration.run_cloud_release_migration(
|
|
ap,
|
|
environ={'TEST_RELEASE_OPERATOR_DSN': postgres_url},
|
|
)
|
|
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} NOBYPASSRLS'))
|
|
await conn.execute(text(f'ALTER TABLE bots OWNER TO {quote(runtime_role)}'))
|
|
with pytest.raises(RuntimeError, match='owns tenant tables'):
|
|
await release_migration.run_cloud_release_migration(
|
|
ap,
|
|
environ={'TEST_RELEASE_OPERATOR_DSN': postgres_url},
|
|
)
|
|
finally:
|
|
async with postgres_engine.connect() as conn:
|
|
table_exists = await conn.scalar(text("SELECT to_regclass('bots') IS NOT NULL"))
|
|
if table_exists:
|
|
await conn.execute(text(f'ALTER TABLE bots OWNER TO {quote(operator_role)}'))
|
|
await conn.execute(text(f'DROP OWNED BY {quote(runtime_role)}'))
|
|
await conn.execute(text(f'DROP ROLE IF EXISTS {quote(runtime_role)}'))
|
|
|
|
|
|
async def test_runtime_role_catalog_validator_rejects_delegation_and_escape_hatches(
|
|
postgres_url: str,
|
|
postgres_engine: AsyncEngine,
|
|
clean_database,
|
|
monkeypatch,
|
|
) -> None:
|
|
_restore_postgres_registry(monkeypatch)
|
|
monkeypatch.setattr(constants, 'instance_id', 'release-runtime-role-catalog-test')
|
|
suffix = uuid.uuid4().hex[:12]
|
|
runtime_role = f'lb_release_runtime_{suffix}'
|
|
delegated_role = f'lb_release_delegate_{suffix}'
|
|
extra_schema = f'lb_release_schema_{suffix}'
|
|
extra_view = f'lb_release_view_{suffix}'
|
|
owned_routine = f'lb_release_owned_routine_{suffix}'
|
|
security_definer = f'lb_release_definer_{suffix}'
|
|
foreign_wrapper = f'lb_release_fdw_{suffix}'
|
|
foreign_server = f'lb_release_server_{suffix}'
|
|
persistent_setting_secret = f'lb_release_secret_{suffix}'
|
|
database_name = sa.engine.make_url(postgres_url).database
|
|
assert database_name
|
|
quote = postgres_engine.dialect.identifier_preparer.quote
|
|
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(text(f"CREATE ROLE {quote(runtime_role)} LOGIN PASSWORD '{_RUNTIME_PASSWORD}'"))
|
|
await conn.execute(text(f'CREATE ROLE {quote(delegated_role)}'))
|
|
|
|
ap = _application(postgres_url, runtime_role=runtime_role)
|
|
try:
|
|
await release_migration.run_cloud_release_migration(
|
|
ap,
|
|
environ={'TEST_RELEASE_OPERATOR_DSN': postgres_url},
|
|
)
|
|
manager = ap.persistence_mgr
|
|
assert isinstance(manager, PersistenceManager)
|
|
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(text(f'GRANT pg_read_all_data TO {quote(runtime_role)}'))
|
|
with pytest.raises(RuntimeError, match='must not participate in role memberships'):
|
|
await manager._validate_configured_runtime_postgres_role()
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(text(f'REVOKE pg_read_all_data FROM {quote(runtime_role)}'))
|
|
|
|
# Reject delegation in the other direction too: no role may be
|
|
# allowed to SET ROLE to the runtime identity or administer it.
|
|
await conn.execute(text(f'GRANT {quote(runtime_role)} TO {quote(delegated_role)} WITH ADMIN OPTION'))
|
|
with pytest.raises(RuntimeError, match='must not participate in role memberships'):
|
|
await manager._validate_configured_runtime_postgres_role()
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(text(f'REVOKE {quote(runtime_role)} FROM {quote(delegated_role)}'))
|
|
|
|
await conn.execute(
|
|
text(f'GRANT SELECT ON TABLE public.metadata TO {quote(runtime_role)} WITH GRANT OPTION')
|
|
)
|
|
with pytest.raises(RuntimeError, match='GRANT OPTION'):
|
|
await manager._validate_configured_runtime_postgres_role()
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(
|
|
text(f'REVOKE GRANT OPTION FOR SELECT ON TABLE public.metadata FROM {quote(runtime_role)}')
|
|
)
|
|
|
|
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} SET search_path TO public'))
|
|
with pytest.raises(RuntimeError, match='persistent session overrides'):
|
|
await manager._validate_configured_runtime_postgres_role()
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} RESET search_path'))
|
|
|
|
await conn.execute(text(f'ALTER DATABASE {quote(database_name)} SET search_path TO public'))
|
|
with pytest.raises(RuntimeError, match='persistent session overrides'):
|
|
await manager._validate_configured_runtime_postgres_role()
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(text(f'ALTER DATABASE {quote(database_name)} RESET search_path'))
|
|
|
|
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} SET session_replication_role TO replica'))
|
|
with pytest.raises(RuntimeError, match='persistent session overrides'):
|
|
await manager._validate_configured_runtime_postgres_role()
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} RESET session_replication_role'))
|
|
|
|
await conn.execute(
|
|
text(f"ALTER ROLE {quote(runtime_role)} SET application_name TO '{persistent_setting_secret}'")
|
|
)
|
|
with pytest.raises(RuntimeError, match='persistent session overrides') as error:
|
|
await manager._validate_configured_runtime_postgres_role()
|
|
assert persistent_setting_secret not in str(error.value)
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} RESET application_name'))
|
|
|
|
await conn.execute(text('CREATE EXTENSION dblink'))
|
|
with pytest.raises(RuntimeError, match='extensions must include vector and be limited'):
|
|
await manager._validate_configured_runtime_postgres_role()
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(text('DROP EXTENSION dblink'))
|
|
await manager._validate_configured_runtime_postgres_role()
|
|
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(text(f'GRANT CREATE ON DATABASE {quote(database_name)} TO {quote(runtime_role)}'))
|
|
await conn.execute(text(f'GRANT CREATE ON SCHEMA public TO {quote(runtime_role)}'))
|
|
runtime_url = (
|
|
sa.engine.make_url(postgres_url)
|
|
.set(username=runtime_role, password=_RUNTIME_PASSWORD)
|
|
.render_as_string(hide_password=False)
|
|
)
|
|
runtime_engine = create_async_engine(runtime_url)
|
|
try:
|
|
async with runtime_engine.begin() as conn:
|
|
# hstore is a trusted extension in the production PG16 image,
|
|
# so this creates a real runtime-owned extension catalog row.
|
|
await conn.execute(text('CREATE EXTENSION hstore'))
|
|
finally:
|
|
await runtime_engine.dispose()
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(text(f'REVOKE CREATE ON SCHEMA public FROM {quote(runtime_role)}'))
|
|
await conn.execute(text(f'REVOKE CREATE ON DATABASE {quote(database_name)} FROM {quote(runtime_role)}'))
|
|
with pytest.raises(RuntimeError, match='must not own extensions'):
|
|
await manager._validate_configured_runtime_postgres_role()
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(text('DROP EXTENSION hstore CASCADE'))
|
|
await manager._validate_configured_runtime_postgres_role()
|
|
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(text(f'CREATE FOREIGN DATA WRAPPER {quote(foreign_wrapper)}'))
|
|
await conn.execute(
|
|
text(f'CREATE SERVER {quote(foreign_server)} FOREIGN DATA WRAPPER {quote(foreign_wrapper)}')
|
|
)
|
|
await conn.execute(text(f'CREATE USER MAPPING FOR {quote(runtime_role)} SERVER {quote(foreign_server)}'))
|
|
with pytest.raises(RuntimeError, match='foreign data wrappers, servers, or user mappings'):
|
|
await manager._validate_configured_runtime_postgres_role()
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(text(f'DROP FOREIGN DATA WRAPPER {quote(foreign_wrapper)} CASCADE'))
|
|
await manager._validate_configured_runtime_postgres_role()
|
|
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(text(f'CREATE SCHEMA {quote(extra_schema)} AUTHORIZATION {quote(runtime_role)}'))
|
|
with pytest.raises(RuntimeError, match='non-business schemas'):
|
|
await manager._validate_configured_runtime_postgres_role()
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(text(f'DROP SCHEMA {quote(extra_schema)} CASCADE'))
|
|
|
|
await conn.execute(text(f'CREATE VIEW public.{quote(extra_view)} AS SELECT key FROM public.metadata'))
|
|
await conn.execute(text(f'GRANT SELECT ON public.{quote(extra_view)} TO {quote(runtime_role)}'))
|
|
with pytest.raises(RuntimeError, match='non-business objects|table privileges are unsafe'):
|
|
await manager._validate_configured_runtime_postgres_role()
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(text(f'DROP VIEW public.{quote(extra_view)}'))
|
|
|
|
await conn.execute(text(f'GRANT SELECT (key) ON public.metadata TO {quote(runtime_role)}'))
|
|
with pytest.raises(RuntimeError, match='column-level ACLs'):
|
|
await manager._validate_configured_runtime_postgres_role()
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(text(f'REVOKE SELECT (key) ON public.metadata FROM {quote(runtime_role)}'))
|
|
|
|
# System file access functions are not SECURITY DEFINER, so the
|
|
# validator must reject their explicit EXECUTE ACL independently.
|
|
await conn.execute(
|
|
text(f'GRANT EXECUTE ON FUNCTION pg_catalog.pg_read_file(text) TO {quote(runtime_role)}')
|
|
)
|
|
runtime_application = _application(postgres_url, runtime_role=runtime_role)
|
|
runtime_manager = PersistenceManager(runtime_application, mode=PersistenceMode.CLOUD_RUNTIME)
|
|
runtime_application.persistence_mgr = runtime_manager
|
|
try:
|
|
with pytest.raises(RuntimeError, match='explicit EXECUTE privileges on routines'):
|
|
await runtime_manager.initialize()
|
|
finally:
|
|
await runtime_manager.get_db_engine().dispose()
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(
|
|
text(f'REVOKE EXECUTE ON FUNCTION pg_catalog.pg_read_file(text) FROM {quote(runtime_role)}')
|
|
)
|
|
await manager._validate_configured_runtime_postgres_role()
|
|
|
|
async with postgres_engine.connect() as conn:
|
|
# replica disables ordinary triggers/rules and foreign-key
|
|
# enforcement; it must never reach the runtime identity.
|
|
await conn.execute(text(f'GRANT SET ON PARAMETER session_replication_role TO {quote(runtime_role)}'))
|
|
with pytest.raises(RuntimeError, match='explicit SET or ALTER SYSTEM parameter privileges'):
|
|
await manager._validate_configured_runtime_postgres_role()
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(text(f'REVOKE SET ON PARAMETER session_replication_role FROM {quote(runtime_role)}'))
|
|
await manager._validate_configured_runtime_postgres_role()
|
|
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(
|
|
text(f"CREATE FUNCTION public.{quote(owned_routine)}() RETURNS integer LANGUAGE sql AS 'SELECT 1'")
|
|
)
|
|
await conn.execute(text(f'ALTER FUNCTION public.{quote(owned_routine)}() OWNER TO {quote(runtime_role)}'))
|
|
with pytest.raises(RuntimeError, match='must not own routines'):
|
|
await manager._validate_configured_runtime_postgres_role()
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(text(f'DROP FUNCTION public.{quote(owned_routine)}()'))
|
|
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(
|
|
text(
|
|
f'CREATE FUNCTION public.{quote(security_definer)}() RETURNS integer '
|
|
"LANGUAGE sql SECURITY DEFINER AS 'SELECT 1'"
|
|
)
|
|
)
|
|
# Extension membership must not exempt an executable definer from
|
|
# the runtime audit, even for an allowlisted extension.
|
|
await conn.execute(text(f'ALTER EXTENSION vector ADD FUNCTION public.{quote(security_definer)}()'))
|
|
with pytest.raises(RuntimeError, match='SECURITY DEFINER'):
|
|
await manager._validate_configured_runtime_postgres_role()
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(text(f'ALTER EXTENSION vector DROP FUNCTION public.{quote(security_definer)}()'))
|
|
finally:
|
|
async with postgres_engine.connect() as conn:
|
|
await conn.execute(text(f'ALTER DATABASE {quote(database_name)} RESET search_path'))
|
|
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} RESET search_path'))
|
|
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} RESET session_replication_role'))
|
|
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} RESET application_name'))
|
|
await conn.execute(text(f'REVOKE pg_read_all_data FROM {quote(runtime_role)}'))
|
|
await conn.execute(text(f'REVOKE {quote(runtime_role)} FROM {quote(delegated_role)}'))
|
|
await conn.execute(text(f'REVOKE CREATE ON SCHEMA public FROM {quote(runtime_role)}'))
|
|
await conn.execute(text(f'REVOKE CREATE ON DATABASE {quote(database_name)} FROM {quote(runtime_role)}'))
|
|
await conn.execute(
|
|
text(f'REVOKE EXECUTE ON FUNCTION pg_catalog.pg_read_file(text) FROM {quote(runtime_role)}')
|
|
)
|
|
await conn.execute(text(f'REVOKE SET ON PARAMETER session_replication_role FROM {quote(runtime_role)}'))
|
|
await conn.execute(text('DROP EXTENSION IF EXISTS dblink CASCADE'))
|
|
await conn.execute(text('DROP EXTENSION IF EXISTS hstore CASCADE'))
|
|
await conn.execute(text(f'DROP FOREIGN DATA WRAPPER IF EXISTS {quote(foreign_wrapper)} CASCADE'))
|
|
await conn.execute(text(f'DROP FUNCTION IF EXISTS public.{quote(owned_routine)}()'))
|
|
security_definer_is_extension_member = await conn.scalar(
|
|
text(
|
|
"""
|
|
SELECT EXISTS (
|
|
SELECT 1
|
|
FROM pg_depend dependency
|
|
JOIN pg_extension extension ON extension.oid = dependency.refobjid
|
|
WHERE dependency.classid = 'pg_proc'::regclass
|
|
AND dependency.objid = to_regprocedure(:routine)
|
|
AND dependency.refclassid = 'pg_extension'::regclass
|
|
AND dependency.deptype = 'e'
|
|
AND extension.extname = 'vector'
|
|
)
|
|
"""
|
|
),
|
|
{'routine': f'public.{security_definer}()'},
|
|
)
|
|
if security_definer_is_extension_member:
|
|
await conn.execute(text(f'ALTER EXTENSION vector DROP FUNCTION public.{quote(security_definer)}()'))
|
|
await conn.execute(text(f'DROP FUNCTION IF EXISTS public.{quote(security_definer)}()'))
|
|
await conn.execute(text(f'DROP VIEW IF EXISTS public.{quote(extra_view)}'))
|
|
await conn.execute(text(f'DROP SCHEMA IF EXISTS {quote(extra_schema)} CASCADE'))
|
|
await conn.execute(text(f'DROP OWNED BY {quote(runtime_role)}'))
|
|
await conn.execute(text(f'DROP ROLE IF EXISTS {quote(delegated_role)}'))
|
|
await conn.execute(text(f'DROP ROLE IF EXISTS {quote(runtime_role)}'))
|
|
|
|
|
|
async def test_direct_postgres_head_stamp_fails_without_business_schema(
|
|
postgres_engine: AsyncEngine,
|
|
clean_database,
|
|
) -> None:
|
|
await run_alembic_stamp(postgres_engine, '0012_plugin_identity')
|
|
|
|
with pytest.raises(RuntimeError, match='requires the knowledge_bases table'):
|
|
await run_alembic_upgrade(postgres_engine)
|
|
|
|
assert await get_alembic_current(postgres_engine) == '0012_plugin_identity'
|
|
|
|
|
|
async def test_release_entrypoint_fails_immediately_when_another_job_holds_lock(
|
|
postgres_url: str,
|
|
postgres_engine: AsyncEngine,
|
|
clean_database,
|
|
monkeypatch,
|
|
) -> None:
|
|
_restore_postgres_registry(monkeypatch)
|
|
ap = _application(postgres_url)
|
|
|
|
async with postgres_engine.connect() as lock_connection:
|
|
assert (
|
|
await lock_connection.scalar(
|
|
text('SELECT pg_try_advisory_lock(:lock_id)'),
|
|
{'lock_id': _RELEASE_MIGRATION_ADVISORY_LOCK_ID},
|
|
)
|
|
is True
|
|
)
|
|
try:
|
|
with pytest.raises(RuntimeError, match='already holds the advisory lock'):
|
|
await release_migration.run_cloud_release_migration(
|
|
ap,
|
|
environ={'TEST_RELEASE_OPERATOR_DSN': postgres_url},
|
|
)
|
|
finally:
|
|
assert (
|
|
await lock_connection.scalar(
|
|
text('SELECT pg_advisory_unlock(:lock_id)'),
|
|
{'lock_id': _RELEASE_MIGRATION_ADVISORY_LOCK_ID},
|
|
)
|
|
is True
|
|
)
|
|
|
|
async with postgres_engine.connect() as conn:
|
|
assert await conn.run_sync(lambda sync_conn: sa.inspect(sync_conn).get_table_names()) == []
|