agnes-the-ai-analyst/src/scheduler.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

260 lines
8.7 KiB
Python

"""
Schedule evaluation for automatic data sync.
Parses sync_schedule strings from table configuration and determines
whether a table is due for synchronization based on its last sync time.
Schedule formats:
"every 15m" - every 15 minutes
"every 1h" - every hour
"daily 05:00" - once per day at 05:00 UTC
"daily 07:00,13:00,18:00" - multiple times per day (UTC)
"""
import logging
import re
from datetime import datetime, timezone
from typing import Optional
logger = logging.getLogger(__name__)
# Pattern: "every 15m", "every 2h"
INTERVAL_PATTERN = re.compile(r"^every (\d+)([mh])$")
# Pattern: "daily 05:00", "daily 17:30", "daily 07:00,13:00,18:00"
DAILY_PATTERN = re.compile(r"^daily ([\d:,]+)$")
def parse_interval_minutes(schedule: str) -> Optional[int]:
"""Parse an interval schedule into minutes.
Args:
schedule: Schedule string like "every 15m" or "every 1h"
Returns:
Interval in minutes, or None if not an interval schedule.
"""
match = INTERVAL_PATTERN.match(schedule)
if not match:
return None
value = int(match.group(1))
unit = match.group(2)
if unit == "h":
return value * 60
return value
def is_table_due(
schedule: str,
last_sync_iso: Optional[str],
now: Optional[datetime] = None,
) -> bool:
"""Determine whether a table is due for sync based on its schedule.
Args:
schedule: Schedule string from table config (e.g., "every 1h", "daily 05:00")
last_sync_iso: ISO timestamp of last sync, or None if never synced
now: Current time (UTC). Defaults to datetime.now(timezone.utc).
Returns:
True if the table should be synced now.
"""
if now is None:
now = datetime.now(timezone.utc)
# Never synced -> always due
if not last_sync_iso:
logger.info("Table never synced, marking as due")
return True
# Parse last_sync timestamp
last_sync = _parse_timestamp(last_sync_iso)
if last_sync is None:
logger.warning(f"Cannot parse last_sync timestamp: {last_sync_iso}, marking as due")
return True
# Ensure timezone-aware comparison
if last_sync.tzinfo is None:
last_sync = last_sync.replace(tzinfo=timezone.utc)
# Check interval schedule: "every Xm" / "every Xh"
interval_minutes = parse_interval_minutes(schedule)
if interval_minutes is not None:
elapsed_minutes = (now - last_sync).total_seconds() / 60
due = elapsed_minutes >= interval_minutes
if due:
logger.debug(
f"Interval schedule: {elapsed_minutes:.0f}m elapsed >= {interval_minutes}m interval"
)
return due
# Check daily schedule: "daily HH:MM" or "daily HH:MM,HH:MM,..."
match = DAILY_PATTERN.match(schedule)
if match:
times_str = match.group(1)
target_times = _parse_daily_times(times_str)
if not target_times:
logger.warning(f"Invalid daily schedule times: {schedule}")
return False
return _is_daily_due(last_sync, now, target_times)
logger.warning(f"Unknown schedule format: {schedule}")
return False
def _parse_daily_times(times_str: str) -> list[tuple[int, int]]:
"""Parse comma-separated HH:MM times into list of (hour, minute) tuples."""
time_pattern = re.compile(r"^(\d{2}):(\d{2})$")
result = []
for part in times_str.split(","):
m = time_pattern.match(part.strip())
if not m:
return []
hour, minute = int(m.group(1)), int(m.group(2))
if hour > 23 or minute > 59:
return []
result.append((hour, minute))
return result
def _is_daily_due(
last_sync: datetime,
now: datetime,
target_times: list[tuple[int, int]],
) -> bool:
"""Check if a daily schedule is due.
Supports multiple target times per day. A target time is due when:
1. Current time is at or past HH:MM today, AND
2. Last sync was before HH:MM today
Returns True if ANY of the target times is due.
"""
for target_hour, target_minute in target_times:
today_target = now.replace(
hour=target_hour, minute=target_minute, second=0, microsecond=0
)
if now >= today_target and last_sync < today_target:
logger.debug(
f"Daily schedule: target {target_hour:02d}:{target_minute:02d} UTC, "
f"last sync {last_sync.isoformat()}, now {now.isoformat()} -> due"
)
return True
return False
def _parse_timestamp(iso_string: str) -> Optional[datetime]:
"""Parse an ISO timestamp string, handling various formats.
Args:
iso_string: ISO 8601 timestamp string
Returns:
datetime object or None if parsing fails
"""
try:
# Python 3.11+ fromisoformat handles most formats
return datetime.fromisoformat(iso_string)
except (ValueError, TypeError):
pass
# Fallback: try common formats
for fmt in ("%Y-%m-%dT%H:%M:%S.%f", "%Y-%m-%dT%H:%M:%S", "%Y-%m-%d %H:%M:%S"):
try:
return datetime.strptime(iso_string, fmt)
except ValueError:
continue
return None
def is_valid_schedule(schedule: Optional[str]) -> bool:
"""Return True iff ``schedule`` parses as a documented schedule string.
Accepted forms (mirroring the rest of this module):
- ``"every Nm"`` / ``"every Nh"`` with N a non-negative integer
(``every 0m`` = always due, useful for force-resync of a row whose
previous attempt errored without recording last_sync — bypasses
the rate limit on the next ``/api/sync/trigger`` tick)
- ``"daily HH:MM"`` (24-h, UTC) optionally comma-separated:
``"daily 07:00,13:00"``
Anything else — including ``None``, empty string, or a parseable-looking
but out-of-range value (``"daily 25:00"``) — returns False. Pydantic
validators on the admin API call this to reject malformed input with
422 instead of accepting it and silently no-op'ing later.
"""
if not schedule or not isinstance(schedule, str):
return False
interval = parse_interval_minutes(schedule)
if interval is not None:
return interval >= 0
match = DAILY_PATTERN.match(schedule)
if not match:
return False
return bool(_parse_daily_times(match.group(1)))
def filter_due_tables(
table_configs: list[dict],
sync_state_repo,
now: Optional[datetime] = None,
) -> list[dict]:
"""Drop table configs whose ``sync_schedule`` says they are not due.
Behaviour:
- ``sync_schedule`` is None / empty / not a valid string → table passes
through (no schedule = "sync on every tick", existing behaviour).
- Valid schedule + last_sync within the cadence → drop.
- Valid schedule + last_sync past cadence (or never) → keep.
- Invalid schedule string → log a warning and let the table through
(do NOT silently skip — operator surprise is worse than a redundant
sync).
``sync_state_repo`` is duck-typed: only ``get_last_sync(table_id)`` is
called, returning a ``datetime`` (tz-aware preferred, naive treated as
UTC) or ``None``.
"""
if now is None:
now = datetime.now(timezone.utc)
out: list[dict] = []
for tc in table_configs:
# sync_state.table_id is populated from _meta.table_name by the
# orchestrator and equals table_registry.name (NOT id). When
# id != name (auto-discovered Keboola rows: id="in_c-crm_company",
# name="company") an id-keyed lookup misses every row and the
# filter degrades to "always sync" — defeating the schedule. The
# same pitfall is documented at app/api/sync.py:244-249.
table_id = tc.get("name") or tc.get("id")
schedule = tc.get("sync_schedule")
if not schedule:
out.append(tc)
continue
if not is_valid_schedule(schedule):
logger.warning(
"Table %s has malformed sync_schedule %r — syncing anyway "
"(fix the schedule string to suppress this message)",
table_id,
schedule,
)
out.append(tc)
continue
last_sync = sync_state_repo.get_last_sync(table_id)
if last_sync is None:
last_sync_iso = None
else:
if last_sync.tzinfo is None:
last_sync = last_sync.replace(tzinfo=timezone.utc)
last_sync_iso = last_sync.isoformat()
if is_table_due(schedule, last_sync_iso, now=now):
out.append(tc)
else:
logger.info(
"Table %s skipped: schedule=%r, last_sync=%s, not due yet",
table_id,
schedule,
last_sync_iso,
)
return out