diff --git a/README.md b/README.md index 59a04c3..98c50c4 100644 --- a/README.md +++ b/README.md @@ -16,7 +16,7 @@ CleanStack cleans tabular files and documents with AI suggestions and human appr - **Pipelines and templates:** save the approved rules of a pipeline as a template and reuse it. A chat builder can draft pipelines and synthetic data. - **Schema drift alerts:** optional Slack webhook per pipeline. - **Usage controls:** monthly row quotas per plan (Free: 50,000 rows; plans are assigned by an admin, with no payment integration), Upstash rate limits, and an AI spend cap checked in Postgres before every Bedrock call. -- **Guest access:** **Try as guest** on the landing and sign-in pages starts a 24-hour session with no account. **Try with sample data** loads a seeded demo template, so it makes no Bedrock call. Guests get 2 MB files, 3 uploads, 5,000 rows per run and 10 AI calls per hour, with $0.25 of AI spend per guest. AI-heavy or irreversible features (chat builder, auto-clean, training export, Slack alerts, templates, account deletion) are blocked for guests. A guest's data is erased after 24 hours. [Full limits](docs/security-model.md#guest-access) +- **Guest access:** **Try as guest** on the landing and sign-in pages starts a 24-hour session with no account. **Try with sample data** loads a seeded demo template, so it makes no Bedrock call. Guests get 2 MB files, 3 uploads, 5,000 rows per run and 10 AI calls per hour, with $0.25 of AI spend per guest. AI-heavy or irreversible features (chat builder, auto-clean, training export, Slack alerts, templates, account deletion) are blocked for guests. A guest's data is erased by the purge that runs every 4 hours, once the 24-hour session has ended. [Full limits](docs/security-model.md#guest-access) ## Architecture @@ -41,7 +41,7 @@ Data PR approval / auto-clean ──► SQS executor queue ──► Lambda exec Other services: Clerk (auth), Upstash Redis (rate limits and quota cache, which fail open when not configured) and Sentry (optional). AWS resources are configured by hand, and there is no infrastructure as code in this repo. -A run moves through `pending → profiling → awaiting_ai → awaiting_approval → queued → running → completed | failed`. Execution is idempotent per run: the executor claims a run with a conditional `UPDATE`, so a duplicate SQS delivery is skipped. A stale run is marked failed after 20 minutes by `/api/cron/reconcile-runs`. +A run moves through `pending → profiling → awaiting_ai → awaiting_approval → queued → running → completed | failed`. Execution is idempotent per run: the executor claims a run with a conditional `UPDATE`, so a duplicate SQS delivery is skipped. A run stuck for 20 minutes is marked failed by `/api/cron/reconcile-runs`, which runs every 4 hours. ## Local development @@ -96,7 +96,7 @@ Vercel builds and deploys `main`. The Lambdas, migrations and AWS resources are - applying migrations 001–003 (`run-migration.mjs`, or the RDS Data API from CloudShell); - the CloudShell zip-swap deploy for `cleanstack-profiler`, `cleanstack-executor`, `cleanstack-ai-trigger` and `cleanstack-drift`; - raw bucket CORS (PUT, GET, POST), the profiler trigger on all object create events, and processed bucket CORS; -- EventBridge rules `cleanstack-reconciler-5min` (every 5 minutes) and `cleanstack-purge-guests-hourly` (every hour), which call the cron routes through API destinations as the `cleanstack-eventbridge-invoker` role; +- EventBridge rules `cleanstack-reconciler-5min` and `cleanstack-purge-guests-hourly`, both now `cron(0 */4 * * ? *)` (every 4 hours, so Aurora can auto-pause), which call the cron routes through API destinations as the `cleanstack-eventbridge-invoker` role; - the S3 lifecycle rule on the `guest_` prefix; - least-privilege IAM for the `cleanstack-vercel` user and for the Lambda roles. diff --git a/docs/deployment.md b/docs/deployment.md index b7aa9c1..67cb789 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -97,8 +97,10 @@ The profiler, ai-trigger and web app must share the same `WEBHOOK_SECRET`. | Rule | Schedule | API destination | Route | |---|---|---|---| -| `cleanstack-reconciler-5min` | `rate(5 minutes)` | `cleanstack-reconciler-destination` | `/api/cron/reconcile-runs`: marks runs stuck for 20 minutes as `failed` | -| `cleanstack-purge-guests-hourly` | `rate(1 hour)` | `cleanstack-purge-guests-destination` | `/api/cron/purge-guests`: erases up to 20 expired guests per call | +| `cleanstack-reconciler-5min` | `cron(0 */4 * * ? *)` | `cleanstack-reconciler-destination` | `/api/cron/reconcile-runs`: marks runs stuck for 20 minutes as `failed` | +| `cleanstack-purge-guests-hourly` | `cron(0 */4 * * ? *)` | `cleanstack-purge-guests-destination` | `/api/cron/purge-guests`: erases up to 20 expired guests per call | + +Both rules run every 4 hours (00:00, 04:00, ... UTC). The names are from the original 5-minute and hourly schedules. The reconciler used to wake the database every 5 minutes, so the Aurora Serverless v2 cluster `database-1` (minimum 0 ACU, auto-pause after 300 s) never paused, which cost about $45 a month. Now the cluster is paused most of the time, and the first request after a pause waits while it resumes (see [Aurora auto-pause](operations.md#aurora-auto-pause)). If you rotate `CRON_SECRET`, update the connection `cleanstack-reconciler-auth` at the same time as Vercel. diff --git a/docs/operations.md b/docs/operations.md index 0b6eca1..65d9e99 100644 --- a/docs/operations.md +++ b/docs/operations.md @@ -4,7 +4,7 @@ | Symptom | What to check or do | |---|---| -| Run stuck in `profiling`, `awaiting_ai`, `queued` or `running` | `reconcile-runs` marks these runs `failed` after 20 minutes, so first confirm that the EventBridge rule `cleanstack-reconciler-5min` is firing. | +| Run stuck in `profiling`, `awaiting_ai`, `queued` or `running` | `reconcile-runs` marks runs stuck for 20 minutes as `failed`. It runs every 4 hours (`cron(0 */4 * * ? *)`), so a stuck run can stay that way for up to about 4 h 20 min. First confirm that the EventBridge rule `cleanstack-reconciler-5min` is firing. | | Run `failed` with "Could not queue this run …" | Sending to SQS failed (log `SQS send failed; marking run failed` or `AI jobs SQS send failed; marking run failed`), so the route failed the run and returned 503. Check the queue URL and the `sqs:SendMessage` permission of `cleanstack-vercel`. The user retries by uploading again. | | Runs stay in `pending` after upload | The profiler trigger is missing `Post`. It must fire on all object create events. | | Run `failed` with `Attempt n/3 failed, retrying` in its history | SQS retries transient errors up to `EXECUTOR_MAX_ATTEMPTS` times, and the last attempt marks the run failed. Look up the run id in the executor's CloudWatch logs. | @@ -16,4 +16,15 @@ | Rotating `WEBHOOK_SECRET` | Update Vercel and the profiler and ai-trigger Lambdas together. Requests fail with 401 while the values differ. | | AI spend | `GET /api/admin/ai-spend` with `x-admin-secret` lists this month's estimated spend per team. | | Guests see "Guest AI capacity is used up for today" | All guests together have reached $5 of estimated Bedrock spend since 00:00 UTC. The cap resets at midnight UTC (`GUEST_LIMITS.aiSpendAllGuestsPerDayUsd`). | -| Guest data not being deleted | Check the rule `cleanstack-purge-guests-hourly` and its `invoke-purge-guests` policy. Each call erases up to 20 expired guests. The `guest_` lifecycle rule removes S3 objects after 1 day even if the purge stops, but the database rows remain until the purge runs. | +| Guest data not being deleted | Check the rule `cleanstack-purge-guests-hourly` (now every 4 hours, `cron(0 */4 * * ? *)`) and its `invoke-purge-guests` policy. Each call erases up to 20 expired guests, so about 120 a day. The `guest_` lifecycle rule removes S3 objects after 1 day even if the purge stops, but the database rows remain until the purge runs. | + +## Aurora auto-pause + +The cluster `database-1` (Aurora Serverless v2, PostgreSQL 17.7) has a minimum of 0 ACU and pauses after 300 s without connections. It used to never pause, because the reconciler called the database every 5 minutes, and that cost about $45 a month. Both EventBridge rules (`cleanstack-reconciler-5min` and `cleanstack-purge-guests-hourly`) now run on `cron(0 */4 * * ? *)`, every 4 hours, so the cluster is paused most of the time. + +The first request after a pause starts the resume, which takes about 15 s, or 30 s or more after a pause longer than 24 h. Callers wait instead of failing: + +- **Web app (RDS Data API):** while the cluster resumes, the Data API answers `DatabaseResumingException`. `src/lib/db-resume.ts` retries only that error, for `ExecuteStatement`, `BeginTransaction`, `CommitTransaction` and `RollbackTransaction`, with backoff of 1, 2, 4 and then 8 s, for up to 35 s. AWS documents the request as cancelled before it ran, so retrying writes is safe. Any other error, including statement timeouts, is never retried. The log line is `database is resuming from auto-pause; retrying Data API call` (warn, with `op`). The cron routes have `maxDuration` 60 (reconcile) and 120 (purge) to leave room for the wait. +- **Lambdas (psycopg2, IAM auth):** profiler, executor and drift connect directly. A connection to a paused cluster normally waits until it has resumed. A failed connect (`OperationalError`, except authentication failures) is retried for up to 35 s. The CloudWatch line is `[db] connect failed, database may be resuming from auto-pause`. ai-trigger does not connect to the database. + +If a request still fails with `DatabaseResumingException` after a long pause, the resume took longer than the 35 s budget. The constants are `RESUME_RETRY_BUDGET_MS` in `src/lib/db-resume.ts` and `DB_CONNECT_RETRY_BUDGET_S` in each Lambda. diff --git a/docs/security-model.md b/docs/security-model.md index b5b28f9..f4aacd6 100644 --- a/docs/security-model.md +++ b/docs/security-model.md @@ -29,6 +29,6 @@ Guest access is off until `GUEST_COOKIE_SECRET` is set. `POST /api/guest` checks | AI spend | $0.25 per guest, and $5 per UTC day for all guests together | `checkAiBudget` | | Sessions | 5 per IP hash, and 200 overall, per 24 h | `POST /api/guest` | -The values live in `src/lib/guest-limits.ts`. Guests cannot use the chat builder or synthetic data, auto-clean, training export, Slack alerts, account deletion, templates or `/api/admin`. The middleware answers 403 for these, or redirects `/templates` to the dashboard, and each route also rejects guest ids (`forbidGuest`). **Try with sample data** (`POST /api/sample-data`) uploads the CSV from `src/lib/sample-data.ts` and attaches the demo template that migration 003 seeds, so it makes no Bedrock call. After 24 h the cookie stops working, and `/api/cron/purge-guests` erases the guest's data. Guest work is not carried over to a new account on sign-up. +The values live in `src/lib/guest-limits.ts`. Guests cannot use the chat builder or synthetic data, auto-clean, training export, Slack alerts, account deletion, templates or `/api/admin`. The middleware answers 403 for these, or redirects `/templates` to the dashboard, and each route also rejects guest ids (`forbidGuest`). **Try with sample data** (`POST /api/sample-data`) uploads the CSV from `src/lib/sample-data.ts` and attaches the demo template that migration 003 seeds, so it makes no Bedrock call. After 24 h the cookie stops working, and the next run of `/api/cron/purge-guests` (every 4 hours) erases the guest's data. Guest work is not carried over to a new account on sign-up. Report vulnerabilities as described in [SECURITY.md](../SECURITY.md). diff --git a/lambdas/drift/handler.py b/lambdas/drift/handler.py index 6b4b044..0eb10a9 100644 --- a/lambdas/drift/handler.py +++ b/lambdas/drift/handler.py @@ -1,5 +1,6 @@ import json import os +import time import boto3 import psycopg2 import requests @@ -7,6 +8,31 @@ secrets = boto3.client("secretsmanager") +# Aurora Serverless v2 auto-pause: a connection to a paused cluster normally waits while it +# resumes (~15s, 30s+ after a pause longer than 24h), but AWS recommends retrying failed +# connections, since a burst of connections during a resume can fail. A failed connect ran no +# SQL, so retrying it is safe. Authentication errors are not retried. Same helper in every +# Lambda that connects (profiler, executor, drift); each Lambda ships as one handler.py. +DB_CONNECT_RETRY_BUDGET_S = 35 + + +def _connect_with_resume_retry(**kwargs): + deadline = time.monotonic() + DB_CONNECT_RETRY_BUDGET_S + delay, logged = 1, False + while True: + try: + return psycopg2.connect(**kwargs) + except psycopg2.OperationalError as e: + if "authentication failed" in str(e) or time.monotonic() + delay > deadline: + raise + if not logged: + print(f"[db] connect failed, database may be resuming from auto-pause; retrying " + f"for up to {DB_CONNECT_RETRY_BUDGET_S}s: {str(e).strip()[:200]}") + logged = True + time.sleep(delay) + delay = min(delay * 2, 8) + + def get_db_conn(): secret = json.loads( secrets.get_secret_value(SecretId=os.environ["DB_SECRET_ARN"])["SecretString"] @@ -17,7 +43,7 @@ def get_db_conn(): dbname = secret.get("dbname", "cleanstack") rds = boto3.client("rds", region_name=os.environ.get("AWS_REGION", "us-east-1")) token = rds.generate_db_auth_token(DBHostname=host, Port=port, DBUsername=user) - return psycopg2.connect(host=host, port=port, user=user, password=token, dbname=dbname, sslmode="require") + return _connect_with_resume_retry(host=host, port=port, user=user, password=token, dbname=dbname, sslmode="require") def compute_diff(old: dict, new: dict) -> dict: diff --git a/lambdas/executor/handler.py b/lambdas/executor/handler.py index 023de57..6a4e4b9 100644 --- a/lambdas/executor/handler.py +++ b/lambdas/executor/handler.py @@ -557,6 +557,31 @@ def _parse_params_simple(raw) -> dict: return raw or {} +# Aurora Serverless v2 auto-pause: a connection to a paused cluster normally waits while it +# resumes (~15s, 30s+ after a pause longer than 24h), but AWS recommends retrying failed +# connections, since a burst of connections during a resume can fail. A failed connect ran no +# SQL, so retrying it is safe. Authentication errors are not retried. Same helper in every +# Lambda that connects (profiler, executor, drift); each Lambda ships as one handler.py. +DB_CONNECT_RETRY_BUDGET_S = 35 + + +def _connect_with_resume_retry(**kwargs): + deadline = time.monotonic() + DB_CONNECT_RETRY_BUDGET_S + delay, logged = 1, False + while True: + try: + return psycopg2.connect(**kwargs) + except psycopg2.OperationalError as e: + if "authentication failed" in str(e) or time.monotonic() + delay > deadline: + raise + if not logged: + print(f"[db] connect failed, database may be resuming from auto-pause; retrying " + f"for up to {DB_CONNECT_RETRY_BUDGET_S}s: {str(e).strip()[:200]}") + logged = True + time.sleep(delay) + delay = min(delay * 2, 8) + + def get_db_conn(): secret = json.loads( secrets.get_secret_value(SecretId=os.environ["DB_SECRET_ARN"])["SecretString"] @@ -567,7 +592,7 @@ def get_db_conn(): dbname = secret.get("dbname", "cleanstack") rds = boto3.client("rds", region_name=os.environ.get("AWS_REGION", "us-east-1")) token = rds.generate_db_auth_token(DBHostname=host, Port=port, DBUsername=user) - return psycopg2.connect(host=host, port=port, user=user, password=token, dbname=dbname, sslmode="require") + return _connect_with_resume_retry(host=host, port=port, user=user, password=token, dbname=dbname, sslmode="require") def load_raw_dataframe(file_bytes: bytes, fmt: str) -> pd.DataFrame: diff --git a/lambdas/profiler/handler.py b/lambdas/profiler/handler.py index b1436f0..9c581ad 100644 --- a/lambdas/profiler/handler.py +++ b/lambdas/profiler/handler.py @@ -15,6 +15,7 @@ import json import io import re +import time import boto3 import psycopg2 import requests @@ -299,6 +300,31 @@ def compute_quality_score(df: pd.DataFrame) -> dict: } +# Aurora Serverless v2 auto-pause: a connection to a paused cluster normally waits while it +# resumes (~15s, 30s+ after a pause longer than 24h), but AWS recommends retrying failed +# connections, since a burst of connections during a resume can fail. A failed connect ran no +# SQL, so retrying it is safe. Authentication errors are not retried. Same helper in every +# Lambda that connects (profiler, executor, drift); each Lambda ships as one handler.py. +DB_CONNECT_RETRY_BUDGET_S = 35 + + +def _connect_with_resume_retry(**kwargs): + deadline = time.monotonic() + DB_CONNECT_RETRY_BUDGET_S + delay, logged = 1, False + while True: + try: + return psycopg2.connect(**kwargs) + except psycopg2.OperationalError as e: + if "authentication failed" in str(e) or time.monotonic() + delay > deadline: + raise + if not logged: + print(f"[db] connect failed, database may be resuming from auto-pause; retrying " + f"for up to {DB_CONNECT_RETRY_BUDGET_S}s: {str(e).strip()[:200]}") + logged = True + time.sleep(delay) + delay = min(delay * 2, 8) + + def get_db_conn(): from urllib.parse import urlparse url = os.environ["DATABASE_URL"] @@ -306,7 +332,7 @@ def get_db_conn(): host, port, user, dbname = p.hostname, p.port or 5432, p.username, p.path.lstrip("/") rds = boto3.client("rds", region_name=os.environ.get("AWS_REGION", "us-east-1")) token = rds.generate_db_auth_token(DBHostname=host, Port=port, DBUsername=user) - return psycopg2.connect(host=host, port=port, user=user, password=token, dbname=dbname, sslmode="require") + return _connect_with_resume_retry(host=host, port=port, user=user, password=token, dbname=dbname, sslmode="require") def detect_format(key: str, content_type: str, file_bytes: bytes = b"") -> str: diff --git a/lambdas/tests/test_db_resume_retry.py b/lambdas/tests/test_db_resume_retry.py new file mode 100644 index 0000000..0e5d98b --- /dev/null +++ b/lambdas/tests/test_db_resume_retry.py @@ -0,0 +1,83 @@ +"""Every Lambda that connects to Aurora retries a failed connect while the cluster resumes +from auto-pause, never retries authentication errors, and gives up after its budget.""" +import psycopg2 +import pytest + + +@pytest.fixture(params=["profiler", "executor", "drift"]) +def handler(request): + return request.getfixturevalue(request.param) + + +class FakeClock: + def __init__(self): + self.t = 0.0 + self.sleeps = [] + + def monotonic(self): + return self.t + + def sleep(self, s): + self.sleeps.append(s) + self.t += s + + +@pytest.fixture +def clock(handler, monkeypatch): + c = FakeClock() + monkeypatch.setattr(handler.time, "monotonic", c.monotonic) + monkeypatch.setattr(handler.time, "sleep", c.sleep) + return c + + +def _connect_failing(monkeypatch, handler, errors, result="conn"): + calls = [] + + def fake_connect(**kwargs): + calls.append(kwargs) + if len(calls) <= len(errors): + raise errors[len(calls) - 1] + return result + + monkeypatch.setattr(handler.psycopg2, "connect", fake_connect) + return calls + + +def test_retries_while_resuming_then_connects(handler, clock, monkeypatch, capsys): + err = psycopg2.OperationalError("could not connect to server: Connection timed out") + calls = _connect_failing(monkeypatch, handler, [err, err, err]) + assert handler._connect_with_resume_retry(host="h", dbname="d") == "conn" + assert len(calls) == 4 and calls[-1] == {"host": "h", "dbname": "d"} + assert clock.sleeps == [1, 2, 4] + out = capsys.readouterr().out + assert out.count("database may be resuming") == 1 # one log line per connect + + +def test_authentication_errors_are_not_retried(handler, clock, monkeypatch): + err = psycopg2.OperationalError('FATAL: PAM authentication failed for user "cleanstack"') + calls = _connect_failing(monkeypatch, handler, [err]) + with pytest.raises(psycopg2.OperationalError, match="authentication failed"): + handler._connect_with_resume_retry(host="h") + assert len(calls) == 1 and clock.sleeps == [] + + +def test_other_exceptions_are_not_retried(handler, clock, monkeypatch): + calls = _connect_failing(monkeypatch, handler, [ValueError("bad dsn")]) + with pytest.raises(ValueError): + handler._connect_with_resume_retry(host="h") + assert len(calls) == 1 and clock.sleeps == [] + + +def test_gives_up_after_the_budget(handler, clock, monkeypatch): + err = psycopg2.OperationalError("could not connect to server") + calls = _connect_failing(monkeypatch, handler, [err] * 100) + with pytest.raises(psycopg2.OperationalError): + handler._connect_with_resume_retry(host="h") + assert clock.sleeps == [1, 2, 4, 8, 8, 8] + assert sum(clock.sleeps) <= handler.DB_CONNECT_RETRY_BUDGET_S + assert len(calls) == len(clock.sleeps) + 1 + + +def test_get_db_conn_uses_the_retry(handler): + import inspect + assert "_connect_with_resume_retry(" in inspect.getsource(handler.get_db_conn) diff --git a/src/app/api/cron/purge-guests/route.ts b/src/app/api/cron/purge-guests/route.ts index 83d886c..c10c068 100644 --- a/src/app/api/cron/purge-guests/route.ts +++ b/src/app/api/cron/purge-guests/route.ts @@ -10,7 +10,9 @@ const log = logger.child({ route: "GET /api/cron/purge-guests" }); const s3 = new S3Client({ region: awsRegion() }); -export const maxDuration = 60; +// Room for up to 35s of Aurora resume wait (src/lib/db-resume.ts) before the batch: +// this rule runs every 4 hours, so the DB is usually paused. +export const maxDuration = 120; /** Guests erased per call; a backlog drains over the next calls. */ export const PURGE_BATCH = 20; diff --git a/src/app/api/cron/reconcile-runs/route.ts b/src/app/api/cron/reconcile-runs/route.ts index fb7aec2..97febae 100644 --- a/src/app/api/cron/reconcile-runs/route.ts +++ b/src/app/api/cron/reconcile-runs/route.ts @@ -6,7 +6,9 @@ import { logger } from "@/lib/logger"; const log = logger.child({ route: "GET /api/cron/reconcile-runs" }); -export const maxDuration = 30; +// Room for the first query to wait while Aurora resumes from auto-pause (up to +// 35s, src/lib/db-resume.ts): this rule runs every 4 hours, so the DB is usually paused. +export const maxDuration = 60; const STUCK_AFTER_MINUTES = 20; const PENDING_ORPHAN_AFTER_MINUTES = 60; diff --git a/src/lib/db-resume.test.ts b/src/lib/db-resume.test.ts new file mode 100644 index 0000000..8f6b780 --- /dev/null +++ b/src/lib/db-resume.test.ts @@ -0,0 +1,74 @@ +import { describe, expect, it, vi } from "vitest"; +import { isDatabaseResuming, withResumeRetry, RESUME_RETRY_BUDGET_MS } from "@/lib/db-resume"; + +const resuming = () => Object.assign(new Error("The database is resuming"), { name: "DatabaseResumingException" }); +const named = (name: string, message = "x") => Object.assign(new Error(message), { name }); + +/** Fake clock: sleep advances time instead of waiting. */ +function clock() { + let t = 0; + const sleeps: number[] = []; + return { + sleeps, + now: () => t, + sleep: async (ms: number) => { sleeps.push(ms); t += ms; }, + }; +} + +describe("isDatabaseResuming", () => { + it("matches only the Data API resuming error", () => { + expect(isDatabaseResuming(resuming())).toBe(true); + expect(isDatabaseResuming(named("BadRequestException", "Communications link failure"))).toBe(false); + expect(isDatabaseResuming(named("StatementTimeoutException"))).toBe(false); + expect(isDatabaseResuming(named("DatabaseUnavailableException"))).toBe(false); + expect(isDatabaseResuming(null)).toBe(false); + expect(isDatabaseResuming("DatabaseResumingException")).toBe(false); + }); +}); + +describe("withResumeRetry", () => { + it("waits with backoff while the database resumes, then returns the result", async () => { + const c = clock(); + const send = vi.fn() + .mockRejectedValueOnce(resuming()) + .mockRejectedValueOnce(resuming()) + .mockRejectedValueOnce(resuming()) + .mockResolvedValueOnce({ records: [] }); + await expect(withResumeRetry("ExecuteStatement", send, c)).resolves.toEqual({ records: [] }); + expect(send).toHaveBeenCalledTimes(4); + expect(c.sleeps).toEqual([1000, 2000, 4000]); + }); + + it("does not retry other errors", async () => { + for (const err of [named("StatementTimeoutException"), named("BadRequestException", "Communications link failure"), + named("DatabaseErrorException", "duplicate key")]) { + const c = clock(); + const send = vi.fn().mockRejectedValue(err); + await expect(withResumeRetry("ExecuteStatement", send, c)).rejects.toBe(err); + expect(send).toHaveBeenCalledTimes(1); + expect(c.sleeps).toEqual([]); + } + }); + + it("gives up after the budget and rethrows the resuming error", async () => { + const c = clock(); + const err = resuming(); + const send = vi.fn().mockRejectedValue(err); + await expect(withResumeRetry("BeginTransaction", send, c)).rejects.toBe(err); + const waited = c.sleeps.reduce((a, b) => a + b, 0); + expect(c.sleeps).toEqual([1000, 2000, 4000, 8000, 8000, 8000]); // capped at 8s + expect(waited).toBeLessThanOrEqual(RESUME_RETRY_BUDGET_MS); + expect(send).toHaveBeenCalledTimes(c.sleeps.length + 1); + }); + + it("logs one warn line per call that had to wait", async () => { + const c = clock(); + const warn = vi.spyOn(console, "warn").mockImplementation(() => {}); + const send = vi.fn().mockRejectedValueOnce(resuming()).mockRejectedValueOnce(resuming()).mockResolvedValueOnce(1); + await withResumeRetry("ExecuteStatement", send, c); + const lines = warn.mock.calls.map((a) => String(a[0])).filter((l) => l.includes("resuming from auto-pause")); + expect(lines).toHaveLength(1); + expect(JSON.parse(lines[0])).toMatchObject({ level: "warn", op: "ExecuteStatement" }); + warn.mockRestore(); + }); +}); diff --git a/src/lib/db-resume.ts b/src/lib/db-resume.ts new file mode 100644 index 0000000..d1e0b64 --- /dev/null +++ b/src/lib/db-resume.ts @@ -0,0 +1,68 @@ +/** + * Wait-and-retry for RDS Data API calls while Aurora Serverless v2 resumes + * from auto-pause (min 0 ACU). The first Data API request to a paused cluster + * fails with DatabaseResumingException and starts the resume, which usually + * takes about 15s (30s or more after a pause longer than 24h, per the Aurora + * auto-pause docs). + * + * Only the resuming error is retried. AWS documents it as "the request was + * cancelled because the DB instance was paused", so the statement never ran + * and retrying a write (INSERT/UPDATE, BeginTransaction, Commit) is safe. Every + * other error is rethrown at once, including StatementTimeoutException, + * DatabaseUnavailableException and connection failures, which can happen after + * a statement was sent, so a non-idempotent statement is never run twice. + * + * The wait is capped (RESUME_RETRY_BUDGET_MS) so a call still fits in short + * routes: the cron routes have maxDuration 60. + */ +import { logger } from "@/lib/logger"; + +const log = logger.child({ module: "db" }); + +export const RESUME_RETRY_BUDGET_MS = 35_000; +const FIRST_DELAY_MS = 1_000; +const MAX_DELAY_MS = 8_000; + +/** + * The Data API's error for this case on Aurora Serverless v2 (HTTP 400, client + * fault, so the SDK's own retry strategy does not retry it). The + * BadRequestException "Communications link failure" seen during resumes is + * documented for Serverless v1 only and is not treated as resuming. + */ +export function isDatabaseResuming(err: unknown): boolean { + return !!err && typeof err === "object" && (err as { name?: unknown }).name === "DatabaseResumingException"; +} + +export interface ResumeRetryOptions { + budgetMs?: number; + sleep?: (ms: number) => Promise; + now?: () => number; +} + +const realSleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms)); + +/** + * Run `send`, retrying with backoff (1s, 2s, 4s, 8s, 8s, ...) only while the + * database is resuming, until the budget would be exceeded; then the last + * resuming error is rethrown. Logs one warn line per call that had to wait. + */ +export async function withResumeRetry( + op: string, + send: () => Promise, + { budgetMs = RESUME_RETRY_BUDGET_MS, sleep = realSleep, now = Date.now }: ResumeRetryOptions = {} +): Promise { + const deadline = now() + budgetMs; + let delay = FIRST_DELAY_MS; + for (let attempt = 1; ; attempt++) { + try { + return await send(); + } catch (err) { + if (!isDatabaseResuming(err) || now() + delay > deadline) throw err; + if (attempt === 1) { + log.warn("database is resuming from auto-pause; retrying Data API call", { op, budget_ms: budgetMs }); + } + await sleep(delay); + delay = Math.min(delay * 2, MAX_DELAY_MS); + } + } +} diff --git a/src/lib/db.test.ts b/src/lib/db.test.ts new file mode 100644 index 0000000..ad8d9c8 --- /dev/null +++ b/src/lib/db.test.ts @@ -0,0 +1,62 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; + +// The Data API path of src/lib/db.ts: every client call goes through the resume retry. +const send = vi.fn(); +vi.mock("@aws-sdk/client-rds-data", () => { + const cmd = (type: string) => class { type = type; constructor(public input: unknown) {} }; + return { + RDSDataClient: class { send = (...a: unknown[]) => send(...a); }, + ExecuteStatementCommand: cmd("ExecuteStatement"), + BeginTransactionCommand: cmd("BeginTransaction"), + CommitTransactionCommand: cmd("CommitTransaction"), + RollbackTransactionCommand: cmd("RollbackTransaction"), + }; +}); +vi.mock("@/lib/db-resume", async (orig) => { + const real = await orig(); + // Same retry logic, no real waiting. + return { ...real, withResumeRetry: (op: string, fn: () => Promise) => real.withResumeRetry(op, fn, { sleep: async () => {} }) }; +}); + +process.env.AURORA_CLUSTER_ARN = "arn:aws:rds:us-east-1:1:cluster:database-1"; +process.env.AURORA_SECRET_ARN = "arn:aws:secretsmanager:us-east-1:1:secret:x"; +delete process.env.DB_DRIVER; +const { query, withTransaction } = await import("@/lib/db"); + +const resuming = () => Object.assign(new Error("resuming"), { name: "DatabaseResumingException" }); + +describe("db (Data API driver)", () => { + beforeEach(() => { + send.mockReset(); + vi.spyOn(console, "warn").mockImplementation(() => {}); + }); + + it("retries ExecuteStatement while the cluster resumes", async () => { + send.mockRejectedValueOnce(resuming()).mockRejectedValueOnce(resuming()).mockResolvedValueOnce({ + columnMetadata: [{ name: "n", typeName: "int4" }], records: [[{ longValue: 1 }]], + }); + await expect(query("SELECT 1 AS n")).resolves.toEqual([{ n: 1 }]); + expect(send).toHaveBeenCalledTimes(3); + }); + + it("does not retry a failed statement", async () => { + send.mockRejectedValueOnce(Object.assign(new Error("duplicate key"), { name: "DatabaseErrorException" })); + await expect(query("INSERT INTO t VALUES (1)")).rejects.toThrow("duplicate key"); + expect(send).toHaveBeenCalledTimes(1); + }); + + it("retries BeginTransaction and Commit", async () => { + send.mockImplementation(async (cmd: { type: string }) => { + if (cmd.type === "BeginTransaction") { + if (send.mock.calls.length === 1) throw resuming(); + return { transactionId: "tx1" }; + } + if (cmd.type === "CommitTransaction" && send.mock.calls.filter((c) => c[0].type === "CommitTransaction").length === 1) { + throw resuming(); + } + return {}; + }); + await expect(withTransaction(async () => "ok")).resolves.toBe("ok"); + expect(send.mock.calls.map((c) => c[0].type)).toEqual(["BeginTransaction", "BeginTransaction", "CommitTransaction", "CommitTransaction"]); + }); +}); diff --git a/src/lib/db.ts b/src/lib/db.ts index a6557f8..9d5aabb 100644 --- a/src/lib/db.ts +++ b/src/lib/db.ts @@ -10,6 +10,7 @@ import { import { ConfigError, awsRegion, optionalEnv, requireEnv } from "@/lib/env"; import { parseDataApiString } from "@/lib/db-shape"; import { isTestMode } from "@/lib/test-mode"; +import { withResumeRetry } from "@/lib/db-resume"; const client = new RDSDataClient({ region: awsRegion(), @@ -171,7 +172,8 @@ export async function query( const { clusterArn, secretArn } = getArns(); const { sql, parameters } = convertQuery(text, params); - const result = await client.send( + // Every Data API call waits while the cluster resumes from auto-pause (src/lib/db-resume.ts). + const result = await withResumeRetry("ExecuteStatement", () => client.send( new ExecuteStatementCommand({ resourceArn: clusterArn, secretArn: secretArn, @@ -181,7 +183,7 @@ export async function query( transactionId, includeResultMetadata: true, }) - ); + )); if (!result.columnMetadata || !result.records) return []; return recordsToRows(result.columnMetadata, result.records); @@ -201,32 +203,32 @@ export async function withTransaction( ): Promise { if (pgDriverSelected()) return (await loadPg()).pgWithTransaction(fn); const { clusterArn, secretArn } = getArns(); - const begin = await client.send( + const begin = await withResumeRetry("BeginTransaction", () => client.send( new BeginTransactionCommand({ resourceArn: clusterArn, secretArn: secretArn, database: DATABASE, }) - ); + )); const txId = begin.transactionId!; try { const result = await fn(txId); - await client.send( + await withResumeRetry("CommitTransaction", () => client.send( new CommitTransactionCommand({ resourceArn: clusterArn, secretArn: secretArn, transactionId: txId, }) - ); + )); return result; } catch (e) { - await client.send( + await withResumeRetry("RollbackTransaction", () => client.send( new RollbackTransactionCommand({ resourceArn: clusterArn, secretArn: secretArn, transactionId: txId, }) - ); + )); throw e; } }