fix(cd): bound workload DB connection preflight

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

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(
+                      async def probe_connection() -> None:
-                      os.environ["REQUIRED_CONCURRENT_CONNECTIONS"]
+                          async with engine.connect() as probe:
-                  )
+                              await probe.execute(
+                                  text("SELECT pg_sleep(0.25)")
+                              )
 
-                  async def probe_connection() -> None:
+                      await asyncio.gather(*(
-                      async with engine.connect() as probe:
+                          probe_connection()
-                          await probe.execute(text("SELECT pg_sleep(0.25)"))
+                          for _ in range(additional_probe_count)
-
+                      ))
-                  await asyncio.gather(*(
-                      probe_connection()
-                      for _ in range(required_concurrency)
-                  ))
               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