agnes-the-ai-analyst/tests/test_sync_trigger_singleton.py
ZdenekSrotyr 28430ced09
Keboola cutover: native parquet path + sync correctness + auto-discover protection (#190)
* fix: cutover regressions + parallel Keboola legacy fallback

Bundled fixes from a fresh-deploy run on a Keboola Storage backend with
the block-shared-snowflake-access feature flag — DuckDB Keboola
extension's per-table scan can't access bucket schemas, so the legacy
kbcstorage Storage-API client is the only working path.

CUTOVER REGRESSIONS

- agnes pull hash mismatch on every Keboola local-mode table —
  src/orchestrator.py:_update_sync_state stored md5(mtime+size)[:12]
  while the CLI compares against full 32-char content MD5. Now stores
  the same content MD5 the materialized SQL path already used.

- Trailing-slash sanitization in connectors/keboola/access.py and
  extractor.py — DuckDB Keboola extension's ATTACH fails when the URL
  ends in / (canonical form).

- src/profiler.py:TableInfo.description becomes optional — two call
  sites instantiated without it, crashing the profiler pass.

- scripts/ops/agnes-auto-upgrade.sh: chown on UID change — older images
  ran as root, current runs as agnes (uid 999). Reads target uid:gid
  from /etc/passwd inside the new image and chowns ${STATE_DIR},
  /data/extracts, /data/analytics when the digest moves.

- POST /api/sync/trigger is now singleton per process — two
  near-simultaneous trigger calls each forked an extractor subprocess,
  fought for extract.duckdb's file lock, starved uvicorn, flipped the
  container to unhealthy. Trigger now returns 409
  (sync_already_in_progress) when held; _run_sync acquires non-blocking.

PARALLEL LEGACY FALLBACK

- Process pool fan-out for the _extract_via_legacy queue (default 8
  workers, override via AGNES_KEBOOLA_PARALLELISM). Process pool, not
  thread pool, because connectors/keboola/client.py:export_table does
  os.chdir(temp_dir) — process-global, so threads raced and slice files
  landed in the wrong directory ("[Errno 2] No such file or directory:
  '<job_id>.csv_X_Y_Z.csv'").

- Extractor subprocess timeout 1800s -> 3600s (configurable via
  AGNES_EXTRACTOR_TIMEOUT_SEC). 28+ tables × multi-minute Keboola export
  jobs need the headroom on telemetry-class projects.

- Process group cleanup on timeout — Popen(start_new_session=True) puts
  the extractor in its own group. On timeout the parent SIGTERMs the
  group (10s grace) then SIGKILLs stragglers. Without this, the pool
  workers were reparented to PID 1 and continued holding open Keboola
  Storage export jobs. Inline extractor script also installs a SIGTERM
  -> sys.exit(143) handler so the with ProcessPoolExecutor(...) block
  __exit__ runs cleanly.

Tests: existing tests that patched subprocess.run updated to patch
subprocess.Popen with a _FakePopen stand-in (same exit-code-injection
contract). Two tests that exercised the parallel path forced
AGNES_KEBOOLA_PARALLELISM=1 to keep mocks alive (mocks don't ride into
ProcessPoolExecutor subprocesses).

Squashed onto current main (was 7 commits + multi-commit CHANGELOG +
agnes-auto-upgrade.sh conflicts; squash avoids per-commit conflict
resolution against main's flat-mount STATE_DIR refactor and 0.38.0
release cut).

* feat(keboola): Storage API direct extract path; drop extension data path

The DuckDB Keboola extension's COPY routes through Keboola QueryService,
which is unreliable on linked-bucket projects (extension v0.1.6 fixes
that case but isn't yet in the community CDN, and pre-fix any project
with the block-shared-snowflake-access feature flag couldn't see bucket
schemas at all). Move the extract path off the extension entirely and
talk to the Storage API directly via signed-URL download — works on any
project, regardless of extension state.

connectors/keboola/storage_api.py (NEW)
  Lightweight client built on requests.Session. Three endpoints:
  - POST /v2/storage/tables/{id}/export-async        (kicks off job)
  - GET  /v2/storage/jobs/{id}                        (poll until done)
  - GET  /v2/storage/files/{id}?federationToken=1     (signed URL detail)
  - GET  <signed_url>                                 (download bytes)
  Supports sliced exports (manifest + per-slice signed URLs) and gzipped
  payloads. ExportFilter dataclass mirrors the Keboola filter spec
  (whereFilters / columns / changedSince / limit) and handles JSON
  round-trip with the registry's source_query column. Token redaction
  in error messages. Bounded exponential backoff on job polling.
  No cloud-SDK dependency on the data path; thread-safe.

connectors/keboola/extractor.py
  - materialize_query() rewritten: takes bucket/source_table/source_query
    (JSON filter spec), exports via KeboolaStorageClient, converts CSV
    to parquet via DuckDB, atomic os.replace. Same return shape so
    sync.py downstream code stays uniform with the BQ branch.
  - _extract_via_legacy() also moved to Storage API direct (kept the
    name for caller compatibility with _legacy_worker / the parallel
    batch extractor). Per-call temp directories — no os.chdir, threads
    don't race.

app/api/sync.py
  _run_materialized_pass for source_type='keboola' rows now constructs a
  KeboolaStorageClient (replaces KeboolaAccess) and passes
  bucket/source_table/source_query to materialize_query. Reuses one
  client across rows for HTTP keep-alive. Sources keboola URL from env
  too (KEBOOLA_STACK_URL) when instance.yaml doesn't have stack_url
  configured.

cli/commands/admin.py
  discover-and-register defaults Keboola rows to query_mode='materialized'
  (NULL source_query = full table), matching the v26 migration's
  unification of the local/materialized split for Keboola. BigQuery and
  Jira keep their per-source defaults.

src/db.py
  Schema bump 25 → 26. Migration: UPDATE table_registry SET
  query_mode='materialized' WHERE source_type='keboola' AND
  query_mode='local'. NULL source_query on those rows means "full table
  export" — same effective behavior the local mode provided, but now
  via Storage API instead of the extension.

pyproject.toml
  kbcstorage dep stays (admin-side bucket/table list still uses the
  SDK in app/api/admin.py / connectors/keboola/client.py); only the
  data path is migrated off the SDK. Comment updated to reflect the
  new boundary.

tests
  - test_keboola_storage_api.py (NEW, 19 tests): ExportFilter parsing,
    HTTP client (token redaction, retry logic, polling), download_file
    (single, gzipped, sliced), end-to-end export_table_to_csv.
  - test_keboola_materialize.py rewritten: mocks KeboolaStorageClient
    instead of FakeAccess; same atomic-write + zero-rows + unsafe-id
    contracts.
  - test_sync_trigger_keboola_materialized.py: registry rows now carry
    bucket+source_table+JSON-shape source_query.

114+ Keboola-impacted tests green locally.

* test: schema version assertion bumped to 26 alongside the keboola query_mode migration

* fix(keboola): cutover hot-patches surfaced on agnes-dev

Five small fixes that were applied as in-container hot-patches during
agnes-dev cutover and need to be on the source-of-truth image so a fresh
upgrade does not undo them.

- app/api/sync.py: auto-discover gate considers the WHOLE registry (any
  source, any mode), not just rows where source matches and query_mode
  is local. After the v25→v26 keboola materialized migration an
  instance can have 30 materialized rows and zero local rows; the
  previous gate kept re-firing _discover_and_register_tables every
  scheduler tick, creating duplicate auto-discovered rows with the
  wrong bucket prefix every time.

- app/api/admin.py: _discover_and_register_tables reassembles the
  bucket as <stage>.<bucket-id> (e.g. in.c-finance) instead of
  dropping the stage prefix; default query_mode for keboola is now
  materialized (the v26 contract); validator allows NULL source_query
  for keboola materialized rows (full-table export via Storage API
  export-async, no SQL needed).

- cli/commands/admin.py: register-table mirrors the server validator
  (NULL source_query allowed for source_type=keboola); --bucket help
  text generalized to cover both BQ dataset and Keboola bucket id.

- connectors/keboola/extractor.py: max_line_size=64 MiB on
  read_csv_auto so embedded JSON / SQL cells (kbc_component_configuration
  in particular) do not trip the default 2 MiB ceiling.

- connectors/keboola/storage_api.py: GCP backend support — when the
  Storage API returns a manifest whose slice URLs are gs://
  references with a gcsCredentials block, rewrite to the JSON REST
  download endpoint and authenticate with the issued OAuth bearer
  token; redact tokens in any surfaced error string.

* test: align with new keboola materialized + auto-discover-gate contracts

- test_admin_keboola_materialized: rename
  test_register_keboola_materialized_rejects_missing_source_query →
  test_register_keboola_materialized_accepts_missing_source_query.
  v25→v26 introduced 'keboola materialized with NULL source_query
  means full-table export via Storage API export-async' as the
  default registration shape; the rejection case is no longer the
  contract.

- test_sync_filter: add list_all() to _StubRegistry. The auto-discover
  gate in _run_sync now keys off the WHOLE registry (not just local
  rows) so materialized-only Keboola instances do not re-trigger
  discovery on every tick.

* feat(keboola): native parquet export — skip CSV roundtrip

Storage API export-async accepts fileType={csv,parquet}. Switching the
materialized sync to parquet eliminates the CSV → DuckDB COPY → parquet
roundtrip that pinned a single uvicorn worker over 4 GiB on multi-GB
tables (read_csv with all_varchar + max_line_size=64MB has to
materialize the whole CSV in memory before COPY can stream out a
parquet). Snowflake UNLOAD on Keboola's side already produces typed,
self-contained parquet files; the extractor downloads them and renames
into place.

Two cases:

- **Single-file** export (small table): file_info.url points at one
  signed URL; download_file streams chunks straight to .parquet.tmp
  and we're done. No DuckDB.

- **Sliced** export (Snowflake UNLOAD respects MAX_FILE_SIZE — 16 MiB
  default — so anything larger arrives as N parquet slices): each
  slice is a complete parquet file with its own footer; naive concat
  would corrupt them. download_file_slices keeps the slices as
  separate files in a tempdir, then DuckDB COPY (SELECT * FROM
  read_parquet([slice0, slice1, ...])) merges them into one
  consolidated parquet. DuckDB streams row groups during this — peak
  memory bounded to one row group (~1 MiB) regardless of source size.

The legacy CSV path stays as the explicit opt-in via source_query=
'{"file_type":"csv"}' for projects whose backend can't UNLOAD
parquet (none known today; cheap escape hatch). Backward-compat alias
KeboolaStorageClient.export_table_to_csv kept.

Also fixes a latent bug in download_file's gzip detection: previous
heuristic flagged any unencrypted file as gzipped, which would have
corrupted parquet downloads at gunzip time. Name-suffix-only now.

* fix: tempdir leak cleanup, every 0m schedule, /sync/trigger body shapes

Three small self-contained fixes uncovered during agnes-dev cutover.

- connectors/keboola/extractor.py: tempfile.TemporaryDirectory now uses
  ignore_cleanup_errors=True so a worker death mid-write doesn't leave
  multi-GiB stale slice trees on the boot disk. (12 GiB seen after a
  disk-full crash where TemporaryDirectory's own cleanup also raised
  and got swallowed.)

- src/scheduler.py: is_valid_schedule accepts 'every 0m' (interval=0
  = always due). Force-resync of an errored row no longer requires
  waiting out the default 'every 1h' interval — admin can flip the
  schedule, trigger, then flip back.

- app/api/sync.py: POST /api/sync/trigger accepts both ['table_id']
  (legacy bare-array body) and {'tables': ['table_id']} (matches the
  response payload shape, more discoverable for clients building
  requests by hand). Malformed bodies return 422 with a structured
  detail; null/missing means 'sync everything' as before.

Tests cover: tempdir cleanup on raise (sliced parquet path),
is_valid_schedule + is_table_due 'every 0m' acceptance, and trigger
body parametrized matrix (8 valid shapes + 6 rejection cases).

* fix: targeted-trigger filter in materialized pass + auto-upgrade defer

Two operational gaps observed during agnes-dev cutover, in the same
sync-routing area.

- _run_materialized_pass now takes a 'tables' arg and skips rows not in
  the target set with reason='not_in_target'. POST /api/sync/trigger
  with a body of tables previously only scoped the legacy extractor
  subprocess — the materialized pass kept iterating every due
  materialized row, so an admin asking to re-sync kbc_job re-ran
  every other due materialized row alongside it. Match on registry id
  OR name (admins commonly pass either form). tables=None preserves
  the no-filter behavior.

- New GET /api/sync/status (public, no auth) returns {locked: bool}
  off _sync_lock.locked(). agnes-auto-upgrade.sh probes this before
  docker compose up -d and exits 0 with a 'deferred recreate' log
  line if a sync is in flight — the next 5-min cron tick retries.
  Pre-fix, an auto-upgrade triggered mid-sync would recreate the
  uvicorn worker and kill the in-flight extractor / Snowflake-UNLOAD
  download (observed when kbc_job's first 7-day retry got SIGKILLed).
  Connection failures in the probe fall through to the upgrade —
  being stuck on a wedged image is worse than interrupting a
  hypothetical sync.

* fix: auto-discover protects admin overrides + surfaces drift

Two real-world incidents on agnes-dev drove this:

1. kbc_job was registered manually with the correct
   (in.c-kbc_telemetry, kbc_job) coordinates. A naive auto-discover
   re-run would have inserted a SECOND kbc_job row at the slugified
   id 'in_c-keboola-storage_kbc_job' (where Keboola's discovery
   places it) — and that row's Storage API export-async 404s.

2. An earlier auto-discover bug stripped the stage prefix from
   bucket ids ('c-finance' instead of 'in.c-finance'), inserting
   137 rows whose syncs all failed.

Fix:

- _discover_and_register_tables now builds a plan first
  (_build_keboola_discovery_plan) classifying each discovered table
  into one of new / existing_match / existing_drift / invalid, then
  executes only the 'new' bucket. Drift rows are reported with both
  sides of the disagreement plus drift_kind:
  - same_id_diff_coords: registry has the same id but different
    bucket / source_table (admin migrated coords inline).
  - name_collision: discovery's slugified id differs from any
    registry id, but the discovered .name matches an existing row's
    .name (case-insensitive). Catches the kbc_job case.

- Bucket detection now prefers the API's authoritative bucket_id
  field (separate field on the Keboola tables.list response,
  normalised by KeboolaClient.discover_all_tables). Falls back to
  id-string parsing only when bucket_id is missing (older fallback
  path inside discover_all_tables).

- Endpoint POST /api/admin/discover-and-register?dry_run=true
  returns the plan without writing — would_register, drift,
  invalid lists. Lets an operator audit before merging discovery
  with a registry that has admin overrides.

Removed 'every 0m' from test_register_request_rejects_malformed_sync_schedule
— the runtime started accepting it in the previous commit (force-resync
override) and the validator follows suit.

* feat(keboola): AGNES_TEMP_DIR routes tempfiles off overlayfs /tmp

The container's /tmp lives on the boot disk's overlayfs (29 GiB on
agnes-dev, shared with /var). Snowflake UNLOAD of a wide table writes
slices into per-call /tmp tempdirs that fill multi-GiB / many-slice
exports long before the dedicated data disk fills. agnes-dev hit
100% boot-disk while the 20 GiB data disk had 15 GiB free.

connectors.keboola.storage_api.get_temp_root() reads AGNES_TEMP_DIR;
mkdirs the target on first use; unset / empty / unwritable falls
back to None (system tempdir, OSS-pre-fix behaviour). Both
materialize_query (parquet path) and _extract_via_legacy (CSV
fallback) and the sliced-CSV concat path in storage_api use the
helper now.

docker-compose.yml defaults AGNES_TEMP_DIR=/data/tmp on app, scheduler,
and extract services. The data volume is the dedicated disk in
production layouts and a plain docker volume in single-disk
dev/laptop setups — same blast radius as the previous /tmp default
on the latter, no regression.
2026-05-07 12:12:14 +02:00

224 lines
8.4 KiB
Python

"""Process-wide singleton guard around POST /api/sync/trigger.
Without it, two near-simultaneous trigger calls each launch their own
extractor subprocess, both write `extract.duckdb`, fight for the file
lock, starve uvicorn, and Docker flips the container to `unhealthy`.
These tests cover the trigger handler's 409 fast-fail (the
operator-visible behavior) and the in-`_run_sync` defense-in-depth
(if something bypasses the handler).
"""
from unittest.mock import patch
import pytest
from app.api import sync as sync_module
@pytest.fixture(autouse=True)
def reset_sync_lock():
"""Make sure each test starts with a free lock — and never leaves one
held even if an assertion fires mid-test."""
if sync_module._sync_lock.locked():
sync_module._sync_lock.release()
yield
if sync_module._sync_lock.locked():
sync_module._sync_lock.release()
def test_run_sync_skips_when_lock_held(capsys):
"""When `_sync_lock` is already held, `_run_sync` no-ops with a log
line instead of starting a second extractor subprocess."""
sync_module._sync_lock.acquire()
# Patch the heavy parts so a successful path would otherwise execute.
# Reaching any of them while the lock is held would be the bug.
with patch("app.api.sync.subprocess.run") as run_mock, \
patch("app.instance_config.get_data_source_type") as src_mock:
sync_module._run_sync(tables=None)
assert not run_mock.called, "extractor subprocess must not run when lock is held"
assert not src_mock.called, "_run_sync must short-circuit before reading config"
captured = capsys.readouterr()
assert "another sync is already in flight" in captured.err
def test_run_sync_releases_lock_on_exception():
"""Even if the body throws, the lock must release so the next sync can
run. Asserts the `finally:` covers all exit paths.
`_run_sync` imports `get_data_source_type` lazily inside the body, so
we patch the source module rather than the re-export in `app.api.sync`.
"""
with patch(
"app.instance_config.get_data_source_type",
side_effect=RuntimeError("boom"),
):
# Should not raise — `_run_sync` catches and logs
sync_module._run_sync(tables=None)
assert not sync_module._sync_lock.locked()
def test_trigger_endpoint_returns_409_when_locked():
"""Handler-level fast-fail: when a sync is already running, the
trigger endpoint returns 409 without scheduling a second background
task."""
from fastapi.testclient import TestClient
from fastapi import FastAPI
# Stand up a minimal app exposing only the sync router. Bypass auth
# by overriding the require_admin dependency.
from app.auth.access import require_admin
app = FastAPI()
app.include_router(sync_module.router)
app.dependency_overrides[require_admin] = lambda: {"id": "test", "email": "t@e"}
client = TestClient(app)
sync_module._sync_lock.acquire()
try:
resp = client.post("/api/sync/trigger")
assert resp.status_code == 409
assert resp.json()["detail"] == "sync_already_in_progress"
finally:
sync_module._sync_lock.release()
def test_trigger_endpoint_succeeds_when_lock_free():
"""When the lock is free, the trigger endpoint schedules the
background task and returns 200. The background task itself doesn't
execute synchronously in TestClient — that's how FastAPI background
tasks work — so we patch `_run_sync` to a no-op and only assert the
handler shape."""
from fastapi.testclient import TestClient
from fastapi import FastAPI
from app.auth.access import require_admin
app = FastAPI()
app.include_router(sync_module.router)
app.dependency_overrides[require_admin] = lambda: {"id": "test", "email": "t@e"}
client = TestClient(app)
with patch("app.api.sync._run_sync") as run_mock:
resp = client.post("/api/sync/trigger")
assert resp.status_code == 200
body = resp.json()
assert body["status"] == "triggered"
# BackgroundTasks runs after response; TestClient awaits them
run_mock.assert_called_once()
# ---- body shape acceptance -------------------------------------------------
def _make_client():
"""Stand up a minimal FastAPI app exposing the sync router with auth
bypassed and `_run_sync` mocked. Returns (client, run_mock)."""
from fastapi.testclient import TestClient
from fastapi import FastAPI
from app.auth.access import require_admin
app = FastAPI()
app.include_router(sync_module.router)
app.dependency_overrides[require_admin] = lambda: {"id": "test", "email": "t@e"}
return TestClient(app)
@pytest.mark.parametrize("body,expected_tables", [
(None, None), # no body
([], []), # empty array
(["kbc_job"], ["kbc_job"]), # bare array
(["a", "b", "c"], ["a", "b", "c"]),
({"tables": None}, None), # explicit null
({"tables": []}, []),
({"tables": ["kbc_job"]}, ["kbc_job"]), # object form
({"tables": ["a", "b"], "extra": "x"}, ["a", "b"]), # extra keys ignored
])
def test_trigger_accepts_both_body_shapes(body, expected_tables):
"""Both ``["x", "y"]`` and ``{"tables": ["x", "y"]}`` (and `null` /
no body) reach `_run_sync` with the same `tables` arg. Lets older
clients (raw array) and newer ones (object matching the response
payload shape) both work."""
client = _make_client()
with patch("app.api.sync._run_sync") as run_mock:
if body is None:
resp = client.post("/api/sync/trigger")
else:
resp = client.post("/api/sync/trigger", json=body)
assert resp.status_code == 200, resp.text
run_mock.assert_called_once_with(expected_tables)
@pytest.mark.parametrize("bad_body", [
"kbc_job", # bare string
42, # number
{"tables": "kbc_job"}, # tables as string, not array
{"tables": [1, 2, 3]}, # tables entries not strings
[1, 2, 3], # array of ints
[{"id": "x"}], # array of objects
])
def test_trigger_rejects_malformed_bodies(bad_body):
"""Anything that isn't a list-of-strings, an object with a
list-of-strings under `tables`, or null/missing returns 422 with a
structured detail — never silently treated as 'sync everything'."""
client = _make_client()
with patch("app.api.sync._run_sync") as run_mock:
resp = client.post("/api/sync/trigger", json=bad_body)
assert resp.status_code == 422, resp.text
assert not run_mock.called
# ---- /api/sync/status (auto-upgrade defer probe) --------------------------
def test_sync_status_unlocked_returns_locked_false():
"""Default state: no sync running → ``{"locked": false}``. No auth
required (host-side cron probes from outside the auth boundary)."""
from fastapi.testclient import TestClient
from fastapi import FastAPI
app = FastAPI()
app.include_router(sync_module.router)
client = TestClient(app)
if sync_module._sync_lock.locked():
sync_module._sync_lock.release()
resp = client.get("/api/sync/status")
assert resp.status_code == 200
assert resp.json() == {"locked": False}
def test_sync_status_locked_returns_locked_true():
"""Held lock → ``{"locked": true}``. agnes-auto-upgrade.sh greps for
`"locked":true` to defer the recreate, so the wire format must be
exactly this shape."""
from fastapi.testclient import TestClient
from fastapi import FastAPI
app = FastAPI()
app.include_router(sync_module.router)
client = TestClient(app)
sync_module._sync_lock.acquire()
try:
resp = client.get("/api/sync/status")
assert resp.status_code == 200
body = resp.json()
assert body == {"locked": True}
finally:
sync_module._sync_lock.release()
def test_sync_status_does_not_require_auth():
"""No `require_admin` / `get_current_user` dependency — the host's
cron has no PAT and shouldn't need one for a status check."""
from fastapi.testclient import TestClient
from fastapi import FastAPI
app = FastAPI()
app.include_router(sync_module.router)
client = TestClient(app)
# No auth headers at all.
resp = client.get("/api/sync/status")
assert resp.status_code == 200