from __future__ import annotations import hashlib import os import subprocess import uuid from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from pathlib import Path from threading import Barrier, Event, Lock import pytest from sqlalchemy import create_engine, inspect, text from sqlalchemy.engine import make_url from sqlalchemy.exc import DBAPIError, IntegrityError from sqlalchemy.orm import Session from app.core.connectors.builtin.oracle import OracleConnector from app.core.connectors.builtin.rest_catalog import RestCatalogConnector from app.core.connectors.builtin.sqlserver import SqlServerConnector from app.core.connectors.errors import ( ConnectorAuthenticationError, ConnectorConfigurationError, ConnectorRateLimitError, ConnectorUpstreamError, ) from app.core.connectors.identity import ConnectorIdentityRepository from app.core.connectors.registry import ConnectorRegistry from app.core.connectors.repository import ConnectorRepository from app.core.connectors.runtime import ( ConnectorRuntime, deterministic_idempotency_key, ) from app.core.connectors.sdk import ( SDK_VERSION, SECRET_REF_PATTERN, Connector, ConnectorManifest, OperationRequest, OperationResult, ) ROOT = Path(__file__).resolve().parents[2] pytestmark = pytest.mark.integration def _uid(): return str(uuid.uuid4()) def _alembic(url, command, target): subprocess.run( [ str(ROOT / ".venv/bin/alembic"), "-c", str(ROOT / "alembic.ini"), command, target, ], cwd=ROOT, env={**os.environ, "SQLALCHEMY_DATABASE_URI": url}, check=True, capture_output=True, text=True, ) def _clear_connector_test_data(engine): tables = set(inspect(engine).get_table_names(schema="public")) with engine.begin() as connection: for table in ( "connector_audit_events", "connector_graph_edges", "connector_evidence", "connector_checkpoints", "connector_run_attempts", "connector_runs", "connector_machine_credentials", "connector_principals", "connector_source_bindings", "connector_manifests", "connector_rate_limits", ): if table in tables: connection.execute(text(f"DELETE FROM public.{table}")) def _create_wp03_rollback_database(base_url: str) -> tuple[str, str]: """Create a uniquely named DB so historical WP03 rollback never crosses WP11 fences.""" database = f"wp03_rollback_{uuid.uuid4().hex}" url = make_url(base_url) control_url = url.set(database="postgres").render_as_string(hide_password=False) control = create_engine(control_url, isolation_level="AUTOCOMMIT", pool_pre_ping=True) try: quoted = control.dialect.identifier_preparer.quote(database) with control.connect() as connection: connection.execute(text(f"CREATE DATABASE {quoted}")) finally: control.dispose() return database, url.set(database=database).render_as_string(hide_password=False) def _drop_wp03_rollback_database(base_url: str, database: str) -> None: url = make_url(base_url) control_url = url.set(database="postgres").render_as_string(hide_password=False) control = create_engine(control_url, isolation_level="AUTOCOMMIT", pool_pre_ping=True) try: quoted = control.dialect.identifier_preparer.quote(database) with control.connect() as connection: connection.execute( text( "SELECT pg_terminate_backend(pid) FROM pg_stat_activity " "WHERE datname=:database AND pid<>pg_backend_pid()" ), {"database": database}, ) connection.execute(text(f"DROP DATABASE IF EXISTS {quoted}")) finally: control.dispose() def test_connector_flask_postgres_security_bindings_and_machine_actions( monkeypatch, ): url = os.environ.get("TEST_DATABASE_URL") if not url: pytest.skip("TEST_DATABASE_URL is required") engine = create_engine(url, pool_pre_ping=True) _clear_connector_test_data(engine) engine.dispose() _alembic(url, "upgrade", "head") monkeypatch.setenv("DATABASE_URL", url) from app import create_app actor = _uid() monkeypatch.setattr( "app.core.system.permissions.authenticate_request", lambda: {"id": actor, "sub": actor, "roles": ["admin"]}, ) class SpyTransport: def __init__(self): self.calls = [] def get_json(self, **values): self.calls.append(values) return { "assets": [ { "key": "approved-asset", "name": "Approved", "namespace": "security", "type": "table", } ] } spy = SpyTransport() registry = ConnectorRegistry() registry.register(RestCatalogConnector(spy)) app = create_app() app.config.update(TESTING=True) app.extensions["connector_registry"] = registry client = app.test_client() source, domain = _uid(), _uid() binding_response = client.post( "/api/datasource/connectors/source-bindings", json={ "connector_id": "rest-catalog", "version": "1.0.0", "source_uid": source, "business_domain_uid": domain, "environment": "staging", "approved_config": { "base_url": "https://approved.example.test", "allowed_host": "approved.example.test", "credential_ref": "env:DATAOPS_CONNECTOR_APPROVED", }, }, ) assert binding_response.status_code == 201 binding = binding_response.get_json()["data"] principal_response = client.post( "/api/datasource/connectors/principals", json={ "connector_id": "rest-catalog", "version": "1.0.0", "source_uid": source, "business_domain_uid": domain, "environment": "staging", "operations": ["discover", "cancel", "resume"], "scopes": {}, "source_binding_uid": binding["binding_uid"], "source_binding_version": binding["binding_version"], }, ) assert principal_response.status_code == 201 principal = principal_response.get_json()["data"]["principal_uid"] rebound_response = client.post( "/api/datasource/connectors/source-bindings", json={ "binding_uid": binding["binding_uid"], "connector_id": "rest-catalog", "version": "1.0.0", "source_uid": source, "business_domain_uid": domain, "environment": "staging", "approved_config": { "base_url": "https://approved.example.test", "allowed_host": "approved.example.test", "credential_ref": "env:DATAOPS_CONNECTOR_APPROVED_V2", }, }, ) assert rebound_response.status_code == 201 rebound = rebound_response.get_json()["data"] assert rebound["binding_version"] == 2 assert rebound["rebound_principals"] == 1 with create_engine(url).connect() as connection: versions = connection.execute( text(""" SELECT p.source_binding_version, ARRAY_AGG(b.status ORDER BY b.binding_version) FROM public.connector_principals p JOIN public.connector_source_bindings b ON b.uid=p.source_binding_uid WHERE p.uid=CAST(:principal AS uuid) GROUP BY p.source_binding_version """), {"principal": principal}, ).one() assert versions == (2, ["revoked", "approved"]) def issue(principal_uid): response = client.post( f"/api/datasource/connectors/principals/{principal_uid}/credentials", json={"ttl_seconds": 300}, ) assert response.status_code == 201 return response.get_json()["data"]["credential"] hint = uuid.uuid4().hex + uuid.uuid4().hex token = issue(principal) empty_config = client.post( "/api/datasource/connectors/machine/runs", headers={"X-Connector-Credential": token}, json={ "connector_id": "rest-catalog", "version": "1.0.0", "source_uid": source, "business_domain_uid": domain, "environment": "staging", "process_key": "security-probe", "operation": "discover", "config": {}, "scope": {}, "idempotency_key": hint, }, ) assert empty_config.status_code == 400 and spy.calls == [] attack = client.post( "/api/datasource/connectors/machine/runs", headers={"X-Connector-Credential": token}, json={ "connector_id": "rest-catalog", "version": "1.0.0", "source_uid": source, "business_domain_uid": domain, "environment": "staging", "process_key": "security-probe", "operation": "discover", "config": { "base_url": "https://attacker.example.test", "allowed_host": "attacker.example.test", "credential_ref": "env:DATAOPS_CONNECTOR_APPROVED", }, "scope": {}, "idempotency_key": hint, }, ) assert attack.status_code == 400 and spy.calls == [] created = client.post( "/api/datasource/connectors/machine/runs", headers={"X-Connector-Credential": token}, json={ "connector_id": "rest-catalog", "version": "1.0.0", "source_uid": source, "business_domain_uid": domain, "environment": "staging", "process_key": "security-probe", "operation": "discover", "scope": {}, "idempotency_key": hint, }, ) assert created.status_code == 201 assert len(spy.calls) == 1 assert spy.calls[0]["url"].startswith("https://approved.example.test/") assert spy.calls[0]["allowed_host"] == "approved.example.test" assert spy.calls[0]["credential_ref"] == "env:DATAOPS_CONNECTOR_APPROVED_V2" assert client.post( f"/api/datasource/connectors/runs/{hint}/cancel" ).status_code == 400 assert client.post( f"/api/datasource/connectors/runs/{hint}/resume" ).status_code == 400 second_source, second_domain = _uid(), _uid() second_binding = client.post( "/api/datasource/connectors/source-bindings", json={ "connector_id": "rest-catalog", "version": "1.0.0", "source_uid": second_source, "business_domain_uid": second_domain, "environment": "production", "approved_config": { "base_url": "https://second.example.test", "allowed_host": "second.example.test", "credential_ref": "env:DATAOPS_CONNECTOR_SECOND", }, }, ).get_json()["data"] second_principal = client.post( "/api/datasource/connectors/principals", json={ "connector_id": "rest-catalog", "version": "1.0.0", "source_uid": second_source, "business_domain_uid": second_domain, "environment": "production", "operations": ["discover", "cancel", "resume"], "scopes": {}, "source_binding_uid": second_binding["binding_uid"], "source_binding_version": second_binding["binding_version"], }, ).get_json()["data"]["principal_uid"] cross = client.post( f"/api/datasource/connectors/machine/runs/{hint}/cancel", headers={"X-Connector-Credential": issue(second_principal)}, ) assert cross.status_code == 403 replay = client.post( f"/api/datasource/connectors/machine/runs/{hint}/cancel", headers={"X-Connector-Credential": token}, ) assert replay.status_code == 401 with create_engine(url).begin() as connection: connection.execute( text(""" UPDATE public.connector_runs SET status='running',cancel_requested=FALSE WHERE client_hint_hash=:hint """), {"hint": hashlib.sha256(hint.encode()).hexdigest()}, ) cancelled = client.post( f"/api/datasource/connectors/machine/runs/{hint}/cancel", headers={"X-Connector-Credential": issue(principal)}, ) assert cancelled.status_code == 200 resumed = client.post( f"/api/datasource/connectors/machine/runs/{hint}/resume", headers={"X-Connector-Credential": issue(principal)}, ) assert resumed.status_code == 200 conflict = client.post( "/api/datasource/connectors/machine/runs", headers={"X-Connector-Credential": issue(second_principal)}, json={ "connector_id": "rest-catalog", "version": "1.0.0", "source_uid": second_source, "business_domain_uid": second_domain, "environment": "production", "process_key": "different-process", "operation": "discover", "scope": {}, "idempotency_key": hint, }, ) assert conflict.status_code == 409 listed = client.get("/api/datasource/connectors/source-bindings") assert listed.status_code == 200 serialized = str(listed.get_json()) assert "credential_ref" not in serialized assert "DATAOPS_CONNECTOR_APPROVED" not in serialized with create_engine(url).connect() as connection: persisted = connection.execute( text(""" SELECT COALESCE(string_agg(payload::text,''),'') FROM public.connector_evidence WHERE run_uid IN ( SELECT uid FROM public.connector_runs WHERE principal_uid=CAST(:principal AS uuid) ) """), {"principal": principal}, ).scalar_one() assert "DATAOPS_CONNECTOR_APPROVED" not in persisted revoked = client.post( f"/api/datasource/connectors/source-bindings/{binding['binding_uid']}/revoke" ) assert revoked.status_code == 200 assert revoked.get_json()["data"] == { "revoked": True, "principals_deactivated": 1, } with create_engine(url).connect() as connection: statuses = connection.execute( text(""" SELECT p.status,COALESCE(MAX(c.status),'none') FROM public.connector_principals p LEFT JOIN public.connector_machine_credentials c ON c.principal_uid=p.uid WHERE p.uid=CAST(:principal AS uuid) GROUP BY p.status """), {"principal": principal}, ).one() assert statuses[0] == "revoked" def test_oracle_sqlserver_bindings_authentication_and_trusted_environment(monkeypatch): url = os.environ.get("TEST_DATABASE_URL") if not url: pytest.skip("TEST_DATABASE_URL is required") engine = create_engine(url, pool_pre_ping=True) _clear_connector_test_data(engine) _alembic(url, "upgrade", "head") monkeypatch.setenv("DATABASE_URL", url) from app import create_app actor = _uid() monkeypatch.setattr( "app.core.system.permissions.authenticate_request", lambda: {"id": actor, "sub": actor, "roles": ["admin"]}, ) provider_calls = [] class Rows: def mappings(self): return self def all(self): return [ { "schema_name": "APP", "asset_name": "ORDERS", "asset_type": "TABLE", "column_name": "ID", "ordinal_position": 1, "data_type": "NUMBER", "is_nullable": "NO", "column_default": None, "column_comment": None, } ] class Connection: def execute(self, _statement, _parameters): return Rows() @contextmanager def provider( source_uid, purpose, *, environment="production", allow_insecure_development=False, ): provider_calls.append( { "source_uid": source_uid, "purpose": purpose, "environment": environment, "allow_insecure_development": allow_insecure_development, } ) yield Connection() class TrackingOracle(OracleConnector): seen_configs = [] def discover(self, request): self.seen_configs.append(dict(request.config)) return super().discover(request) class TrackingSqlServer(SqlServerConnector): seen_configs = [] def discover(self, request): self.seen_configs.append(dict(request.config)) return super().discover(request) registry = ConnectorRegistry() oracle = TrackingOracle(provider) sqlserver = TrackingSqlServer(provider) registry.register(oracle) registry.register(sqlserver) app = create_app() app.config.update(TESTING=True) app.extensions["connector_registry"] = registry client = app.test_client() rejected = client.post( "/api/datasource/connectors/principals", json={ "connector_id": "oracle", "version": "1.0.0", "source_uid": _uid(), "business_domain_uid": _uid(), "environment": "staging", "operations": ["discover"], "scopes": {}, }, ) assert rejected.status_code == 400 for connector_id, environment, approved_config in ( ( "oracle", "staging", {"credential_ref": "env:DATAOPS_CONNECTOR_ORACLE"}, ), ( "sqlserver", "production", { "credential_ref": "env:DATAOPS_CONNECTOR_SQLSERVER", "allow_insecure_development": True, }, ), ): source, domain = _uid(), _uid() binding_response = client.post( "/api/datasource/connectors/source-bindings", json={ "connector_id": connector_id, "version": "1.0.0", "source_uid": source, "business_domain_uid": domain, "environment": environment, "approved_config": approved_config, }, ) assert binding_response.status_code == 201 binding = binding_response.get_json()["data"] principal_response = client.post( "/api/datasource/connectors/principals", json={ "connector_id": connector_id, "version": "1.0.0", "source_uid": source, "business_domain_uid": domain, "environment": environment, "operations": ["discover"], "scopes": {}, "source_binding_uid": binding["binding_uid"], "source_binding_version": binding["binding_version"], }, ) assert principal_response.status_code == 201 principal = principal_response.get_json()["data"]["principal_uid"] credential_response = client.post( f"/api/datasource/connectors/principals/{principal}/credentials" ) assert credential_response.status_code == 201 token = credential_response.get_json()["data"]["credential"] payload = { "connector_id": connector_id, "version": "1.0.0", "source_uid": source, "business_domain_uid": domain, "environment": environment, "process_key": f"{connector_id}-catalog", "operation": "discover", "scope": {}, } if connector_id == "sqlserver": forged = client.post( "/api/datasource/connectors/machine/runs", headers={"X-Connector-Credential": token}, json={**payload, "environment": "development"}, ) assert forged.status_code == 403 created = client.post( "/api/datasource/connectors/machine/runs", headers={"X-Connector-Credential": token}, json=payload, ) assert created.status_code == 201 assert created.get_json()["data"]["records"][0]["name"] == "ORDERS" assert oracle.seen_configs == [ {"credential_ref": "env:DATAOPS_CONNECTOR_ORACLE"} ] assert sqlserver.seen_configs == [ { "credential_ref": "env:DATAOPS_CONNECTOR_SQLSERVER", "allow_insecure_development": True, } ] assert provider_calls[0]["environment"] == "production" assert provider_calls[0]["allow_insecure_development"] is False assert provider_calls[1]["environment"] == "production" assert provider_calls[1]["allow_insecure_development"] is True with engine.connect() as connection: configs = connection.execute( text(""" SELECT connector_id,safe_config FROM public.connector_runs WHERE connector_id IN ('oracle','sqlserver') ORDER BY connector_id """) ).all() assert configs == [ ("oracle", {"credential_ref": "env:DATAOPS_CONNECTOR_ORACLE"}), ( "sqlserver", { "credential_ref": "env:DATAOPS_CONNECTOR_SQLSERVER", "allow_insecure_development": True, }, ), ] engine.dispose() def test_connector_postgres_attempt_lease_rejects_late_terminal_writes(): url = os.environ.get("TEST_DATABASE_URL") if not url: pytest.skip("TEST_DATABASE_URL is required") engine = create_engine(url, pool_pre_ping=True) _clear_connector_test_data(engine) _alembic(url, "upgrade", "head") actor, source = _uid(), _uid() connector_id = f"lease-test-{uuid.uuid4().hex[:8]}" with engine.begin() as connection: connection.execute( text(""" INSERT INTO public.connector_manifests (uid,connector_id,connector_version,sdk_version,display_name, capabilities,config_schema,status,created_by) VALUES(CAST(:uid AS uuid),:connector,'1.0.0','1.0','lease-test', '["discover","cancel","resume"]'::jsonb, CAST(:schema AS jsonb), 'active',CAST(:actor AS uuid)) """), { "uid": _uid(), "connector": connector_id, "actor": actor, "schema": '{"type":"object","additionalProperties":false,"properties":{}}', }, ) key = uuid.uuid4().hex + uuid.uuid4().hex first = Session(engine) second = Session(engine) try: repository = ConnectorRepository(first) repository.claim( key, { "idempotency_key": key, "request_hash": key, "client_hint_hash": None, "connector_id": connector_id, "connector_version": "1.0.0", "source_uid": source, "operation": "discover", "config": {}, "scope": {}, "checkpoint": {}, "cursor": {}, "dry_run": True, "actor_uid": actor, }, ) lease_one, lease_two = _uid(), _uid() assert repository.update( key, status="running", attempt_count=1, lease_token=lease_one ).acquired assert repository.cancel(key)["status"] == "cancelled" with Session(engine) as observer: observed = ConnectorRepository(observer).get(key) assert observed["status"] == "cancelled" assert observed["cancel_requested"] is True assert ConnectorRepository(observer).is_cancel_requested(key) is True resumed = ConnectorRepository(second).update( key, status="resumable", cancel_requested=False ) assert resumed.acquired assert ConnectorRepository(second).update( key, status="running", attempt_count=2, lease_token=lease_two ).acquired stale_success = ConnectorRepository(first).update( key, status="succeeded", expected_attempt=1, lease_token=lease_one, ) stale_failure = ConnectorRepository(first).update( key, status="failed", error_category="upstream", error_code="late", expected_attempt=1, lease_token=lease_one, ) assert stale_success.acquired is False assert stale_failure.acquired is False winner = ConnectorRepository(second).update( key, status="succeeded", expected_attempt=2, lease_token=lease_two, ) assert winner.acquired final = ConnectorRepository(second).get(key) assert final["status"] == "succeeded" assert final["attempt_count"] == 2 attempts = second.execute( text(""" SELECT attempt_number,status,lease_token::text FROM public.connector_run_attempts WHERE run_uid=CAST(:run AS uuid) ORDER BY attempt_number """), {"run": final["uid"]}, ).all() assert attempts == [ (1, "cancelled", lease_one), (2, "succeeded", lease_two), ] finally: first.close() second.close() engine.dispose() def test_connector_postgres_repository_identity_graph_cas_and_473_rollback(): shared_url = os.environ.get("TEST_DATABASE_URL") if not shared_url: pytest.skip("TEST_DATABASE_URL is required") rollback_database, url = _create_wp03_rollback_database(shared_url) engine = create_engine(url, pool_pre_ping=True) actor, source, domain = _uid(), _uid(), _uid() connector_id, version = f"pg-test-{uuid.uuid4().hex[:8]}", "1.0.0" _alembic(url, "upgrade", "20260802_476") try: tables = inspect(engine).get_table_names(schema="public") assert {"connector_runs", "connector_rate_limits"}.issubset(tables) binding_uid = _uid() with engine.begin() as connection: connection.execute( text(""" INSERT INTO public.connector_manifests (uid,connector_id,connector_version,sdk_version,display_name,capabilities,config_schema,status,created_by) VALUES(CAST(:uid AS uuid),:connector,:version,'1.0','test','["discover","cancel","resume"]'::jsonb, '{"type":"object"}'::jsonb,'active',CAST(:actor AS uuid)) """), { "uid": _uid(), "connector": connector_id, "version": version, "actor": actor, }, ) connection.execute( text(""" INSERT INTO public.connector_source_bindings (uid,binding_version,connector_id,connector_version,source_uid, business_domain_uid,environment,approved_config,status,approved_by) VALUES(CAST(:uid AS uuid),1,:connector,:version,CAST(:source AS uuid), CAST(:domain AS uuid),'staging', '{"credential_ref":"env:DATAOPS_CONNECTOR_TEST"}'::jsonb, 'approved',CAST(:actor AS uuid)) """), { "uid": binding_uid, "connector": connector_id, "version": version, "source": source, "domain": domain, "actor": actor, }, ) idem = uuid.uuid4().hex + uuid.uuid4().hex def claim(_index): with engine.begin() as connection: return connection.execute( text(""" INSERT INTO public.connector_runs (uid,idempotency_key,connector_id,connector_version,source_uid,operation,status,attempt_count, checkpoint,cursor,dry_run,actor_uid,safe_config,scope,request_hash) VALUES(CAST(:uid AS uuid),:key,:connector,:version,CAST(:source AS uuid),'discover','running',0, '{}'::jsonb,'{}'::jsonb,TRUE,CAST(:actor AS uuid),'{}'::jsonb,'{}'::jsonb,:key) ON CONFLICT(idempotency_key) DO NOTHING RETURNING uid """), { "uid": _uid(), "key": idem, "connector": connector_id, "version": version, "source": source, "actor": actor, }, ).scalar_one_or_none() with ThreadPoolExecutor(max_workers=4) as executor: assert sum(item is not None for item in executor.map(claim, range(4))) == 1 principal = _uid() with engine.begin() as connection: connection.execute( text(""" INSERT INTO public.connector_principals (uid,connector_id,connector_version,source_uid,business_domain_uid,environment, allowed_operations,allowed_scopes,status,created_by, source_binding_uid,source_binding_version) VALUES(CAST(:uid AS uuid),:connector,:version,CAST(:source AS uuid),CAST(:domain AS uuid), 'staging',ARRAY['discover'],'{}'::jsonb,'active',CAST(:actor AS uuid), CAST(:binding AS uuid),1) """), { "uid": principal, "connector": connector_id, "version": version, "source": source, "domain": domain, "actor": actor, "binding": binding_uid, }, ) with pytest.raises(IntegrityError), engine.begin() as connection: connection.execute( text(""" INSERT INTO public.connector_machine_credentials(uid,principal_uid,token_hash,status,issued_by,expires_at) VALUES(CAST(:uid AS uuid),CAST(:principal AS uuid),:hash,'active',CAST(:actor AS uuid), CURRENT_TIMESTAMP+INTERVAL '901 seconds') """), { "uid": _uid(), "principal": principal, "hash": "a" * 64, "actor": actor, }, ) session = Session(engine) identity = ConnectorIdentityRepository(session) issued = identity.issue(principal, ttl_seconds=300, actor_uid=actor) authenticated = identity.authenticate( issued["credential"], connector_id=connector_id, connector_version=version, source_uid=source, business_domain_uid=domain, environment="staging", operation="discover", scope={}, ) assert authenticated["principal_uid"] == principal with pytest.raises(ConnectorAuthenticationError): identity.authenticate( issued["credential"], connector_id=connector_id, connector_version=version, source_uid=source, business_domain_uid=domain, environment="staging", operation="discover", scope={}, ) second = identity.issue(principal, ttl_seconds=300, actor_uid=actor) rotated = identity.rotate( second["credential_uid"], ttl_seconds=300, actor_uid=actor ) assert identity.revoke(rotated["credential_uid"], actor) is True limiter_key = f"{connector_id}:{source}" for _ in range(30): with Session(engine) as limiter_session: ConnectorRepository(limiter_session).acquire_rate_limit(limiter_key) with Session(engine) as limiter_session, pytest.raises(ConnectorRateLimitError): ConnectorRepository(limiter_session).acquire_rate_limit(limiter_key) run_key = uuid.uuid4().hex + uuid.uuid4().hex repository = ConnectorRepository(session, principal) repository.claim( run_key, { "connector_id": connector_id, "connector_version": version, "source_uid": source, "operation": "discover", "checkpoint": {}, "cursor": {}, "dry_run": False, "principal_uid": principal, "business_domain_uid": domain, "environment": "staging", "process_key": "catalog-sync", "config": {"credential_ref": "secret:test/key"}, "scope": {}, }, ) running_graph = repository.graph(run_uid=repository.get(run_key)["uid"]) assert running_graph["summary"]["edge_count"] == 3 assert {node["type"] for node in running_graph["nodes"]} == { "source", "process", "business_domain", "run", } assert {edge["run_status"] for edge in running_graph["edges"]} == {"running"} run_lease = _uid() repository.update( run_key, status="running", attempt_count=1, lease_token=run_lease ) repository.update( run_key, status="succeeded", result=OperationResult( records=({"asset_key": f"{source}:APP.ORDERS"},), cursor={"page": 2}, checkpoint={"cursor": {"page": 2}, "snapshot": [{"key": "orders"}]}, evidence={"safe": True}, ), checkpoint={"cursor": {"page": 2}}, cursor={"page": 2}, error_category=None, expected_attempt=1, lease_token=run_lease, ) graph = repository.graph(source_uid=source, business_domain_uid=domain) assert {node["type"] for node in graph["nodes"]} == { "source", "asset", "process", "business_domain", "run", } assert graph["summary"]["edge_count"] == 5 process_graph = repository.graph(process_key="catalog-sync") assert process_graph["summary"]["edge_count"] >= 1 assert {"process", "run"}.issubset( {node["type"] for node in process_graph["nodes"]} ) failed_key = uuid.uuid4().hex + uuid.uuid4().hex repository.claim( failed_key, { "connector_id": connector_id, "connector_version": version, "source_uid": source, "operation": "discover", "checkpoint": {"page": 4}, "cursor": {"page": 4}, "dry_run": False, "principal_uid": principal, "business_domain_uid": domain, "environment": "staging", "process_key": "failed-catalog-sync", "config": {"credential_ref": "env:DATAOPS_CONNECTOR_TEST"}, "scope": {}, }, ) failed_lease = _uid() repository.update( failed_key, status="running", attempt_count=1, lease_token=failed_lease, ) failed = repository.update( failed_key, status="failed", error_category="upstream", error_code="ConnectorUpstreamError", expected_attempt=1, lease_token=failed_lease, ).record assert failed["checkpoint"] == {"page": 4} failed_graph = repository.graph(run_uid=failed["uid"]) assert failed_graph["summary"]["edge_count"] == 3 assert {edge["run_status"] for edge in failed_graph["edges"]} == {"failed"} cancelled_key = uuid.uuid4().hex + uuid.uuid4().hex repository.claim( cancelled_key, { "connector_id": connector_id, "connector_version": version, "source_uid": source, "operation": "discover", "checkpoint": {"page": 1}, "cursor": {"page": 1}, "dry_run": False, "principal_uid": principal, "business_domain_uid": domain, "environment": "staging", "process_key": "catalog-sync", "config": {"credential_ref": "secret:test/key"}, "scope": {}, }, ) cancelled = Event() def cancel_worker(): with Session(engine) as worker_session: result = ConnectorRepository(worker_session, principal).cancel( cancelled_key ) cancelled.set() return result def completion_worker(): assert cancelled.wait(timeout=5) with Session(engine) as worker_session: return ConnectorRepository(worker_session, principal).update( cancelled_key, status="succeeded", result=OperationResult(), checkpoint={}, cursor={}, error_category=None, ) with ThreadPoolExecutor(max_workers=2) as executor: cancel_future = executor.submit(cancel_worker) complete_future = executor.submit(completion_worker) assert cancel_future.result()["status"] == "cancelled" completion_outcome = complete_future.result() assert completion_outcome.acquired is False after = completion_outcome.record assert after["status"] == "cancelled" cancelled_graph = repository.graph(run_uid=after["uid"]) assert cancelled_graph["summary"]["edge_count"] == 3 assert {edge["run_status"] for edge in cancelled_graph["edges"]} == { "cancelled" } with pytest.raises(IntegrityError), engine.begin() as connection: connection.execute( text(""" INSERT INTO public.connector_runs(uid,idempotency_key,connector_id,connector_version,source_uid, operation,status,attempt_count,checkpoint,cursor,dry_run,actor_uid,safe_config,scope) VALUES(CAST(:uid AS uuid),:key,:connector,:version,CAST(:source AS uuid),'discover','running',0, '{}'::jsonb,'{}'::jsonb,FALSE,CAST(:actor AS uuid),'{}'::jsonb,'{}'::jsonb) """), { "uid": _uid(), "key": uuid.uuid4().hex + uuid.uuid4().hex, "connector": connector_id, "version": version, "source": source, "actor": actor, }, ) class RuntimeProbeConnector(Connector): manifest = ConnectorManifest( connector_id=connector_id, version=version, sdk_version=SDK_VERSION, display_name="PostgreSQL Runtime Probe", capabilities=( "discover", "snapshot", "incremental", "lineage", "profile", "cancel", "resume", "evidence", ), config_schema={ "type": "object", "additionalProperties": False, "required": ["credential_ref"], "properties": { "credential_ref": { "type": "string", "pattern": SECRET_REF_PATTERN.pattern, } }, }, ) def __init__(self): self.discover_calls = 0 self.resume_calls = 0 self.discover_succeeds = False self.resume_failures = 2 self.call_lock = Lock() def discover(self, request): with self.call_lock: self.discover_calls += 1 if self.discover_succeeds: return OperationResult(evidence={"concurrent": True}) raise ConnectorUpstreamError() snapshot = incremental = lineage = profile = evidence = discover def cancel(self, request): return OperationResult(status="cancelled") def resume(self, request): with self.call_lock: self.resume_calls += 1 call = self.resume_calls if call <= self.resume_failures: raise ConnectorUpstreamError() return OperationResult( checkpoint={"page": 8}, cursor={"page": 9}, evidence={"resumed": True}, ) probe = RuntimeProbeConnector() registry = ConnectorRegistry() registry.register(probe) capped_source = _uid() capped_request = OperationRequest( source_uid=capped_source, operation="discover", config={"credential_ref": "env:DATAOPS_CONNECTOR_TEST"}, principal_uid=principal, business_domain_uid=domain, environment="staging", process_key="attempt-cap", ) capped_runtime = ConnectorRuntime( registry, store=ConnectorRepository(session, principal), sleeper=lambda _seconds: None, ) with pytest.raises(ConnectorUpstreamError): capped_runtime.execute(connector_id, version, capped_request) with pytest.raises(ConnectorUpstreamError): capped_runtime.execute(connector_id, version, capped_request) with pytest.raises(ConnectorConfigurationError): capped_runtime.execute(connector_id, version, capped_request) capped_key = capped_request.idempotency_key or deterministic_idempotency_key( connector_id, version, capped_request ) capped = ConnectorRepository(session).get(capped_key) assert capped["attempt_count"] == 5 and probe.discover_calls == 5 attempts = session.execute( text(""" SELECT COUNT(*),MAX(attempt_number) FROM public.connector_run_attempts WHERE run_uid=CAST(:run AS uuid) """), {"run": capped["uid"]}, ).one() assert attempts == (5, 5) rejected_source = _uid() rejected_key = f"{connector_id}:{rejected_source}" for _ in range(30): ConnectorRepository(session).acquire_rate_limit(rejected_key) rejected_request = OperationRequest( source_uid=rejected_source, operation="discover", config={"credential_ref": "env:DATAOPS_CONNECTOR_TEST"}, dry_run=True, ) rejected_idempotency = deterministic_idempotency_key( connector_id, version, rejected_request ) with pytest.raises(ConnectorRateLimitError): ConnectorRuntime( registry, store=ConnectorRepository(session), sleeper=lambda _: None ).execute(connector_id, version, rejected_request) assert ConnectorRepository(session).get(rejected_idempotency) is None resume_source = _uid() resume_key = uuid.uuid4().hex + uuid.uuid4().hex resume_repository = ConnectorRepository(session, principal) resume_repository.claim( resume_key, { "connector_id": connector_id, "connector_version": version, "source_uid": resume_source, "operation": "discover", "checkpoint": {"page": 7}, "cursor": {"page": 7}, "dry_run": False, "principal_uid": principal, "business_domain_uid": domain, "environment": "staging", "process_key": "resume-state-machine", "config": {"credential_ref": "env:DATAOPS_CONNECTOR_TEST"}, "scope": {}, }, ) resume_repository.update( resume_key, status="running", attempt_count=1, lease_token=_uid() ) resume_repository.cancel(resume_key) resume_runtime = ConnectorRuntime( registry, store=resume_repository, max_attempts=2, sleeper=lambda _seconds: None, ) with pytest.raises(ConnectorUpstreamError): resume_runtime.resume(resume_key) resume_failed = resume_repository.get(resume_key) assert resume_failed["status"] == "failed" assert resume_failed["attempt_count"] == 3 assert resume_failed["checkpoint"] == {"page": 7} resumed_result = resume_runtime.resume(resume_key) assert resumed_result.checkpoint == {"page": 8} resumed = resume_repository.get(resume_key) assert resumed["status"] == "succeeded" and resumed["attempt_count"] == 4 persisted = session.execute( text(""" SELECT (SELECT COUNT(*) FROM public.connector_run_attempts WHERE run_uid=CAST(:run AS uuid)), (SELECT COUNT(*) FROM public.connector_evidence WHERE run_uid=CAST(:run AS uuid)), (SELECT COUNT(*) FROM public.connector_checkpoints WHERE run_uid=CAST(:run AS uuid)) """), {"run": resumed["uid"]}, ).one() assert persisted == (4, 1, 1) probe.discover_succeeds = True probe.discover_calls = 0 retry_hint = uuid.uuid4().hex + uuid.uuid4().hex retry_source = _uid() retry_request = OperationRequest( source_uid=retry_source, operation="discover", config={"credential_ref": "env:DATAOPS_CONNECTOR_TEST"}, idempotency_key=retry_hint, principal_uid=principal, business_domain_uid=domain, environment="staging", process_key="concurrent-retry", ) retry_key = deterministic_idempotency_key( connector_id, version, retry_request ) retry_seed = ConnectorRepository(session, principal) retry_seed.claim( retry_key, { "connector_id": connector_id, "connector_version": version, "source_uid": retry_source, "operation": "discover", "checkpoint": {"page": 1}, "cursor": {"page": 1}, "dry_run": False, "principal_uid": principal, "business_domain_uid": domain, "environment": "staging", "process_key": "concurrent-retry", "config": {"credential_ref": "env:DATAOPS_CONNECTOR_TEST"}, "scope": {}, }, ) retry_lease = _uid() retry_seed.update( retry_key, status="running", attempt_count=1, lease_token=retry_lease ) retry_seed.update( retry_key, status="failed", error_category="upstream", error_code="ConnectorUpstreamError", expected_attempt=1, lease_token=retry_lease, ) retry_gate = Barrier(2) class RetryBarrierRepository(ConnectorRepository): def claim(self, key, record): claimed, created = super().claim(key, record) if not created: retry_gate.wait(timeout=5) return claimed, created def competing_retry(): with Session(engine) as worker_session: worker_session.execute(text("SET statement_timeout='5s'")) worker_session.commit() try: result = ConnectorRuntime( registry, store=RetryBarrierRepository(worker_session, principal), max_attempts=1, sleeper=lambda _seconds: None, ).execute(connector_id, version, retry_request) return "success", result.status except Exception as error: return "error", type(error).__name__ with ThreadPoolExecutor(max_workers=2) as executor: retry_futures = [executor.submit(competing_retry) for _index in range(2)] retry_results = [future.result(timeout=10) for future in retry_futures] assert sorted(item[0] for item in retry_results) == ["error", "success"] assert [item[1] for item in retry_results if item[0] == "error"] == [ "ConnectorConfigurationError" ] assert probe.discover_calls == 1 retried = ConnectorRepository(session).get(retry_key) assert retried["status"] == "succeeded" and retried["attempt_count"] == 2 retry_attempts = session.execute( text(""" SELECT ARRAY_AGG(attempt_number ORDER BY attempt_number) FROM public.connector_run_attempts WHERE run_uid=CAST(:run AS uuid) """), {"run": retried["uid"]}, ).scalar_one() assert retry_attempts == [1, 2] probe.resume_failures = 0 probe.resume_calls = 0 competing_resume_key = uuid.uuid4().hex + uuid.uuid4().hex competing_resume_source = _uid() resume_seed = ConnectorRepository(session, principal) resume_seed.claim( competing_resume_key, { "connector_id": connector_id, "connector_version": version, "source_uid": competing_resume_source, "operation": "discover", "checkpoint": {"page": 5}, "cursor": {"page": 5}, "dry_run": False, "principal_uid": principal, "business_domain_uid": domain, "environment": "staging", "process_key": "concurrent-resume", "config": {"credential_ref": "env:DATAOPS_CONNECTOR_TEST"}, "scope": {}, }, ) resume_seed.update( competing_resume_key, status="running", attempt_count=1, lease_token=_uid(), ) resume_seed.cancel(competing_resume_key) resume_gate = Barrier(2) class ResumeBarrierRepository(ConnectorRepository): def __init__(self, worker_session, actor_uid): super().__init__(worker_session, actor_uid) self.waited = False def get(self, key): record = super().get(key) if not self.waited and record and record.get("status") == "cancelled": self.waited = True resume_gate.wait(timeout=5) return record def competing_resume(): with Session(engine) as worker_session: worker_session.execute(text("SET statement_timeout='5s'")) worker_session.commit() try: result = ConnectorRuntime( registry, store=ResumeBarrierRepository(worker_session, principal), max_attempts=1, sleeper=lambda _seconds: None, ).resume(competing_resume_key) return "success", result.status except Exception as error: return "error", type(error).__name__ with ThreadPoolExecutor(max_workers=2) as executor: resume_futures = [executor.submit(competing_resume) for _index in range(2)] resume_results = [future.result(timeout=10) for future in resume_futures] assert sorted(item[0] for item in resume_results) == ["error", "success"] assert [item[1] for item in resume_results if item[0] == "error"] == [ "ConnectorConfigurationError" ] assert probe.resume_calls == 1 concurrently_resumed = ConnectorRepository(session).get(competing_resume_key) assert concurrently_resumed["status"] == "succeeded" assert concurrently_resumed["attempt_count"] == 2 resume_attempts = session.execute( text(""" SELECT ARRAY_AGG(attempt_number ORDER BY attempt_number) FROM public.connector_run_attempts WHERE run_uid=CAST(:run AS uuid) """), {"run": concurrently_resumed["uid"]}, ).scalar_one() assert resume_attempts == [1, 2] session.close() with engine.begin() as connection: for table in ( "connector_audit_events", "connector_graph_edges", "connector_evidence", "connector_checkpoints", "connector_run_attempts", "connector_runs", "connector_machine_credentials", "connector_principals", "connector_source_bindings", "connector_manifests", "connector_rate_limits", ): connection.execute(text(f"DELETE FROM public.{table}")) _alembic(url, "downgrade", "20260802_472") ambiguous_connector = f"ambiguous-{uuid.uuid4().hex[:8]}" ambiguous_source, ambiguous_domain = _uid(), _uid() with engine.begin() as connection: for ambiguous_version in ("1.0.0", "2.0.0"): connection.execute( text(""" INSERT INTO public.connector_manifests (uid,connector_id,connector_version,sdk_version,display_name, capabilities,config_schema,status,created_by) VALUES(CAST(:uid AS uuid),:connector,:version,'1.0','ambiguous', '["discover"]'::jsonb,'{"type":"object"}'::jsonb, 'active',CAST(:actor AS uuid)) """), { "uid": _uid(), "connector": ambiguous_connector, "version": ambiguous_version, "actor": actor, }, ) connection.execute( text(""" INSERT INTO public.connector_principals (uid,connector_id,source_uid,business_domain_uid,environment, allowed_operations,allowed_scopes,status,created_by) VALUES(CAST(:uid AS uuid),:connector,CAST(:source AS uuid), CAST(:domain AS uuid),'staging',ARRAY['discover'], '{}'::jsonb,'active',CAST(:actor AS uuid)) """), { "uid": _uid(), "connector": ambiguous_connector, "source": ambiguous_source, "domain": ambiguous_domain, "actor": actor, }, ) with pytest.raises(subprocess.CalledProcessError) as ambiguous_upgrade: _alembic(url, "upgrade", "20260802_473") assert "rejected before backfill" in ambiguous_upgrade.value.stderr with engine.connect() as connection: assert connection.execute( text("SELECT version_num FROM alembic_version") ).scalar_one() == "20260802_472" assert "connector_version" not in { item["name"] for item in inspect(engine).get_columns( "connector_principals", schema="public" ) } single_connector = f"single-{uuid.uuid4().hex[:8]}" single_principal, single_source, single_domain = _uid(), _uid(), _uid() with engine.begin() as connection: connection.execute( text("DELETE FROM public.connector_principals") ) connection.execute(text("DELETE FROM public.connector_manifests")) connection.execute( text(""" INSERT INTO public.connector_manifests (uid,connector_id,connector_version,sdk_version,display_name, capabilities,config_schema,status,created_by) VALUES(CAST(:uid AS uuid),:connector,'1.0.0','1.0','single', '["discover"]'::jsonb,'{"type":"object"}'::jsonb, 'active',CAST(:actor AS uuid)) """), {"uid": _uid(), "connector": single_connector, "actor": actor}, ) connection.execute( text(""" INSERT INTO public.connector_principals (uid,connector_id,source_uid,business_domain_uid,environment, allowed_operations,allowed_scopes,status,created_by) VALUES(CAST(:uid AS uuid),:connector,CAST(:source AS uuid), CAST(:domain AS uuid),'staging',ARRAY['discover'], '{}'::jsonb,'active',CAST(:actor AS uuid)) """), { "uid": single_principal, "connector": single_connector, "source": single_source, "domain": single_domain, "actor": actor, }, ) _alembic(url, "upgrade", "20260802_475") with pytest.raises(subprocess.CalledProcessError) as unbound_upgrade: _alembic(url, "upgrade", "20260802_476") assert "bind every active enterprise principal first" in ( unbound_upgrade.value.stderr ) binding_uid = _uid() revoked_principal = _uid() revoked_source, revoked_domain, revoked_binding = _uid(), _uid(), _uid() with engine.begin() as connection: connection.execute( text(""" INSERT INTO public.connector_source_bindings (uid,binding_version,connector_id,connector_version,source_uid, business_domain_uid,environment,approved_config,status,approved_by) VALUES(CAST(:uid AS uuid),1,:connector,'1.0.0',CAST(:source AS uuid), CAST(:domain AS uuid),'staging','{}'::jsonb,'approved',CAST(:actor AS uuid)) """), { "uid": binding_uid, "connector": single_connector, "source": single_source, "domain": single_domain, "actor": actor, }, ) connection.execute( text(""" UPDATE public.connector_principals SET source_binding_uid=CAST(:binding AS uuid),source_binding_version=1 WHERE uid=CAST(:principal AS uuid) """), {"binding": binding_uid, "principal": single_principal}, ) connection.execute( text(""" INSERT INTO public.connector_principals (uid,connector_id,connector_version,source_uid,business_domain_uid, environment,allowed_operations,allowed_scopes,status,created_by) VALUES(CAST(:uid AS uuid),:connector,'1.0.0',CAST(:source AS uuid), CAST(:domain AS uuid),'staging',ARRAY['discover'], '{}'::jsonb,'revoked',CAST(:actor AS uuid)) """), { "uid": revoked_principal, "connector": single_connector, "source": revoked_source, "domain": revoked_domain, "actor": actor, }, ) _alembic(url, "upgrade", "20260802_476") with pytest.raises(DBAPIError), engine.begin() as connection: connection.execute( text(""" UPDATE public.connector_principals SET status='active' WHERE uid=CAST(:principal AS uuid) """), {"principal": revoked_principal}, ) with engine.begin() as connection: connection.execute( text(""" INSERT INTO public.connector_source_bindings (uid,binding_version,connector_id,connector_version,source_uid, business_domain_uid,environment,approved_config,status,approved_by) VALUES(CAST(:uid AS uuid),1,:connector,'1.0.0',CAST(:source AS uuid), CAST(:domain AS uuid),'staging','{}'::jsonb,'approved',CAST(:actor AS uuid)) """), { "uid": revoked_binding, "connector": single_connector, "source": revoked_source, "domain": revoked_domain, "actor": actor, }, ) connection.execute( text(""" UPDATE public.connector_principals SET source_binding_uid=CAST(:binding AS uuid),source_binding_version=1 WHERE uid=CAST(:principal AS uuid) """), {"binding": revoked_binding, "principal": revoked_principal}, ) connection.execute( text(""" UPDATE public.connector_principals SET status='active' WHERE uid=CAST(:principal AS uuid) """), {"principal": revoked_principal}, ) assert connection.execute( text(""" SELECT status FROM public.connector_principals WHERE uid=CAST(:principal AS uuid) """), {"principal": revoked_principal}, ).scalar_one() == "active" with pytest.raises(DBAPIError), engine.begin() as connection: connection.execute( text(""" INSERT INTO public.connector_manifests (uid,connector_id,connector_version,sdk_version,display_name, capabilities,config_schema,status,created_by) VALUES(CAST(:uid AS uuid),:connector,'2.0.0','1.0','ambiguous', '["discover"]'::jsonb,'{"type":"object"}'::jsonb, 'active',CAST(:actor AS uuid)) """), {"uid": _uid(), "connector": single_connector, "actor": actor}, ) with pytest.raises(subprocess.CalledProcessError) as guarded_downgrade: _alembic(url, "downgrade", "20260802_474") assert "explicitly remove source bindings" in guarded_downgrade.value.stderr _alembic(url, "downgrade", "20260802_475") with engine.begin() as connection: connection.execute( text(""" UPDATE public.connector_principals SET source_binding_uid=NULL,source_binding_version=NULL WHERE uid IN (CAST(:principal AS uuid),CAST(:revoked AS uuid)) """), {"principal": single_principal, "revoked": revoked_principal}, ) connection.execute( text( "DELETE FROM public.connector_source_bindings WHERE uid=CAST(:uid AS uuid)" ), {"uid": binding_uid}, ) connection.execute( text( "DELETE FROM public.connector_source_bindings WHERE uid=CAST(:uid AS uuid)" ), {"uid": revoked_binding}, ) _alembic(url, "downgrade", "20260802_474") assert "connector_source_bindings" not in inspect(engine).get_table_names( schema="public" ) _clear_connector_test_data(engine) _alembic(url, "upgrade", "20260802_476") finally: try: _clear_connector_test_data(engine) finally: engine.dispose() _drop_wp03_rollback_database(shared_url, rollback_database) with create_engine(shared_url, pool_pre_ping=True).connect() as connection: assert connection.execute(text("SELECT version_num FROM alembic_version")).scalar_one() == "20260818_545"