diff --git a/src/langbot/pkg/persistence/alembic/versions/0032_certification_artifact_digest_backfill.py b/src/langbot/pkg/persistence/alembic/versions/0032_certification_artifact_digest_backfill.py new file mode 100644 index 000000000..1dc4e613f --- /dev/null +++ b/src/langbot/pkg/persistence/alembic/versions/0032_certification_artifact_digest_backfill.py @@ -0,0 +1,138 @@ +"""Backfill raw artifact digests for legacy shared certificates. + +Revision ID: 0032_cert_artifact_digest +Revises: 0031_merge_totp_assistant + +Older certified-plugin rows persisted every shared-admission fact except the raw +archive digest. Placement now requires that digest to match the installation's +immutable artifact digest, so this migration fills only records whose remaining +facts already prove the exact legacy shared-admission shape. +""" + +from __future__ import annotations + +from collections.abc import Mapping + +import sqlalchemy as sa +from alembic import op + +revision = '0032_cert_artifact_digest' +down_revision = '0031_merge_totp_assistant' +branch_labels = None +depends_on = None + +_TABLE = 'plugin_settings' +_SHARED_RUNTIME = 'shared-runtime-v1' +_SHARED_ADMISSION_CODE = 'CERTIFIED_PLUGIN_SHARED_ELIGIBLE' +_LOWER_HEX = frozenset('0123456789abcdef') +_HEX = frozenset('0123456789abcdefABCDEF') + + +def _is_digest(value: object, *, lowercase: bool) -> bool: + allowed = _LOWER_HEX if lowercase else _HEX + return isinstance(value, str) and len(value) == 64 and all(character in allowed for character in value) + + +def _eligible_certification(install_info: object, artifact_digest: object) -> Mapping[object, object] | None: + if not _is_digest(artifact_digest, lowercase=True) or not isinstance(install_info, Mapping): + return None + certification = install_info.get('_certification') + if not isinstance(certification, Mapping) or 'artifact_digest' in certification: + return None + certificate_id = certification.get('certificate_id') + if not isinstance(certificate_id, str) or not certificate_id.strip(): + return None + if not _is_digest(certification.get('normalized_digest'), lowercase=False): + return None + required_facts = { + 'verification': 'valid', + 'certificate_runtime_profile': _SHARED_RUNTIME, + 'runtime_profile': _SHARED_RUNTIME, + 'admission_code': _SHARED_ADMISSION_CODE, + } + if not all(certification.get(key) == value for key, value in required_facts.items()): + return None + return certification + + +def _suspend_postgres_rls(conn: sa.Connection) -> tuple[bool, bool]: + if conn.dialect.name != 'postgresql': + return False, False + state = conn.execute( + sa.text('SELECT relrowsecurity, relforcerowsecurity FROM pg_class WHERE oid = to_regclass(:table_name)'), + {'table_name': _TABLE}, + ).one() + rls_enabled, rls_forced = bool(state.relrowsecurity), bool(state.relforcerowsecurity) + if rls_forced: + conn.execute(sa.text(f'ALTER TABLE {_TABLE} NO FORCE ROW LEVEL SECURITY')) + if rls_enabled: + conn.execute(sa.text(f'ALTER TABLE {_TABLE} DISABLE ROW LEVEL SECURITY')) + return rls_enabled, rls_forced + + +def _restore_postgres_rls(conn: sa.Connection, state: tuple[bool, bool]) -> None: + if conn.dialect.name != 'postgresql': + return + rls_enabled, rls_forced = state + if rls_enabled: + conn.execute(sa.text(f'ALTER TABLE {_TABLE} ENABLE ROW LEVEL SECURITY')) + if rls_forced: + conn.execute(sa.text(f'ALTER TABLE {_TABLE} FORCE ROW LEVEL SECURITY')) + + +def backfill_certification_artifact_digests(conn: sa.Connection) -> None: + inspector = sa.inspect(conn) + if _TABLE not in inspector.get_table_names(): + return + columns = {column['name'] for column in inspector.get_columns(_TABLE)} + required_columns = { + 'workspace_uuid', + 'plugin_author', + 'plugin_name', + 'artifact_digest', + 'install_info', + } + if not required_columns <= columns: + return + + plugin_settings = sa.table( + _TABLE, + sa.column('workspace_uuid', sa.String(36)), + sa.column('plugin_author', sa.String(255)), + sa.column('plugin_name', sa.String(255)), + sa.column('artifact_digest', sa.String(64)), + sa.column('install_info', sa.JSON()), + ) + rows = conn.execute(sa.select(plugin_settings)).mappings().all() + for row in rows: + certification = _eligible_certification(row['install_info'], row['artifact_digest']) + if certification is None: + continue + updated_certification = dict(certification) + updated_certification['artifact_digest'] = row['artifact_digest'] + updated_install_info = dict(row['install_info']) + updated_install_info['_certification'] = updated_certification + conn.execute( + plugin_settings.update() + .where(plugin_settings.c.workspace_uuid == row['workspace_uuid']) + .where(plugin_settings.c.plugin_author == row['plugin_author']) + .where(plugin_settings.c.plugin_name == row['plugin_name']) + .values(install_info=updated_install_info) + ) + + +def upgrade() -> None: + conn = op.get_bind() + if _TABLE not in sa.inspect(conn).get_table_names(): + return + rls_state = _suspend_postgres_rls(conn) + try: + backfill_certification_artifact_digests(conn) + finally: + _restore_postgres_rls(conn, rls_state) + + +def downgrade() -> None: + # The source digest cannot be distinguished from a digest persisted by + # current Core, so downgrade intentionally preserves the safe enrichment. + pass diff --git a/tests/integration/persistence/test_migration_branch_convergence.py b/tests/integration/persistence/test_migration_branch_convergence.py index 5b03caa29..759b46b03 100644 --- a/tests/integration/persistence/test_migration_branch_convergence.py +++ b/tests/integration/persistence/test_migration_branch_convergence.py @@ -16,6 +16,7 @@ import uuid import pytest import sqlalchemy as sa from sqlalchemy.ext.asyncio import create_async_engine +from langbot_plugin.entities.io.context import PluginExecutionMode from langbot.pkg.entity import persistence from langbot.pkg.entity.persistence.base import Base @@ -25,6 +26,7 @@ from langbot.pkg.persistence.alembic_runner import ( run_alembic_downgrade, run_alembic_upgrade, ) +from langbot.pkg.plugin.certification import execution_mode_for_persisted_installation from langbot.pkg.utils import importutil @@ -353,6 +355,143 @@ async def test_empty_database_startup_schema_then_real_migrations(convergence_en assert await get_alembic_current(engine) == get_alembic_head() +def _legacy_shared_certification(**overrides): + certification = { + 'normalized_digest': 'B' * 64, + 'verification': 'valid', + 'certificate_runtime_profile': 'shared-runtime-v1', + 'certificate_id': 'ed25519:trusted-issuer', + 'runtime_profile': 'shared-runtime-v1', + 'admission_code': 'CERTIFIED_PLUGIN_SHARED_ELIGIBLE', + 'preserved_certificate_key': {'nested': True}, + } + certification.update(overrides) + return certification + + +@pytest.mark.asyncio +async def test_certification_artifact_digest_backfill_is_safe_and_enables_shared_placement(convergence_engine): + engine = convergence_engine + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + await run_alembic_upgrade(engine, '0031_merge_totp_assistant') + + valid_digest = 'a' * 64 + rows = { + 'eligible-a': (_legacy_shared_certification(), valid_digest), + 'eligible-b': (_legacy_shared_certification(), valid_digest), + 'present-matching': (_legacy_shared_certification(artifact_digest=valid_digest), valid_digest), + 'present-mismatched': (_legacy_shared_certification(artifact_digest='c' * 64), valid_digest), + 'invalid': (_legacy_shared_certification(verification='invalid'), valid_digest), + 'incomplete': (_legacy_shared_certification(certificate_id=' '), valid_digest), + 'missing-fact': ( + { + key: value + for key, value in _legacy_shared_certification().items() + if key != 'certificate_runtime_profile' + }, + valid_digest, + ), + 'malformed-normalized': (_legacy_shared_certification(normalized_digest='g' * 64), valid_digest), + 'dedicated-certificate': ( + _legacy_shared_certification(certificate_runtime_profile='dedicated'), + valid_digest, + ), + 'dedicated': (_legacy_shared_certification(runtime_profile='dedicated'), valid_digest), + 'wrong-admission': (_legacy_shared_certification(admission_code='SHARED_ELIGIBLE'), valid_digest), + 'uppercase-row-digest': (_legacy_shared_certification(), 'A' * 64), + 'nonhex-row-digest': (_legacy_shared_certification(), 'g' * 64), + } + plugin_settings = Base.metadata.tables['plugin_settings'] + workspaces = Base.metadata.tables['workspaces'] + async with engine.begin() as conn: + for index, (name, (certification, artifact_digest)) in enumerate(rows.items(), start=1): + workspace_uuid = f'41100000-0000-4000-8000-{index:012d}' + await conn.execute( + workspaces.insert().values( + uuid=workspace_uuid, + instance_uuid=f'cert-backfill-{index}', + name=name, + slug=name, + ) + ) + await conn.execute( + plugin_settings.insert().values( + workspace_uuid=workspace_uuid, + plugin_author='langbot', + plugin_name=name, + installation_uuid=f'51100000-0000-4000-8000-{index:012d}', + artifact_digest=artifact_digest, + runtime_revision=1, + install_info={ + '_certification': certification, + 'preserved_install_key': ['keep', {'nested': True}], + }, + ) + ) + + await run_alembic_upgrade(engine, 'head') + # Re-running the data revision itself must also be harmless. + migration = __import__( + 'langbot.pkg.persistence.alembic.versions.0032_certification_artifact_digest_backfill', + fromlist=['upgrade'], + ) + async with engine.begin() as conn: + await conn.run_sync(lambda sync: migration.backfill_certification_artifact_digests(sync)) + stored = { + name: (artifact_digest, install_info) + for name, artifact_digest, install_info in ( + await conn.execute( + sa.select( + plugin_settings.c.plugin_name, + plugin_settings.c.artifact_digest, + plugin_settings.c.install_info, + ) + ) + ).all() + } + + for name in ('eligible-a', 'eligible-b'): + eligible_digest, eligible_info = stored[name] + assert eligible_info['preserved_install_key'] == ['keep', {'nested': True}] + assert eligible_info['_certification']['preserved_certificate_key'] == {'nested': True} + assert eligible_info['_certification']['artifact_digest'] == eligible_digest == valid_digest + assert ( + execution_mode_for_persisted_installation( + artifact_digest=eligible_digest, + install_info=eligible_info, + ) + is PluginExecutionMode.SHARED_CERTIFIED + ) + + for name, (original_certification, artifact_digest) in rows.items(): + if name in {'eligible-a', 'eligible-b'}: + continue + stored_digest, stored_info = stored[name] + assert stored_digest == artifact_digest + assert stored_info['_certification'] == original_certification + if name != 'present-matching': + assert ( + execution_mode_for_persisted_installation( + artifact_digest=stored_digest, + install_info=stored_info, + ) + is PluginExecutionMode.DEDICATED + ) + + +@pytest.mark.asyncio +async def test_certification_artifact_digest_backfill_accepts_fresh_current_schema(convergence_engine): + engine = convergence_engine + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + + await run_alembic_upgrade(engine, 'head') + await run_alembic_upgrade(engine, 'head') + + assert await get_alembic_current(engine) == get_alembic_head() + + @pytest.mark.asyncio async def test_merge_only_downgrade_preserves_both_branch_schemas(convergence_engine): engine = convergence_engine diff --git a/tests/integration/persistence/test_migrations_postgres.py b/tests/integration/persistence/test_migrations_postgres.py index db3625848..babe2c207 100644 --- a/tests/integration/persistence/test_migrations_postgres.py +++ b/tests/integration/persistence/test_migrations_postgres.py @@ -329,6 +329,76 @@ class TestPostgreSQLMigrationBaseline: rev = await get_alembic_current(postgres_engine) assert rev == '0001_baseline' + @pytest.mark.asyncio + async def test_certification_artifact_digest_backfill_updates_only_complete_legacy_shared_rows( + self, + postgres_engine, + clean_tables, + clean_alembic_version, + ): + async with postgres_engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + await run_alembic_stamp(postgres_engine, '0031_merge_totp_assistant') + + plugin_settings = Base.metadata.tables['plugin_settings'] + workspaces = Base.metadata.tables['workspaces'] + digest = 'a' * 64 + complete = { + 'normalized_digest': 'B' * 64, + 'verification': 'valid', + 'certificate_runtime_profile': 'shared-runtime-v1', + 'certificate_id': 'ed25519:trusted-issuer', + 'runtime_profile': 'shared-runtime-v1', + 'admission_code': 'CERTIFIED_PLUGIN_SHARED_ELIGIBLE', + 'preserved': {'key': True}, + } + rows = { + 'eligible': complete, + 'invalid': {**complete, 'verification': 'invalid'}, + 'present': {**complete, 'artifact_digest': 'c' * 64}, + } + async with postgres_engine.begin() as conn: + for index, (name, certification) in enumerate(rows.items(), start=1): + workspace_uuid = f'61100000-0000-4000-8000-{index:012d}' + await conn.execute( + workspaces.insert().values( + uuid=workspace_uuid, + instance_uuid=f'postgres-cert-backfill-{index}', + name=name, + slug=name, + ) + ) + await conn.execute( + plugin_settings.insert().values( + workspace_uuid=workspace_uuid, + plugin_author='langbot', + plugin_name=name, + installation_uuid=f'71100000-0000-4000-8000-{index:012d}', + artifact_digest=digest, + runtime_revision=1, + install_info={'_certification': certification, 'preserved': ['install-info']}, + ) + ) + + await run_alembic_upgrade(postgres_engine, 'head') + + async with postgres_engine.connect() as conn: + stored = dict( + (await conn.execute(sa.select(plugin_settings.c.plugin_name, plugin_settings.c.install_info))).all() + ) + assert stored['eligible'] == { + '_certification': {**complete, 'artifact_digest': digest}, + 'preserved': ['install-info'], + } + assert stored['invalid'] == { + '_certification': rows['invalid'], + 'preserved': ['install-info'], + } + assert stored['present'] == { + '_certification': rows['present'], + 'preserved': ['install-info'], + } + @pytest.mark.asyncio async def test_fresh_postgres_schema_accepts_application_casefold_identity( self,