diff --git a/apps/api/tests/test_workload_database_identity_projection.py b/apps/api/tests/test_workload_database_identity_projection.py index 0962a796e..9fb230ff0 100644 --- a/apps/api/tests/test_workload_database_identity_projection.py +++ b/apps/api/tests/test_workload_database_identity_projection.py @@ -104,6 +104,12 @@ def test_workload_database_bootstrap_keeps_secret_values_off_output() -> None: assert "sequence_privilege_gap_count" in script assert "REQUIRED_CONCURRENT_CONNECTIONS" in script assert "concurrent_connection_probe_count" in script + assert "existing_role_connection_count" in script + assert "additional_connection_probe_count" in script + assert "FROM pg_stat_activity" in script + assert "required_concurrency\n - existing_role_connection_count" in script + assert "for preflight_attempt in 1 2 3" in script + assert "preflight_attempts_exhausted=3" in script assert 'text("SELECT pg_sleep(0.25)")' in script assert "api) printf '12'" in script assert "worker) printf '8'" in script diff --git a/scripts/ops/awoooi-workload-db-identity-bootstrap.sh b/scripts/ops/awoooi-workload-db-identity-bootstrap.sh index 189187d1a..f5d9152a4 100755 --- a/scripts/ops/awoooi-workload-db-identity-bootstrap.sh +++ b/scripts/ops/awoooi-workload-db-identity-bootstrap.sh @@ -217,7 +217,7 @@ SQL unset database_url password } -preflight_workload() { +preflight_workload_once() { workload="$1" secret="$2" connection_limit="$3" @@ -332,20 +332,33 @@ spec: OR (has_sequence_privilege('awoooi', sequence.oid, 'UPDATE') AND NOT has_sequence_privilege(current_user, sequence.oid, 'UPDATE')) ) - """))).scalar_one()) + """))).scalar_one()) + required_concurrency = int( + os.environ["REQUIRED_CONCURRENT_CONNECTIONS"] + ) + existing_role_connection_count = int(( + await conn.execute(text(""" + SELECT count(*) + FROM pg_stat_activity + WHERE usename = current_user + """)) + ).scalar_one()) + additional_probe_count = max( + 0, + required_concurrency + - existing_role_connection_count, + ) - required_concurrency = int( - os.environ["REQUIRED_CONCURRENT_CONNECTIONS"] - ) + async def probe_connection() -> None: + async with engine.connect() as probe: + await probe.execute( + text("SELECT pg_sleep(0.25)") + ) - async def probe_connection() -> None: - async with engine.connect() as probe: - await probe.execute(text("SELECT pg_sleep(0.25)")) - - await asyncio.gather(*( - probe_connection() - for _ in range(required_concurrency) - )) + await asyncio.gather(*( + probe_connection() + for _ in range(additional_probe_count) + )) finally: await engine.dispose() @@ -363,6 +376,12 @@ spec: ).hexdigest(), "connection_limit": observed_limit, "concurrent_connection_probe_count": required_concurrency, + "existing_role_connection_count": ( + existing_role_connection_count + ), + "additional_connection_probe_count": ( + additional_probe_count + ), "representative_read_count": 3, "table_privilege_gap_count": table_gap_count, "sequence_privilege_gap_count": sequence_gap_count, @@ -407,6 +426,23 @@ YAML cleanup_preflight } +preflight_workload() { + workload="$1" + for preflight_attempt in 1 2 3; do + if preflight_workload_once "$@"; then + echo "workload=$workload preflight_attempt=$preflight_attempt preflight_verified=true" + return 0 + fi + if [ "$preflight_attempt" -lt 3 ]; then + backoff_seconds=$((preflight_attempt * 5)) + echo "workload=$workload preflight_attempt=$preflight_attempt retry_in_seconds=$backoff_seconds" + sleep "$backoff_seconds" + fi + done + echo "workload=$workload preflight_attempts_exhausted=3" >&2 + return 1 +} + require_control_path if [ "$MODE" = "disable" ]; then