From f2b172f0efe3cd3723c01e78eb522d3d8cf26fcf Mon Sep 17 00:00:00 2001 From: Petr Date: Sun, 30 Aug 2026 07:25:47 +0200 Subject: [PATCH 1/7] fix(storage): budget create-table/swap-tables job waits, add --timeout (#713) `storage create-table` and `storage swap-tables` -- the two halves of the BigQuery repartition path -- both move real data, yet both inherited the 60s STORAGE_JOB_MAX_WAIT meant for metadata jobs. Measured on an 800 MB BigQuery table: create ~15s, swap ~31s, so a table a few times larger blew the budget and reported STORAGE_JOB_TIMEOUT for a job that was still running and would succeed. Giving up locally never cancels the Storage job, so a short budget buys nothing and costs a false negative on the only in-place repartition path available today (the Storage API's PUT .../definition has no clustering field, so the copy+swap workaround cannot be replaced). - New TABLE_DATA_JOB_MAX_WAIT (300s) default for both commands. - New --timeout SECONDS on both, plus an optional `timeout` on the two REST routes (gt=0, so a bad budget is a 422 at the boundary). - The `None`-sentinel/reject-non-positive guard is now one shared normalize_job_timeout() in services/base.py; workspace_service's private copy is migrated onto it. The client-level default is unchanged (max_wait=None still means 60s), so the raw SDK client and the upload-table auto-create path keep their existing behaviour. --- CLAUDE.md | 11 +- plugins/kbagent/agents/keboola-expert.md | 2 +- .../kbagent/references/commands-reference.md | 4 +- .../skills/kbagent/references/gotchas.md | 23 ++ .../references/storage-types-workflow.md | 22 ++ .../client/storage_tables.py | 14 +- src/keboola_agent_cli/commands/context.py | 13 +- src/keboola_agent_cli/commands/storage.py | 29 +++ src/keboola_agent_cli/constants.py | 9 + .../server/routers/storage.py | 10 +- src/keboola_agent_cli/services/base.py | 33 ++- .../services/storage_service.py | 23 +- .../services/workspace_service.py | 26 +-- tests/test_base_service.py | 26 +++ tests/test_server_router_calls.py | 93 ++++++++ tests/test_storage_swap.py | 207 +++++++++++++++++- tests/test_storage_write.py | 105 +++++++++ 17 files changed, 612 insertions(+), 38 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 928454d07..275b87c97 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -549,7 +549,7 @@ kbagent storage table-detail --project NAME --table-id ID [--branch ID] # -- untyped tables get one too -- so null means the stack omitted the key, NOT "untyped". # `storage tables` (the LIST endpoint) is unaffected: the API has no `definition` include. kbagent storage create-bucket --project NAME --stage STAGE --name NAME [--description D] [--backend B] [--branch ID] -kbagent storage create-table --project NAME --bucket-id ID --name NAME [--column COL:TYPE[(length)] ...] [--primary-key COL] [--not-null COL ...] [--default NAME=VALUE ...] [--source-table-id ID] [--source-branch-id N] [--time-partitioning-type DAY|HOUR|MONTH|YEAR] [--time-partitioning-field COL] [--time-partitioning-expiration-ms MS] [--range-partitioning-field COL --range-partitioning-start S --range-partitioning-end E --range-partitioning-interval I] [--clustering-field COL ...] [--branch ID] [--if-not-exists] +kbagent storage create-table --project NAME --bucket-id ID --name NAME [--column COL:TYPE[(length)] ...] [--primary-key COL] [--not-null COL ...] [--default NAME=VALUE ...] [--source-table-id ID] [--source-branch-id N] [--time-partitioning-type DAY|HOUR|MONTH|YEAR] [--time-partitioning-field COL] [--time-partitioning-expiration-ms MS] [--range-partitioning-field COL --range-partitioning-start S --range-partitioning-end E --range-partitioning-interval I] [--clustering-field COL ...] [--branch ID] [--if-not-exists] [--timeout SECONDS] # --column XOR --source-table-id (0.66.0+, BigQuery only): --source-table-id copies an existing table's data into the requested partition/clustering layout (schema derived from source) -> swap into place with swap-tables. Partition/clustering flags work in both modes (BigQuery only); time vs range partitioning are mutually exclusive. A non-BigQuery project fails fast (pre-flight backend check). kbagent storage upload-table --project NAME --table-id ID --file PATH [--incremental] [--branch ID] kbagent storage download-table --project NAME --table-id ID [--output FILE] [--columns COL ...] [--limit N] [--where-column COL --where-value VAL ... [--where-operator eq|neq]] [--changed-since WHEN] [--changed-until WHEN] [--branch ID] @@ -558,7 +558,14 @@ kbagent storage delete-table --project NAME --table-id ID [--table-id ...] [--fo kbagent storage truncate-table --project NAME --table-id ID [--table-id ...] [--dry-run] [--yes] [--branch ID] kbagent storage delete-column --project NAME --table-id ID --column COL [--column ...] [--force] [--dry-run] [--yes] [--branch ID] kbagent storage delete-bucket --project NAME --bucket-id ID [--bucket-id ...] [--force] [--dry-run] [--yes] [--branch ID] -kbagent storage swap-tables --project NAME --table-id ID --target-table-id ID --branch ID [--dry-run] [--yes] +kbagent storage swap-tables --project NAME --table-id ID --target-table-id ID --branch ID [--dry-run] [--yes] [--timeout SECONDS] + # --timeout SECONDS on create-table and swap-tables (since vNEXT, #713): both halves + # of the repartition path move real data yet used to inherit the 60s metadata-job + # budget, so a large BigQuery table reported STORAGE_JOB_TIMEOUT for a job that was + # still running and would succeed. Default is now 300s on both. A local timeout + # NEVER cancels the job -- it keeps running server-side; poll + # GET /v2/storage/jobs/{id} or verify with `storage table-detail` rather than + # assuming the operation failed. kbagent storage clone-table --project NAME --table-id ID --branch ID [--dry-run] kbagent storage snapshot-create --project NAME --table-id ID [--description D] [--branch ID] kbagent storage snapshots --project NAME --table-id ID [--limit N] [--branch ID] diff --git a/plugins/kbagent/agents/keboola-expert.md b/plugins/kbagent/agents/keboola-expert.md index fd9380bcd..7430065a7 100644 --- a/plugins/kbagent/agents/keboola-expert.md +++ b/plugins/kbagent/agents/keboola-expert.md @@ -120,7 +120,7 @@ been retired, so its absence is NOT a promise (see ยง1 Rule 6). | Create typed table with native types | `kbagent storage create-table --column pk:VARCHAR(40) --column amount:NUMBER(18,2) --not-null pk --default amount=0` | -- | re-creating via raw REST to `tables-definition` | | Add one column to an existing table | `kbagent storage add-column --project P --table-id in.c-foo.data --column status:VARCHAR(20) [--not-null] [--default active]` -- synchronous, same `name:TYPE(length)` grammar as `create-table` | -- | re-creating the whole table to add a field (loses data / PK / dependents) | | Promote a rebuilt table into the original name | `kbagent storage swap-tables --project P --table-id in.c-foo.data --target-table-id in.c-foo.data_new --branch --yes` -- async `tableSwap`; the service refuses without a branch (any branch, prod included) | -- | renaming or delete + re-upload (loses history; every downstream config needs rewriting) | -| Repartition / recluster a populated BigQuery table | `kbagent storage create-table --source-table-id --time-partitioning-type DAY --time-partitioning-field created_at --clustering-field tenant_id` (0.66.0+, BigQuery only) to copy the data into the new layout, then `swap-tables` to flip it in. `--source-table-id` derives the schema, so `--column` is forbidden. Range partitioning: all four `--range-partitioning-*` flags together. VERIFY the result with `storage table-detail --json` -> `.definition.timePartitioning` / `.clustering` (0.88.0+) -- `create-table` only echoes the layout you REQUESTED, so it proves nothing. See [storage-types-workflow.md](../skills/kbagent/references/storage-types-workflow.md) | -- | `CREATE TABLE ... AS SELECT` in a workspace (drops NOT NULL + primary key) | +| Repartition / recluster a populated BigQuery table | `kbagent storage create-table --source-table-id --time-partitioning-type DAY --time-partitioning-field created_at --clustering-field tenant_id` (0.66.0+, BigQuery only) to copy the data into the new layout, then `swap-tables` to flip it in. `--source-table-id` derives the schema, so `--column` is forbidden. Range partitioning: all four `--range-partitioning-*` flags together. VERIFY the result with `storage table-detail --json` -> `.definition.timePartitioning` / `.clustering` (0.88.0+) -- `create-table` only echoes the layout you REQUESTED, so it proves nothing. Both steps take `--timeout SECONDS` (vNEXT+, default 300); a `STORAGE_JOB_TIMEOUT` means kbagent stopped WATCHING, not that the job failed -- verify before retrying, because re-running a swap that succeeded swaps the tables back. See [storage-types-workflow.md](../skills/kbagent/references/storage-types-workflow.md) | -- | `CREATE TABLE ... AS SELECT` in a workspace (drops NOT NULL + primary key) | | Re-seed a table without losing schema / PK / dependents | `kbagent storage truncate-table --project P --table-id in.c-foo.data [--dry-run] [--yes]` -- rows only, uniformly async-via-job on every branch; batch via repeated `--table-id` | -- | drop + recreate (loses descriptions, PK, sharing edges, and breaks every downstream reference); deleting rows via raw SQL in a workspace (bypasses the Storage audit trail) | | Back up / restore a table around a risky change | `kbagent storage snapshot-create --table-id ...` then, to restore, `kbagent storage table-from-snapshot --snapshot-id ID --bucket-id B --name NEW` -- restore is always a NEW table (`--name` REQUIRED, no overwrite): verify it, then `swap-tables`. See [snapshot-workflow.md](../skills/kbagent/references/snapshot-workflow.md) | `storage snapshots` / `snapshot-detail` to find one | exporting to CSV as a "backup" (loses column types + PK); `create-table --snapshot-id` (not a thing) | | Debug a failed job | `kbagent job detail --project P --job-id J --json` + `kbagent job run ... --log-tail-lines 200` | `kbagent workspace from-transformation` for SQL repro | "I think the issue is..." without reading logs | diff --git a/plugins/kbagent/skills/kbagent/references/commands-reference.md b/plugins/kbagent/skills/kbagent/references/commands-reference.md index 71784fe2a..09ef829cb 100644 --- a/plugins/kbagent/skills/kbagent/references/commands-reference.md +++ b/plugins/kbagent/skills/kbagent/references/commands-reference.md @@ -163,7 +163,7 @@ Requires a **super-admin** Manage API token (same kind as `org setup`). Same def - `storage tables [--project NAME ...] [--bucket-id ID] [--branch ID] [--include-usage]` -- list tables across all connected projects in parallel (multi-project by default, same as `storage buckets`); repeat `--project` to target a subset; `--bucket-id` is applied independently per project (missing buckets become per-project errors); `--branch` requires exactly one `--project`. `--include-usage` (0.88.0+) adds `used_by` per table: the configurations naming it in their **storage input/output mapping** only -- a table id inside a transformation's SQL is NOT a reference. Costs one extra component listing per project (not per table), which is the slow call in a big project; unreadable components degrade to an empty `used_by` - `storage table-detail --project NAME --table-id ID [--branch ID]` -- table detail with columns, types, primary key, row count (branch-aware). Since 0.88.0 (#621) also returns the raw Storage API `definition`: on BigQuery that carries `timePartitioning` / `rangePartitioning` / `clustering` / `requirePartitionFilter` / `partitions[]`, and it is the only way to verify a repartition landed. Human mode prints the layout and a partition COUNT; `--json` passes `definition` through verbatim. Present on every response (untyped tables too), so `null` means the stack omitted the key, not "untyped". Since 0.88.0 (#624) it also RESOLVES column descriptions written by anyone -- the UI, a component, or kbagent -- into `column_details[].description`, with precedence native definition -> `columnMetadata` `KBC.description` -> legacy flat `KBC.column.*` (an alias table falls back to the source table's `columnMetadata`, matching the MCP server). The response always carries `legacy_column_descriptions`, naming the columns still backed by the pre-0.88.0 convention -- human mode warns and points at `storage describe-migrate`. *(since v0.89.0)* Human mode also renders a `Description` column in the Columns table (only when at least one column has one; long text wraps rather than truncating); on 0.88.0 descriptions were visible in `--json` only. Reading never writes, so it is safe under a read-only token or `--deny-writes` - `storage create-bucket --project NAME --stage STAGE --name NAME [--description D] [--backend B] [--branch ID]` -- create bucket (branch-aware). With `--branch ID` on a project lacking the `storage-branches` feature (legacy fake-branch), response carries `legacy_branch_storage: true` and human mode prints a warning -- the runner will create a parallel `out.c--*` bucket at job time. See `storage-types-workflow.md` -- `storage create-table --project NAME --bucket-id ID --name NAME [--column col:TYPE[(length)] ...] [--primary-key COL] [--not-null COL ...] [--default NAME=VALUE ...] [--source-table-id ID] [--source-branch-id N] [--time-partitioning-type DAY|HOUR|MONTH|YEAR] [--time-partitioning-field COL] [--time-partitioning-expiration-ms MS] [--range-partitioning-field COL --range-partitioning-start S --range-partitioning-end E --range-partitioning-interval I] [--clustering-field COL ...] [--branch ID] [--if-not-exists]` -- create typed table. Base types `STRING/INTEGER/NUMERIC/FLOAT/BOOLEAN/DATE/TIMESTAMP` plus native backend types with length (`VARCHAR(40)`, `NUMBER(18,2)`, `TIMESTAMP_TZ`, `VARIANT`, etc.) -- type/length validation delegated to the Storage API. `--not-null` marks a column `nullable=false`; `--default NAME=VALUE` sets a DEFAULT expression (booleans must be lowercase `true`/`false`). In a dev branch, the target bucket is auto-materialized if it has not yet been written to there -- response surfaces this via `auto_created_bucket: bool`. On legacy fake-branch projects (no `storage-branches` feature), `legacy_branch_storage: true` flags that the runner will use a separate `out.c--*` bucket at job time. `--if-not-exists` (0.47.0+) turns a duplicate-display-name failure into `action: skipped` when the table really exists at the expected id (safe for parallel workers). Since 0.47.1 the skipped envelope reports the EXISTING table's actual `columns`/`primary_key`/`name`, mirrors the request under `requested_columns`/`requested_primary_key`, and sets `schema_drift: true` when they diverge. **`--source-table-id` (0.66.0+, BigQuery only)** copies an existing table's data into the requested partition/clustering layout instead of building from `--column` (schema derived from source -> `--column`/`--not-null`/`--default` forbidden; the two are mutually exclusive). This is the supported way to repartition a populated BigQuery table -- then promote it with `storage swap-tables`. Partition/clustering flags (`--time-partitioning-*`, `--range-partitioning-*`, `--clustering-field`) also work on a plain `--column` create (BigQuery only); time vs range partitioning are mutually exclusive and range bounds are strings. When any source/partition/clustering flag is used, a one-call backend pre-flight rejects non-BigQuery projects (exit 2) before the create. See `storage-types-workflow.md` +- `storage create-table --project NAME --bucket-id ID --name NAME [--column col:TYPE[(length)] ...] [--primary-key COL] [--not-null COL ...] [--default NAME=VALUE ...] [--source-table-id ID] [--source-branch-id N] [--time-partitioning-type DAY|HOUR|MONTH|YEAR] [--time-partitioning-field COL] [--time-partitioning-expiration-ms MS] [--range-partitioning-field COL --range-partitioning-start S --range-partitioning-end E --range-partitioning-interval I] [--clustering-field COL ...] [--branch ID] [--if-not-exists] [--timeout SECONDS]` -- create typed table. Base types `STRING/INTEGER/NUMERIC/FLOAT/BOOLEAN/DATE/TIMESTAMP` plus native backend types with length (`VARCHAR(40)`, `NUMBER(18,2)`, `TIMESTAMP_TZ`, `VARIANT`, etc.) -- type/length validation delegated to the Storage API. `--not-null` marks a column `nullable=false`; `--default NAME=VALUE` sets a DEFAULT expression (booleans must be lowercase `true`/`false`). In a dev branch, the target bucket is auto-materialized if it has not yet been written to there -- response surfaces this via `auto_created_bucket: bool`. On legacy fake-branch projects (no `storage-branches` feature), `legacy_branch_storage: true` flags that the runner will use a separate `out.c--*` bucket at job time. `--if-not-exists` (0.47.0+) turns a duplicate-display-name failure into `action: skipped` when the table really exists at the expected id (safe for parallel workers). Since 0.47.1 the skipped envelope reports the EXISTING table's actual `columns`/`primary_key`/`name`, mirrors the request under `requested_columns`/`requested_primary_key`, and sets `schema_drift: true` when they diverge. **`--source-table-id` (0.66.0+, BigQuery only)** copies an existing table's data into the requested partition/clustering layout instead of building from `--column` (schema derived from source -> `--column`/`--not-null`/`--default` forbidden; the two are mutually exclusive). This is the supported way to repartition a populated BigQuery table -- then promote it with `storage swap-tables`. Partition/clustering flags (`--time-partitioning-*`, `--range-partitioning-*`, `--clustering-field`) also work on a plain `--column` create (BigQuery only); time vs range partitioning are mutually exclusive and range bounds are strings. When any source/partition/clustering flag is used, a one-call backend pre-flight rejects non-BigQuery projects (exit 2) before the create. See `storage-types-workflow.md`. **`--timeout SECONDS` (vNEXT+, default 300)** budgets the async create job; a `--source-table-id` copy moves real data and used to inherit the 60s metadata-job default, so a large table reported `STORAGE_JOB_TIMEOUT` for a copy that was still running and would succeed. A local timeout NEVER cancels the job -- poll `GET /v2/storage/jobs/{id}` or verify with `table-detail`; exit code is 4 (retryable) - `storage upload-table --project NAME --table-id ID --file PATH [--incremental] [--branch ID]` -- upload CSV (branch-aware) - `storage download-table --project NAME --table-id ID [--output FILE] [--columns COL ...] [--limit N] [--where-column COL --where-value VAL ... [--where-operator eq|neq]] [--changed-since WHEN] [--changed-until WHEN] [--branch ID]` -- export table to CSV (branch-aware). `--where-column` + `--where-value` (repeatable, OR within the set) + `--where-operator eq|neq` filter rows server-side; `--changed-since`/`--changed-until` (unix ts or strtotime like `-2 days`) filter by import time -- the credential-only, no-workspace way to pull a filtered/incremental slice (0.62.0+) - `storage add-column --project NAME --table-id ID --column COL:TYPE[(length)] [--not-null] [--default VALUE] [--branch ID]` -- add a single column to an existing table (0.62.0+). Same `name:TYPE(length)` grammar as `create-table --column`; a bare `name` adds an untyped STRING column. Synchronous endpoint (no job to wait on). `--not-null` needs an empty table or a `--default`. Mirror of `delete-column` @@ -171,7 +171,7 @@ Requires a **super-admin** Manage API token (same kind as `org setup`). Same def - `storage truncate-table --project NAME --table-id ID [--table-id ...] [--dry-run] [--yes] [--branch ID]` -- delete all rows while preserving table schema, primary key, descriptions, sharing edges, and downstream dependents. Batch via repeated `--table-id`. Endpoint is uniformly async-via-job on every branch (returns a queued `tableRowsDelete` job; client polls via `_wait_for_storage_job` before returning). Idempotent (truncating an empty table is a no-op). Use when re-seeding a table without losing the schema contract - `storage delete-column --project NAME --table-id ID --column COL [--column ...] [--force] [--dry-run] [--yes] [--branch ID]` -- delete columns from a table (branch-aware) - `storage delete-bucket --project NAME --bucket-id ID [--bucket-id ...] [--force] [--dry-run] [--yes] [--branch ID]` -- delete buckets (branch-aware) -- `storage swap-tables --project NAME --table-id ID --target-table-id ID --branch ID [--dry-run] [--yes]` -- swap two storage tables in any branch, including the default/production branch (POST `/tables/{id}/swap`). Both tables exchange physical positions; aliases are NOT transferred (they keep pointing at the same physical position and therefore expose the OTHER table's data after the swap). Service refuses without a branch (active branch via `branch use` works too). Use to flip a typed rebuild ("data_change_log") into the original name ("data") without touching downstream config references +- `storage swap-tables --project NAME --table-id ID --target-table-id ID --branch ID [--dry-run] [--yes] [--timeout SECONDS]` -- swap two storage tables in any branch, including the default/production branch (POST `/tables/{id}/swap`). Both tables exchange physical positions; aliases are NOT transferred (they keep pointing at the same physical position and therefore expose the OTHER table's data after the swap). Service refuses without a branch (active branch via `branch use` works too). Use to flip a typed rebuild ("data_change_log") into the original name ("data") without touching downstream config references. **`--timeout SECONDS` (vNEXT+, default 300)** budgets the async swap job -- a swap on a large BigQuery table outlasts the old 60s default. A local timeout NEVER cancels the swap; it usually lands anyway, so verify with `table-detail` rather than re-issuing it - `storage clone-table --project NAME --table-id ID --branch ID [--dry-run]` -- pull (clone) a production table into a dev branch (POST `/tables/{id}/pull`, operationName `devBranchTablePull`). On `storage-branches` projects a dev branch reads prod tables transparently until the first write, so an in-branch schema mutation (`swap-tables`, dropping a column) fails with a misleading "bucket not found" until the table is materialized branch-local; `clone-table` does that. One-way (default -> branch). Service refuses without a branch (active branch via `branch use` works too). Permission class `write` - `storage snapshot-create --project NAME --table-id ID [--description D] [--branch ID]` -- create a point-in-time snapshot of a table (data + columns + primary key; async `tableSnapshotCreate` job, client polls to completion). The receipt carries `snapshot_id` -- restores are addressed by it. Permission class `write` - `storage snapshots --project NAME --table-id ID [--limit N] [--branch ID]` -- list a table's snapshots (`id`, `createdTime`, `description`, `creatorToken`). Read command: production endpoint by default diff --git a/plugins/kbagent/skills/kbagent/references/gotchas.md b/plugins/kbagent/skills/kbagent/references/gotchas.md index 48891fa68..0fa2a5afd 100644 --- a/plugins/kbagent/skills/kbagent/references/gotchas.md +++ b/plugins/kbagent/skills/kbagent/references/gotchas.md @@ -1644,6 +1644,29 @@ config, the retry fires, and the retry destroys it for good. branch and run the real build + swap in the production (default) branch. Full procedure: `typify-table-workflow.md`. +## A `STORAGE_JOB_TIMEOUT` is not a failure -- the job is still running + +*(since vNEXT, #713)* + +- `storage create-table` and `storage swap-tables` -- the two halves of the + BigQuery repartition path -- both move real data, yet both used to inherit + the **60s** `STORAGE_JOB_MAX_WAIT` budget meant for metadata jobs. Measured + on an 800 MB BigQuery table: create ~15s, swap ~31s. A table a few times + larger blew the budget, so a perfectly healthy operation surfaced as + `STORAGE_JOB_TIMEOUT`. +- **Since vNEXT both default to 300s and accept `--timeout SECONDS`.** The + REST routes (`POST /storage/tables/{p}`, `POST /storage/tables/{p}/{tid}/swap`) + take an optional `timeout` field; a non-positive value is a 422 at the + boundary and `INVALID_ARGUMENT` from the service. +- **Never read the timeout as "nothing happened".** kbagent only stops + *watching*; the Storage job keeps running server-side (and keeps consuming + backend resources). It usually lands. That is why the code is + `retryable: true` and maps to **exit 4**, not exit 1. +- **Do not re-issue the command on a timeout.** Poll + `GET /v2/storage/jobs/{id}` (the id is named in the error message) or check + the physical result with `storage table-detail --json` -> `.definition`. + Re-running a swap that actually succeeded swaps the tables *back*. + ## `storage clone-table` materializes a prod table into a dev branch - `kbagent storage clone-table --project P --table-id T --branch ` diff --git a/plugins/kbagent/skills/kbagent/references/storage-types-workflow.md b/plugins/kbagent/skills/kbagent/references/storage-types-workflow.md index 420ad6a51..e586dc0e3 100644 --- a/plugins/kbagent/skills/kbagent/references/storage-types-workflow.md +++ b/plugins/kbagent/skills/kbagent/references/storage-types-workflow.md @@ -333,6 +333,28 @@ kbagent --json storage table-detail --project prod --table-id in.c-main.events \ | jq '.data.definition | {timePartitioning, clustering, requirePartitionFilter}' ``` +### Budget both steps on a large table + +*(since vNEXT, #713)* + +Steps 1 and 2 both move real data. Until vNEXT they inherited the 60s +`STORAGE_JOB_MAX_WAIT` meant for metadata jobs -- measured on an 800 MB table, +create takes ~15s and the swap ~31s, so a table a few times larger reported +`STORAGE_JOB_TIMEOUT` on a job that was still running and would succeed. Both +now default to **300s** and accept `--timeout SECONDS`: + +```bash +kbagent storage create-table --project prod --bucket-id in.c-main \ + --name events_repart --source-table-id in.c-main.events \ + --time-partitioning-type DAY --time-partitioning-field created_at \ + --timeout 3600 +``` + +**A timeout is not a failure.** kbagent stops watching; the Storage job keeps +running server-side. The error names the job id -- poll +`GET /v2/storage/jobs/{id}`, or just run step 3. Never blindly re-issue the +swap: repeating one that already landed swaps the tables straight back. + On 0.87.0 and earlier `table-detail` dropped `definition`, so step 3 was not possible from kbagent at all. `create-table`'s own output is not a substitute: it echoes the layout you *requested*, and its `--if-not-exists` skip path diff --git a/src/keboola_agent_cli/client/storage_tables.py b/src/keboola_agent_cli/client/storage_tables.py index ca2c05875..16f54d108 100644 --- a/src/keboola_agent_cli/client/storage_tables.py +++ b/src/keboola_agent_cli/client/storage_tables.py @@ -512,6 +512,7 @@ def create_table( time_partitioning: dict[str, Any] | None = None, range_partitioning: dict[str, Any] | None = None, clustering: dict[str, Any] | None = None, + max_wait: float | None = None, ) -> dict[str, Any]: """Create a new table with typed columns (async, waits for completion). @@ -537,6 +538,10 @@ def create_table( range_partitioning: Optional ``{"field": str, "range": {"start": str, "end": str, "interval": str}}`` (BigQuery). clustering: Optional ``{"fields": list[str]}`` (BigQuery). + max_wait: Seconds to wait for the create job. ``None`` defers to + ``_wait_for_storage_job``'s own default. A ``source`` copy + moves real data and can far outlast the metadata-job default, + so callers doing one pass their own budget. Returns: Completed storage job results dict. @@ -560,7 +565,7 @@ def create_table( if clustering is not None: body["clustering"] = clustering response = self._request("POST", f"{prefix}/buckets/{safe_id}/tables-definition", json=body) - job = self._wait_for_storage_job(response.json()) + job = self._wait_for_storage_job(response.json(), max_wait=max_wait) return job.get("results", {}) def prepare_file_upload( @@ -882,6 +887,7 @@ def swap_tables( table_id: str, target_table_id: str, branch_id: int, + max_wait: float | None = None, ) -> dict[str, Any]: """Swap two storage tables (async, waits for completion; branch-scoped). @@ -902,6 +908,10 @@ def swap_tables( table_id: Full ID of the first table (e.g. "in.c-bucket.table"). target_table_id: Full ID of the second table to swap with. branch_id: Development branch ID. Required by the API. + max_wait: Seconds to wait for the swap job. ``None`` defers to + ``_wait_for_storage_job``'s own default. A swap rewrites + table metadata but is measurably slower than the copy that + produced the replacement, so callers pass their own budget. Returns: Completed storage job dict. @@ -910,7 +920,7 @@ def swap_tables( safe_id = quote(table_id, safe="") body = {"targetTableId": target_table_id} response = self._request("POST", f"{prefix}/tables/{safe_id}/swap", json=body) - return self._wait_for_storage_job(response.json()) + return self._wait_for_storage_job(response.json(), max_wait=max_wait) def pull_table(self, table_id: str, branch_id: int) -> dict[str, Any]: """Pull (clone) a table from the default branch into a dev branch. diff --git a/src/keboola_agent_cli/commands/context.py b/src/keboola_agent_cli/commands/context.py index d3b6e99e6..14cadd5f2 100644 --- a/src/keboola_agent_cli/commands/context.py +++ b/src/keboola_agent_cli/commands/context.py @@ -651,8 +651,14 @@ time. Response includes `legacy_branch_storage: true` and human mode prints a warning when this applies. See storage-types-workflow.md. - kbagent storage create-table --project NAME --bucket-id BUCKET_ID --name TABLE_NAME [--column col:TYPE[(length)] ...] [--primary-key COL] [--not-null COL ...] [--default NAME=VALUE ...] [--source-table-id ID] [--source-branch-id N] [--time-partitioning-type DAY|HOUR|MONTH|YEAR] [--time-partitioning-field COL] [--time-partitioning-expiration-ms MS] [--range-partitioning-field COL --range-partitioning-start S --range-partitioning-end E --range-partitioning-interval I] [--clustering-field COL ...] [--branch ID] [--if-not-exists] + kbagent storage create-table --project NAME --bucket-id BUCKET_ID --name TABLE_NAME [--column col:TYPE[(length)] ...] [--primary-key COL] [--not-null COL ...] [--default NAME=VALUE ...] [--source-table-id ID] [--source-branch-id N] [--time-partitioning-type DAY|HOUR|MONTH|YEAR] [--time-partitioning-field COL] [--time-partitioning-expiration-ms MS] [--range-partitioning-field COL --range-partitioning-start S --range-partitioning-end E --range-partitioning-interval I] [--clustering-field COL ...] [--branch ID] [--if-not-exists] [--timeout SECONDS] Create a typed table. --column repeatable. + - --timeout SECONDS (since vNEXT, default 300): budget for the async create job. A + --source-table-id copy moves real data and used to inherit the 60s metadata-job + default, so a large BigQuery table reported STORAGE_JOB_TIMEOUT for a copy that was + still running and would succeed. A local timeout NEVER cancels the job -- it keeps + running server-side; poll GET /v2/storage/jobs/{{id}} or verify with table-detail + instead of assuming failure. Exit code is 4 (retryable), not 1. - --if-not-exists: opt-in idempotency. On a duplicate-display-name failure, probe get-table-detail at the expected id and, if the table really exists, return `action: "skipped", skip_reason: "table already exists"` instead of raising. A different @@ -724,7 +730,7 @@ kbagent storage delete-bucket --project NAME --bucket-id ID [--bucket-id ...] [--force] [--dry-run] [--yes] [--branch ID] Delete one or more buckets. --force cascade-deletes tables. Linked/shared buckets protected. Branch-aware. - kbagent storage swap-tables --project NAME --table-id ID --target-table-id ID --branch ID [--dry-run] [--yes] + kbagent storage swap-tables --project NAME --table-id ID --target-table-id ID --branch ID [--dry-run] [--yes] [--timeout SECONDS] Swap two storage tables in any branch, including the default/production branch (POST /tables/{{id}}/swap). Both tables exchange physical positions; aliases are NOT transferred (they keep pointing at the same physical position and therefore expose the OTHER table's data after the swap). Use to promote a typed rebuild back into the original name without @@ -732,6 +738,9 @@ branch use'); service guards before any HTTP call when none is set. Any branch works, INCLUDING the default/production branch -- a default-branch swap is how a typed rebuild reaches prod (dev-branch merge does not carry storage schema). + --timeout SECONDS (since vNEXT, default 300): budget for the async swap job. A swap on a large + BigQuery table outlasts the old 60s metadata-job default. A local timeout NEVER cancels the swap -- + it keeps running server-side and usually lands; verify with table-detail rather than re-issuing it. kbagent storage clone-table --project NAME --table-id ID --branch ID [--dry-run] Clone (pull) a production table into a dev branch (POST /tables/{{id}}/pull). On storage-branches projects a diff --git a/src/keboola_agent_cli/commands/storage.py b/src/keboola_agent_cli/commands/storage.py index f69faf042..2f4df93a0 100644 --- a/src/keboola_agent_cli/commands/storage.py +++ b/src/keboola_agent_cli/commands/storage.py @@ -10,6 +10,7 @@ from rich.markup import escape from ..config_store import ConfigStore +from ..constants import TABLE_DATA_JOB_MAX_WAIT from ..errors import ConfigError, ErrorCode, KeboolaApiError from ._helpers import ( check_cli_permission, @@ -603,6 +604,15 @@ def storage_create_table( "with the same display name still surfaces the original error." ), ), + timeout: float = typer.Option( + TABLE_DATA_JOB_MAX_WAIT, + "--timeout", + help=( + "Seconds to wait for the create job. A --source-table-id copy moves " + "real data; on timeout the job KEEPS RUNNING server-side -- kbagent " + "stops watching, it does not cancel." + ), + ), ) -> None: """Create a new storage table with typed columns. @@ -661,6 +671,7 @@ def storage_create_table( range_partitioning_end=range_partitioning_end, range_partitioning_interval=range_partitioning_interval, clustering_fields=clustering_field, + timeout=timeout, ) except ValueError as exc: formatter.error(message=str(exc), error_code=ErrorCode.INVALID_ARGUMENT) @@ -1401,6 +1412,14 @@ def storage_swap_tables( "-y", help="Skip confirmation prompt", ), + timeout: float = typer.Option( + TABLE_DATA_JOB_MAX_WAIT, + "--timeout", + help=( + "Seconds to wait for the storage job. On timeout the job KEEPS " + "RUNNING server-side -- kbagent stops watching, it does not cancel." + ), + ), ) -> None: """Swap two storage tables (any branch, including the default/production branch). @@ -1437,10 +1456,19 @@ def storage_swap_tables( target_table_id=target_table_id, branch_id=effective_branch, dry_run=True, + timeout=timeout, ) except ConfigError as exc: formatter.error(message=exc.message, error_code=ErrorCode.CONFIG_ERROR) raise typer.Exit(code=5) from None + except KeboolaApiError as exc: + formatter.error( + message=exc.message, + error_code=exc.error_code, + project=project, + retryable=exc.retryable, + ) + raise typer.Exit(code=map_error_to_exit_code(exc)) from None if formatter.json_mode: formatter.output(result) @@ -1468,6 +1496,7 @@ def storage_swap_tables( target_table_id=target_table_id, branch_id=effective_branch, dry_run=False, + timeout=timeout, ) except ConfigError as exc: formatter.error(message=exc.message, error_code=ErrorCode.CONFIG_ERROR) diff --git a/src/keboola_agent_cli/constants.py b/src/keboola_agent_cli/constants.py index ffc809af4..7badbd447 100644 --- a/src/keboola_agent_cli/constants.py +++ b/src/keboola_agent_cli/constants.py @@ -161,6 +161,15 @@ def _resolve_app_name() -> str: IMPORT_JOB_MAX_WAIT: float = 600.0 # 10 min for table import jobs (large files) MERGE_JOB_MAX_WAIT: float = 600.0 # 10 min for merge-request merge jobs (many-config branches) +# A create-table-from-source (BigQuery repartition copy) and the swap that +# puts it in place both move real data, yet both used to inherit the 60s +# STORAGE_JOB_MAX_WAIT meant for metadata operations. Measured on an 800 MB +# BigQuery table: create 15s, swap 31s -- so a table a few times larger +# blows the budget and reports STORAGE_JOB_TIMEOUT for a job that is still +# running and will succeed. Giving up locally never cancels the job, so a +# too-short budget buys nothing and costs a false negative. +TABLE_DATA_JOB_MAX_WAIT: float = 300.0 # 5 min default for data-moving table jobs + # --- Workspace Table Loading --- # A workspace load is NOT fire-and-forget: when the local poller gives up, the # server-side job keeps running (and keeps consuming warehouse resources). The diff --git a/src/keboola_agent_cli/server/routers/storage.py b/src/keboola_agent_cli/server/routers/storage.py index 832416fc6..20f362b97 100644 --- a/src/keboola_agent_cli/server/routers/storage.py +++ b/src/keboola_agent_cli/server/routers/storage.py @@ -8,7 +8,7 @@ from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, UploadFile from fastapi.responses import FileResponse -from pydantic import BaseModel, model_validator +from pydantic import BaseModel, Field, model_validator from ..dependencies import ServiceRegistry, get_registry @@ -47,6 +47,9 @@ class CreateTable(BaseModel): range_partitioning_end: str | None = None range_partitioning_interval: str | None = None clustering_fields: list[str] | None = None + # `gt=0` mirrors the CLI guard so a bad budget is a 422 at the boundary + # rather than an INVALID_ARGUMENT raised from the service. + timeout: float | None = Field(default=None, gt=0) @model_validator(mode="after") def _columns_xor_source(self) -> CreateTable: @@ -110,6 +113,9 @@ class LoadFileToTable(BaseModel): class SwapTables(BaseModel): target_table_id: str branch_id: int + # `gt=0` mirrors the CLI guard so a bad budget is a 422 at the boundary + # rather than an INVALID_ARGUMENT raised from the service. + timeout: float | None = Field(default=None, gt=0) class CloneTable(BaseModel): @@ -343,6 +349,7 @@ def create_table( range_partitioning_end=body.range_partitioning_end, range_partitioning_interval=body.range_partitioning_interval, clustering_fields=body.clustering_fields, + timeout=body.timeout, ) @@ -443,6 +450,7 @@ def swap_tables( target_table_id=body.target_table_id, branch_id=body.branch_id, dry_run=dry_run, + timeout=body.timeout, ) diff --git a/src/keboola_agent_cli/services/base.py b/src/keboola_agent_cli/services/base.py index b8e4e5af9..343aac00b 100644 --- a/src/keboola_agent_cli/services/base.py +++ b/src/keboola_agent_cli/services/base.py @@ -15,7 +15,7 @@ from ..client import KeboolaClient from ..config_store import ConfigError, ConfigStore, project_not_found_error from ..constants import ENV_MAX_PARALLEL_WORKERS, UNEXPECTED_ERROR_MAX_MESSAGE_LEN -from ..errors import ErrorCode +from ..errors import ErrorCode, KeboolaApiError from ..models import ProjectConfig logger = logging.getLogger(__name__) @@ -79,6 +79,37 @@ def sanitize_unexpected_error(exc: BaseException) -> str: return raw +def normalize_job_timeout(timeout: float | None, default: float) -> float: + """Resolve an async-job wait budget, rejecting a non-positive override. + + ``None`` is the documented "use the default" sentinel and branches + explicitly -- deliberately NOT ``timeout or default``, which would + silently promote a falsy-but-real ``0.0`` to the default instead of + rejecting it. Every CLI command already constrains its own ``--timeout``, + but a direct service or SDK caller goes through neither guard, so the + check belongs here rather than in each command. + + Args: + timeout: Caller-supplied budget in seconds, or ``None`` for ``default``. + default: Budget to use when ``timeout`` is ``None``. + + Returns: + The resolved budget in seconds. + + Raises: + KeboolaApiError: ``INVALID_ARGUMENT`` when ``timeout`` is <= 0. + """ + if timeout is None: + return default + if timeout <= 0: + raise KeboolaApiError( + message=f"Invalid timeout {timeout}. Must be greater than 0.", + status_code=0, + error_code=ErrorCode.INVALID_ARGUMENT, + ) + return timeout + + def project_error_entry( alias: str, exc: BaseException, diff --git a/src/keboola_agent_cli/services/storage_service.py b/src/keboola_agent_cli/services/storage_service.py index 80d4c5bf9..25610e322 100644 --- a/src/keboola_agent_cli/services/storage_service.py +++ b/src/keboola_agent_cli/services/storage_service.py @@ -11,12 +11,13 @@ from pathlib import Path from typing import Any -from ..constants import STORAGE_BRANCHES_FEATURE +from ..constants import STORAGE_BRANCHES_FEATURE, TABLE_DATA_JOB_MAX_WAIT from ..errors import ConfigError, ErrorCode, KeboolaApiError from ..models import ProjectConfig from ._column_descriptions import ColumnDescriptionsMixin from ._storage_tables import normalize_table_rows from ._table_detail import build_table_detail +from .base import normalize_job_timeout from .table_usage import collect_table_usage, fetch_usage_components logger = logging.getLogger(__name__) @@ -792,6 +793,7 @@ def create_table( range_partitioning_end: str | None = None, range_partitioning_interval: str | None = None, clustering_fields: list[str] | None = None, + timeout: float | None = None, ) -> dict[str, Any]: """Create a new table with typed columns, or by copying a source table. @@ -839,6 +841,10 @@ def create_table( partitioning (all four required together). Mutually exclusive with time partitioning. clustering_fields: BigQuery clustering columns. + timeout: Seconds to wait for the create job. ``None`` means + ``TABLE_DATA_JOB_MAX_WAIT``. A ``source_table_id`` copy moves + real data, so a large table needs more than the metadata-job + default; giving up locally never cancels the job. Returns: Dict with table details and ``auto_created_bucket`` flag. @@ -851,6 +857,7 @@ def create_table( ``schema_drift`` is ``True`` when the two diverge. Raises: + KeboolaApiError: ``INVALID_ARGUMENT`` for a non-positive ``timeout``. ValueError: Malformed column spec or ``--default`` assignment; ``--not-null`` / ``--default`` references an unknown column; ``columns`` and ``source`` both/neither given; @@ -858,6 +865,7 @@ def create_table( incomplete or conflicting partitioning flags; or BigQuery-only features requested on a non-BigQuery backend. """ + max_wait = normalize_job_timeout(timeout, TABLE_DATA_JOB_MAX_WAIT) not_null_set = set(not_null_columns or []) defaults_map = _parse_default_assignments(defaults) @@ -965,6 +973,7 @@ def create_table( time_partitioning=time_partitioning, range_partitioning=range_partitioning, clustering=clustering, + max_wait=max_wait, ) except KeboolaApiError as exc: # IF-NOT-EXISTS: if the create failed because the table @@ -1602,6 +1611,7 @@ def swap_tables( target_table_id: str, branch_id: int | None, dry_run: bool = False, + timeout: float | None = None, ) -> dict[str, Any]: """Swap two storage tables (branch-scoped; branch_id mandatory). @@ -1623,6 +1633,11 @@ def swap_tables( target_table_id: Full ID of the second table. branch_id: Branch ID (must not be None; any branch accepted, including the default/production branch). dry_run: If True, only report what would be swapped. + timeout: Seconds to wait for the swap job. ``None`` means + ``TABLE_DATA_JOB_MAX_WAIT``. A swap on a large BigQuery table + outlasts the metadata-job default, and giving up locally does + not cancel it -- the swap still lands, it is just no longer + being watched. Returns: Dict with 'project_alias', 'branch_id', 'table_id', @@ -1630,8 +1645,11 @@ def swap_tables( Raises: ConfigError: If branch_id is None. - KeboolaApiError: If the API call fails. + KeboolaApiError: ``INVALID_ARGUMENT`` for a non-positive + ``timeout``, or if the API call fails. """ + max_wait = normalize_job_timeout(timeout, TABLE_DATA_JOB_MAX_WAIT) + if branch_id is None: raise ConfigError( "swap-tables requires a branch. Set one with " @@ -1664,6 +1682,7 @@ def swap_tables( table_id=table_id, target_table_id=target_table_id, branch_id=branch_id, + max_wait=max_wait, ) finally: client.close() diff --git a/src/keboola_agent_cli/services/workspace_service.py b/src/keboola_agent_cli/services/workspace_service.py index f41cd48fd..f4968bfc9 100644 --- a/src/keboola_agent_cli/services/workspace_service.py +++ b/src/keboola_agent_cli/services/workspace_service.py @@ -38,7 +38,7 @@ coerce_data_size_bytes, plan_auto_load_type, ) -from .base import BaseService +from .base import BaseService, normalize_job_timeout logger = logging.getLogger(__name__) @@ -873,7 +873,7 @@ def load_tables( size guard trips and is neither forced nor approved. """ requested = self._normalize_load_type(load_type) - max_wait = self._normalize_timeout(timeout) + max_wait = normalize_job_timeout(timeout, WORKSPACE_LOAD_JOB_MAX_WAIT) projects = self.resolve_projects([alias]) project = projects[alias] @@ -952,28 +952,6 @@ def _normalize_load_type(load_type: str | None) -> str | None: ) return normalized - @staticmethod - def _normalize_timeout(timeout: float | None) -> float: - """Resolve the load-job wait budget, rejecting a non-positive override. - - ``None`` is the documented "use the default" sentinel and branches - explicitly to ``WORKSPACE_LOAD_JOB_MAX_WAIT`` -- deliberately NOT - ``timeout or WORKSPACE_LOAD_JOB_MAX_WAIT``, which would silently - promote a falsy-but-real ``0.0`` to the 300s default instead of - rejecting it. The CLI and the ``kbagent serve`` router both already - guard a non-positive timeout before calling in, but a direct - service/SDK caller would not go through either guard. - """ - if timeout is None: - return WORKSPACE_LOAD_JOB_MAX_WAIT - if timeout <= 0: - raise KeboolaApiError( - message=f"Invalid timeout {timeout}. Must be greater than 0.", - status_code=0, - error_code=ErrorCode.INVALID_ARGUMENT, - ) - return timeout - def _plan_table_loads( self, client: Any, diff --git a/tests/test_base_service.py b/tests/test_base_service.py index dec4bfe6b..a3a1a05c4 100644 --- a/tests/test_base_service.py +++ b/tests/test_base_service.py @@ -26,6 +26,7 @@ ENV_MAX_PARALLEL_WORKERS, UNEXPECTED_ERROR_CODE, BaseService, + normalize_job_timeout, project_error_entry, ) @@ -544,3 +545,28 @@ def test_custom_client_factory_is_used(self, tmp_config_dir: Path) -> None: service = _TestService(config_store=store, client_factory=mock_factory) assert service._client_factory is mock_factory + + +class TestNormalizeJobTimeout: + """Shared async-job budget guard used by every service that waits on a job.""" + + def test_none_uses_the_default(self) -> None: + assert normalize_job_timeout(None, 300.0) == 300.0 + + def test_positive_override_wins(self) -> None: + assert normalize_job_timeout(900.0, 300.0) == 900.0 + + @pytest.mark.parametrize("bad", [0.0, -1.0, -900.0]) + def test_non_positive_is_rejected(self, bad: float) -> None: + """0.0 must raise, not fall back to the default. + + This is the whole reason the guard branches on ``is None`` rather + than writing ``timeout or default``: 0.0 is falsy, so the shorter + spelling would silently accept an invalid budget as "unset" and every + caller would look correct while validating nothing. + """ + with pytest.raises(KeboolaApiError) as exc_info: + normalize_job_timeout(bad, 300.0) + + assert exc_info.value.error_code == ErrorCode.INVALID_ARGUMENT + assert str(bad) in exc_info.value.message diff --git a/tests/test_server_router_calls.py b/tests/test_server_router_calls.py index 51aa27ff4..09bc64a75 100644 --- a/tests/test_server_router_calls.py +++ b/tests/test_server_router_calls.py @@ -2668,3 +2668,96 @@ def test_workspace_load_invalid_load_type_answers_400(tmp_path: Path) -> None: assert res.status_code == 400, res.text assert res.json()["error"]["code"] == "INVALID_ARGUMENT" + + +# --------------------------------------------------------------------------- +# storage.py POST /tables/{p} and POST /tables/{p}/{tid}/swap -- job budget +# Service: storage.create_table(timeout=...) / storage.swap_tables(timeout=...) +# Issue #713: the data-moving half of the BigQuery repartition path needs a +# budget the caller can raise; a bad one must be a 422 at the boundary. +# --------------------------------------------------------------------------- + + +def test_create_table_route_defaults_timeout_to_none(tmp_path: Path) -> None: + """Omitting `timeout` sends None so the service applies its own default.""" + storage_svc = MagicMock() + storage_svc.create_table.return_value = {"table_id": "in.c-b.t"} + app = _make_app_with_registry(tmp_path, _mock_registry(storage=storage_svc)) + + with TestClient(app) as client: + res = client.post( + f"/storage/tables/{PROJECT}", + headers=AUTH, + json={"bucket_id": "in.c-b", "name": "t", "columns": ["id:INTEGER"]}, + ) + + assert res.status_code == 200, res.text + assert storage_svc.create_table.call_args.kwargs["timeout"] is None + + +def test_create_table_route_forwards_timeout(tmp_path: Path) -> None: + storage_svc = MagicMock() + storage_svc.create_table.return_value = {"table_id": "in.c-b.t"} + app = _make_app_with_registry(tmp_path, _mock_registry(storage=storage_svc)) + + with TestClient(app) as client: + res = client.post( + f"/storage/tables/{PROJECT}", + headers=AUTH, + json={ + "bucket_id": "in.c-b", + "name": "t", + "source_table_id": "in.c-b.src", + "timeout": 1800, + }, + ) + + assert res.status_code == 200, res.text + assert storage_svc.create_table.call_args.kwargs["timeout"] == 1800 + + +def test_create_table_route_rejects_non_positive_timeout(tmp_path: Path) -> None: + """A zero budget would make every create 'time out' instantly.""" + storage_svc = MagicMock() + app = _make_app_with_registry(tmp_path, _mock_registry(storage=storage_svc)) + + with TestClient(app) as client: + res = client.post( + f"/storage/tables/{PROJECT}", + headers=AUTH, + json={"bucket_id": "in.c-b", "name": "t", "columns": ["id:INTEGER"], "timeout": 0}, + ) + + assert res.status_code == 422, res.text + storage_svc.create_table.assert_not_called() + + +def test_swap_tables_route_forwards_timeout(tmp_path: Path) -> None: + storage_svc = MagicMock() + storage_svc.swap_tables.return_value = {"dry_run": False} + app = _make_app_with_registry(tmp_path, _mock_registry(storage=storage_svc)) + + with TestClient(app) as client: + res = client.post( + f"/storage/tables/{PROJECT}/in.c-b.a/swap", + headers=AUTH, + json={"target_table_id": "in.c-b.b", "branch_id": 42, "timeout": 900}, + ) + + assert res.status_code == 200, res.text + assert storage_svc.swap_tables.call_args.kwargs["timeout"] == 900 + + +def test_swap_tables_route_rejects_non_positive_timeout(tmp_path: Path) -> None: + storage_svc = MagicMock() + app = _make_app_with_registry(tmp_path, _mock_registry(storage=storage_svc)) + + with TestClient(app) as client: + res = client.post( + f"/storage/tables/{PROJECT}/in.c-b.a/swap", + headers=AUTH, + json={"target_table_id": "in.c-b.b", "branch_id": 42, "timeout": -1}, + ) + + assert res.status_code == 422, res.text + storage_svc.swap_tables.assert_not_called() diff --git a/tests/test_storage_swap.py b/tests/test_storage_swap.py index f4d1941cc..7536847dc 100644 --- a/tests/test_storage_swap.py +++ b/tests/test_storage_swap.py @@ -10,7 +10,8 @@ from keboola_agent_cli.cli import app from keboola_agent_cli.client import KeboolaClient from keboola_agent_cli.config_store import ConfigStore -from keboola_agent_cli.errors import ConfigError, KeboolaApiError +from keboola_agent_cli.constants import TABLE_DATA_JOB_MAX_WAIT +from keboola_agent_cli.errors import ConfigError, ErrorCode, KeboolaApiError from keboola_agent_cli.models import AppConfig, ProjectConfig from keboola_agent_cli.services.storage_service import StorageService @@ -198,6 +199,7 @@ def test_success(self, tmp_path: Path) -> None: table_id="in.c-foo.data", target_table_id="in.c-foo.data_change_log", branch_id=9999, + max_wait=TABLE_DATA_JOB_MAX_WAIT, ) mock_client.close.assert_called_once() @@ -339,6 +341,7 @@ def test_swap_json(self, tmp_path: Path) -> None: target_table_id="in.c-foo.data_change_log", branch_id=9999, dry_run=False, + timeout=TABLE_DATA_JOB_MAX_WAIT, ) def test_swap_dry_run(self, tmp_path: Path) -> None: @@ -454,3 +457,205 @@ def test_swap_missing_branch_fails_clearly(self, tmp_path: Path) -> None: payload = json.loads(result.output) assert payload["status"] == "error" assert "requires a branch" in payload["error"]["message"] + + +# --------------------------------------------------------------------------- +# Job wait budget (issue #713) +# --------------------------------------------------------------------------- + + +class TestSwapTablesTimeout: + """The swap job budget is caller-controllable and defaults to 5 minutes. + + A swap moves a populated table into place; it used to inherit the 60s + ``STORAGE_JOB_MAX_WAIT`` meant for metadata jobs, so a large BigQuery + table reported ``STORAGE_JOB_TIMEOUT`` for a job that was still running + and would succeed. Giving up locally never cancels the job. + """ + + def test_service_default_is_table_data_budget(self, tmp_path: Path) -> None: + """No --timeout means TABLE_DATA_JOB_MAX_WAIT, not the 60s metadata default.""" + from keboola_agent_cli.constants import STORAGE_JOB_MAX_WAIT + + store = _make_store(tmp_path) + mock_client = MagicMock() + mock_client.swap_tables.return_value = {"status": "ok"} + service = _make_service(store, mock_client) + + service.swap_tables( + alias="test", + table_id="in.c-foo.a", + target_table_id="in.c-foo.b", + branch_id=42, + ) + + max_wait = mock_client.swap_tables.call_args.kwargs["max_wait"] + assert max_wait == TABLE_DATA_JOB_MAX_WAIT + assert max_wait > STORAGE_JOB_MAX_WAIT + + def test_service_forwards_explicit_timeout(self, tmp_path: Path) -> None: + store = _make_store(tmp_path) + mock_client = MagicMock() + mock_client.swap_tables.return_value = {"status": "ok"} + service = _make_service(store, mock_client) + + service.swap_tables( + alias="test", + table_id="in.c-foo.a", + target_table_id="in.c-foo.b", + branch_id=42, + timeout=900.0, + ) + + assert mock_client.swap_tables.call_args.kwargs["max_wait"] == 900.0 + + @pytest.mark.parametrize("bad", [0.0, -1.0]) + def test_service_rejects_non_positive_timeout(self, tmp_path: Path, bad: float) -> None: + """0.0 must be rejected, not silently promoted to the default. + + ``timeout or DEFAULT`` would treat a falsy-but-real 0 as "unset"; + the guard branches on ``is None`` instead. + """ + store = _make_store(tmp_path) + mock_client = MagicMock() + service = _make_service(store, mock_client) + + with pytest.raises(KeboolaApiError) as exc_info: + service.swap_tables( + alias="test", + table_id="in.c-foo.a", + target_table_id="in.c-foo.b", + branch_id=42, + timeout=bad, + ) + + assert exc_info.value.error_code == ErrorCode.INVALID_ARGUMENT + mock_client.swap_tables.assert_not_called() + + def test_service_rejects_non_positive_timeout_on_dry_run(self, tmp_path: Path) -> None: + """The guard runs before the dry-run short-circuit, so --dry-run validates too.""" + store = _make_store(tmp_path) + mock_client = MagicMock() + service = _make_service(store, mock_client) + + with pytest.raises(KeboolaApiError): + service.swap_tables( + alias="test", + table_id="in.c-foo.a", + target_table_id="in.c-foo.b", + branch_id=42, + dry_run=True, + timeout=0.0, + ) + + def test_client_forwards_budget_to_the_poller(self, httpx_mock) -> None: + """swap_tables passes max_wait through to _wait_for_storage_job. + + Asserted on the kwarg, not on elapsed time: an httpx mock never sees + the budget, so dropping the argument would leave every wire test green + while the swap silently fell back to the 60s metadata default. + """ + httpx_mock.add_response( + url="https://connection.keboola.com/v2/storage/branch/42/tables/in.c-foo.a/swap", + method="POST", + json={"id": 777, "status": "waiting", "operationName": "tableSwap"}, + status_code=200, + ) + + client = KeboolaClient(stack_url="https://connection.keboola.com", token=TEST_TOKEN) + with patch.object( + KeboolaClient, + "_wait_for_storage_job", + return_value={"id": 777, "status": "success"}, + ) as poller: + client.swap_tables( + table_id="in.c-foo.a", + target_table_id="in.c-foo.b", + branch_id=42, + max_wait=900.0, + ) + client.close() + + assert poller.call_args.kwargs["max_wait"] == 900.0 + + def test_expired_budget_reports_the_job_as_still_running(self, httpx_mock) -> None: + """An exhausted budget raises STORAGE_JOB_TIMEOUT naming the live job. + + The message must not read as "nothing happened": kbagent stops + watching, the swap keeps running server-side. That is also why the + code is retryable (exit 4), not a plain failure. + """ + httpx_mock.add_response( + url="https://connection.keboola.com/v2/storage/branch/42/tables/in.c-foo.a/swap", + method="POST", + json={"id": 777, "status": "waiting", "operationName": "tableSwap"}, + status_code=200, + ) + httpx_mock.add_response( + url="https://connection.keboola.com/v2/storage/jobs/777", + method="GET", + json={"id": 777, "status": "waiting", "operationName": "tableSwap"}, + status_code=200, + ) + + client = KeboolaClient(stack_url="https://connection.keboola.com", token=TEST_TOKEN) + # The deadline is only checked at the top of each iteration, so the + # budget cannot expire before the first poll however small it is. + # Skipping the real sleep keeps the test instant. + with ( + patch("keboola_agent_cli.client._core.time.sleep"), + pytest.raises(KeboolaApiError) as exc_info, + ): + client.swap_tables( + table_id="in.c-foo.a", + target_table_id="in.c-foo.b", + branch_id=42, + max_wait=0.0001, + ) + client.close() + + exc = exc_info.value + assert exc.error_code == ErrorCode.STORAGE_JOB_TIMEOUT + assert exc.retryable is True + assert "777" in exc.message + assert "continues running" in exc.message + + def test_cli_forwards_timeout(self, tmp_path: Path) -> None: + store = _make_store(tmp_path) + + with ( + patch("keboola_agent_cli.cli.ConfigStore") as MockStore, + patch("keboola_agent_cli.cli.StorageService") as MockSvc, + ): + MockStore.return_value = store + svc = MockSvc.return_value + svc.swap_tables.return_value = { + "project_alias": "test", + "branch_id": 42, + "table_id": "in.c-foo.a", + "target_table_id": "in.c-foo.b", + "dry_run": False, + "response": {"status": "success"}, + } + result = runner.invoke( + app, + [ + "--json", + "storage", + "swap-tables", + "--project", + "test", + "--table-id", + "in.c-foo.a", + "--target-table-id", + "in.c-foo.b", + "--branch", + "42", + "--timeout", + "900", + "--yes", + ], + ) + + assert result.exit_code == 0 + assert svc.swap_tables.call_args.kwargs["timeout"] == 900.0 diff --git a/tests/test_storage_write.py b/tests/test_storage_write.py index d30004738..abc4fb341 100644 --- a/tests/test_storage_write.py +++ b/tests/test_storage_write.py @@ -9,6 +9,7 @@ from keboola_agent_cli.cli import app from keboola_agent_cli.config_store import ConfigStore +from keboola_agent_cli.constants import TABLE_DATA_JOB_MAX_WAIT from keboola_agent_cli.errors import ErrorCode, KeboolaApiError from keboola_agent_cli.models import AppConfig, ProjectConfig from keboola_agent_cli.services.storage_service import StorageService @@ -220,6 +221,7 @@ def test_success(self, tmp_path: Path) -> None: time_partitioning=None, range_partitioning=None, clustering=None, + max_wait=TABLE_DATA_JOB_MAX_WAIT, ) mock_client.close.assert_called_once() @@ -266,6 +268,7 @@ def test_no_primary_key(self, tmp_path: Path) -> None: time_partitioning=None, range_partitioning=None, clustering=None, + max_wait=TABLE_DATA_JOB_MAX_WAIT, ) def test_unknown_type_is_passed_through_to_api(self, tmp_path: Path) -> None: @@ -1426,6 +1429,7 @@ def test_create_table_json(self, tmp_path: Path) -> None: range_partitioning_end=None, range_partitioning_interval=None, clustering_fields=None, + timeout=TABLE_DATA_JOB_MAX_WAIT, ) def test_create_table_native_types_and_attributes(self, tmp_path: Path) -> None: @@ -1883,6 +1887,7 @@ def test_service_passes_branch_id(self, tmp_path: Path) -> None: time_partitioning=None, range_partitioning=None, clustering=None, + max_wait=TABLE_DATA_JOB_MAX_WAIT, ) # In a branch we check bucket existence first (auto-materialize). mock_client.get_bucket_detail.assert_called_once_with("in.c-b", branch_id=77) @@ -2443,3 +2448,103 @@ def test_download_table_with_branch(self, tmp_path: Path) -> None: assert result.exit_code == 0 call_kwargs = svc.download_table.call_args.kwargs assert call_kwargs["branch_id"] == 33 + + +class TestCreateTableJobTimeout: + """The create job budget is caller-controllable (issue #713). + + A ``--source-table-id`` create copies a populated table into a new + BigQuery layout -- the first half of the only available repartition + path. It used to inherit the 60s ``STORAGE_JOB_MAX_WAIT`` meant for + metadata jobs, so a large table reported ``STORAGE_JOB_TIMEOUT`` for a + copy that was still running and would succeed. Giving up locally never + cancels the job, so a short budget buys nothing. + """ + + def test_default_is_table_data_budget(self, tmp_path: Path) -> None: + from keboola_agent_cli.constants import STORAGE_JOB_MAX_WAIT + + store = _make_store(tmp_path) + mock_client = MagicMock() + mock_client.create_table.return_value = {"id": "in.c-b.t"} + service = _make_service(store, mock_client) + + service.create_table(alias="test", bucket_id="in.c-b", name="t", columns=["id:INTEGER"]) + + max_wait = mock_client.create_table.call_args.kwargs["max_wait"] + assert max_wait == TABLE_DATA_JOB_MAX_WAIT + assert max_wait > STORAGE_JOB_MAX_WAIT + + def test_forwards_explicit_timeout(self, tmp_path: Path) -> None: + store = _make_store(tmp_path) + mock_client = MagicMock() + mock_client.create_table.return_value = {"id": "in.c-b.t"} + service = _make_service(store, mock_client) + + service.create_table( + alias="test", + bucket_id="in.c-b", + name="t", + columns=["id:INTEGER"], + timeout=1800.0, + ) + + assert mock_client.create_table.call_args.kwargs["max_wait"] == 1800.0 + + @pytest.mark.parametrize("bad", [0.0, -5.0]) + def test_rejects_non_positive_timeout(self, tmp_path: Path, bad: float) -> None: + """0.0 is rejected rather than silently promoted to the default.""" + store = _make_store(tmp_path) + mock_client = MagicMock() + service = _make_service(store, mock_client) + + with pytest.raises(KeboolaApiError) as exc_info: + service.create_table( + alias="test", + bucket_id="in.c-b", + name="t", + columns=["id:INTEGER"], + timeout=bad, + ) + + assert exc_info.value.error_code == ErrorCode.INVALID_ARGUMENT + mock_client.create_table.assert_not_called() + + def test_cli_forwards_timeout(self, tmp_path: Path) -> None: + store = _make_store(tmp_path) + + with ( + patch("keboola_agent_cli.cli.ConfigStore") as MockStore, + patch("keboola_agent_cli.cli.StorageService") as MockSvc, + ): + MockStore.return_value = store + svc = MockSvc.return_value + svc.create_table.return_value = { + "table_id": "in.c-b.t", + "action": "created", + "name": "t", + "columns": ["id"], + "primary_key": [], + "auto_created_bucket": False, + } + result = runner.invoke( + app, + [ + "--json", + "storage", + "create-table", + "--project", + "test", + "--bucket-id", + "in.c-b", + "--name", + "t", + "--column", + "id:INTEGER", + "--timeout", + "1800", + ], + ) + + assert result.exit_code == 0 + assert svc.create_table.call_args.kwargs["timeout"] == 1800.0 From d713c7eee430489515425dede8b2ceb3b22635a8 Mon Sep 17 00:00:00 2001 From: Petr Date: Sun, 30 Aug 2026 07:35:34 +0200 Subject: [PATCH 2/7] test(storage): make the swap timeout test clock-independent (#713) The new test armed a 0.0001s budget and patched out time.sleep, so it depended on real elapsed time between two loop iterations. time.monotonic() advances in ~15.6ms ticks on Windows, so the deadline never registered as expired there: the poller kept polling and the failure surfaced as ErrorCode.TIMEOUT plus unmatched GET /jobs/777 requests. It passed on macOS only because one HTTP roundtrip happens to exceed 100us. Use an already-expired budget (max_wait=-1.0) instead: the deadline is behind the first check, so the timeout branch is reached with no poll, no sleep and no dependence on the clock. Reproduced both shapes under a simulated 15.6ms tick -- old: TIMEOUT after 8 polls; new: STORAGE_JOB_TIMEOUT after 0. Stubbing time.monotonic was rejected as the fix: patching it on the client module patches the global time module, which httpx reads too. --- tests/test_storage_swap.py | 26 ++++++++++++-------------- 1 file changed, 12 insertions(+), 14 deletions(-) diff --git a/tests/test_storage_swap.py b/tests/test_storage_swap.py index 7536847dc..f31d0a0ee 100644 --- a/tests/test_storage_swap.py +++ b/tests/test_storage_swap.py @@ -584,6 +584,16 @@ def test_expired_budget_reports_the_job_as_still_running(self, httpx_mock) -> No The message must not read as "nothing happened": kbagent stops watching, the swap keeps running server-side. That is also why the code is retryable (exit 4), not a plain failure. + + A budget already in the past puts the deadline behind the very first + check, so the timeout branch is reached without a poll, a sleep, or + any dependence on elapsed time. A tiny positive budget does NOT work: + ``time.monotonic()`` advances in ~15.6 ms ticks on Windows, so a + sub-millisecond deadline never registers as expired there and the + poller keeps polling (this test failed exactly that way in CI). + Stubbing the clock instead would patch the global ``time`` module, + which httpx also reads. Non-positive budgets are rejected for real + callers one layer up, in ``StorageService.swap_tables``. """ httpx_mock.add_response( url="https://connection.keboola.com/v2/storage/branch/42/tables/in.c-foo.a/swap", @@ -591,26 +601,14 @@ def test_expired_budget_reports_the_job_as_still_running(self, httpx_mock) -> No json={"id": 777, "status": "waiting", "operationName": "tableSwap"}, status_code=200, ) - httpx_mock.add_response( - url="https://connection.keboola.com/v2/storage/jobs/777", - method="GET", - json={"id": 777, "status": "waiting", "operationName": "tableSwap"}, - status_code=200, - ) client = KeboolaClient(stack_url="https://connection.keboola.com", token=TEST_TOKEN) - # The deadline is only checked at the top of each iteration, so the - # budget cannot expire before the first poll however small it is. - # Skipping the real sleep keeps the test instant. - with ( - patch("keboola_agent_cli.client._core.time.sleep"), - pytest.raises(KeboolaApiError) as exc_info, - ): + with pytest.raises(KeboolaApiError) as exc_info: client.swap_tables( table_id="in.c-foo.a", target_table_id="in.c-foo.b", branch_id=42, - max_wait=0.0001, + max_wait=-1.0, ) client.close() From 87825c797bcd1d66cf2b1c7dbe82e148e3f8c7fd Mon Sep 17 00:00:00 2001 From: soustruh Date: Sat, 10 Oct 2026 03:51:36 +0200 Subject: [PATCH 3/7] test(storage): cover the create-table job timeout, move the timeout note in the docs --- .../references/storage-types-workflow.md | 44 +++++++++---------- tests/test_storage_create_table.py | 30 +++++++++++++ 2 files changed, 52 insertions(+), 22 deletions(-) diff --git a/plugins/kbagent/skills/kbagent/references/storage-types-workflow.md b/plugins/kbagent/skills/kbagent/references/storage-types-workflow.md index e586dc0e3..63579580e 100644 --- a/plugins/kbagent/skills/kbagent/references/storage-types-workflow.md +++ b/plugins/kbagent/skills/kbagent/references/storage-types-workflow.md @@ -333,28 +333,6 @@ kbagent --json storage table-detail --project prod --table-id in.c-main.events \ | jq '.data.definition | {timePartitioning, clustering, requirePartitionFilter}' ``` -### Budget both steps on a large table - -*(since vNEXT, #713)* - -Steps 1 and 2 both move real data. Until vNEXT they inherited the 60s -`STORAGE_JOB_MAX_WAIT` meant for metadata jobs -- measured on an 800 MB table, -create takes ~15s and the swap ~31s, so a table a few times larger reported -`STORAGE_JOB_TIMEOUT` on a job that was still running and would succeed. Both -now default to **300s** and accept `--timeout SECONDS`: - -```bash -kbagent storage create-table --project prod --bucket-id in.c-main \ - --name events_repart --source-table-id in.c-main.events \ - --time-partitioning-type DAY --time-partitioning-field created_at \ - --timeout 3600 -``` - -**A timeout is not a failure.** kbagent stops watching; the Storage job keeps -running server-side. The error names the job id -- poll -`GET /v2/storage/jobs/{id}`, or just run step 3. Never blindly re-issue the -swap: repeating one that already landed swaps the tables straight back. - On 0.87.0 and earlier `table-detail` dropped `definition`, so step 3 was not possible from kbagent at all. `create-table`'s own output is not a substitute: it echoes the layout you *requested*, and its `--if-not-exists` skip path @@ -382,3 +360,25 @@ Rules: - The partition/clustering flags also work on a plain `--column` create (same BigQuery-only rule) when you want a fresh empty table in a specific layout rather than a copy. + +### Budget both steps on a large table + +*(since vNEXT, #713)* + +Steps 1 and 2 both move real data. Until vNEXT they inherited the 60s +`STORAGE_JOB_MAX_WAIT` meant for metadata jobs -- measured on an 800 MB table, +create takes ~15s and the swap ~31s, so a table a few times larger reported +`STORAGE_JOB_TIMEOUT` on a job that was still running and would succeed. Both +now default to **300s** and accept `--timeout SECONDS`: + +```bash +kbagent storage create-table --project prod --bucket-id in.c-main \ + --name events_repart --source-table-id in.c-main.events \ + --time-partitioning-type DAY --time-partitioning-field created_at \ + --timeout 3600 +``` + +**A timeout is not a failure.** kbagent stops watching; the Storage job keeps +running server-side. The error names the job id -- poll +`GET /v2/storage/jobs/{id}`, or just run step 3. Never blindly re-issue the +swap: repeating one that already landed swaps the tables straight back. diff --git a/tests/test_storage_create_table.py b/tests/test_storage_create_table.py index 4dfdf530f..c681df890 100644 --- a/tests/test_storage_create_table.py +++ b/tests/test_storage_create_table.py @@ -146,6 +146,36 @@ def test_columns_mode_unchanged(self, httpx_mock) -> None: assert "clustering" not in body client.close() + def test_forwards_budget_to_the_poller(self, httpx_mock) -> None: + """create_table passes max_wait through to _wait_for_storage_job (issue #713). + + Asserted on the kwarg, not on elapsed time: an httpx mock never sees + the budget, so dropping the argument would leave every wire test green + while a --source-table-id copy silently fell back to the 60s metadata + default. + """ + httpx_mock.add_response( + url="https://connection.keboola.com/v2/storage/buckets/in.c-main/tables-definition", + method="POST", + json={"id": 1, "status": "waiting"}, + status_code=200, + ) + client = KeboolaClient(stack_url="https://connection.keboola.com", token=TEST_TOKEN) + with patch.object( + KeboolaClient, + "_wait_for_storage_job", + return_value={"id": 1, "status": "success", "results": {}}, + ) as poller: + client.create_table( + bucket_id="in.c-main", + name="events_repart", + source={"tableId": "in.c-main.events"}, + max_wait=900.0, + ) + client.close() + + assert poller.call_args.kwargs["max_wait"] == 900.0 + # --------------------------------------------------------------------------- # Service layer -- validation From a63613f0e1f5941ea3c74fb3fa5cda36618601f3 Mon Sep 17 00:00:00 2001 From: soustruh Date: Sat, 10 Oct 2026 03:51:36 +0200 Subject: [PATCH 4/7] fix(storage): a create-table or swap-tables wait timeout is not retryable, name the job --- CLAUDE.md | 8 +- plugins/kbagent/agents/keboola-expert.md | 2 +- .../kbagent/references/commands-reference.md | 4 +- .../skills/kbagent/references/gotchas.md | 17 +-- .../references/storage-types-workflow.md | 14 ++- .../client/storage_tables.py | 103 +++++++++++++++--- src/keboola_agent_cli/commands/context.py | 9 +- src/keboola_agent_cli/commands/storage.py | 8 +- .../services/_storage_jobs.py | 57 ++++++++-- .../services/storage_service.py | 15 ++- tests/test_server_router_calls.py | 42 +++++++ tests/test_storage_create_table.py | 32 ++++++ tests/test_storage_swap.py | 76 ++++++++++++- tests/test_storage_write.py | 60 ++++++++++ 14 files changed, 392 insertions(+), 55 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 548ad0a29..9d9c2e2e9 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -668,9 +668,11 @@ kbagent storage swap-tables --project NAME --table-id ID --target-table-id ID -- # of the repartition path move real data yet used to inherit the 60s metadata-job # budget, so a large BigQuery table reported STORAGE_JOB_TIMEOUT for a job that was # still running and would succeed. Default is now 300s on both. A local timeout - # NEVER cancels the job -- it keeps running server-side; poll - # GET /v2/storage/jobs/{id} or verify with `storage table-detail` rather than - # assuming the operation failed. + # NEVER cancels the job -- it keeps running server-side. The error is + # STORAGE_JOB_TIMEOUT, exit 4, retryable=false, details.job_id, and the message + # names `storage job-detail --job-id ID --wait` (create-table also names + # `storage table-detail`). Never re-run: a repeated swap swaps the tables back, + # a repeated --source-table-id create starts a second copy. kbagent storage clone-table --project NAME --table-id ID --branch ID [--dry-run] kbagent storage snapshot-create --project NAME --table-id ID [--description D] [--branch ID] kbagent storage snapshots --project NAME --table-id ID [--limit N] [--branch ID] diff --git a/plugins/kbagent/agents/keboola-expert.md b/plugins/kbagent/agents/keboola-expert.md index 10cbc1784..30b54d8a5 100644 --- a/plugins/kbagent/agents/keboola-expert.md +++ b/plugins/kbagent/agents/keboola-expert.md @@ -121,7 +121,7 @@ been retired, so its absence is NOT a promise (see ยง1 Rule 6). | Create typed table with native types | `kbagent storage create-table --column pk:VARCHAR(40) --column amount:NUMBER(18,2) --not-null pk --default amount=0` | -- | re-creating via raw REST to `tables-definition` | | Add one column to an existing table | `kbagent storage add-column --project P --table-id in.c-foo.data --column status:VARCHAR(20) [--not-null] [--default active]` -- synchronous, same `name:TYPE(length)` grammar as `create-table` | -- | re-creating the whole table to add a field (loses data / PK / dependents) | | Promote a rebuilt table into the original name | `kbagent storage swap-tables --project P --table-id in.c-foo.data --target-table-id in.c-foo.data_new --branch --yes` -- async `tableSwap`; the service refuses without a branch (any branch, prod included) | -- | renaming or delete + re-upload (loses history; every downstream config needs rewriting) | -| Repartition / recluster a populated BigQuery table | `kbagent storage create-table --source-table-id --time-partitioning-type DAY --time-partitioning-field created_at --clustering-field tenant_id` (BigQuery only) to copy the data into the new layout, then `swap-tables` to flip it in. `--source-table-id` derives the schema, so `--column` is forbidden. Range partitioning: all four `--range-partitioning-*` flags together. VERIFY the result with `storage table-detail --json` -> `.definition.timePartitioning` / `.clustering` (0.88.0+) -- `create-table` only echoes the layout you REQUESTED, so it proves nothing. Both steps take `--timeout SECONDS` (vNEXT+, default 300); a `STORAGE_JOB_TIMEOUT` means kbagent stopped WATCHING, not that the job failed -- verify before retrying, because re-running a swap that succeeded swaps the tables back. See [storage-types-workflow.md](../skills/kbagent/references/storage-types-workflow.md) | -- | `CREATE TABLE ... AS SELECT` in a workspace (drops NOT NULL + primary key) | +| Repartition / recluster a populated BigQuery table | `kbagent storage create-table --source-table-id --time-partitioning-type DAY --time-partitioning-field created_at --clustering-field tenant_id` (BigQuery only) to copy the data into the new layout, then `swap-tables` to flip it in. `--source-table-id` derives the schema, so `--column` is forbidden. Range partitioning: all four `--range-partitioning-*` flags together. VERIFY the result with `storage table-detail --json` -> `.definition.timePartitioning` / `.clustering` (0.88.0+) -- `create-table` only echoes the layout you REQUESTED, so it proves nothing. Both steps take `--timeout SECONDS` (vNEXT+, default 300); a `STORAGE_JOB_TIMEOUT` (exit 4, `retryable: false`) means kbagent stopped WATCHING, not that the job failed -- follow it with `storage job-detail --job-id ID --wait`, never re-run: a repeated swap swaps the tables back, a repeated create starts a second copy. See [storage-types-workflow.md](../skills/kbagent/references/storage-types-workflow.md) | -- | `CREATE TABLE ... AS SELECT` in a workspace (drops NOT NULL + primary key) | | Re-seed a table without losing schema / PK / dependents | `kbagent storage truncate-table --project P --table-id in.c-foo.data [--dry-run] [--yes]` -- rows only, uniformly async-via-job on every branch; batch via repeated `--table-id` | -- | drop + recreate (loses descriptions, PK, sharing edges, and breaks every downstream reference); deleting rows via raw SQL in a workspace (bypasses the Storage audit trail) | | Back up / restore a table around a risky change | `kbagent storage snapshot-create --table-id ...` then, to restore, `kbagent storage table-from-snapshot --snapshot-id ID --bucket-id B --name NEW` -- restore is always a NEW table (`--name` REQUIRED, no overwrite): verify it, then `swap-tables`. See [snapshot-workflow.md](../skills/kbagent/references/snapshot-workflow.md) | `storage snapshots` / `snapshot-detail` to find one | exporting to CSV as a "backup" (loses column types + PK); `create-table --snapshot-id` (not a thing) | | Debug a failed job | `kbagent job detail --project P --job-id J --log-tail-lines 200 --json` | `kbagent workspace from-transformation` for SQL repro (on a dev branch it reads the branch's transformation only since 0.96.1, #807; older versions load production's input mapping) | "I think the issue is..." without reading logs; a new job run only to see its logs (a writer can send its data again) | diff --git a/plugins/kbagent/skills/kbagent/references/commands-reference.md b/plugins/kbagent/skills/kbagent/references/commands-reference.md index b3b4cdc30..362e72100 100644 --- a/plugins/kbagent/skills/kbagent/references/commands-reference.md +++ b/plugins/kbagent/skills/kbagent/references/commands-reference.md @@ -178,7 +178,7 @@ Requires a **super-admin** Manage API token (same kind as `org setup`). Same def - `storage tables [--project NAME ...] [--bucket-id ID] [--branch ID] [--include-usage]` -- list tables across all connected projects in parallel (multi-project by default, same as `storage buckets`); repeat `--project` to target a subset; `--bucket-id` is applied independently per project (missing buckets become per-project errors); `--branch` requires exactly one `--project`. `--include-usage` (0.88.0+) adds `used_by` per table: the configurations naming it in their **storage input/output mapping** only -- a table id inside a transformation's SQL is NOT a reference. Costs one extra component listing per project (not per table), which is the slow call in a big project; unreadable components degrade to an empty `used_by` - `storage table-detail --project NAME --table-id ID [--branch ID]` -- table detail with columns, types, primary key, row count (branch-aware). Since 0.88.0 (#621) also returns the raw Storage API `definition`: on BigQuery that carries `timePartitioning` / `rangePartitioning` / `clustering` / `requirePartitionFilter` / `partitions[]`, and it is the only way to verify a repartition landed. Human mode prints the layout and a partition COUNT; `--json` passes `definition` through verbatim. Present on every response (untyped tables too), so `null` means the stack omitted the key, not "untyped". Since 0.88.0 (#624) it also RESOLVES column descriptions written by anyone -- the UI, a component, or kbagent -- into `column_details[].description`, with precedence native definition -> `columnMetadata` `KBC.description` -> legacy flat `KBC.column.*` (an alias table falls back to the source table's `columnMetadata`, matching the MCP server). The response always carries `legacy_column_descriptions`, naming the columns still backed by the pre-0.88.0 convention -- human mode warns and points at `storage describe-migrate`. *(since v0.89.0)* Human mode also renders a `Description` column in the Columns table (only when at least one column has one; long text wraps rather than truncating); on 0.88.0 descriptions were visible in `--json` only. Reading never writes, so it is safe under a read-only token or `--deny-writes`. *(since 0.95.0, #761)* The response also carries `backend_path` (the owning bucket's Storage `backendPath`, verbatim: Snowflake `[database, schema]`, BigQuery `[dataset]`) and `sql_path` (the quoted, directly queryable table path, same convention as `storage bucket-detail`'s per-table `sql_path`; `null` when Storage reports no location or the backend is not Snowflake/BigQuery) - `storage create-bucket --project NAME --stage STAGE --name NAME [--description D] [--backend B] [--branch ID]` -- create bucket (branch-aware). With `--branch ID` on a project lacking the `storage-branches` feature (legacy fake-branch), response carries `legacy_branch_storage: true` and human mode prints a warning -- the runner will create a parallel `out.c--*` bucket at job time. See `storage-types-workflow.md` -- `storage create-table --project NAME --bucket-id ID --name NAME [--column col:TYPE[(length)] ...] [--primary-key COL] [--not-null COL ...] [--default NAME=VALUE ...] [--source-table-id ID] [--source-branch-id N] [--time-partitioning-type DAY|HOUR|MONTH|YEAR] [--time-partitioning-field COL] [--time-partitioning-expiration-ms MS] [--range-partitioning-field COL --range-partitioning-start S --range-partitioning-end E --range-partitioning-interval I] [--clustering-field COL ...] [--branch ID] [--if-not-exists] [--timeout SECONDS]` -- create typed table. Base types `STRING/INTEGER/NUMERIC/FLOAT/BOOLEAN/DATE/TIMESTAMP` plus native backend types with length (`VARCHAR(40)`, `NUMBER(18,2)`, `TIMESTAMP_TZ`, `VARIANT`, etc.) -- type/length validation delegated to the Storage API. `--not-null` marks a column `nullable=false`; `--default NAME=VALUE` sets a DEFAULT expression (booleans must be lowercase `true`/`false`). In a dev branch, the target bucket is auto-materialized if it has not yet been written to there -- response surfaces this via `auto_created_bucket: bool`. On legacy fake-branch projects (no `storage-branches` feature), `legacy_branch_storage: true` flags that the runner will use a separate `out.c--*` bucket at job time. `--if-not-exists` (0.47.0+) turns a duplicate-display-name failure into `action: skipped` when the table really exists at the expected id (safe for parallel workers). Since 0.47.1 the skipped envelope reports the EXISTING table's actual `columns`/`primary_key`/`name`, mirrors the request under `requested_columns`/`requested_primary_key`, and sets `schema_drift: true` when they diverge. **`--source-table-id` (0.66.0+, BigQuery only)** copies an existing table's data into the requested partition/clustering layout instead of building from `--column` (schema derived from source -> `--column`/`--not-null`/`--default` forbidden; the two are mutually exclusive). This is the supported way to repartition a populated BigQuery table -- then promote it with `storage swap-tables`. Partition/clustering flags (`--time-partitioning-*`, `--range-partitioning-*`, `--clustering-field`) also work on a plain `--column` create (BigQuery only); time vs range partitioning are mutually exclusive and range bounds are strings. When any source/partition/clustering flag is used, a one-call backend pre-flight rejects non-BigQuery projects (exit 2) before the create. See `storage-types-workflow.md`. **`--timeout SECONDS` (vNEXT+, default 300)** budgets the async create job; a `--source-table-id` copy moves real data and used to inherit the 60s metadata-job default, so a large table reported `STORAGE_JOB_TIMEOUT` for a copy that was still running and would succeed. A local timeout NEVER cancels the job -- poll `GET /v2/storage/jobs/{id}` or verify with `table-detail`; exit code is 4 (retryable) +- `storage create-table --project NAME --bucket-id ID --name NAME [--column col:TYPE[(length)] ...] [--primary-key COL] [--not-null COL ...] [--default NAME=VALUE ...] [--source-table-id ID] [--source-branch-id N] [--time-partitioning-type DAY|HOUR|MONTH|YEAR] [--time-partitioning-field COL] [--time-partitioning-expiration-ms MS] [--range-partitioning-field COL --range-partitioning-start S --range-partitioning-end E --range-partitioning-interval I] [--clustering-field COL ...] [--branch ID] [--if-not-exists] [--timeout SECONDS]` -- create typed table. Base types `STRING/INTEGER/NUMERIC/FLOAT/BOOLEAN/DATE/TIMESTAMP` plus native backend types with length (`VARCHAR(40)`, `NUMBER(18,2)`, `TIMESTAMP_TZ`, `VARIANT`, etc.) -- type/length validation delegated to the Storage API. `--not-null` marks a column `nullable=false`; `--default NAME=VALUE` sets a DEFAULT expression (booleans must be lowercase `true`/`false`). In a dev branch, the target bucket is auto-materialized if it has not yet been written to there -- response surfaces this via `auto_created_bucket: bool`. On legacy fake-branch projects (no `storage-branches` feature), `legacy_branch_storage: true` flags that the runner will use a separate `out.c--*` bucket at job time. `--if-not-exists` (0.47.0+) turns a duplicate-display-name failure into `action: skipped` when the table really exists at the expected id (safe for parallel workers). Since 0.47.1 the skipped envelope reports the EXISTING table's actual `columns`/`primary_key`/`name`, mirrors the request under `requested_columns`/`requested_primary_key`, and sets `schema_drift: true` when they diverge. **`--source-table-id` (0.66.0+, BigQuery only)** copies an existing table's data into the requested partition/clustering layout instead of building from `--column` (schema derived from source -> `--column`/`--not-null`/`--default` forbidden; the two are mutually exclusive). This is the supported way to repartition a populated BigQuery table -- then promote it with `storage swap-tables`. Partition/clustering flags (`--time-partitioning-*`, `--range-partitioning-*`, `--clustering-field`) also work on a plain `--column` create (BigQuery only); time vs range partitioning are mutually exclusive and range bounds are strings. When any source/partition/clustering flag is used, a one-call backend pre-flight rejects non-BigQuery projects (exit 2) before the create. See `storage-types-workflow.md`. **`--timeout SECONDS` (vNEXT+, default 300)** budgets the async create job; a `--source-table-id` copy moves real data and used to inherit the 60s metadata-job default, so a large table reported `STORAGE_JOB_TIMEOUT` for a copy that was still running and would succeed. A local timeout NEVER cancels the job: the error is `STORAGE_JOB_TIMEOUT` (exit 4, `retryable: false`, `details.job_id`) -- follow it with `storage job-detail --job-id ID --wait`, verify the layout with `table-detail`, and never re-run the create (a second copy) - `storage upload-table --project NAME --table-id ID --file PATH [--incremental] [--delimiter D] [--enclosure E] [--no-auto-create] [--wait/--no-wait] [--timeout SECONDS] [--branch ID] [--progress]` -- upload a CSV (or a `.csv.gz`, uploaded as-is) into a table; auto-creates the bucket/table from the header unless `--no-auto-create` (branch-aware). Two phases: cloud upload, then an async Storage import job. *(since 0.98.0, #834)* On AWS stacks a file above 64 MiB is sent as an S3 **multipart** upload (64 MiB parts, auto-scaled to fit 10,000 parts, 4 in parallel, per-part retry; peak memory ~4 x part size) -- the old 5 GiB single-PUT ceiling and whole-file-in-RAM read are gone; Azure/GCP unchanged. Human mode shows an upload progress bar on stderr. *(since 0.98.0)* `--wait` (default) waits `--timeout SECONDS` (default 600) for the import; `--no-wait` still waits for the upload, then returns `file_id`, `job_id`, `job_status` (`imported_rows` null while pending) -- follow it with `storage job-detail --job-id ID --wait`. A wait timeout is `STORAGE_JOB_TIMEOUT` (exit 4, `retryable: false`): the import keeps running server-side, so never re-run the upload -- re-import a failed one from the uploaded file with `storage load-file --file-id`. Big files: see `large-upload-workflow.md`. *(since 0.98.0)* `--progress` always reports progress on stderr -- also with `--json` and without a terminal (a bar on a terminal, otherwise one line every 10 s plus a final line: percent, transferred/total, speed over a trailing 30 s window, elapsed, ETA); stdout is untouched, so `--json ... --progress 2>progress.log` keeps a machine-readable result and a readable log - `storage download-table --project NAME --table-id ID [--output FILE] [--columns COL ...] [--limit N] [--where-column COL --where-value VAL ... [--where-operator eq|neq]] [--changed-since WHEN] [--changed-until WHEN] [--branch ID] [--progress]` -- export table to CSV (branch-aware). `--where-column` + `--where-value` (repeatable, OR within the set) + `--where-operator eq|neq` filter rows server-side; `--changed-since`/`--changed-until` (unix ts or strtotime like `-2 days`) filter by import time -- the credential-only, no-workspace way to pull a filtered/incremental slice (0.62.0+). *(since 0.98.0)* `--progress` as on `upload-table`; counts network (compressed) bytes and starts the clock when the download starts (the export job wait is excluded); a sliced export whose manifest has no sizes shows bytes + speed without percent/ETA - `storage add-column --project NAME --table-id ID --column COL:TYPE[(length)] [--not-null] [--default VALUE] [--branch ID]` -- add a single column to an existing table (0.62.0+). Same `name:TYPE(length)` grammar as `create-table --column`; a bare `name` adds an untyped STRING column. Synchronous endpoint (no job to wait on). `--not-null` needs an empty table or a `--default`. Mirror of `delete-column` @@ -186,7 +186,7 @@ Requires a **super-admin** Manage API token (same kind as `org setup`). Same def - `storage truncate-table --project NAME --table-id ID [--table-id ...] [--dry-run] [--yes] [--branch ID]` -- delete all rows while preserving table schema, primary key, descriptions, sharing edges, and downstream dependents. Batch via repeated `--table-id`. Endpoint is uniformly async-via-job on every branch (returns a queued `tableRowsDelete` job; client polls via `_wait_for_storage_job` before returning). Idempotent (truncating an empty table is a no-op). Use when re-seeding a table without losing the schema contract - `storage delete-column --project NAME --table-id ID --column COL [--column ...] [--force] [--dry-run] [--yes] [--branch ID]` -- delete columns from a table (branch-aware) - `storage delete-bucket --project NAME --bucket-id ID [--bucket-id ...] [--force] [--dry-run] [--yes] [--branch ID]` -- delete buckets (branch-aware) -- `storage swap-tables --project NAME --table-id ID --target-table-id ID --branch ID [--dry-run] [--yes] [--timeout SECONDS]` -- swap two storage tables in any branch, including the default/production branch (POST `/tables/{id}/swap`). Both tables exchange physical positions; aliases are NOT transferred (they keep pointing at the same physical position and therefore expose the OTHER table's data after the swap). Service refuses without a branch (active branch via `branch use` works too). Use to flip a typed rebuild ("data_change_log") into the original name ("data") without touching downstream config references. **`--timeout SECONDS` (vNEXT+, default 300)** budgets the async swap job -- a swap on a large BigQuery table outlasts the old 60s default. A local timeout NEVER cancels the swap; it usually lands anyway, so verify with `table-detail` rather than re-issuing it +- `storage swap-tables --project NAME --table-id ID --target-table-id ID --branch ID [--dry-run] [--yes] [--timeout SECONDS]` -- swap two storage tables in any branch, including the default/production branch (POST `/tables/{id}/swap`). Both tables exchange physical positions; aliases are NOT transferred (they keep pointing at the same physical position and therefore expose the OTHER table's data after the swap). Service refuses without a branch (active branch via `branch use` works too). Use to flip a typed rebuild ("data_change_log") into the original name ("data") without touching downstream config references. **`--timeout SECONDS` (vNEXT+, default 300)** budgets the async swap job -- a swap on a large BigQuery table outlasts the old 60s default. A local timeout NEVER cancels the swap: the error is `STORAGE_JOB_TIMEOUT` (exit 4, `retryable: false`, `details.job_id`) -- follow it with `storage job-detail --job-id ID --wait` and never re-run the swap (a repeat after it lands swaps the tables back) - `storage clone-table --project NAME --table-id ID --branch ID [--dry-run]` -- pull (clone) a production table into a dev branch (POST `/tables/{id}/pull`, operationName `devBranchTablePull`). On `storage-branches` projects a dev branch reads prod tables transparently until the first write, so an in-branch schema mutation (`swap-tables`, dropping a column) fails with a misleading "bucket not found" until the table is materialized branch-local; `clone-table` does that. One-way (default -> branch). Service refuses without a branch (active branch via `branch use` works too). Permission class `write` - `storage snapshot-create --project NAME --table-id ID [--description D] [--branch ID]` -- create a point-in-time snapshot of a table (data + columns + primary key; async `tableSnapshotCreate` job, client polls to completion). The receipt carries `snapshot_id` -- restores are addressed by it. Permission class `write` - `storage snapshots --project NAME --table-id ID [--limit N] [--branch ID]` -- list a table's snapshots (`id`, `createdTime`, `description`, `creatorToken`). Read command: production endpoint by default diff --git a/plugins/kbagent/skills/kbagent/references/gotchas.md b/plugins/kbagent/skills/kbagent/references/gotchas.md index 2675ca70d..08b47edd9 100644 --- a/plugins/kbagent/skills/kbagent/references/gotchas.md +++ b/plugins/kbagent/skills/kbagent/references/gotchas.md @@ -2154,7 +2154,7 @@ config, the retry fires, and the retry destroys it for good. BigQuery repartition path -- both move real data, yet both used to inherit the **60s** `STORAGE_JOB_MAX_WAIT` budget meant for metadata jobs. Measured on an 800 MB BigQuery table: create ~15s, swap ~31s. A table a few times - larger blew the budget, so a perfectly healthy operation surfaced as + larger blew the budget, so a healthy operation was reported as `STORAGE_JOB_TIMEOUT`. - **Since vNEXT both default to 300s and accept `--timeout SECONDS`.** The REST routes (`POST /storage/tables/{p}`, `POST /storage/tables/{p}/{tid}/swap`) @@ -2162,12 +2162,15 @@ config, the retry fires, and the retry destroys it for good. boundary and `INVALID_ARGUMENT` from the service. - **Never read the timeout as "nothing happened".** kbagent only stops *watching*; the Storage job keeps running server-side (and keeps consuming - backend resources). It usually lands. That is why the code is - `retryable: true` and maps to **exit 4**, not exit 1. -- **Do not re-issue the command on a timeout.** Poll - `GET /v2/storage/jobs/{id}` (the id is named in the error message) or check - the physical result with `storage table-detail --json` -> `.definition`. - Re-running a swap that actually succeeded swaps the tables *back*. + backend resources). It usually lands. The error is **exit 4** with + `retryable: false` and `details.job_id`, the same contract as an + `upload-table` import timeout (0.98.0). +- **Do not re-issue the command on a timeout.** Follow the job with + `kbagent storage job-detail --project P --job-id ID --wait` (the error + message gives the full command), then check the physical result with + `storage table-detail --json` -> `.definition`. Re-running a swap that + succeeded swaps the tables *back*; re-running a `--source-table-id` create + starts a second copy. ## `storage clone-table` materializes a prod table into a dev branch diff --git a/plugins/kbagent/skills/kbagent/references/storage-types-workflow.md b/plugins/kbagent/skills/kbagent/references/storage-types-workflow.md index 63579580e..8844fd10f 100644 --- a/plugins/kbagent/skills/kbagent/references/storage-types-workflow.md +++ b/plugins/kbagent/skills/kbagent/references/storage-types-workflow.md @@ -379,6 +379,14 @@ kbagent storage create-table --project prod --bucket-id in.c-main \ ``` **A timeout is not a failure.** kbagent stops watching; the Storage job keeps -running server-side. The error names the job id -- poll -`GET /v2/storage/jobs/{id}`, or just run step 3. Never blindly re-issue the -swap: repeating one that already landed swaps the tables straight back. +running server-side. The error is `STORAGE_JOB_TIMEOUT`, exit 4, +`retryable: false`, with the job id in `details.job_id` and the follow-up +command in the message: + +```bash +kbagent storage job-detail --project prod --job-id --wait +``` + +Then run step 3. Never re-issue the command: repeating a swap that already +landed swaps the tables straight back, and repeating the create starts a +second copy. diff --git a/src/keboola_agent_cli/client/storage_tables.py b/src/keboola_agent_cli/client/storage_tables.py index 19e7fa164..9bbaa8417 100644 --- a/src/keboola_agent_cli/client/storage_tables.py +++ b/src/keboola_agent_cli/client/storage_tables.py @@ -18,6 +18,7 @@ FILE_UPLOAD_TIMEOUT, IMPORT_JOB_MAX_WAIT, S3_MULTIPART_THRESHOLD, + STORAGE_JOB_MAX_WAIT, ) from ..errors import ErrorCode, KeboolaApiError from ._core import _CoreClient @@ -606,6 +607,11 @@ def create_table( Returns: Completed storage job results dict. + + Raises: + KeboolaApiError: ``STORAGE_JOB_TIMEOUT`` (``retryable=False``, + details ``job_id``) when the wait budget runs out -- the + create keeps running server-side. """ prefix = f"/v2/storage/branch/{branch_id}" if branch_id else "/v2/storage" safe_id = quote(bucket_id, safe="") @@ -626,7 +632,15 @@ def create_table( if clustering is not None: body["clustering"] = clustering response = self._request("POST", f"{prefix}/buckets/{safe_id}/tables-definition", json=body) - job = self._wait_for_storage_job(response.json(), max_wait=max_wait) + job = self._wait_keeps_running( + response.json(), + budget=STORAGE_JOB_MAX_WAIT if max_wait is None else max_wait, + operation="create-table", + warning=( + ". Re-running the command starts a second create (with a source " + "table, a second copy) -- check the job before retrying." + ), + ) return job.get("results", {}) def prepare_file_upload( @@ -912,30 +926,74 @@ def _wait_for_import_job( the job (``storage job-detail``) or the file (``storage load-file``). """ budget = IMPORT_JOB_MAX_WAIT if max_wait is None else max_wait + try: + return self._wait_keeps_running( + job, + budget=budget, + operation="import", + warning=( + f" (source: Storage file {file_id}). Re-running an incremental " + "upload can duplicate rows -- check the job before retrying." + ), + details={"file_id": file_id}, + ) + except KeboolaApiError as exc: + if exc.error_code == ErrorCode.STORAGE_JOB_TIMEOUT: + raise + raise KeboolaApiError( + message=exc.message, + status_code=exc.status_code, + error_code=exc.error_code, + retryable=exc.retryable, + details={"job_id": job.get("id"), **exc.details, "file_id": file_id}, + ) from exc + + def _wait_keeps_running( + self, + job: dict[str, Any], + *, + budget: float, + operation: str, + warning: str, + details: dict[str, Any] | None = None, + ) -> dict[str, Any]: + """Poll a job whose wait timeout must not be retried. + + Giving up locally does not stop a Storage job. For an operation that + is not safe to repeat (an import, a create-from-source copy, a swap) + a retry would run it a second time next to the first, so the timeout + is re-raised NOT retryable, with ``job_id`` in the details. A caller + that knows the project alias appends the ``storage job-detail`` + command (``services._storage_jobs``). + + Args: + job: Initial job response from the enqueue request. + budget: Seconds to wait. + operation: What the job does, named in the message. + warning: Text that follows "keeps running server-side" in the + message: what a repeat of the operation would do. + details: Extra keys for the timeout's ``details``. + + Raises: + KeboolaApiError: ``STORAGE_JOB_TIMEOUT`` (``retryable=False``) when + the budget runs out; any other wait failure unchanged. + """ try: return self._wait_for_storage_job(job, max_wait=budget) except KeboolaApiError as exc: - job_id = exc.details.get("job_id", job.get("id")) - details = {**exc.details, "job_id": job_id, "file_id": file_id} if exc.error_code != ErrorCode.STORAGE_JOB_TIMEOUT: - raise KeboolaApiError( - message=exc.message, - status_code=exc.status_code, - error_code=exc.error_code, - retryable=exc.retryable, - details=details, - ) from exc + raise + job_id = exc.details.get("job_id", job.get("id")) raise KeboolaApiError( message=( - f"Storage import job {job_id} did not finish within {budget:g}s. " - "Waiting stopped locally; the import keeps running server-side " - f"(source: Storage file {file_id}). Re-running an incremental " - "upload can duplicate rows -- check the job before retrying." + f"Storage {operation} job {job_id} did not finish within {budget:g}s. " + f"Waiting stopped locally; the {operation} keeps running server-side" + f"{warning}" ), status_code=exc.status_code, error_code=ErrorCode.STORAGE_JOB_TIMEOUT, retryable=False, - details=details, + details={**exc.details, "job_id": job_id, **(details or {})}, ) from exc def upload_table( @@ -1173,12 +1231,25 @@ def swap_tables( Returns: Completed storage job dict. + + Raises: + KeboolaApiError: ``STORAGE_JOB_TIMEOUT`` (``retryable=False``, + details ``job_id``) when the wait budget runs out -- the swap + keeps running server-side. """ prefix = f"/v2/storage/branch/{branch_id}" safe_id = quote(table_id, safe="") body = {"targetTableId": target_table_id} response = self._request("POST", f"{prefix}/tables/{safe_id}/swap", json=body) - return self._wait_for_storage_job(response.json(), max_wait=max_wait) + return self._wait_keeps_running( + response.json(), + budget=STORAGE_JOB_MAX_WAIT if max_wait is None else max_wait, + operation="swap-tables", + warning=( + ". Re-running the command after the swap lands swaps the tables " + "back -- check the job before retrying." + ), + ) def pull_table(self, table_id: str, branch_id: int) -> dict[str, Any]: """Pull (clone) a table from the default branch into a dev branch. diff --git a/src/keboola_agent_cli/commands/context.py b/src/keboola_agent_cli/commands/context.py index c3a4a132c..c06323ad2 100644 --- a/src/keboola_agent_cli/commands/context.py +++ b/src/keboola_agent_cli/commands/context.py @@ -725,8 +725,9 @@ --source-table-id copy moves real data and used to inherit the 60s metadata-job default, so a large BigQuery table reported STORAGE_JOB_TIMEOUT for a copy that was still running and would succeed. A local timeout NEVER cancels the job -- it keeps - running server-side; poll GET /v2/storage/jobs/{{id}} or verify with table-detail - instead of assuming failure. Exit code is 4 (retryable), not 1. + running server-side. The error is STORAGE_JOB_TIMEOUT, exit 4, retryable: false, + details.job_id; follow it with `storage job-detail --job-id ID --wait`, verify with + table-detail, and never re-run the create (a second copy). - --if-not-exists: opt-in idempotency. On a duplicate-display-name failure, probe get-table-detail at the expected id and, if the table really exists, return `action: "skipped", skip_reason: "table already exists"` instead of raising. A different @@ -826,7 +827,9 @@ does not carry storage schema). --timeout SECONDS (since vNEXT, default 300): budget for the async swap job. A swap on a large BigQuery table outlasts the old 60s metadata-job default. A local timeout NEVER cancels the swap -- - it keeps running server-side and usually lands; verify with table-detail rather than re-issuing it. + it keeps running server-side and usually lands. The error is STORAGE_JOB_TIMEOUT, exit 4, + retryable: false, details.job_id; follow it with `storage job-detail --job-id ID --wait` and never + re-issue the swap (a repeat after it lands swaps the tables back). kbagent storage clone-table --project NAME --table-id ID --branch ID [--dry-run] Clone (pull) a production table into a dev branch (POST /tables/{{id}}/pull). On storage-branches projects a diff --git a/src/keboola_agent_cli/commands/storage.py b/src/keboola_agent_cli/commands/storage.py index c14cff62d..eec440f63 100644 --- a/src/keboola_agent_cli/commands/storage.py +++ b/src/keboola_agent_cli/commands/storage.py @@ -685,7 +685,12 @@ def storage_create_table( formatter.error(message=exc.message, error_code=ErrorCode.CONFIG_ERROR) raise typer.Exit(code=5) from None except KeboolaApiError as exc: - formatter.error(message=exc.message, error_code=exc.error_code, retryable=exc.retryable) + formatter.error( + message=exc.message, + error_code=exc.error_code, + retryable=exc.retryable, + details=exc.details, + ) raise typer.Exit(code=map_error_to_exit_code(exc)) from None if formatter.json_mode: @@ -1531,6 +1536,7 @@ def storage_swap_tables( error_code=exc.error_code, project=project, retryable=exc.retryable, + details=exc.details, ) raise typer.Exit(code=exit_code) from None diff --git a/src/keboola_agent_cli/services/_storage_jobs.py b/src/keboola_agent_cli/services/_storage_jobs.py index 02206220b..149364886 100644 --- a/src/keboola_agent_cli/services/_storage_jobs.py +++ b/src/keboola_agent_cli/services/_storage_jobs.py @@ -180,6 +180,48 @@ def _table_detail_command(alias: str, table_id: str, branch_id: int | None) -> s return shlex.join(args + _branch_args(branch_id)) +def _job_detail_command(alias: str, job_id: Any) -> str: + """The ``storage job-detail --wait`` command that follows a running job.""" + args = ["kbagent", "storage", "job-detail", "--project", alias, "--job-id", str(job_id)] + return shlex.join([*args, "--wait"]) + + +def _with_hint(exc: KeboolaApiError, hint: str) -> KeboolaApiError: + """An equal copy of ``exc`` with ``hint`` appended to its message.""" + return KeboolaApiError( + message=exc.message + hint, + status_code=exc.status_code, + error_code=exc.error_code, + retryable=exc.retryable, + details=exc.details, + ) + + +def with_job_timeout_hint( + exc: KeboolaApiError, + alias: str, + table_id: str | None = None, + branch_id: int | None = None, +) -> KeboolaApiError: + """Append the follow-up commands to a wait timeout (create-table, swap-tables). + + The client cannot name them (it knows no project alias). A + ``STORAGE_JOB_TIMEOUT`` with a ``job_id`` gets ``storage job-detail + --wait``, plus ``storage table-detail`` for ``table_id`` when given. + Anything else comes back as an equal copy, so a caller can always + ``raise with_job_timeout_hint(exc, ...) from exc``. + """ + job_id = exc.details.get("job_id") + if exc.error_code != ErrorCode.STORAGE_JOB_TIMEOUT or job_id is None: + return _with_hint(exc, "") + hint = f" Follow it with: {_job_detail_command(alias, job_id)}" + if table_id is not None: + hint += ( + f" -- then verify the table with: {_table_detail_command(alias, table_id, branch_id)}" + ) + return _with_hint(exc, hint) + + def with_import_hint( exc: KeboolaApiError, alias: str, table_id: str, options: ImportOptions ) -> KeboolaApiError: @@ -203,11 +245,8 @@ def with_import_hint( job_id = exc.details.get("job_id") file_id = exc.details.get("file_id") if exc.error_code == ErrorCode.STORAGE_JOB_TIMEOUT and job_id is not None: - hint = ( - f" Follow it with: kbagent storage job-detail --project {shlex.quote(alias)} " - f"--job-id {job_id} --wait" - ) - elif exc.details.get("import_may_be_running") and job_id is None: + return with_job_timeout_hint(exc, alias) + if exc.details.get("import_may_be_running") and job_id is None: hint = ( " The import request may have reached Storage, so the import may already be " "running -- do NOT re-import yet. Check the table first: " @@ -223,13 +262,7 @@ def with_import_hint( hint = f" Import it with: {_load_file_command(alias, file_id, table_id, options)}" else: hint = "" - return KeboolaApiError( - message=exc.message + hint, - status_code=exc.status_code, - error_code=exc.error_code, - retryable=exc.retryable, - details=exc.details, - ) + return _with_hint(exc, hint) def read_csv_header(file_path: str, delimiter: str = ",", enclosure: str = '"') -> list[str]: diff --git a/src/keboola_agent_cli/services/storage_service.py b/src/keboola_agent_cli/services/storage_service.py index aad6a4aca..5e43f7d83 100644 --- a/src/keboola_agent_cli/services/storage_service.py +++ b/src/keboola_agent_cli/services/storage_service.py @@ -28,6 +28,7 @@ summarize_storage_job, validate_wait_timeout, with_import_hint, + with_job_timeout_hint, ) from ._storage_tables import normalize_table_rows from ._table_detail import build_table_detail @@ -783,7 +784,10 @@ def create_table( ``schema_drift`` is ``True`` when the two diverge. Raises: - KeboolaApiError: ``INVALID_ARGUMENT`` for a non-positive ``timeout``. + KeboolaApiError: ``INVALID_ARGUMENT`` for a non-positive ``timeout``; + ``STORAGE_JOB_TIMEOUT`` (``retryable=False``, details ``job_id``, + message names ``storage job-detail`` and ``storage table-detail``) + when the wait runs out -- the create keeps running server-side. ValueError: Malformed column spec or ``--default`` assignment; ``--not-null`` / ``--default`` references an unknown column; ``columns`` and ``source`` both/neither given; @@ -957,7 +961,7 @@ def create_table( "range_partitioning": None, "clustering": None, } - raise + raise with_job_timeout_hint(exc, alias, target_table_id, branch_id) from exc legacy_branch_storage = _detect_legacy_branch_storage(client, branch_id) finally: client.close() @@ -1603,7 +1607,10 @@ def swap_tables( Raises: ConfigError: If branch_id is None. KeboolaApiError: ``INVALID_ARGUMENT`` for a non-positive - ``timeout``, or if the API call fails. + ``timeout``, or if the API call fails. A wait timeout is + ``STORAGE_JOB_TIMEOUT`` (``retryable=False``, details + ``job_id``, message names ``storage job-detail``): the swap + keeps running, and repeating it would swap the tables back. """ max_wait = normalize_job_timeout(timeout, TABLE_DATA_JOB_MAX_WAIT) @@ -1641,6 +1648,8 @@ def swap_tables( branch_id=branch_id, max_wait=max_wait, ) + except KeboolaApiError as exc: + raise with_job_timeout_hint(exc, alias) from exc finally: client.close() diff --git a/tests/test_server_router_calls.py b/tests/test_server_router_calls.py index 8d910427f..40d426eda 100644 --- a/tests/test_server_router_calls.py +++ b/tests/test_server_router_calls.py @@ -2913,6 +2913,48 @@ def test_swap_tables_route_rejects_non_positive_timeout(tmp_path: Path) -> None: storage_svc.swap_tables.assert_not_called() +@pytest.mark.parametrize( + ("service_method", "path", "body"), + [ + ( + "create_table", + f"/storage/tables/{PROJECT}", + {"bucket_id": "in.c-b", "name": "t", "source_table_id": "in.c-b.src"}, + ), + ( + "swap_tables", + f"/storage/tables/{PROJECT}/in.c-b.a/swap", + {"target_table_id": "in.c-b.b", "branch_id": 42}, + ), + ], +) +def test_table_job_timeout_route_keeps_code_and_follow_up( + tmp_path: Path, service_method: str, path: str, body: dict +) -> None: + """A wait timeout reaches the REST caller as STORAGE_JOB_TIMEOUT with the job-detail command.""" + storage_svc = MagicMock() + getattr(storage_svc, service_method).side_effect = KeboolaApiError( + message=( + "Storage job 777 did not finish within 300s. Waiting stopped locally; " + "the job keeps running server-side. Follow it with: kbagent storage " + f"job-detail --project {PROJECT} --job-id 777 --wait" + ), + status_code=0, + error_code=ErrorCode.STORAGE_JOB_TIMEOUT, + retryable=False, + details={"job_id": 777}, + ) + app = _make_app_with_registry(tmp_path, _mock_registry(storage=storage_svc)) + + with TestClient(app) as client: + res = client.post(path, headers=AUTH, json=body) + + error = res.json()["error"] + assert error["code"] == ErrorCode.STORAGE_JOB_TIMEOUT + assert "storage job-detail" in error["message"] + assert "--job-id 777" in error["message"] + + # --------------------------------------------------------------------------- # merge_requests.py -- the REST mirror of `kbagent merge-request *` (DMD-1900) # Service: MergeRequestService.(project, merge_request_id, ...) diff --git a/tests/test_storage_create_table.py b/tests/test_storage_create_table.py index c681df890..5b071302b 100644 --- a/tests/test_storage_create_table.py +++ b/tests/test_storage_create_table.py @@ -18,6 +18,7 @@ from keboola_agent_cli.cli import app from keboola_agent_cli.client import KeboolaClient from keboola_agent_cli.config_store import ConfigStore +from keboola_agent_cli.errors import ErrorCode, KeboolaApiError from keboola_agent_cli.models import AppConfig, ProjectConfig from keboola_agent_cli.services.storage_service import StorageService @@ -176,6 +177,37 @@ def test_forwards_budget_to_the_poller(self, httpx_mock) -> None: assert poller.call_args.kwargs["max_wait"] == 900.0 + def test_timeout_is_not_retryable(self, httpx_mock) -> None: + """A --source-table-id copy that outlives the budget keeps running. + + A retry would start a second copy next to it, so the timeout is NOT + retryable and carries the job id. A budget already in the past + reaches the timeout branch without a poll (see the swap-tables twin + in test_storage_swap.py for why a tiny positive budget is flaky). + """ + httpx_mock.add_response( + url="https://connection.keboola.com/v2/storage/buckets/in.c-main/tables-definition", + method="POST", + json={"id": 555, "status": "waiting"}, + status_code=200, + ) + client = KeboolaClient(stack_url="https://connection.keboola.com", token=TEST_TOKEN) + with pytest.raises(KeboolaApiError) as exc_info: + client.create_table( + bucket_id="in.c-main", + name="events_repart", + source={"tableId": "in.c-main.events"}, + max_wait=-1.0, + ) + client.close() + + exc = exc_info.value + assert exc.error_code == ErrorCode.STORAGE_JOB_TIMEOUT + assert exc.retryable is False + assert exc.details["job_id"] == 555 + assert "keeps running server-side" in exc.message + assert "second copy" in exc.message + # --------------------------------------------------------------------------- # Service layer -- validation diff --git a/tests/test_storage_swap.py b/tests/test_storage_swap.py index f31d0a0ee..a698bd4b8 100644 --- a/tests/test_storage_swap.py +++ b/tests/test_storage_swap.py @@ -459,6 +459,17 @@ def test_swap_missing_branch_fails_clearly(self, tmp_path: Path) -> None: assert "requires a branch" in payload["error"]["message"] +def _swap_timeout() -> KeboolaApiError: + """The client's swap timeout: not retryable, ``job_id`` in the details.""" + return KeboolaApiError( + message="Storage swap-tables job 777 did not finish within 300s.", + status_code=0, + error_code=ErrorCode.STORAGE_JOB_TIMEOUT, + retryable=False, + details={"job_id": 777}, + ) + + # --------------------------------------------------------------------------- # Job wait budget (issue #713) # --------------------------------------------------------------------------- @@ -582,8 +593,8 @@ def test_expired_budget_reports_the_job_as_still_running(self, httpx_mock) -> No """An exhausted budget raises STORAGE_JOB_TIMEOUT naming the live job. The message must not read as "nothing happened": kbagent stops - watching, the swap keeps running server-side. That is also why the - code is retryable (exit 4), not a plain failure. + watching, the swap keeps running server-side. It is NOT retryable: + a repeat after the swap lands swaps the tables back. Exit code 4. A budget already in the past puts the deadline behind the very first check, so the timeout branch is reached without a poll, a sleep, or @@ -614,9 +625,66 @@ def test_expired_budget_reports_the_job_as_still_running(self, httpx_mock) -> No exc = exc_info.value assert exc.error_code == ErrorCode.STORAGE_JOB_TIMEOUT - assert exc.retryable is True + assert exc.retryable is False + assert exc.details["job_id"] == 777 assert "777" in exc.message - assert "continues running" in exc.message + assert "keeps running server-side" in exc.message + assert "swaps the tables back" in exc.message + + def test_service_timeout_names_job_detail(self, tmp_path: Path) -> None: + """The service appends the follow-up command; the client knows no alias.""" + store = _make_store(tmp_path) + mock_client = MagicMock() + mock_client.swap_tables.side_effect = _swap_timeout() + service = _make_service(store, mock_client) + + with pytest.raises(KeboolaApiError) as exc_info: + service.swap_tables( + alias="test", + table_id="in.c-foo.a", + target_table_id="in.c-foo.b", + branch_id=42, + ) + + exc = exc_info.value + assert exc.error_code == ErrorCode.STORAGE_JOB_TIMEOUT + assert exc.retryable is False + assert exc.details["job_id"] == 777 + assert "kbagent storage job-detail --project test --job-id 777 --wait" in exc.message + mock_client.close.assert_called_once() + + def test_cli_timeout_is_exit_4_not_retryable(self, tmp_path: Path) -> None: + store = _make_store(tmp_path) + + with ( + patch("keboola_agent_cli.cli.ConfigStore") as MockStore, + patch("keboola_agent_cli.cli.StorageService") as MockSvc, + ): + MockStore.return_value = store + MockSvc.return_value.swap_tables.side_effect = _swap_timeout() + result = runner.invoke( + app, + [ + "--json", + "storage", + "swap-tables", + "--project", + "test", + "--table-id", + "in.c-foo.a", + "--target-table-id", + "in.c-foo.b", + "--branch", + "42", + "--yes", + ], + ) + + assert result.exit_code == 4 + error = json.loads(result.output)["error"] + assert error["code"] == ErrorCode.STORAGE_JOB_TIMEOUT + assert error["retryable"] is False + assert error["details"]["job_id"] == 777 def test_cli_forwards_timeout(self, tmp_path: Path) -> None: store = _make_store(tmp_path) diff --git a/tests/test_storage_write.py b/tests/test_storage_write.py index 43570755a..e9a9596dc 100644 --- a/tests/test_storage_write.py +++ b/tests/test_storage_write.py @@ -2599,3 +2599,63 @@ def test_cli_forwards_timeout(self, tmp_path: Path) -> None: assert result.exit_code == 0 assert svc.create_table.call_args.kwargs["timeout"] == 1800.0 + + def test_timeout_names_job_detail_and_table_detail(self, tmp_path: Path) -> None: + """The service appends the follow-up commands; the client knows no alias.""" + store = _make_store(tmp_path) + mock_client = MagicMock() + mock_client.create_table.side_effect = _create_timeout() + service = _make_service(store, mock_client) + + with pytest.raises(KeboolaApiError) as exc_info: + service.create_table(alias="test", bucket_id="in.c-b", name="t", columns=["id:INTEGER"]) + + exc = exc_info.value + assert exc.error_code == ErrorCode.STORAGE_JOB_TIMEOUT + assert exc.retryable is False + assert exc.details["job_id"] == 555 + assert "kbagent storage job-detail --project test --job-id 555 --wait" in exc.message + assert "kbagent storage table-detail --project test --table-id in.c-b.t" in exc.message + + def test_cli_timeout_is_exit_4_not_retryable(self, tmp_path: Path) -> None: + store = _make_store(tmp_path) + + with ( + patch("keboola_agent_cli.cli.ConfigStore") as MockStore, + patch("keboola_agent_cli.cli.StorageService") as MockSvc, + ): + MockStore.return_value = store + MockSvc.return_value.create_table.side_effect = _create_timeout() + result = runner.invoke( + app, + [ + "--json", + "storage", + "create-table", + "--project", + "test", + "--bucket-id", + "in.c-b", + "--name", + "t", + "--column", + "id:INTEGER", + ], + ) + + assert result.exit_code == 4 + error = json.loads(result.output)["error"] + assert error["code"] == ErrorCode.STORAGE_JOB_TIMEOUT + assert error["retryable"] is False + assert error["details"]["job_id"] == 555 + + +def _create_timeout() -> KeboolaApiError: + """The client's create-table timeout: not retryable, ``job_id`` in the details.""" + return KeboolaApiError( + message="Storage create-table job 555 did not finish within 300s.", + status_code=0, + error_code=ErrorCode.STORAGE_JOB_TIMEOUT, + retryable=False, + details={"job_id": 555}, + ) From 60e5a03c24949a74f9a5ad3e9b99e70d6bdfa0e4 Mon Sep 17 00:00:00 2001 From: soustruh Date: Sat, 10 Oct 2026 03:55:42 +0200 Subject: [PATCH 5/7] fix(storage): name a second copy only when create-table has a source table --- src/keboola_agent_cli/client/storage_tables.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/keboola_agent_cli/client/storage_tables.py b/src/keboola_agent_cli/client/storage_tables.py index 9bbaa8417..5a46fe58a 100644 --- a/src/keboola_agent_cli/client/storage_tables.py +++ b/src/keboola_agent_cli/client/storage_tables.py @@ -637,9 +637,11 @@ def create_table( budget=STORAGE_JOB_MAX_WAIT if max_wait is None else max_wait, operation="create-table", warning=( - ". Re-running the command starts a second create (with a source " - "table, a second copy) -- check the job before retrying." - ), + ". Re-running the command starts a second copy of the source table" + if source is not None + else ". Re-running the command starts a second create job" + ) + + " -- check the job before retrying.", ) return job.get("results", {}) From e0b382eb8a47d0eb554ea81892ef430bf4cdee70 Mon Sep 17 00:00:00 2001 From: soustruh Date: Sat, 10 Oct 2026 03:56:24 +0200 Subject: [PATCH 6/7] fix(storage): reject a NaN or infinite job timeout --- src/keboola_agent_cli/services/base.py | 7 ++++--- tests/test_base_service.py | 8 ++++++++ 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/src/keboola_agent_cli/services/base.py b/src/keboola_agent_cli/services/base.py index e756db80a..7b7a0dd5a 100644 --- a/src/keboola_agent_cli/services/base.py +++ b/src/keboola_agent_cli/services/base.py @@ -5,6 +5,7 @@ """ import logging +import math import os from collections.abc import Callable from concurrent.futures import ThreadPoolExecutor, as_completed @@ -107,13 +108,13 @@ def normalize_job_timeout(timeout: float | None, default: float) -> float: The resolved budget in seconds. Raises: - KeboolaApiError: ``INVALID_ARGUMENT`` when ``timeout`` is <= 0. + KeboolaApiError: ``INVALID_ARGUMENT`` when ``timeout`` is <= 0, NaN or infinite. """ if timeout is None: return default - if timeout <= 0: + if not math.isfinite(timeout) or timeout <= 0: raise KeboolaApiError( - message=f"Invalid timeout {timeout}. Must be greater than 0.", + message=f"Invalid timeout {timeout}. Must be a finite number greater than 0.", status_code=0, error_code=ErrorCode.INVALID_ARGUMENT, ) diff --git a/tests/test_base_service.py b/tests/test_base_service.py index 951bd0dee..c93985c7b 100644 --- a/tests/test_base_service.py +++ b/tests/test_base_service.py @@ -568,3 +568,11 @@ def test_non_positive_is_rejected(self, bad: float) -> None: assert exc_info.value.error_code == ErrorCode.INVALID_ARGUMENT assert str(bad) in exc_info.value.message + + @pytest.mark.parametrize("bad", [float("nan"), float("inf")]) + def test_non_finite_is_rejected(self, bad: float) -> None: + """NaN and infinity would make the poller's deadline never expire.""" + with pytest.raises(KeboolaApiError) as exc_info: + normalize_job_timeout(bad, 300.0) + + assert exc_info.value.error_code == ErrorCode.INVALID_ARGUMENT From 5045a6f9d43fa7d7055106e995b0494447f99e84 Mon Sep 17 00:00:00 2001 From: soustruh Date: Sat, 10 Oct 2026 08:58:51 +0200 Subject: [PATCH 7/7] fix(storage): one timeout validator for all storage waits, reject Infinity and NaN with a 422 --- .../skills/kbagent/references/gotchas.md | 6 +- src/keboola_agent_cli/commands/storage.py | 17 +++-- src/keboola_agent_cli/commands/workspace.py | 12 +-- src/keboola_agent_cli/server/app.py | 18 +++++ .../server/routers/storage.py | 12 +-- .../server/routers/workspaces.py | 3 +- .../services/_storage_jobs.py | 13 ++++ src/keboola_agent_cli/services/base.py | 34 +-------- .../services/storage_service.py | 21 ++--- .../services/workspace_service.py | 15 ++-- tests/test_base_service.py | 34 --------- tests/test_server_router_calls.py | 76 ++++++------------- tests/test_storage_import_jobs.py | 22 +++++- tests/test_storage_swap.py | 53 +++++++++++-- tests/test_storage_write.py | 45 +++++++++-- tests/test_workspace_cli.py | 11 ++- tests/test_workspace_service.py | 18 ++--- 17 files changed, 227 insertions(+), 183 deletions(-) diff --git a/plugins/kbagent/skills/kbagent/references/gotchas.md b/plugins/kbagent/skills/kbagent/references/gotchas.md index 08b47edd9..59245efc7 100644 --- a/plugins/kbagent/skills/kbagent/references/gotchas.md +++ b/plugins/kbagent/skills/kbagent/references/gotchas.md @@ -2158,8 +2158,10 @@ config, the retry fires, and the retry destroys it for good. `STORAGE_JOB_TIMEOUT`. - **Since vNEXT both default to 300s and accept `--timeout SECONDS`.** The REST routes (`POST /storage/tables/{p}`, `POST /storage/tables/{p}/{tid}/swap`) - take an optional `timeout` field; a non-positive value is a 422 at the - boundary and `INVALID_ARGUMENT` from the service. + take an optional `timeout` field; a zero, negative, NaN or infinite value + is a 422. On the CLI the same value is a usage error: exit 2, + `INVALID_ARGUMENT`, before `swap-tables` asks for confirmation -- the + same as `upload-table`, `load-file` and `workspace load`. - **Never read the timeout as "nothing happened".** kbagent only stops *watching*; the Storage job keeps running server-side (and keeps consuming backend resources). It usually lands. The error is **exit 4** with diff --git a/src/keboola_agent_cli/commands/storage.py b/src/keboola_agent_cli/commands/storage.py index eec440f63..d017c9716 100644 --- a/src/keboola_agent_cli/commands/storage.py +++ b/src/keboola_agent_cli/commands/storage.py @@ -14,6 +14,7 @@ from ..constants import TABLE_DATA_JOB_MAX_WAIT from ..effective_branch import record_branch, resolve_branch from ..errors import ConfigError, ErrorCode, KeboolaApiError +from ..services._storage_jobs import validate_wait_timeout from ._helpers import ( check_cli_permission, emit_project_warnings, @@ -1476,6 +1477,14 @@ def storage_swap_tables( config_store: ConfigStore = ctx.obj["config_store"] effective_branch = resolve_branch(config_store, project, branch, required=True) + try: + # Before the confirmation prompt: a bad budget is a usage error, and + # the user must not confirm a destructive swap only to be refused. + validate_wait_timeout(timeout) + except ValueError as exc: + formatter.error(message=str(exc), error_code=ErrorCode.INVALID_ARGUMENT) + raise typer.Exit(code=2) from None + if dry_run: try: result = service.swap_tables( @@ -1489,14 +1498,6 @@ def storage_swap_tables( except ConfigError as exc: formatter.error(message=exc.message, error_code=ErrorCode.CONFIG_ERROR) raise typer.Exit(code=5) from None - except KeboolaApiError as exc: - formatter.error( - message=exc.message, - error_code=exc.error_code, - project=project, - retryable=exc.retryable, - ) - raise typer.Exit(code=map_error_to_exit_code(exc)) from None if formatter.json_mode: formatter.output(result) diff --git a/src/keboola_agent_cli/commands/workspace.py b/src/keboola_agent_cli/commands/workspace.py index 35f63b6e9..9120267c9 100644 --- a/src/keboola_agent_cli/commands/workspace.py +++ b/src/keboola_agent_cli/commands/workspace.py @@ -19,6 +19,7 @@ from ..effective_branch import resolve_branch from ..errors import ConfigError, ErrorCode, KeboolaApiError from ..output import OutputFormatter, format_query_results, format_workspaces_table +from ..services._storage_jobs import validate_wait_timeout from ..services._workspace_load_plan import LoadTablePlan from ._helpers import ( check_cli_permission, @@ -425,14 +426,13 @@ def workspace_load( ) raise typer.Exit(code=2) - if timeout <= 0: + try: # Falling back to the default here would silently ignore what the # caller asked for; a zero budget cannot mean "wait forever" either. - formatter.error( - error_code=ErrorCode.INVALID_ARGUMENT, - message=f"Invalid --timeout {timeout}. Must be greater than 0.", - ) - raise typer.Exit(code=2) + validate_wait_timeout(timeout) + except ValueError as exc: + formatter.error(error_code=ErrorCode.INVALID_ARGUMENT, message=str(exc)) + raise typer.Exit(code=2) from None # No prompt in --json mode: there is nobody to answer it, and a machine # caller must get the structured refusal instead of a silent large COPY. diff --git a/src/keboola_agent_cli/server/app.py b/src/keboola_agent_cli/server/app.py index 07354ac02..331fb451d 100644 --- a/src/keboola_agent_cli/server/app.py +++ b/src/keboola_agent_cli/server/app.py @@ -16,6 +16,7 @@ import asyncio import contextlib import logging +import math import secrets import time from contextlib import asynccontextmanager @@ -23,6 +24,8 @@ from urllib.parse import parse_qs from fastapi import Depends, FastAPI +from fastapi.encoders import jsonable_encoder +from fastapi.exceptions import RequestValidationError from fastapi.middleware.cors import CORSMiddleware from fastapi.openapi.utils import get_openapi from fastapi.responses import JSONResponse @@ -549,6 +552,12 @@ def _format_error( {ErrorCode.MR_MERGE_CONFLICT, ErrorCode.MR_NOT_READY_TO_MERGE} ) + +def _finite_or_text(value: float) -> float | str: + """``value`` unchanged when finite, else its text (``"inf"``, ``"nan"``).""" + return value if math.isfinite(value) else str(value) + + _SESSION_REMEDY_ON_HOST = ( "Complete `kbagent auth login` on the host running `kbagent serve` -- this server " "cannot open a browser login for a remote caller." @@ -953,6 +962,15 @@ async def _permission_denied_handler(_request, exc: PermissionDeniedError): # a caller can branch on one value across both surfaces. return _format_error(exc.message, ErrorCode.PERMISSION_DENIED, http_status=403) + @app.exception_handler(RequestValidationError) + async def _validation_error_handler(_request, exc: RequestValidationError): + # FastAPI's own handler, except for one case: a JSON body may carry + # Infinity or NaN (Python's json accepts both), the error echoes it as + # `input`, and JSONResponse cannot encode a non-finite float -- so the + # 422 crashed into a 500. Same body shape, the bad float as text. + detail = jsonable_encoder(exc.errors(), custom_encoder={float: _finite_or_text}) + return JSONResponse(status_code=422, content={"detail": detail}) + @app.exception_handler(StarletteHTTPException) async def _starlette_handler(_request, exc: StarletteHTTPException): return _format_error( diff --git a/src/keboola_agent_cli/server/routers/storage.py b/src/keboola_agent_cli/server/routers/storage.py index 2f789dcbe..79e6c53d5 100644 --- a/src/keboola_agent_cli/server/routers/storage.py +++ b/src/keboola_agent_cli/server/routers/storage.py @@ -49,9 +49,9 @@ class CreateTable(BaseModel): range_partitioning_end: str | None = None range_partitioning_interval: str | None = None clustering_fields: list[str] | None = None - # `gt=0` mirrors the CLI guard so a bad budget is a 422 at the boundary - # rather than an INVALID_ARGUMENT raised from the service. - timeout: float | None = Field(default=None, gt=0) + # `gt=0` + `allow_inf_nan=False` mirror the CLI guard, so a bad budget is + # a 422 at the boundary rather than a ValueError raised from the service. + timeout: float | None = Field(default=None, gt=0, allow_inf_nan=False) @model_validator(mode="after") def _columns_xor_source(self) -> CreateTable: @@ -119,9 +119,9 @@ class LoadFileToTable(BaseModel): class SwapTables(BaseModel): target_table_id: str branch_id: int - # `gt=0` mirrors the CLI guard so a bad budget is a 422 at the boundary - # rather than an INVALID_ARGUMENT raised from the service. - timeout: float | None = Field(default=None, gt=0) + # `gt=0` + `allow_inf_nan=False` mirror the CLI guard, so a bad budget is + # a 422 at the boundary rather than a ValueError raised from the service. + timeout: float | None = Field(default=None, gt=0, allow_inf_nan=False) class CloneTable(BaseModel): diff --git a/src/keboola_agent_cli/server/routers/workspaces.py b/src/keboola_agent_cli/server/routers/workspaces.py index 563e5c41a..bd518d69f 100644 --- a/src/keboola_agent_cli/server/routers/workspaces.py +++ b/src/keboola_agent_cli/server/routers/workspaces.py @@ -33,7 +33,8 @@ class WorkspaceLoad(BaseModel): # WORKSPACE_LOAD_COPY_TOO_LARGE) rather than started unannounced. force: bool = False # gt=0: a zero/negative budget would make every load "time out" instantly. - timeout: float | None = Field(default=None, gt=0) + # allow_inf_nan=False: NaN or infinity would make the wait never end. + timeout: float | None = Field(default=None, gt=0, allow_inf_nan=False) class WorkspaceQuery(BaseModel): diff --git a/src/keboola_agent_cli/services/_storage_jobs.py b/src/keboola_agent_cli/services/_storage_jobs.py index 149364886..62dfaf984 100644 --- a/src/keboola_agent_cli/services/_storage_jobs.py +++ b/src/keboola_agent_cli/services/_storage_jobs.py @@ -64,6 +64,19 @@ def validate_wait_timeout(timeout: float | None) -> None: raise ValueError(f"--timeout must be a positive number of seconds, got {timeout}.") +def resolve_wait_timeout(timeout: float | None, default: float) -> float: + """Validate a wait budget, then resolve ``None`` to ``default``. + + Branches on ``is None`` -- NOT ``timeout or default``, which would turn + a falsy-but-invalid ``0.0`` into the default instead of rejecting it. + + Raises: + ValueError: ``timeout`` is zero, negative, NaN or infinite. + """ + validate_wait_timeout(timeout) + return default if timeout is None else timeout + + def import_job_fields(job: dict[str, Any], file_id: Any) -> dict[str, Any]: """The keys an import result adds for its job (upload-table, load-file). diff --git a/src/keboola_agent_cli/services/base.py b/src/keboola_agent_cli/services/base.py index 7b7a0dd5a..d2aebdcb1 100644 --- a/src/keboola_agent_cli/services/base.py +++ b/src/keboola_agent_cli/services/base.py @@ -5,7 +5,6 @@ """ import logging -import math import os from collections.abc import Callable from concurrent.futures import ThreadPoolExecutor, as_completed @@ -25,7 +24,7 @@ ENV_MAX_PARALLEL_WORKERS, UNEXPECTED_ERROR_MAX_MESSAGE_LEN, ) -from ..errors import ErrorCode, KeboolaApiError +from ..errors import ErrorCode from ..models import ProjectConfig from ..project_ref import resolve_project_ref @@ -90,37 +89,6 @@ def sanitize_unexpected_error(exc: BaseException) -> str: return raw -def normalize_job_timeout(timeout: float | None, default: float) -> float: - """Resolve an async-job wait budget, rejecting a non-positive override. - - ``None`` is the documented "use the default" sentinel and branches - explicitly -- deliberately NOT ``timeout or default``, which would - silently promote a falsy-but-real ``0.0`` to the default instead of - rejecting it. Every CLI command already constrains its own ``--timeout``, - but a direct service or SDK caller goes through neither guard, so the - check belongs here rather than in each command. - - Args: - timeout: Caller-supplied budget in seconds, or ``None`` for ``default``. - default: Budget to use when ``timeout`` is ``None``. - - Returns: - The resolved budget in seconds. - - Raises: - KeboolaApiError: ``INVALID_ARGUMENT`` when ``timeout`` is <= 0, NaN or infinite. - """ - if timeout is None: - return default - if not math.isfinite(timeout) or timeout <= 0: - raise KeboolaApiError( - message=f"Invalid timeout {timeout}. Must be a finite number greater than 0.", - status_code=0, - error_code=ErrorCode.INVALID_ARGUMENT, - ) - return timeout - - def project_error_entry( alias: str, exc: BaseException, diff --git a/src/keboola_agent_cli/services/storage_service.py b/src/keboola_agent_cli/services/storage_service.py index 5e43f7d83..eac33fd93 100644 --- a/src/keboola_agent_cli/services/storage_service.py +++ b/src/keboola_agent_cli/services/storage_service.py @@ -25,6 +25,7 @@ ImportOptions, import_job_fields, read_csv_header, + resolve_wait_timeout, summarize_storage_job, validate_wait_timeout, with_import_hint, @@ -32,7 +33,6 @@ ) from ._storage_tables import normalize_table_rows from ._table_detail import build_table_detail -from .base import normalize_job_timeout from .table_usage import collect_table_usage, fetch_usage_components logger = logging.getLogger(__name__) @@ -784,18 +784,19 @@ def create_table( ``schema_drift`` is ``True`` when the two diverge. Raises: - KeboolaApiError: ``INVALID_ARGUMENT`` for a non-positive ``timeout``; - ``STORAGE_JOB_TIMEOUT`` (``retryable=False``, details ``job_id``, - message names ``storage job-detail`` and ``storage table-detail``) - when the wait runs out -- the create keeps running server-side. - ValueError: Malformed column spec or ``--default`` assignment; + KeboolaApiError: ``STORAGE_JOB_TIMEOUT`` (``retryable=False``, + details ``job_id``, message names ``storage job-detail`` and + ``storage table-detail``) when the wait runs out -- the create + keeps running server-side. + ValueError: ``timeout`` is zero, negative, NaN or infinite; + malformed column spec or ``--default`` assignment; ``--not-null`` / ``--default`` references an unknown column; ``columns`` and ``source`` both/neither given; ``source_branch_id`` given without ``source_table_id``; incomplete or conflicting partitioning flags; or BigQuery-only features requested on a non-BigQuery backend. """ - max_wait = normalize_job_timeout(timeout, TABLE_DATA_JOB_MAX_WAIT) + max_wait = resolve_wait_timeout(timeout, TABLE_DATA_JOB_MAX_WAIT) not_null_set = set(not_null_columns or []) defaults_map = _parse_default_assignments(defaults) @@ -1606,13 +1607,13 @@ def swap_tables( Raises: ConfigError: If branch_id is None. - KeboolaApiError: ``INVALID_ARGUMENT`` for a non-positive - ``timeout``, or if the API call fails. A wait timeout is + ValueError: ``timeout`` is zero, negative, NaN or infinite. + KeboolaApiError: If the API call fails. A wait timeout is ``STORAGE_JOB_TIMEOUT`` (``retryable=False``, details ``job_id``, message names ``storage job-detail``): the swap keeps running, and repeating it would swap the tables back. """ - max_wait = normalize_job_timeout(timeout, TABLE_DATA_JOB_MAX_WAIT) + max_wait = resolve_wait_timeout(timeout, TABLE_DATA_JOB_MAX_WAIT) if branch_id is None: raise ConfigError( diff --git a/src/keboola_agent_cli/services/workspace_service.py b/src/keboola_agent_cli/services/workspace_service.py index bf308768c..c2e91af59 100644 --- a/src/keboola_agent_cli/services/workspace_service.py +++ b/src/keboola_agent_cli/services/workspace_service.py @@ -31,6 +31,7 @@ from ..effective_branch import record_branch, resolve_branch from ..errors import ConfigError, ErrorCode, KeboolaApiError from ..models import ProjectConfig +from ._storage_jobs import resolve_wait_timeout from ._workspace_load_plan import ( LOAD_TYPE_CLONE, LOAD_TYPE_COPY, @@ -39,7 +40,7 @@ coerce_data_size_bytes, plan_auto_load_type, ) -from .base import BaseService, find_default_branch_id, normalize_job_timeout +from .base import BaseService, find_default_branch_id logger = logging.getLogger(__name__) @@ -856,7 +857,8 @@ def load_tables( or None for the auto decision. force: Skip the large-COPY size guard. timeout: Seconds to wait for the load job. ``None`` defaults to - WORKSPACE_LOAD_JOB_MAX_WAIT; any other value must be > 0. + WORKSPACE_LOAD_JOB_MAX_WAIT; any other value must be a finite + number > 0. on_copy_guard: Called with the oversized COPY plans; return True to proceed, False to refuse. None means "refuse without asking". @@ -864,12 +866,13 @@ def load_tables( Dict with load job results, including a per-table ``tables`` list. Raises: - KeboolaApiError: INVALID_ARGUMENT for an unknown load_type or a - non-positive timeout; WORKSPACE_LOAD_COPY_TOO_LARGE when the - size guard trips and is neither forced nor approved. + ValueError: ``timeout`` is zero, negative, NaN or infinite. + KeboolaApiError: INVALID_ARGUMENT for an unknown load_type; + WORKSPACE_LOAD_COPY_TOO_LARGE when the size guard trips and is + neither forced nor approved. """ requested = self._normalize_load_type(load_type) - max_wait = normalize_job_timeout(timeout, WORKSPACE_LOAD_JOB_MAX_WAIT) + max_wait = resolve_wait_timeout(timeout, WORKSPACE_LOAD_JOB_MAX_WAIT) projects = self.resolve_projects([alias]) project = projects[alias] diff --git a/tests/test_base_service.py b/tests/test_base_service.py index c93985c7b..21f06b822 100644 --- a/tests/test_base_service.py +++ b/tests/test_base_service.py @@ -26,7 +26,6 @@ ENV_MAX_PARALLEL_WORKERS, UNEXPECTED_ERROR_CODE, BaseService, - normalize_job_timeout, project_error_entry, ) @@ -543,36 +542,3 @@ def test_custom_client_factory_is_used(self, tmp_config_dir: Path) -> None: service = _TestService(config_store=store, client_factory=mock_factory) assert service._client_factory is mock_factory - - -class TestNormalizeJobTimeout: - """Shared async-job budget guard used by every service that waits on a job.""" - - def test_none_uses_the_default(self) -> None: - assert normalize_job_timeout(None, 300.0) == 300.0 - - def test_positive_override_wins(self) -> None: - assert normalize_job_timeout(900.0, 300.0) == 900.0 - - @pytest.mark.parametrize("bad", [0.0, -1.0, -900.0]) - def test_non_positive_is_rejected(self, bad: float) -> None: - """0.0 must raise, not fall back to the default. - - This is the whole reason the guard branches on ``is None`` rather - than writing ``timeout or default``: 0.0 is falsy, so the shorter - spelling would silently accept an invalid budget as "unset" and every - caller would look correct while validating nothing. - """ - with pytest.raises(KeboolaApiError) as exc_info: - normalize_job_timeout(bad, 300.0) - - assert exc_info.value.error_code == ErrorCode.INVALID_ARGUMENT - assert str(bad) in exc_info.value.message - - @pytest.mark.parametrize("bad", [float("nan"), float("inf")]) - def test_non_finite_is_rejected(self, bad: float) -> None: - """NaN and infinity would make the poller's deadline never expire.""" - with pytest.raises(KeboolaApiError) as exc_info: - normalize_job_timeout(bad, 300.0) - - assert exc_info.value.error_code == ErrorCode.INVALID_ARGUMENT diff --git a/tests/test_server_router_calls.py b/tests/test_server_router_calls.py index 40d426eda..073d9d327 100644 --- a/tests/test_server_router_calls.py +++ b/tests/test_server_router_calls.py @@ -2754,16 +2754,23 @@ def test_workspace_load_route_passes_load_type_force_timeout(tmp_path: Path) -> assert kwargs["timeout"] == 900 -def test_workspace_load_route_rejects_non_positive_timeout(tmp_path: Path) -> None: - """A zero/negative budget would make every load 'time out' instantly.""" +# A JSON number Python's json module accepts: 0 and -1 fail `gt=0`, Infinity +# and NaN fail `allow_inf_nan=False`. Sent as raw text because httpx refuses +# to encode a non-finite float. +_BAD_TIMEOUTS = pytest.mark.parametrize("bad", ["0", "-1", "Infinity", "NaN"]) + + +@_BAD_TIMEOUTS +def test_workspace_load_route_rejects_bad_timeout(tmp_path: Path, bad: str) -> None: + """A zero/negative budget times out at once, a non-finite one never does.""" ws_svc = MagicMock() app = _make_app_with_registry(tmp_path, _mock_registry(workspace=ws_svc)) with TestClient(app) as client: res = client.post( f"/workspaces/{PROJECT}/42/load", - headers=AUTH, - json={"tables": [TABLE_ID], "timeout": 0}, + headers={**AUTH, "Content-Type": "application/json"}, + content=f'{{"tables": ["{TABLE_ID}"], "timeout": {bad}}}', ) assert res.status_code == 422, res.text @@ -2866,16 +2873,20 @@ def test_create_table_route_forwards_timeout(tmp_path: Path) -> None: assert storage_svc.create_table.call_args.kwargs["timeout"] == 1800 -def test_create_table_route_rejects_non_positive_timeout(tmp_path: Path) -> None: - """A zero budget would make every create 'time out' instantly.""" +@_BAD_TIMEOUTS +def test_create_table_route_rejects_bad_timeout(tmp_path: Path, bad: str) -> None: + """A zero budget times out at once, a non-finite one never does.""" storage_svc = MagicMock() app = _make_app_with_registry(tmp_path, _mock_registry(storage=storage_svc)) with TestClient(app) as client: res = client.post( f"/storage/tables/{PROJECT}", - headers=AUTH, - json={"bucket_id": "in.c-b", "name": "t", "columns": ["id:INTEGER"], "timeout": 0}, + headers={**AUTH, "Content-Type": "application/json"}, + content=( + '{"bucket_id": "in.c-b", "name": "t", "columns": ["id:INTEGER"], ' + f'"timeout": {bad}}}' + ), ) assert res.status_code == 422, res.text @@ -2898,63 +2909,22 @@ def test_swap_tables_route_forwards_timeout(tmp_path: Path) -> None: assert storage_svc.swap_tables.call_args.kwargs["timeout"] == 900 -def test_swap_tables_route_rejects_non_positive_timeout(tmp_path: Path) -> None: +@_BAD_TIMEOUTS +def test_swap_tables_route_rejects_bad_timeout(tmp_path: Path, bad: str) -> None: storage_svc = MagicMock() app = _make_app_with_registry(tmp_path, _mock_registry(storage=storage_svc)) with TestClient(app) as client: res = client.post( f"/storage/tables/{PROJECT}/in.c-b.a/swap", - headers=AUTH, - json={"target_table_id": "in.c-b.b", "branch_id": 42, "timeout": -1}, + headers={**AUTH, "Content-Type": "application/json"}, + content=f'{{"target_table_id": "in.c-b.b", "branch_id": 42, "timeout": {bad}}}', ) assert res.status_code == 422, res.text storage_svc.swap_tables.assert_not_called() -@pytest.mark.parametrize( - ("service_method", "path", "body"), - [ - ( - "create_table", - f"/storage/tables/{PROJECT}", - {"bucket_id": "in.c-b", "name": "t", "source_table_id": "in.c-b.src"}, - ), - ( - "swap_tables", - f"/storage/tables/{PROJECT}/in.c-b.a/swap", - {"target_table_id": "in.c-b.b", "branch_id": 42}, - ), - ], -) -def test_table_job_timeout_route_keeps_code_and_follow_up( - tmp_path: Path, service_method: str, path: str, body: dict -) -> None: - """A wait timeout reaches the REST caller as STORAGE_JOB_TIMEOUT with the job-detail command.""" - storage_svc = MagicMock() - getattr(storage_svc, service_method).side_effect = KeboolaApiError( - message=( - "Storage job 777 did not finish within 300s. Waiting stopped locally; " - "the job keeps running server-side. Follow it with: kbagent storage " - f"job-detail --project {PROJECT} --job-id 777 --wait" - ), - status_code=0, - error_code=ErrorCode.STORAGE_JOB_TIMEOUT, - retryable=False, - details={"job_id": 777}, - ) - app = _make_app_with_registry(tmp_path, _mock_registry(storage=storage_svc)) - - with TestClient(app) as client: - res = client.post(path, headers=AUTH, json=body) - - error = res.json()["error"] - assert error["code"] == ErrorCode.STORAGE_JOB_TIMEOUT - assert "storage job-detail" in error["message"] - assert "--job-id 777" in error["message"] - - # --------------------------------------------------------------------------- # merge_requests.py -- the REST mirror of `kbagent merge-request *` (DMD-1900) # Service: MergeRequestService.(project, merge_request_id, ...) diff --git a/tests/test_storage_import_jobs.py b/tests/test_storage_import_jobs.py index c8e623e7d..d454bac44 100644 --- a/tests/test_storage_import_jobs.py +++ b/tests/test_storage_import_jobs.py @@ -33,7 +33,7 @@ from keboola_agent_cli.config_store import ConfigStore from keboola_agent_cli.errors import ErrorCode, KeboolaApiError from keboola_agent_cli.models import AppConfig, ProjectConfig -from keboola_agent_cli.services._storage_jobs import read_csv_header +from keboola_agent_cli.services._storage_jobs import read_csv_header, resolve_wait_timeout from keboola_agent_cli.services.storage_service import StorageService runner = CliRunner() @@ -370,6 +370,26 @@ def test_wait_keeps_failed_job_details(self, tmp_path: Path) -> None: client.follow_storage_job.assert_called_once_with(55, max_wait=10) +class TestResolveWaitTimeout: + """The one wait-budget guard behind every storage and workspace job wait.""" + + def test_none_uses_the_default(self) -> None: + assert resolve_wait_timeout(None, 300.0) == 300.0 + + def test_positive_override_wins(self) -> None: + assert resolve_wait_timeout(900.0, 300.0) == 900.0 + + @pytest.mark.parametrize("bad", [0.0, -1.0, float("nan"), float("inf")]) + def test_bad_budget_is_rejected(self, bad: float) -> None: + """0.0 must raise, not fall back to the default. + + ``timeout or default`` would accept the falsy 0.0 as "unset". NaN and + infinity would make the poller's deadline never expire. + """ + with pytest.raises(ValueError, match="--timeout must be a positive number"): + resolve_wait_timeout(bad, 300.0) + + class TestReadCsvHeader: def test_plain_csv_with_enclosure_and_delimiter(self, tmp_path: Path) -> None: f = tmp_path / "a.csv" diff --git a/tests/test_storage_swap.py b/tests/test_storage_swap.py index a698bd4b8..2e8f97fae 100644 --- a/tests/test_storage_swap.py +++ b/tests/test_storage_swap.py @@ -520,8 +520,8 @@ def test_service_forwards_explicit_timeout(self, tmp_path: Path) -> None: assert mock_client.swap_tables.call_args.kwargs["max_wait"] == 900.0 - @pytest.mark.parametrize("bad", [0.0, -1.0]) - def test_service_rejects_non_positive_timeout(self, tmp_path: Path, bad: float) -> None: + @pytest.mark.parametrize("bad", [0.0, -1.0, float("nan"), float("inf")]) + def test_service_rejects_bad_timeout(self, tmp_path: Path, bad: float) -> None: """0.0 must be rejected, not silently promoted to the default. ``timeout or DEFAULT`` would treat a falsy-but-real 0 as "unset"; @@ -531,7 +531,7 @@ def test_service_rejects_non_positive_timeout(self, tmp_path: Path, bad: float) mock_client = MagicMock() service = _make_service(store, mock_client) - with pytest.raises(KeboolaApiError) as exc_info: + with pytest.raises(ValueError, match="--timeout must be a positive number"): service.swap_tables( alias="test", table_id="in.c-foo.a", @@ -540,7 +540,6 @@ def test_service_rejects_non_positive_timeout(self, tmp_path: Path, bad: float) timeout=bad, ) - assert exc_info.value.error_code == ErrorCode.INVALID_ARGUMENT mock_client.swap_tables.assert_not_called() def test_service_rejects_non_positive_timeout_on_dry_run(self, tmp_path: Path) -> None: @@ -549,7 +548,7 @@ def test_service_rejects_non_positive_timeout_on_dry_run(self, tmp_path: Path) - mock_client = MagicMock() service = _make_service(store, mock_client) - with pytest.raises(KeboolaApiError): + with pytest.raises(ValueError): service.swap_tables( alias="test", table_id="in.c-foo.a", @@ -686,6 +685,50 @@ def test_cli_timeout_is_exit_4_not_retryable(self, tmp_path: Path) -> None: assert error["retryable"] is False assert error["details"]["job_id"] == 777 + @pytest.mark.parametrize("json_mode", [True, False]) + @pytest.mark.parametrize("bad", ["0", "nan"]) + def test_cli_bad_timeout_is_usage_error_before_prompt( + self, tmp_path: Path, bad: str, json_mode: bool + ) -> None: + """A bad --timeout exits 2 with INVALID_ARGUMENT before the confirmation. + + No --yes: the user must not confirm a destructive swap and only then + learn that the budget is invalid. + """ + store = _make_store(tmp_path) + + with ( + patch("keboola_agent_cli.cli.ConfigStore") as MockStore, + patch("keboola_agent_cli.cli.StorageService") as MockSvc, + ): + MockStore.return_value = store + result = runner.invoke( + app, + [ + *(["--json"] if json_mode else []), + "storage", + "swap-tables", + "--project", + "test", + "--table-id", + "in.c-foo.a", + "--target-table-id", + "in.c-foo.b", + "--branch", + "42", + "--timeout", + bad, + ], + input="y\n", + ) + + assert result.exit_code == 2, result.output + assert "Swap 'in.c-foo.a'" not in result.output + assert "--timeout must be a positive number" in result.output + if json_mode: + assert json.loads(result.output)["error"]["code"] == ErrorCode.INVALID_ARGUMENT + MockSvc.return_value.swap_tables.assert_not_called() + def test_cli_forwards_timeout(self, tmp_path: Path) -> None: store = _make_store(tmp_path) diff --git a/tests/test_storage_write.py b/tests/test_storage_write.py index e9a9596dc..fc96d8598 100644 --- a/tests/test_storage_write.py +++ b/tests/test_storage_write.py @@ -2542,14 +2542,18 @@ def test_forwards_explicit_timeout(self, tmp_path: Path) -> None: assert mock_client.create_table.call_args.kwargs["max_wait"] == 1800.0 - @pytest.mark.parametrize("bad", [0.0, -5.0]) - def test_rejects_non_positive_timeout(self, tmp_path: Path, bad: float) -> None: - """0.0 is rejected rather than silently promoted to the default.""" + @pytest.mark.parametrize("bad", [0.0, -5.0, float("nan"), float("inf")]) + def test_rejects_bad_timeout(self, tmp_path: Path, bad: float) -> None: + """0.0 is rejected rather than silently promoted to the default. + + The ValueError is the same one upload-table and load-file raise, so + every storage command reports a bad --timeout the same way. + """ store = _make_store(tmp_path) mock_client = MagicMock() service = _make_service(store, mock_client) - with pytest.raises(KeboolaApiError) as exc_info: + with pytest.raises(ValueError, match="--timeout must be a positive number"): service.create_table( alias="test", bucket_id="in.c-b", @@ -2558,9 +2562,40 @@ def test_rejects_non_positive_timeout(self, tmp_path: Path, bad: float) -> None: timeout=bad, ) - assert exc_info.value.error_code == ErrorCode.INVALID_ARGUMENT mock_client.create_table.assert_not_called() + @pytest.mark.parametrize("bad", ["0", "nan"]) + def test_cli_bad_timeout_is_usage_error(self, tmp_path: Path, bad: str) -> None: + """A bad --timeout exits 2 with INVALID_ARGUMENT, like upload-table. + + Runs the real service: the exit code depends on which exception type + the service raises, and a mocked service would hide a mismatch. + """ + store = _make_store(tmp_path) + + with patch("keboola_agent_cli.cli.ConfigStore", return_value=store): + result = runner.invoke( + app, + [ + "--json", + "storage", + "create-table", + "--project", + "test", + "--bucket-id", + "in.c-b", + "--name", + "t", + "--column", + "id:INTEGER", + "--timeout", + bad, + ], + ) + + assert result.exit_code == 2, result.output + assert json.loads(result.output)["error"]["code"] == ErrorCode.INVALID_ARGUMENT + def test_cli_forwards_timeout(self, tmp_path: Path) -> None: store = _make_store(tmp_path) diff --git a/tests/test_workspace_cli.py b/tests/test_workspace_cli.py index 4cc5cf0d2..7021ad8c4 100644 --- a/tests/test_workspace_cli.py +++ b/tests/test_workspace_cli.py @@ -11,6 +11,7 @@ from pathlib import Path from unittest.mock import MagicMock, patch +import pytest from typer.testing import CliRunner from keboola_agent_cli.cli import app @@ -880,8 +881,12 @@ def test_workspace_load_invalid_load_type_exits_2(self, tmp_path: Path) -> None: assert output["error"]["code"] == "INVALID_ARGUMENT" mock_ws.load_tables.assert_not_called() - def test_workspace_load_rejects_non_positive_timeout(self, tmp_path: Path) -> None: - """--timeout 0 is rejected, not silently swapped for the default.""" + @pytest.mark.parametrize("bad", ["0", "nan", "inf"]) + def test_workspace_load_rejects_bad_timeout(self, tmp_path: Path, bad: str) -> None: + """--timeout 0 is rejected, not silently swapped for the default. + + NaN and infinity would make the wait never end. + """ config_dir = tmp_path / "config" config_dir.mkdir() store = _setup_config(config_dir, {"prod": {"token": TEST_TOKEN}}) @@ -902,7 +907,7 @@ def test_workspace_load_rejects_non_positive_timeout(self, tmp_path: Path) -> No "--tables", "in.c-main.orders", "--timeout", - "0", + bad, ], ) diff --git a/tests/test_workspace_service.py b/tests/test_workspace_service.py index 1ff1a58af..474171829 100644 --- a/tests/test_workspace_service.py +++ b/tests/test_workspace_service.py @@ -1262,25 +1262,23 @@ def test_load_tables_timeout_zero_rejected(self, tmp_config_dir: Path) -> None: store = setup_single_project(tmp_config_dir) svc = WorkspaceService(config_store=store, client_factory=lambda url, token: mock_client) - with pytest.raises(KeboolaApiError) as exc_info: + with pytest.raises(ValueError, match="--timeout must be a positive number"): svc.load_tables(alias="prod", workspace_id=42, tables=["in.c-main.orders"], timeout=0.0) - - assert exc_info.value.error_code == ErrorCode.INVALID_ARGUMENT mock_client.load_workspace_tables.assert_not_called() - def test_load_tables_timeout_negative_rejected(self, tmp_config_dir: Path) -> None: - """A negative timeout is rejected the same way as zero.""" + @pytest.mark.parametrize("bad", [-5.0, float("nan"), float("inf")]) + def test_load_tables_timeout_negative_or_non_finite_rejected( + self, tmp_config_dir: Path, bad: float + ) -> None: + """A negative or non-finite timeout is rejected the same way as zero.""" mock_client = _make_load_client() store = setup_single_project(tmp_config_dir) svc = WorkspaceService(config_store=store, client_factory=lambda url, token: mock_client) - with pytest.raises(KeboolaApiError) as exc_info: - svc.load_tables( - alias="prod", workspace_id=42, tables=["in.c-main.orders"], timeout=-5.0 - ) + with pytest.raises(ValueError, match="--timeout must be a positive number"): + svc.load_tables(alias="prod", workspace_id=42, tables=["in.c-main.orders"], timeout=bad) - assert exc_info.value.error_code == ErrorCode.INVALID_ARGUMENT mock_client.load_workspace_tables.assert_not_called() def test_load_tables_timeout_explicit_positive_passthrough(self, tmp_config_dir: Path) -> None: