From c9451571a377a9db651f5a806a280d9968090ae1 Mon Sep 17 00:00:00 2001 From: Petr Date: Wed, 10 Jun 2026 12:02:07 +0200 Subject: [PATCH 1/3] feat(workspace): fast inline results for `workspace query`, `--full` for complete CSV (0.59.0) `workspace query` materialized a CSV file via the warehouse UNLOAD path (GET .../export?fileType=csv) for every query, paying a slow warehouse -> object-storage -> download round-trip even for tiny result sets. The default now reads results inline via the Query Service GET /api/v1/queries/{job}/{stmt}/results endpoint (JSON columns+rows, no file materialization) -- ~2.5x faster on a 25-row table, more on larger overhead-dominated queries. - Each statement carries structured columns + rows (+ row_count, total_rows, truncated) and a synthesized csv_data, so the CLI preview, web UI table, and any --json consumer keep working unchanged. - --limit N (default 500) caps the fast path; it pages by offset and marks the result truncated when the warehouse has more. The /results endpoint enforces 100 <= pageSize <= 100000, so pageSize is a fixed valid value and --limit only trims locally (a small --limit no longer 400s). - --full opts back into the complete CSV export (slower, uncapped) for bulk extraction. - serve /workspaces/{p}/{w}/query defaults to full=True to preserve the web UI's complete-CSV download until the frontend paginates. Verified live against the demo "shop" project (Snowflake workspace, 25-row table): structured output, NULL handling, --limit truncation, --full export, and the ~2.5x speedup. --- .claude-plugin/marketplace.json | 2 +- CLAUDE.md | 5 +- plugins/kbagent/.claude-plugin/plugin.json | 2 +- .../kbagent/references/commands-reference.md | 2 +- .../skills/kbagent/references/gotchas.md | 28 +++ .../kbagent/references/workspace-workflow.md | 55 +++++ pyproject.toml | 2 +- src/keboola_agent_cli/changelog.py | 15 ++ src/keboola_agent_cli/client.py | 38 ++++ src/keboola_agent_cli/commands/context.py | 4 +- src/keboola_agent_cli/commands/workspace.py | 17 ++ src/keboola_agent_cli/constants.py | 8 + src/keboola_agent_cli/output.py | 76 ++++--- .../server/routers/workspaces.py | 10 +- .../services/workspace_service.py | 137 +++++++++++- tests/test_client.py | 47 +++++ tests/test_e2e.py | 20 ++ tests/test_output.py | 85 ++++++++ tests/test_workspace_cli.py | 100 +++++++++ tests/test_workspace_service.py | 199 ++++++++++++++++-- uv.lock | 2 +- 21 files changed, 794 insertions(+), 60 deletions(-) diff --git a/.claude-plugin/marketplace.json b/.claude-plugin/marketplace.json index 65f6cb2c..82f1a36d 100644 --- a/.claude-plugin/marketplace.json +++ b/.claude-plugin/marketplace.json @@ -10,7 +10,7 @@ "plugins": [ { "name": "kbagent", - "version": "0.58.0", + "version": "0.59.0", "source": "./plugins/kbagent", "description": "AI-friendly interface to Keboola Connection projects — explore configs, jobs, lineage, call MCP tools, manage dev branches, and debug SQL in workspaces", "category": "development" diff --git a/CLAUDE.md b/CLAUDE.md index dc416893..441e44c0 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -402,8 +402,11 @@ kbagent workspace detail --project ALIAS --workspace-id ID [--branch ID] kbagent workspace delete --project ALIAS --workspace-id ID kbagent workspace password --project ALIAS --workspace-id ID kbagent workspace load --project ALIAS --workspace-id ID --tables TABLE_ID [--tables ...] [--preserve] -kbagent workspace query --project ALIAS --workspace-id ID --sql "SELECT ..." [--transactional] +kbagent workspace query --project ALIAS --workspace-id ID --sql "SELECT ..." [--transactional] [--full] [--limit N] kbagent workspace query --project ALIAS --workspace-id ID --file query.sql +# query: default reads results inline via Query Service `GET .../results` (fast, JSON columns+rows), +# capped at --limit rows (default 500); pass --full for the complete CSV export (slower, uncapped). +# Each statement carries structured columns+rows + a synthesized csv_data (back-compat) since 0.59.0. kbagent workspace gc [--project NAME ...] [--dry-run] [--yes] kbagent workspace from-transformation --project ALIAS --component-id ID --config-id ID [--row-id ID] diff --git a/plugins/kbagent/.claude-plugin/plugin.json b/plugins/kbagent/.claude-plugin/plugin.json index 6a9fb2c7..bdb9cec8 100644 --- a/plugins/kbagent/.claude-plugin/plugin.json +++ b/plugins/kbagent/.claude-plugin/plugin.json @@ -1,6 +1,6 @@ { "name": "kbagent", - "version": "0.58.0", + "version": "0.59.0", "description": "AI-friendly interface to Keboola Connection projects — explore configs, jobs, lineage, call MCP tools, manage dev branches, and debug SQL in workspaces", "author": { "name": "Keboola", diff --git a/plugins/kbagent/skills/kbagent/references/commands-reference.md b/plugins/kbagent/skills/kbagent/references/commands-reference.md index e24da331..cfa52a0f 100644 --- a/plugins/kbagent/skills/kbagent/references/commands-reference.md +++ b/plugins/kbagent/skills/kbagent/references/commands-reference.md @@ -169,7 +169,7 @@ Bucket sharing + linking across projects in the same organization. `sharing edge - `workspace delete --project ALIAS --workspace-id ID` -- delete workspace - `workspace password --project ALIAS --workspace-id ID` -- reset and return new password - `workspace load --project ALIAS --workspace-id ID --tables TABLE_ID [...] [--preserve]` -- load storage tables -- `workspace query --project ALIAS --workspace-id ID --sql "..." [--file F] [--transactional]` -- run SQL via Query Service. **Backend-agnostic since v0.58.0**: runs against both Snowflake and BigQuery workspaces (the path was always identical; BigQuery just needed the classification fix). Mind the dialect: Snowflake quotes identifiers with `"..."`, BigQuery with backticks `` `...` `` +- `workspace query --project ALIAS --workspace-id ID --sql "..." [--file F] [--transactional] [--full] [--limit N]` -- run SQL via Query Service. **Fast inline results since v0.59.0**: default reads the result set inline via `GET /api/v1/queries/{job}/{stmt}/results` (JSON `columns`+`rows`, no CSV-file materialization), capped at `--limit` rows (default 500) and marked `truncated` when there are more; pass `--full` for the complete CSV export (slower, uncapped). Each statement still carries `csv_data` (synthesized from the inline rows) so older parsers keep working. **Backend-agnostic since v0.58.0**: runs against both Snowflake and BigQuery workspaces. Mind the dialect: Snowflake quotes identifiers with `"..."`, BigQuery with backticks `` `...` `` - `workspace gc [--project NAME ...] [--dry-run] [--yes]` -- garbage-collect orphaned workspaces (and any lingering `keboola.sandboxes` configs). `--dry-run` previews without deleting; `--project` repeatable, omit to GC across all connected projects - `workspace from-transformation --project ALIAS --component-id ID --config-id ID [--row-id ID]` -- workspace from existing transform diff --git a/plugins/kbagent/skills/kbagent/references/gotchas.md b/plugins/kbagent/skills/kbagent/references/gotchas.md index c2b8e511..dc1f0536 100644 --- a/plugins/kbagent/skills/kbagent/references/gotchas.md +++ b/plugins/kbagent/skills/kbagent/references/gotchas.md @@ -289,6 +289,34 @@ confirmed-good whitelist". For an unknown loginType, `workspace list` renders it as `?` (yellow) in the QS column so callers know the policy is uncertain rather than confirmed-bad. +## `workspace query`: fast inline results vs `--full` CSV export (since v0.59.0) + +By default `workspace query` now reads the result set inline via the Query +Service `GET /api/v1/queries/{job}/{stmt}/results` endpoint (JSON `columns` + +`rows`) instead of materializing a CSV file through the warehouse UNLOAD path +(`.../export?fileType=csv`). The inline path skips the file round-trip, so +interactive queries are markedly faster. + +- The inline path is **paginated**: it fetches at most `--limit` rows (default + 500), walking `offset` in pages. When the warehouse has more rows than + fetched, the statement is marked `truncated: true` (with `total_rows` = + the full count) and the CLI prints `Showing first N of TOTAL rows. Use --full`. +- The `/results` endpoint enforces **`100 <= pageSize <= 100000`** (a smaller + `pageSize` 400s with `Invalid pageSize parameter, must be between 100 and + 100000`). kbagent therefore requests a fixed valid page size and trims the + result to `--limit` locally -- `pageSize` is NOT derived from `--limit`, so a + `--limit 5` still works (fetches one valid page, returns 5). +- Each statement carries structured `columns`, `rows`, `row_count`, + `total_rows`, `truncated`, **and** a synthesized `csv_data` string. Parsers + that read `csv_data` (the pre-0.59.0 shape) keep working unchanged. +- `--full` opts back into the complete CSV export -- slower (warehouse UNLOAD), + but **uncapped**. Use it when you need every row, e.g. a bulk extract + (`workspace query --full --json`). Under `--full` the statement carries only + `csv_data` (no structured `columns`/`rows`). +- The `kbagent serve` `/workspaces/{p}/{w}/query` REST endpoint defaults to + `full=True` so the web UI's "Download CSV" stays complete; REST clients can + pass `full=false` (+ `limit`) in the JSON body to opt into the fast path. + ## Snowflake `workspace create` returns `private_key`, not password (since v0.47.1) Headless `workspace create` on Snowflake requests diff --git a/plugins/kbagent/skills/kbagent/references/workspace-workflow.md b/plugins/kbagent/skills/kbagent/references/workspace-workflow.md index 23448827..c361f11e 100644 --- a/plugins/kbagent/skills/kbagent/references/workspace-workflow.md +++ b/plugins/kbagent/skills/kbagent/references/workspace-workflow.md @@ -20,6 +20,9 @@ kbagent --json workspace from-transformation \ ```bash # Step 2: Run the original SQL to reproduce the error +# Default (0.59.0+): results come back inline as JSON columns+rows (fast), +# capped at --limit rows (default 500). Add --full for the complete CSV export +# (slower, uncapped) when you need every row. kbagent --json workspace query \ --project ALIAS \ --workspace-id WS_ID \ @@ -94,6 +97,58 @@ kbagent --json workspace query \ --file query.sql ``` +## Fast inline results vs `--full` -- mind the result-set volume (since v0.59.0) + +`workspace query` has two ways to retrieve results. **Pick based on how many rows +you actually need**, not by habit. + +**Default (fast, inline):** reads the result set straight from the Query Service +as JSON via `GET /api/v1/queries/{job}/{stmt}/results`. No file is produced. Each +statement comes back with structured `columns` + `rows` (plus `row_count`, +`total_rows`, `truncated`) and a synthesized `csv_data` for back-compat. + +- **Paginated / capped** at `--limit` rows (default 500). When the warehouse has + more rows than were fetched, the statement is flagged `truncated: true` (with + `total_rows` = the real count) and the CLI prints + `Showing first N of TOTAL rows. Use --full for the complete result set.` +- This is the right default for **inspection, row counts, sampling, schema + checks, and iterating on a fix** -- exactly the workspace-debugging loop. + +**`--full` (complete CSV export, slower):** materializes the *entire* result set +as a CSV file through the warehouse UNLOAD path +(`GET .../export?fileType=csv`), then downloads it. Uncapped -- you get every +row -- but it pays a file-export round-trip (warehouse -> object storage -> +download) on every call. Under `--full` the statement carries only `csv_data` +(no structured `columns`/`rows`). + +```bash +# Fast: first 500 rows inline (default). Add --limit to widen/narrow the page. +kbagent --json workspace query --project ALIAS --workspace-id WS_ID \ + --sql 'SELECT * FROM "in.c-main"."events"' --limit 1000 + +# Complete: every row via CSV export (slower -- use only when you need them all). +kbagent --json workspace query --project ALIAS --workspace-id WS_ID \ + --sql 'SELECT * FROM "in.c-main"."events"' --full +``` + +**Decision guide:** + +- **Just looking / counting / sampling?** Use the default. Faster, and the + `truncated` flag tells you whether there is more. +- **Need a complete extract?** Use `--full` -- but **think about the volume + first**. `--full` pulls the whole result set into a single CSV string in + memory; a `SELECT *` over millions of rows is slow and memory-hungry. Narrow + the query (`SELECT` only the columns you need, add a `WHERE`/`LIMIT`) before + reaching for `--full`. +- **Bulk-exporting an actual Storage table** (not an arbitrary query)? Prefer + `storage unload-table` / `storage download-table` -- they stream sliced files + and are built for volume, whereas `workspace query --full` is for ad-hoc SQL. + +**API floor:** the `/results` endpoint requires `100 <= pageSize <= 100000`. +kbagent always requests a valid page size and trims to `--limit` locally, so a +small `--limit` (e.g. `--limit 5`) works fine -- it does not shrink the wire +`pageSize` below the API minimum. + ## Shared/linked buckets -- different database/dataset Linked buckets (shared from another project) live in a **different database diff --git a/pyproject.toml b/pyproject.toml index 1c70639a..7d9eac9f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "keboola-agent-cli" -version = "0.58.0" +version = "0.59.0" description = "AI-friendly CLI for managing Keboola projects" readme = "README.md" requires-python = ">=3.12" diff --git a/src/keboola_agent_cli/changelog.py b/src/keboola_agent_cli/changelog.py index a0e7dd1a..800cc273 100644 --- a/src/keboola_agent_cli/changelog.py +++ b/src/keboola_agent_cli/changelog.py @@ -24,6 +24,21 @@ # Ordered newest-first. Each value is a list of brief one-line descriptions. CHANGELOG: dict[str, list[str]] = { + "0.59.0": [ + "Faster: `kbagent workspace query` now reads results via the Query Service's inline " + "`GET /api/v1/queries/{job}/{stmt}/results` endpoint by default instead of materializing a " + "CSV file through the warehouse UNLOAD path (`.../export?fileType=csv`). The inline path " + "returns the already-computed result set as JSON -- no file export round-trip -- so interactive " + "queries come back markedly faster. Each statement now carries structured `columns` + `rows` " + "(plus `row_count`, `total_rows`, `truncated`) alongside a synthesized `csv_data`, so the CLI " + "preview, web UI table, and any `--json` consumer keep working unchanged.", + "New: `--limit N` (default 500) caps how many rows the fast inline path fetches; it pages " + "through the result set by `offset` until the limit is reached, marking the result `truncated` " + "when the warehouse has more. `--full` opts back into the complete CSV export (slower, " + "uncapped) when you need every row -- e.g. piping `workspace query --full --json` for a bulk " + "extract. The `kbagent serve` `/workspaces/{p}/{w}/query` REST endpoint defaults to `full=True` " + "to preserve the web UI's complete-CSV download until the frontend learns to paginate.", + ], "0.58.0": [ "New: `kbagent workspace query` runs SQL against BigQuery workspaces, not just Snowflake. " "The Query Service path was always backend-agnostic (`POST " diff --git a/src/keboola_agent_cli/client.py b/src/keboola_agent_cli/client.py index 1cef3682..0f207e77 100644 --- a/src/keboola_agent_cli/client.py +++ b/src/keboola_agent_cli/client.py @@ -35,6 +35,7 @@ OAUTH_PATH, QUERY_JOB_MAX_WAIT, QUERY_JOB_POLL_INTERVAL, + QUERY_RESULTS_PAGE_SIZE, STORAGE_JOB_MAX_WAIT, STORAGE_JOB_POLL_INTERVAL, VALID_POLL_STRATEGIES, @@ -2697,6 +2698,43 @@ def export_query_results( ) return response.text + def get_query_results( + self, + query_job_id: str, + statement_id: str, + offset: int = 0, + page_size: int = QUERY_RESULTS_PAGE_SIZE, + ) -> dict[str, Any]: + """Fetch a page of inline statement results from the Query Service. + + Unlike :meth:`export_query_results`, which materializes a CSV file via the + warehouse UNLOAD path (slow), this reads the already-computed result set + inline as JSON -- much faster for interactive queries. The endpoint is + paginated; ``offset``/``page_size`` walk the result set. + + Args: + query_job_id: The query job ID. + statement_id: The statement ID within the job. + offset: Row offset to start from (for pagination). + page_size: Maximum rows to return in this page. + + Returns: + Raw QueryResult dict, e.g.:: + + { + "status": "completed", + "columns": [{"name": "id", "type": "INTEGER", "nullable": false}], + "data": [[1, "a"], [2, "b"]], + "numberOfRows": 2, + } + """ + response = self._query_request( + "GET", + f"/api/v1/queries/{query_job_id}/{statement_id}/results", + params={"offset": offset, "pageSize": page_size}, + ) + return response.json() + def get_query_history( self, branch_id: int, diff --git a/src/keboola_agent_cli/commands/context.py b/src/keboola_agent_cli/commands/context.py index 4a301511..e6d9fc1b 100644 --- a/src/keboola_agent_cli/commands/context.py +++ b/src/keboola_agent_cli/commands/context.py @@ -698,8 +698,10 @@ kbagent workspace load --project ALIAS --workspace-id ID --tables TABLE_ID [...] [--preserve] Load storage tables into workspace. --preserve keeps existing tables. - kbagent workspace query --project ALIAS --workspace-id ID --sql "SQL" [--file F] [--transactional] + kbagent workspace query --project ALIAS --workspace-id ID --sql "SQL" [--file F] [--transactional] [--full] [--limit N] Execute SQL via Query Service. No Snowflake credentials needed. + Default reads results inline (fast JSON columns+rows), capped at --limit (default 500). + --full uses the complete CSV export instead (slower, uncapped). kbagent workspace from-transformation --project ALIAS --component-id ID --config-id ID [--row-id ID] Create workspace from transformation config. Loads input tables automatically. diff --git a/src/keboola_agent_cli/commands/workspace.py b/src/keboola_agent_cli/commands/workspace.py index 35141c5b..45b15439 100644 --- a/src/keboola_agent_cli/commands/workspace.py +++ b/src/keboola_agent_cli/commands/workspace.py @@ -10,6 +10,7 @@ from rich.markup import escape from ..config_store import ConfigStore +from ..constants import QUERY_RESULTS_DEFAULT_LIMIT from ..errors import ConfigError, ErrorCode, KeboolaApiError from ..output import format_query_results, format_workspaces_table from ._helpers import ( @@ -414,10 +415,24 @@ def workspace_query( "--transactional", help="Wrap query in a transaction", ), + full: bool = typer.Option( + False, + "--full", + help="Fetch the complete result set via CSV export (slower). " + "Default fetches a fast inline page capped by --limit.", + ), + limit: int = typer.Option( + QUERY_RESULTS_DEFAULT_LIMIT, + "--limit", + help="Max rows to fetch via the fast inline path (ignored with --full).", + ), ) -> None: """Execute SQL query in a workspace via Query Service. Provide SQL via --sql or --file (exactly one required). + + By default kbagent reads results inline (fast). For a result set larger than + --limit, pass --full to export the complete CSV instead. """ formatter = get_formatter(ctx) service = get_service(ctx, "workspace_service") @@ -450,6 +465,8 @@ def workspace_query( workspace_id=workspace_id, sql=effective_sql, transactional=transactional, + full=full, + limit=limit, ) if formatter.json_mode: formatter.output(result) diff --git a/src/keboola_agent_cli/constants.py b/src/keboola_agent_cli/constants.py index 288c0fbe..62b916c7 100644 --- a/src/keboola_agent_cli/constants.py +++ b/src/keboola_agent_cli/constants.py @@ -414,6 +414,14 @@ # --- Query Service --- QUERY_JOB_POLL_INTERVAL: float = 1.0 # seconds between polls for query job status QUERY_JOB_MAX_WAIT: float = 120.0 # max seconds to wait for a query job +# Fast inline result path (GET .../results) -- reads the already-computed result +# set as JSON instead of materializing a CSV file via the warehouse UNLOAD path +# (GET .../export). The default path fetches at most QUERY_RESULTS_DEFAULT_LIMIT +# rows, accumulated in pages of QUERY_RESULTS_PAGE_SIZE and trimmed to the limit. +# The endpoint enforces 100 <= pageSize <= 100000, so QUERY_RESULTS_PAGE_SIZE is a +# fixed valid page size -- it is NOT derived from --limit (a small limit would 400). +QUERY_RESULTS_DEFAULT_LIMIT: int = 500 # default --limit for `workspace query` fast path +QUERY_RESULTS_PAGE_SIZE: int = 500 # rows per /results page (within the API's 100..100000) # --- Workspace Defaults --- DEFAULT_WORKSPACE_BACKEND: str = "snowflake" diff --git a/src/keboola_agent_cli/output.py b/src/keboola_agent_cli/output.py index ba38799f..ac0117de 100644 --- a/src/keboola_agent_cli/output.py +++ b/src/keboola_agent_cli/output.py @@ -962,7 +962,10 @@ def format_workspaces_table(console: Console, data: dict[str, Any]) -> None: def format_query_results(console: Console, data: dict[str, Any]) -> None: """Render SQL query results. - Shows query status and CSV results for each statement. + Shows query status plus, per statement, a Rich table built from the + structured ``columns``/``rows`` returned by the fast inline path. Falls back + to a CSV preview when only ``csv_data`` is present (the ``--full`` export + path, which returns a CSV string without structured columns). Args: console: Rich Console instance. @@ -972,32 +975,57 @@ def format_query_results(console: Console, data: dict[str, Any]) -> None: workspace_id = data.get("workspace_id", "") status = data.get("status", "unknown") - lines = [ - f"[bold]Project:[/bold] {alias}", - f"[bold]Workspace:[/bold] {workspace_id}", - f"[bold]Status:[/bold] {status}", - ] + console.print( + Panel( + f"[bold]Project:[/bold] {alias}\n" + f"[bold]Workspace:[/bold] {workspace_id}\n" + f"[bold]Status:[/bold] {status}", + title=f"Query Results - Workspace {workspace_id}", + expand=False, + ) + ) statements = data.get("statements", []) for i, stmt in enumerate(statements): - lines.append(f"\n[bold]Statement {i + 1}:[/bold]") - lines.append(f" Status: {stmt.get('status', 'unknown')}") - rows = stmt.get("rows_affected", 0) - lines.append(f" Rows: {rows}") - - csv_data = stmt.get("csv_data", "") - if csv_data: - # Show first few lines of CSV - csv_lines = csv_data.strip().split("\n") - preview_count = min(len(csv_lines), 11) # header + 10 rows - lines.append(" [bold]Results:[/bold]") - for csv_line in csv_lines[:preview_count]: - lines.append(f" {csv_line}") - if len(csv_lines) > preview_count: - lines.append(f" ... ({len(csv_lines) - preview_count} more rows)") - - panel = Panel("\n".join(lines), title=f"Query Results - Workspace {workspace_id}", expand=False) - console.print(panel) + console.print( + f"\n[bold]Statement {i + 1}:[/bold] " + f"{stmt.get('status', 'unknown')} ・ {stmt.get('rows_affected', 0)} rows" + ) + _render_statement_result(console, stmt) + + +def _render_statement_result(console: Console, stmt: dict[str, Any]) -> None: + """Render a single statement's result set (structured table or CSV preview).""" + columns = stmt.get("columns") + rows = stmt.get("rows") + if columns and rows is not None: + table = Table(show_lines=False) + for col in columns: + table.add_column(str(col.get("name", ""))) + for row in rows: + table.add_row(*["" if value is None else str(value) for value in row]) + console.print(table) + if stmt.get("truncated"): + total = stmt.get("total_rows") + shown = stmt.get("row_count", len(rows)) + suffix = f" of {total}" if total is not None else "" + console.print( + f" [dim]Showing first {shown}{suffix} rows. " + f"Use --full for the complete result set.[/dim]" + ) + return + + # Fallback: --full export path returns a CSV string with no structured columns. + csv_data = stmt.get("csv_data", "") + if not csv_data: + return + csv_lines = csv_data.strip().split("\n") + preview_count = min(len(csv_lines), 11) # header + 10 rows + console.print(" [bold]Results:[/bold]") + for csv_line in csv_lines[:preview_count]: + console.print(f" {csv_line}") + if len(csv_lines) > preview_count: + console.print(f" ... ({len(csv_lines) - preview_count} more rows)") def format_search_results(console: Console, data: dict[str, Any]) -> None: diff --git a/src/keboola_agent_cli/server/routers/workspaces.py b/src/keboola_agent_cli/server/routers/workspaces.py index 3d4f0bd2..fc4e2b53 100644 --- a/src/keboola_agent_cli/server/routers/workspaces.py +++ b/src/keboola_agent_cli/server/routers/workspaces.py @@ -10,7 +10,7 @@ from fastapi.responses import StreamingResponse from pydantic import BaseModel -from ...constants import AI_SQL_HELPER_TIMEOUT +from ...constants import AI_SQL_HELPER_TIMEOUT, QUERY_RESULTS_DEFAULT_LIMIT from ..dependencies import ServiceRegistry, get_registry router = APIRouter(prefix="/workspaces", tags=["workspaces"]) @@ -31,6 +31,12 @@ class WorkspaceLoad(BaseModel): class WorkspaceQuery(BaseModel): sql: str transactional: bool = False + # Default True preserves the current web-UI behavior: the Workspaces page + # renders csv_data and offers a "Download CSV" button that expects the + # complete result set. The fast inline path (full=False) is paginated, so a + # REST client must opt in explicitly until the frontend learns to paginate. + full: bool = True + limit: int = QUERY_RESULTS_DEFAULT_LIMIT class FromTransformation(BaseModel): @@ -177,6 +183,8 @@ def query( workspace_id=workspace_id, sql=body.sql, transactional=body.transactional, + full=body.full, + limit=body.limit, ) diff --git a/src/keboola_agent_cli/services/workspace_service.py b/src/keboola_agent_cli/services/workspace_service.py index 78122766..9446d2e8 100644 --- a/src/keboola_agent_cli/services/workspace_service.py +++ b/src/keboola_agent_cli/services/workspace_service.py @@ -5,6 +5,8 @@ single-project operations. """ +import csv +import io import logging from dataclasses import dataclass from typing import Any @@ -14,6 +16,8 @@ from ..constants import ( BIGQUERY_WORKSPACE_LOGIN_TYPE, + QUERY_RESULTS_DEFAULT_LIMIT, + QUERY_RESULTS_PAGE_SIZE, QUERY_SERVICE_COMPATIBLE_LOGIN_TYPES, QUERY_SERVICE_COMPATIBLE_LOGIN_TYPES_BIGQUERY, SNOWFLAKE_WORKSPACE_LOGIN_TYPE, @@ -151,6 +155,76 @@ def _is_orphaned_workspace(ws: dict[str, Any], config_names: dict[str, str]) -> return not config_id or config_id not in config_names +@dataclass(frozen=True) +class InlineQueryResult: + """One statement's result fetched via the fast inline `/results` path.""" + + columns: list[dict[str, Any]] # [{"name", "type", "nullable"}] + rows: list[list[Any]] # row values, row-major; capped at the requested limit + total_rows: int | None # numberOfRows reported by the warehouse (full count) + truncated: bool # True when the warehouse has more rows than we fetched + + +def _rows_to_csv(columns: list[dict[str, Any]], rows: list[list[Any]]) -> str: + """Render structured columns+rows as an RFC-4180 CSV string. + + Synthesized so the inline `/results` payload stays drop-in compatible with + consumers that still read ``csv_data`` (CLI preview, web UI table + export + buttons, REST). ``csv.writer`` handles quoting/escaping; ``None`` becomes an + empty field, matching the warehouse CSV export semantics. + """ + buffer = io.StringIO() + writer = csv.writer(buffer, lineterminator="\n") + writer.writerow([col.get("name", "") for col in columns]) + for row in rows: + writer.writerow(["" if value is None else value for value in row]) + return buffer.getvalue() + + +def _collect_inline_results( + client: Any, + query_job_id: str, + statement_id: str, + limit: int, +) -> InlineQueryResult: + """Page through `GET .../results`, accumulating up to ``limit`` rows. + + The endpoint enforces ``100 <= pageSize <= 100000``, so we always request a + fixed, valid ``QUERY_RESULTS_PAGE_SIZE`` page and cap the accumulated rows at + ``limit`` locally -- deriving ``pageSize`` from a small ``--limit`` (e.g. 5) + would trip the API's minimum with a 400. A ``limit`` larger than one page is + satisfied by walking ``offset``; we stop once the limit is reached (marking + the result truncated) or when the warehouse runs out of rows. + """ + collected: list[list[Any]] = [] + columns: list[dict[str, Any]] = [] + total_rows: int | None = None + offset = 0 + while len(collected) < limit: + payload = client.get_query_results( + query_job_id, statement_id, offset=offset, page_size=QUERY_RESULTS_PAGE_SIZE + ) + if not columns: + columns = payload.get("columns", []) or [] + if total_rows is None: + total_rows = payload.get("numberOfRows") + page_rows = payload.get("data", []) or [] + collected.extend(page_rows) + # Last page: the warehouse returned fewer rows than a full page. + if len(page_rows) < QUERY_RESULTS_PAGE_SIZE: + break + offset += len(page_rows) + + rows = collected[:limit] + truncated = total_rows is not None and total_rows > len(rows) + return InlineQueryResult( + columns=columns, + rows=rows, + total_rows=total_rows, + truncated=truncated, + ) + + class WorkspaceService(BaseService): """Business logic for managing Keboola workspaces. @@ -822,17 +896,32 @@ def execute_query( workspace_id: int, sql: str, transactional: bool = False, + full: bool = False, + limit: int = QUERY_RESULTS_DEFAULT_LIMIT, ) -> dict[str, Any]: """Execute SQL query in a workspace via Query Service. - Submits the query, polls until complete, and exports CSV results - for each statement. + Submits the query, polls until complete, and fetches results for each + statement. Two retrieval paths: + + * Default (``full=False``): the fast inline ``GET .../results`` path. + Reads the already-computed result set as JSON (no warehouse UNLOAD / + CSV-file materialization), paginated up to ``limit`` rows. Each + statement carries structured ``columns`` + ``rows`` and a synthesized + ``csv_data`` (drop-in for the legacy CSV-string consumers). + * ``full=True``: the legacy ``GET .../export?fileType=csv`` path, which + materializes the *complete* result set as a CSV file. Slower, but not + capped at ``limit`` -- use it when you need every row. Args: alias: Project alias. workspace_id: Workspace ID. sql: SQL statement(s) to execute. transactional: Whether to wrap in a transaction. + full: Fetch the complete result set via the CSV export path instead + of the fast inline path. + limit: Max rows to fetch via the fast inline path (ignored when + ``full`` is True). Returns: Dict with query results. @@ -856,7 +945,7 @@ def execute_query( # Wait for completion completed_job = client.wait_for_query_job(query_job_id) - # Export results for each statement + # Fetch results for each statement results: list[dict[str, Any]] = [] statements = completed_job.get("statements", []) for stmt in statements: @@ -869,13 +958,11 @@ def execute_query( "rows_affected": num_rows, } - # Try to export results if there are rows + # Only statements that produced a result set carry rows. if status == "completed" and num_rows > 0: - try: - csv_data = client.export_query_results(query_job_id, stmt_id) - result_entry["csv_data"] = csv_data - except KeboolaApiError: - logger.debug("Could not export results for statement %s", stmt_id) + self._attach_statement_results( + result_entry, client, query_job_id, stmt_id, full=full, limit=limit + ) results.append(result_entry) @@ -891,6 +978,38 @@ def execute_query( finally: client.close() + @staticmethod + def _attach_statement_results( + result_entry: dict[str, Any], + client: Any, + query_job_id: str, + stmt_id: str, + *, + full: bool, + limit: int, + ) -> None: + """Populate a statement's result entry, fast inline path or full export. + + Failures are swallowed (debug-logged): a result-fetch error must not + sink an otherwise-successful query -- the statement still reports its + status and row count, just without the data payload. + """ + try: + if full: + result_entry["csv_data"] = client.export_query_results(query_job_id, stmt_id) + return + inline = _collect_inline_results(client, query_job_id, stmt_id, limit) + result_entry["columns"] = inline.columns + result_entry["rows"] = inline.rows + result_entry["row_count"] = len(inline.rows) + result_entry["total_rows"] = inline.total_rows + result_entry["truncated"] = inline.truncated + # Synthesize csv_data so legacy consumers (web UI table/export, CLI + # preview) keep working without a format change. + result_entry["csv_data"] = _rows_to_csv(inline.columns, inline.rows) + except KeboolaApiError: + logger.debug("Could not fetch results for statement %s", stmt_id) + def create_from_transformation( self, alias: str, diff --git a/tests/test_client.py b/tests/test_client.py index 2afc7208..603704fc 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -2911,6 +2911,53 @@ def test_fixed_yields_storage_interval_forever(self) -> None: assert seq == [STORAGE_JOB_POLL_INTERVAL] * 5 +class TestGetQueryResults: + """Tests for KeboolaClient.get_query_results -- the fast inline /results path. + + Hits ``GET query./api/v1/queries/{job}/{stmt}/results`` and returns the + raw QueryResult dict (columns + data) without materializing a CSV file. + """ + + def test_fetches_inline_results_with_default_pagination(self, httpx_mock) -> None: + payload = { + "status": "completed", + "columns": [{"name": "id", "type": "INTEGER", "nullable": False}], + "data": [[1], [2]], + "numberOfRows": 2, + } + httpx_mock.add_response( + url=( + "https://query.keboola.com/api/v1/queries/qj-1/stmt-1/results?offset=0&pageSize=500" + ), + json=payload, + status_code=200, + ) + client = KeboolaClient( + stack_url="https://connection.keboola.com", + token="901-55555-fakeTestTokenDoNotUseXXXXXXXX", + ) + result = client.get_query_results("qj-1", "stmt-1") + assert result == payload + client.close() + + def test_passes_explicit_offset_and_page_size(self, httpx_mock) -> None: + httpx_mock.add_response( + url=( + "https://query.keboola.com/api/v1/queries/qj-1/stmt-2/results" + "?offset=500&pageSize=200" + ), + json={"status": "completed", "columns": [], "data": [], "numberOfRows": 0}, + status_code=200, + ) + client = KeboolaClient( + stack_url="https://connection.keboola.com", + token="901-55555-fakeTestTokenDoNotUseXXXXXXXX", + ) + result = client.get_query_results("qj-1", "stmt-2", offset=500, page_size=200) + assert result["numberOfRows"] == 0 + client.close() + + class TestExtractQueryJobError: """Tests for _extract_query_job_error -- pulls the real warehouse error out of a failed Query Service job payload. diff --git a/tests/test_e2e.py b/tests/test_e2e.py index 26398da6..7f54ef3c 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -1953,6 +1953,7 @@ def _test_workspace_query(self, workspace_id: int, table_id: str) -> None: )["data"] quote = "`" if detail.get("backend") == "bigquery" else '"' sql = f"SELECT COUNT(*) AS cnt FROM {quote}{ws_table_name}{quote}" + # Default (fast) path: reads inline /results -- structured columns+rows. data = self._run_ok( "workspace", "query", @@ -1964,6 +1965,25 @@ def _test_workspace_query(self, workspace_id: int, table_id: str) -> None: sql, ) assert data["status"] == "ok" + stmt = data["data"]["statements"][0] + assert stmt["columns"], "fast inline path must return structured columns" + assert stmt["rows"], "fast inline path must return structured rows" + assert "csv_data" in stmt, "csv_data must stay populated for legacy consumers" + + # Full (export) path: complete result set via the CSV export endpoint. + full_data = self._run_ok( + "workspace", + "query", + "--project", + self.alias, + "--workspace-id", + str(workspace_id), + "--sql", + sql, + "--full", + ) + assert full_data["status"] == "ok" + assert "csv_data" in full_data["data"]["statements"][0] def _test_workspace_delete(self, workspace_id: int) -> None: """Delete the workspace.""" diff --git a/tests/test_output.py b/tests/test_output.py index 9b0ee6d4..923b5cf6 100644 --- a/tests/test_output.py +++ b/tests/test_output.py @@ -17,6 +17,7 @@ format_job_detail, format_jobs_table, format_lineage_table, + format_query_results, format_tool_result, format_tools_table, ) @@ -989,3 +990,87 @@ def test_doctor_panel_with_failures(self) -> None: assert "WARN" in output assert "1 failed" in output assert "1 warnings" in output + + +class TestFormatQueryResults: + """Tests for format_query_results -- structured table vs. CSV fallback.""" + + def _console(self) -> Console: + return Console(file=StringIO(), no_color=True, force_terminal=False, width=200) + + def test_renders_structured_columns_and_rows_as_table(self) -> None: + """The fast-path payload (columns+rows) renders as a Rich table.""" + console = self._console() + data = { + "project_alias": "prod", + "workspace_id": 42, + "status": "completed", + "statements": [ + { + "statement_id": "stmt-1", + "status": "completed", + "rows_affected": 2, + "columns": [{"name": "id"}, {"name": "name"}], + "rows": [[1, "alice"], [2, None]], + "row_count": 2, + "total_rows": 2, + "truncated": False, + "csv_data": "id,name\n1,alice\n2,\n", + } + ], + } + format_query_results(console, data) + output = cast(StringIO, console.file).getvalue() + assert "id" in output + assert "name" in output + assert "alice" in output + # No truncation hint when truncated is False. + assert "Use --full" not in output + + def test_shows_truncation_hint(self) -> None: + """When the warehouse has more rows than fetched, hint at --full.""" + console = self._console() + data = { + "project_alias": "prod", + "workspace_id": 42, + "status": "completed", + "statements": [ + { + "statement_id": "stmt-1", + "status": "completed", + "rows_affected": 100, + "columns": [{"name": "id"}], + "rows": [[1], [2]], + "row_count": 2, + "total_rows": 100, + "truncated": True, + "csv_data": "id\n1\n2\n", + } + ], + } + format_query_results(console, data) + output = cast(StringIO, console.file).getvalue() + assert "of 100" in output + assert "--full" in output + + def test_falls_back_to_csv_preview_for_full_export(self) -> None: + """The --full export path carries only csv_data (no structured columns).""" + console = self._console() + data = { + "project_alias": "prod", + "workspace_id": 42, + "status": "completed", + "statements": [ + { + "statement_id": "stmt-1", + "status": "completed", + "rows_affected": 2, + "csv_data": "id,name\n1,alice\n2,bob\n", + } + ], + } + format_query_results(console, data) + output = cast(StringIO, console.file).getvalue() + assert "Results:" in output + assert "id,name" in output + assert "alice" in output diff --git a/tests/test_workspace_cli.py b/tests/test_workspace_cli.py index f76c9c16..5c6bc5da 100644 --- a/tests/test_workspace_cli.py +++ b/tests/test_workspace_cli.py @@ -13,6 +13,7 @@ from keboola_agent_cli.cli import app from keboola_agent_cli.config_store import ConfigStore +from keboola_agent_cli.constants import QUERY_RESULTS_DEFAULT_LIMIT from keboola_agent_cli.errors import ConfigError, KeboolaApiError from keboola_agent_cli.models import ProjectConfig from keboola_agent_cli.services.config_service import ConfigService @@ -994,6 +995,105 @@ def test_workspace_query_api_error(self, tmp_path: Path) -> None: assert output["status"] == "error" assert output["error"]["code"] == "QUERY_JOB_FAILED" + def test_workspace_query_defaults_to_fast_inline_path(self, tmp_path: Path) -> None: + """Without --full/--limit the command passes the fast-path defaults.""" + config_dir = tmp_path / "config" + config_dir.mkdir() + store = _setup_config(config_dir, {"prod": {"token": TEST_TOKEN}}) + + mock_ws = _make_workspace_mock() + mock_ws.execute_query.return_value = { + "project_alias": "prod", + "workspace_id": 42, + "status": "completed", + "statements": [], + "message": "ok", + } + + with ( + patch("keboola_agent_cli.cli.ConfigStore") as MockStore, + patch("keboola_agent_cli.cli.ProjectService") as MockProjService, + patch("keboola_agent_cli.cli.ConfigService") as MockCfgService, + patch("keboola_agent_cli.cli.JobService") as MockJobService, + patch("keboola_agent_cli.cli.WorkspaceService") as MockWsService, + ): + MockStore.return_value = store + MockProjService.return_value = ProjectService(config_store=store) + MockCfgService.return_value = ConfigService(config_store=store) + MockJobService.return_value = JobService(config_store=store) + MockWsService.return_value = mock_ws + + result = runner.invoke( + app, + [ + "--json", + "workspace", + "query", + "--project", + "prod", + "--workspace-id", + "42", + "--sql", + "SELECT 1", + ], + ) + + assert result.exit_code == 0, f"Exit {result.exit_code}: {result.output}" + kwargs = mock_ws.execute_query.call_args.kwargs + assert kwargs["full"] is False + assert kwargs["limit"] == QUERY_RESULTS_DEFAULT_LIMIT + + def test_workspace_query_full_and_limit_flags(self, tmp_path: Path) -> None: + """--full and --limit are forwarded to the service.""" + config_dir = tmp_path / "config" + config_dir.mkdir() + store = _setup_config(config_dir, {"prod": {"token": TEST_TOKEN}}) + + mock_ws = _make_workspace_mock() + mock_ws.execute_query.return_value = { + "project_alias": "prod", + "workspace_id": 42, + "status": "completed", + "statements": [], + "message": "ok", + } + + with ( + patch("keboola_agent_cli.cli.ConfigStore") as MockStore, + patch("keboola_agent_cli.cli.ProjectService") as MockProjService, + patch("keboola_agent_cli.cli.ConfigService") as MockCfgService, + patch("keboola_agent_cli.cli.JobService") as MockJobService, + patch("keboola_agent_cli.cli.WorkspaceService") as MockWsService, + ): + MockStore.return_value = store + MockProjService.return_value = ProjectService(config_store=store) + MockCfgService.return_value = ConfigService(config_store=store) + MockJobService.return_value = JobService(config_store=store) + MockWsService.return_value = mock_ws + + result = runner.invoke( + app, + [ + "--json", + "workspace", + "query", + "--project", + "prod", + "--workspace-id", + "42", + "--sql", + "SELECT 1", + "--full", + "--limit", + "2000", + ], + ) + + assert result.exit_code == 0, f"Exit {result.exit_code}: {result.output}" + kwargs = mock_ws.execute_query.call_args.kwargs + assert kwargs["full"] is True + assert kwargs["limit"] == 2000 + class TestWorkspaceFromTransformation: """Tests for `kbagent workspace from-transformation` command.""" diff --git a/tests/test_workspace_service.py b/tests/test_workspace_service.py index dbfa02cb..a8784eb2 100644 --- a/tests/test_workspace_service.py +++ b/tests/test_workspace_service.py @@ -14,6 +14,7 @@ from keboola_agent_cli.config_store import ConfigStore from keboola_agent_cli.errors import ConfigError, KeboolaApiError from keboola_agent_cli.models import ProjectConfig, TokenVerifyResponse +from keboola_agent_cli.services import workspace_service as workspace_service_module from keboola_agent_cli.services.workspace_service import WorkspaceService SAMPLE_TOKEN_VERIFY = TokenVerifyResponse( @@ -931,7 +932,7 @@ class TestExecuteQuery: """Tests for WorkspaceService.execute_query().""" def test_execute_query_success(self, tmp_config_dir: Path) -> None: - """execute_query submits, polls, and exports CSV results.""" + """Default path reads inline /results: structured columns+rows + csv_data.""" mock_client = MagicMock() mock_client.list_dev_branches.return_value = SAMPLE_BRANCHES mock_client.submit_query.return_value = {"id": "qj-abc123"} @@ -941,11 +942,19 @@ def test_execute_query_success(self, tmp_config_dir: Path) -> None: { "id": "stmt-1", "status": "completed", - "resultRows": 5, + "numberOfRows": 2, }, ], } - mock_client.export_query_results.return_value = "col1,col2\na,b\nc,d\n" + mock_client.get_query_results.return_value = { + "status": "completed", + "columns": [ + {"name": "col1", "type": "VARCHAR", "nullable": True}, + {"name": "col2", "type": "VARCHAR", "nullable": True}, + ], + "data": [["a", "b"], ["c", "d"]], + "numberOfRows": 2, + } store = setup_single_project(tmp_config_dir) svc = WorkspaceService( @@ -965,10 +974,29 @@ def test_execute_query_success(self, tmp_config_dir: Path) -> None: assert result["query_job_id"] == "qj-abc123" assert result["status"] == "completed" assert len(result["statements"]) == 1 - assert result["statements"][0]["statement_id"] == "stmt-1" - assert result["statements"][0]["status"] == "completed" - assert result["statements"][0]["rows_affected"] == 5 - assert result["statements"][0]["csv_data"] == "col1,col2\na,b\nc,d\n" + stmt = result["statements"][0] + assert stmt["statement_id"] == "stmt-1" + assert stmt["status"] == "completed" + assert stmt["rows_affected"] == 2 + assert stmt["columns"] == [ + {"name": "col1", "type": "VARCHAR", "nullable": True}, + {"name": "col2", "type": "VARCHAR", "nullable": True}, + ] + assert stmt["rows"] == [["a", "b"], ["c", "d"]] + assert stmt["row_count"] == 2 + assert stmt["total_rows"] == 2 + assert stmt["truncated"] is False + # csv_data synthesized from columns+rows for legacy consumers. + assert stmt["csv_data"] == "col1,col2\na,b\nc,d\n" + # Default path uses the fast inline endpoint, not the CSV export. + # page_size = min(QUERY_RESULTS_PAGE_SIZE, limit) with the default limit. + mock_client.get_query_results.assert_called_once_with( + "qj-abc123", + "stmt-1", + offset=0, + page_size=workspace_service_module.QUERY_RESULTS_PAGE_SIZE, + ) + mock_client.export_query_results.assert_not_called() mock_client.submit_query.assert_called_once_with( branch_id=100, @@ -979,6 +1007,115 @@ def test_execute_query_success(self, tmp_config_dir: Path) -> None: # close() called twice: once in _resolve_branch_id, once in execute_query assert mock_client.close.call_count == 2 + def test_execute_query_full_uses_csv_export(self, tmp_config_dir: Path) -> None: + """full=True takes the legacy CSV export path (complete result set).""" + mock_client = MagicMock() + mock_client.list_dev_branches.return_value = SAMPLE_BRANCHES + mock_client.submit_query.return_value = {"id": "qj-full"} + mock_client.wait_for_query_job.return_value = { + "status": "completed", + "statements": [{"id": "stmt-1", "status": "completed", "numberOfRows": 5}], + } + mock_client.export_query_results.return_value = "col1,col2\na,b\nc,d\n" + + store = setup_single_project(tmp_config_dir) + svc = WorkspaceService( + config_store=store, + client_factory=lambda url, token: mock_client, + ) + + result = svc.execute_query( + alias="prod", workspace_id=42, sql="SELECT * FROM orders", full=True + ) + + stmt = result["statements"][0] + assert stmt["csv_data"] == "col1,col2\na,b\nc,d\n" + # Full export path carries no structured columns/rows. + assert "columns" not in stmt + assert "rows" not in stmt + mock_client.export_query_results.assert_called_once_with("qj-full", "stmt-1") + mock_client.get_query_results.assert_not_called() + + def test_execute_query_inline_pagination( + self, tmp_config_dir: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """A --limit larger than one page walks offset until the limit is reached.""" + # Shrink the page size so a 4-row limit needs two /results calls. + monkeypatch.setattr(workspace_service_module, "QUERY_RESULTS_PAGE_SIZE", 2) + mock_client = MagicMock() + mock_client.list_dev_branches.return_value = SAMPLE_BRANCHES + mock_client.submit_query.return_value = {"id": "qj-page"} + mock_client.wait_for_query_job.return_value = { + "status": "completed", + "statements": [{"id": "stmt-1", "status": "completed", "numberOfRows": 10}], + } + cols = [{"name": "id", "type": "INTEGER", "nullable": False}] + mock_client.get_query_results.side_effect = [ + {"status": "completed", "columns": cols, "data": [[1], [2]], "numberOfRows": 10}, + {"status": "completed", "columns": cols, "data": [[3], [4]], "numberOfRows": 10}, + ] + + store = setup_single_project(tmp_config_dir) + svc = WorkspaceService( + config_store=store, + client_factory=lambda url, token: mock_client, + ) + + result = svc.execute_query(alias="prod", workspace_id=42, sql="SELECT id FROM big", limit=4) + + stmt = result["statements"][0] + assert stmt["rows"] == [[1], [2], [3], [4]] + assert stmt["row_count"] == 4 + assert stmt["total_rows"] == 10 + assert stmt["truncated"] is True # 4 fetched < 10 total + assert mock_client.get_query_results.call_count == 2 + mock_client.get_query_results.assert_any_call("qj-page", "stmt-1", offset=0, page_size=2) + mock_client.get_query_results.assert_any_call("qj-page", "stmt-1", offset=2, page_size=2) + + def test_execute_query_small_limit_keeps_valid_page_size(self, tmp_config_dir: Path) -> None: + """A small --limit must NOT shrink pageSize below the API floor (100..100000). + + Regression: deriving pageSize from --limit (e.g. 5) made the Query Service + reject the /results call with 400 'Invalid pageSize parameter, must be + between 100 and 100000'. pageSize is now a fixed valid value; --limit only + trims the accumulated rows locally. + """ + mock_client = MagicMock() + mock_client.list_dev_branches.return_value = SAMPLE_BRANCHES + mock_client.submit_query.return_value = {"id": "qj-small"} + mock_client.wait_for_query_job.return_value = { + "status": "completed", + "statements": [{"id": "stmt-1", "status": "completed", "numberOfRows": 25}], + } + # One full page returns all 25 rows; the service must trim to --limit. + mock_client.get_query_results.return_value = { + "status": "completed", + "columns": [{"name": "id", "type": "INTEGER", "nullable": False}], + "data": [[i] for i in range(25)], + "numberOfRows": 25, + } + + store = setup_single_project(tmp_config_dir) + svc = WorkspaceService( + config_store=store, + client_factory=lambda url, token: mock_client, + ) + + result = svc.execute_query(alias="prod", workspace_id=42, sql="SELECT id FROM t", limit=5) + + stmt = result["statements"][0] + assert stmt["row_count"] == 5 # trimmed locally + assert stmt["total_rows"] == 25 + assert stmt["truncated"] is True + # pageSize stays at the fixed valid value, NOT the small --limit. + mock_client.get_query_results.assert_called_once_with( + "qj-small", + "stmt-1", + offset=0, + page_size=workspace_service_module.QUERY_RESULTS_PAGE_SIZE, + ) + assert workspace_service_module.QUERY_RESULTS_PAGE_SIZE >= 100 + def test_execute_query_with_active_branch(self, tmp_config_dir: Path) -> None: """execute_query uses active_branch_id when set.""" mock_client = MagicMock() @@ -1063,22 +1200,48 @@ def test_execute_query_no_result_rows(self, tmp_config_dir: Path) -> None: assert result["statements"][0]["rows_affected"] == 0 assert "csv_data" not in result["statements"][0] + mock_client.get_query_results.assert_not_called() mock_client.export_query_results.assert_not_called() - def test_execute_query_export_fails_gracefully(self, tmp_config_dir: Path) -> None: - """execute_query handles export failure gracefully (no csv_data in result).""" + def test_execute_query_inline_fetch_fails_gracefully(self, tmp_config_dir: Path) -> None: + """A failed inline /results fetch must not sink the whole query.""" + mock_client = MagicMock() + mock_client.list_dev_branches.return_value = SAMPLE_BRANCHES + mock_client.submit_query.return_value = {"id": "qj-fetch-fail"} + mock_client.wait_for_query_job.return_value = { + "status": "completed", + "statements": [{"id": "stmt-1", "status": "completed", "numberOfRows": 10}], + } + mock_client.get_query_results.side_effect = KeboolaApiError( + message="Results unavailable", + error_code="QUERY_JOB_FAILED", + status_code=500, + retryable=False, + ) + + store = setup_single_project(tmp_config_dir) + svc = WorkspaceService( + config_store=store, + client_factory=lambda url, token: mock_client, + ) + + result = svc.execute_query(alias="prod", workspace_id=42, sql="SELECT * FROM big") + + # Statement still reports status + row count, just without a data payload. + stmt = result["statements"][0] + assert result["status"] == "completed" + assert stmt["rows_affected"] == 10 + assert "csv_data" not in stmt + assert "rows" not in stmt + + def test_execute_query_full_export_fails_gracefully(self, tmp_config_dir: Path) -> None: + """A failed CSV export (full=True) must not sink the whole query.""" mock_client = MagicMock() mock_client.list_dev_branches.return_value = SAMPLE_BRANCHES mock_client.submit_query.return_value = {"id": "qj-export-fail"} mock_client.wait_for_query_job.return_value = { "status": "completed", - "statements": [ - { - "id": "stmt-1", - "status": "completed", - "resultRows": 10, - }, - ], + "statements": [{"id": "stmt-1", "status": "completed", "numberOfRows": 10}], } mock_client.export_query_results.side_effect = KeboolaApiError( message="Export unavailable", @@ -1094,9 +1257,7 @@ def test_execute_query_export_fails_gracefully(self, tmp_config_dir: Path) -> No ) result = svc.execute_query( - alias="prod", - workspace_id=42, - sql="SELECT * FROM big_table", + alias="prod", workspace_id=42, sql="SELECT * FROM big_table", full=True ) # Should still succeed, just without csv_data diff --git a/uv.lock b/uv.lock index b5617c2a..7bd53f6b 100644 --- a/uv.lock +++ b/uv.lock @@ -580,7 +580,7 @@ wheels = [ [[package]] name = "keboola-agent-cli" -version = "0.58.0" +version = "0.59.0" source = { editable = "." } dependencies = [ { name = "croniter" }, From 79be60002a94b32d3951d5c258ae5eb97940d232 Mon Sep 17 00:00:00 2001 From: Petr Date: Wed, 10 Jun 2026 12:40:37 +0200 Subject: [PATCH 2/3] fix(workspace): address PR review findings on `workspace query` fast path Devin Review + kbagent-pr-reviewer follow-ups (all non-blocking): - csv_data complex types: VARIANT/ARRAY/OBJECT (Snowflake) and STRUCT/ARRAY (BigQuery) cells that arrive as native dict/list are now serialized as compact JSON ({"k":"v"}) in the synthesized csv_data instead of Python repr ({'k': 'v'}), matching the warehouse CSV export. New _csv_cell helper. - truncated robustness: when the Query Service omits numberOfRows, fall back to how the pagination loop ended -- stopping at the --limit cap with a full last page (not exhausted) now reports truncated=True instead of silently False. - --limit validation: reject a non-positive --limit. min=1 on the Typer option (CLI usage error) and Field(ge=1) on the REST WorkspaceQuery model, so a zero/negative limit no longer silently yields an empty result. - docs: keboola-expert.md tool-selection matrix + workspace SQL-debugging workflow now warn that `workspace query` is capped at --limit (default 500) -- use COUNT(*) for counts, check truncated/total_rows, --full for the complete set. gotchas.md notes the complex-type JSON serialization. - tests: + close.call_count assertions on the new service tests, + complex-type csv, truncated-without-count, and non-positive --limit coverage. The client.py file-size finding (3232 LOC > 2000 hard cap) is pre-existing tech debt and out of scope for this feature PR; tracked as a follow-up split. --- plugins/kbagent/agents/keboola-expert.md | 9 ++- .../skills/kbagent/references/gotchas.md | 3 + src/keboola_agent_cli/commands/workspace.py | 1 + .../server/routers/workspaces.py | 5 +- .../services/workspace_service.py | 34 ++++++++-- tests/test_workspace_cli.py | 40 +++++++++++ tests/test_workspace_service.py | 67 +++++++++++++++++++ 7 files changed, 152 insertions(+), 7 deletions(-) diff --git a/plugins/kbagent/agents/keboola-expert.md b/plugins/kbagent/agents/keboola-expert.md index 6d408a8f..40d25b0f 100644 --- a/plugins/kbagent/agents/keboola-expert.md +++ b/plugins/kbagent/agents/keboola-expert.md @@ -100,7 +100,7 @@ a critical failure. | Promote typed rebuild back into the original name | `kbagent storage swap-tables --project P --table-id in.c-foo.data --target-table-id in.c-foo.data_change_log --branch --yes` (0.28.0+) -- async storage job (`tableSwap`); client polls to completion. Service refuses without a branch; any branch incl. prod | -- | renaming or deleting + re-uploading (loses history; downstream configs need to be rewritten) | | Re-seed a table without losing its schema / PK / dependents | `kbagent storage truncate-table --project P --table-id in.c-foo.data [--branch ID] [--dry-run] [--yes]` (0.32.0+) -- DELETE `/tables/{id}/rows?allowTruncate=1`; endpoint is uniformly async on every branch (returns a queued `tableRowsDelete` job; client polls via `_wait_for_storage_job`). Do NOT pass `async=true` -- the API rejects it. Batch via repeated `--table-id`. Returns `{truncated[], failed[], dry_run, project_alias}` with `truncated[]` entries carrying `{table_id, rows_before, rows_after, branch_id}`. Permission class: `destructive` | `tool call delete_table_rows` if the upstream MCP exposes it | drop + recreate the table (loses descriptions, PK, sharing edges, and breaks every downstream config reference); deleting rows via raw SQL in a workspace (bypasses the Storage API audit trail) | | 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 | -| Ad-hoc SQL / row-count / type audit | `kbagent workspace create` + `kbagent workspace load` + `kbagent workspace query --sql "..."` | `kbagent workspace from-transformation` for existing transform debugging; `workspace list --qs-compatible` (0.42.0+, #304) for data-app reuse | querying Keboola Storage directly via Snowflake credentials outside the workspace abstraction | +| Ad-hoc SQL / row-count / type audit | `kbagent workspace create` + `kbagent workspace load` + `kbagent workspace query --sql "..."` (0.59.0+: results come back inline+fast but **capped at `--limit`, default 500** -- check `statements[].truncated`/`total_rows`, use `COUNT(*)` for counts, `--full` for the complete set) | `kbagent workspace from-transformation` for existing transform debugging; `workspace list --qs-compatible` (0.42.0+, #304) for data-app reuse | trusting a default `SELECT *` as the full result (it is truncated at 500); querying Storage via raw Snowflake credentials outside the workspace abstraction | | Inspect dev branch | `kbagent branch list --project P`, `kbagent branch use --project P --branch ID` | `tool call get_branch` | acting on `main` when a dev branch exists | | Audit project capabilities / features | `kbagent project info --project P` (0.30.0+) -- returns project ID, name, backend, enabled features, quota limits, and metrics | `tool call verify_token` (returns less structured info; no feature list) | inspecting the UI project settings manually | | Manage feature flags (stack catalogue / project / user) | `kbagent feature list\|project-show\|project-add\|project-remove\|user-show\|user-add\|user-remove --project P [--email E] [--feature NAME] [--dry-run] [--yes]` (0.48.0+) -- Manage API; needs a SUPER-ADMIN manage token (interactive prompt; `--allow-env-manage-token`+`KBC_MANAGE_API_TOKEN` for CI); `--project` resolves the stack URL (+project_id for `project-*`); add=admin, remove=destructive; add body is `{"feature":NAME}` | `kbagent project info` for a project's *enabled* features (read-only, no super-admin) | raw `/manage/...` calls; manage token via a CLI flag | @@ -358,6 +358,13 @@ kbagent workspace delete --project P --workspace-id W Use this for TYPE AUDITS before planning retypes, ROW COUNT COMPARISONS between branches, and SQL DEBUGGING of failing transformations. +(0.59.0+) `workspace query` returns results inline and fast, but **capped at +`--limit` rows (default 500)**. For ROW COUNTS use `SELECT COUNT(*)` (one row, +never truncated), NOT `len(rows)` of a `SELECT *`. For exact comparisons of a +result set bigger than the cap, raise `--limit` or pass `--full` (complete CSV +export, slower). Always check `statements[].truncated` / `total_rows` in `--json` +before treating the rows as complete. + ### 4.5 Cross-project migration (high-risk) Preconditions (REFUSE if not met): diff --git a/plugins/kbagent/skills/kbagent/references/gotchas.md b/plugins/kbagent/skills/kbagent/references/gotchas.md index dc1f0536..4cc7851c 100644 --- a/plugins/kbagent/skills/kbagent/references/gotchas.md +++ b/plugins/kbagent/skills/kbagent/references/gotchas.md @@ -309,6 +309,9 @@ interactive queries are markedly faster. - Each statement carries structured `columns`, `rows`, `row_count`, `total_rows`, `truncated`, **and** a synthesized `csv_data` string. Parsers that read `csv_data` (the pre-0.59.0 shape) keep working unchanged. + VARIANT/ARRAY/OBJECT (Snowflake) and STRUCT/ARRAY (BigQuery) cells are + emitted in `csv_data` as compact JSON (`{"k":"v"}`) to match the warehouse + CSV export, not Python `repr`. - `--full` opts back into the complete CSV export -- slower (warehouse UNLOAD), but **uncapped**. Use it when you need every row, e.g. a bulk extract (`workspace query --full --json`). Under `--full` the statement carries only diff --git a/src/keboola_agent_cli/commands/workspace.py b/src/keboola_agent_cli/commands/workspace.py index 45b15439..b07247de 100644 --- a/src/keboola_agent_cli/commands/workspace.py +++ b/src/keboola_agent_cli/commands/workspace.py @@ -424,6 +424,7 @@ def workspace_query( limit: int = typer.Option( QUERY_RESULTS_DEFAULT_LIMIT, "--limit", + min=1, help="Max rows to fetch via the fast inline path (ignored with --full).", ), ) -> None: diff --git a/src/keboola_agent_cli/server/routers/workspaces.py b/src/keboola_agent_cli/server/routers/workspaces.py index fc4e2b53..13a2f201 100644 --- a/src/keboola_agent_cli/server/routers/workspaces.py +++ b/src/keboola_agent_cli/server/routers/workspaces.py @@ -8,7 +8,7 @@ from fastapi import APIRouter, Depends, HTTPException, Query from fastapi.responses import StreamingResponse -from pydantic import BaseModel +from pydantic import BaseModel, Field from ...constants import AI_SQL_HELPER_TIMEOUT, QUERY_RESULTS_DEFAULT_LIMIT from ..dependencies import ServiceRegistry, get_registry @@ -36,7 +36,8 @@ class WorkspaceQuery(BaseModel): # complete result set. The fast inline path (full=False) is paginated, so a # REST client must opt in explicitly until the frontend learns to paginate. full: bool = True - limit: int = QUERY_RESULTS_DEFAULT_LIMIT + # ge=1: a zero/negative limit would otherwise silently yield an empty result. + limit: int = Field(default=QUERY_RESULTS_DEFAULT_LIMIT, ge=1) class FromTransformation(BaseModel): diff --git a/src/keboola_agent_cli/services/workspace_service.py b/src/keboola_agent_cli/services/workspace_service.py index 9446d2e8..8c1a5ba1 100644 --- a/src/keboola_agent_cli/services/workspace_service.py +++ b/src/keboola_agent_cli/services/workspace_service.py @@ -7,6 +7,7 @@ import csv import io +import json import logging from dataclasses import dataclass from typing import Any @@ -165,19 +166,35 @@ class InlineQueryResult: truncated: bool # True when the warehouse has more rows than we fetched +def _csv_cell(value: Any) -> Any: + """Coerce one `/results` JSON cell to its CSV representation. + + ``None`` -> empty field (matches the warehouse CSV export). VARIANT/ARRAY/ + OBJECT (Snowflake) and STRUCT/ARRAY (BigQuery) columns arrive as native + Python ``dict``/``list``; ``csv.writer`` would otherwise emit their Python + ``repr`` (``{'k': 'v'}``), so we serialize them as compact JSON + (``{"k":"v"}``) to match the warehouse's CSV serialization. Scalars pass + through and are stringified by ``csv.writer``. + """ + if value is None: + return "" + if isinstance(value, (dict, list)): + return json.dumps(value, separators=(",", ":"), ensure_ascii=False) + return value + + def _rows_to_csv(columns: list[dict[str, Any]], rows: list[list[Any]]) -> str: """Render structured columns+rows as an RFC-4180 CSV string. Synthesized so the inline `/results` payload stays drop-in compatible with consumers that still read ``csv_data`` (CLI preview, web UI table + export - buttons, REST). ``csv.writer`` handles quoting/escaping; ``None`` becomes an - empty field, matching the warehouse CSV export semantics. + buttons, REST). ``csv.writer`` handles quoting/escaping. """ buffer = io.StringIO() writer = csv.writer(buffer, lineterminator="\n") writer.writerow([col.get("name", "") for col in columns]) for row in rows: - writer.writerow(["" if value is None else value for value in row]) + writer.writerow([_csv_cell(value) for value in row]) return buffer.getvalue() @@ -200,6 +217,7 @@ def _collect_inline_results( columns: list[dict[str, Any]] = [] total_rows: int | None = None offset = 0 + exhausted = False while len(collected) < limit: payload = client.get_query_results( query_job_id, statement_id, offset=offset, page_size=QUERY_RESULTS_PAGE_SIZE @@ -212,11 +230,19 @@ def _collect_inline_results( collected.extend(page_rows) # Last page: the warehouse returned fewer rows than a full page. if len(page_rows) < QUERY_RESULTS_PAGE_SIZE: + exhausted = True break offset += len(page_rows) rows = collected[:limit] - truncated = total_rows is not None and total_rows > len(rows) + if total_rows is not None: + truncated = total_rows > len(rows) + else: + # The Query Service normally reports numberOfRows, but if it omits the + # count we fall back to *how* the loop ended: stopping at the limit cap + # without exhausting a full last page means there may be more rows. Bias + # toward over-warning ("use --full") when the true count is unknown. + truncated = not exhausted and len(collected) >= limit return InlineQueryResult( columns=columns, rows=rows, diff --git a/tests/test_workspace_cli.py b/tests/test_workspace_cli.py index 5c6bc5da..fb31d7be 100644 --- a/tests/test_workspace_cli.py +++ b/tests/test_workspace_cli.py @@ -1094,6 +1094,46 @@ def test_workspace_query_full_and_limit_flags(self, tmp_path: Path) -> None: assert kwargs["full"] is True assert kwargs["limit"] == 2000 + def test_workspace_query_rejects_non_positive_limit(self, tmp_path: Path) -> None: + """--limit 0 is rejected by Typer (min=1) before the service is called.""" + config_dir = tmp_path / "config" + config_dir.mkdir() + store = _setup_config(config_dir, {"prod": {"token": TEST_TOKEN}}) + + mock_ws = _make_workspace_mock() + + with ( + patch("keboola_agent_cli.cli.ConfigStore") as MockStore, + patch("keboola_agent_cli.cli.ProjectService") as MockProjService, + patch("keboola_agent_cli.cli.ConfigService") as MockCfgService, + patch("keboola_agent_cli.cli.JobService") as MockJobService, + patch("keboola_agent_cli.cli.WorkspaceService") as MockWsService, + ): + MockStore.return_value = store + MockProjService.return_value = ProjectService(config_store=store) + MockCfgService.return_value = ConfigService(config_store=store) + MockJobService.return_value = JobService(config_store=store) + MockWsService.return_value = mock_ws + + result = runner.invoke( + app, + [ + "workspace", + "query", + "--project", + "prod", + "--workspace-id", + "42", + "--sql", + "SELECT 1", + "--limit", + "0", + ], + ) + + assert result.exit_code == 2 + mock_ws.execute_query.assert_not_called() + class TestWorkspaceFromTransformation: """Tests for `kbagent workspace from-transformation` command.""" diff --git a/tests/test_workspace_service.py b/tests/test_workspace_service.py index a8784eb2..3acb1569 100644 --- a/tests/test_workspace_service.py +++ b/tests/test_workspace_service.py @@ -1035,6 +1035,8 @@ def test_execute_query_full_uses_csv_export(self, tmp_config_dir: Path) -> None: assert "rows" not in stmt mock_client.export_query_results.assert_called_once_with("qj-full", "stmt-1") mock_client.get_query_results.assert_not_called() + # close() twice: once in _resolve_branch_id, once in execute_query. + assert mock_client.close.call_count == 2 def test_execute_query_inline_pagination( self, tmp_config_dir: Path, monkeypatch: pytest.MonkeyPatch @@ -1071,6 +1073,7 @@ def test_execute_query_inline_pagination( assert mock_client.get_query_results.call_count == 2 mock_client.get_query_results.assert_any_call("qj-page", "stmt-1", offset=0, page_size=2) mock_client.get_query_results.assert_any_call("qj-page", "stmt-1", offset=2, page_size=2) + assert mock_client.close.call_count == 2 def test_execute_query_small_limit_keeps_valid_page_size(self, tmp_config_dir: Path) -> None: """A small --limit must NOT shrink pageSize below the API floor (100..100000). @@ -1115,6 +1118,7 @@ def test_execute_query_small_limit_keeps_valid_page_size(self, tmp_config_dir: P page_size=workspace_service_module.QUERY_RESULTS_PAGE_SIZE, ) assert workspace_service_module.QUERY_RESULTS_PAGE_SIZE >= 100 + assert mock_client.close.call_count == 2 def test_execute_query_with_active_branch(self, tmp_config_dir: Path) -> None: """execute_query uses active_branch_id when set.""" @@ -1233,6 +1237,7 @@ def test_execute_query_inline_fetch_fails_gracefully(self, tmp_config_dir: Path) assert stmt["rows_affected"] == 10 assert "csv_data" not in stmt assert "rows" not in stmt + assert mock_client.close.call_count == 2 def test_execute_query_full_export_fails_gracefully(self, tmp_config_dir: Path) -> None: """A failed CSV export (full=True) must not sink the whole query.""" @@ -1263,6 +1268,68 @@ def test_execute_query_full_export_fails_gracefully(self, tmp_config_dir: Path) # Should still succeed, just without csv_data assert result["status"] == "completed" assert "csv_data" not in result["statements"][0] + assert mock_client.close.call_count == 2 + + def test_execute_query_truncated_when_number_of_rows_missing( + self, tmp_config_dir: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """If the Query Service omits numberOfRows, fall back to how the loop ended. + + Stopping at the --limit cap with a full last page (not exhausted) means + there may be more rows, so `truncated` must be True even without a count. + """ + monkeypatch.setattr(workspace_service_module, "QUERY_RESULTS_PAGE_SIZE", 2) + mock_client = MagicMock() + mock_client.list_dev_branches.return_value = SAMPLE_BRANCHES + mock_client.submit_query.return_value = {"id": "qj-nocount"} + mock_client.wait_for_query_job.return_value = { + "status": "completed", + "statements": [{"id": "stmt-1", "status": "completed", "numberOfRows": 2}], + } + # A full page (== page_size) with NO numberOfRows in the payload. + mock_client.get_query_results.return_value = { + "status": "completed", + "columns": [{"name": "id"}], + "data": [[1], [2]], + } + + store = setup_single_project(tmp_config_dir) + svc = WorkspaceService( + config_store=store, + client_factory=lambda url, token: mock_client, + ) + + result = svc.execute_query(alias="prod", workspace_id=42, sql="SELECT id FROM t", limit=2) + + stmt = result["statements"][0] + assert stmt["total_rows"] is None + assert stmt["truncated"] is True # full last page, capped at limit -> maybe more + + +class TestRowsToCsv: + """Tests for the module-level _rows_to_csv / _csv_cell helpers.""" + + def test_none_becomes_empty_field(self) -> None: + csv_str = workspace_service_module._rows_to_csv( + [{"name": "a"}, {"name": "b"}], [["x", None]] + ) + assert csv_str == "a,b\nx,\n" + + def test_dict_and_list_cells_serialize_as_compact_json(self) -> None: + """VARIANT/ARRAY/OBJECT cells arrive as Python dict/list -- emit compact + JSON (``{"k":"v"}``), not Python repr (``{'k': 'v'}``), to match the + warehouse CSV export. + """ + csv_str = workspace_service_module._rows_to_csv( + [{"name": "payload"}, {"name": "tags"}], + [[{"k": "v", "n": 1}, [1, 2, 3]]], + ) + # csv.writer quotes fields containing commas. + assert '"{""k"":""v"",""n"":1}"' in csv_str + assert '"[1,2,3]"' in csv_str + # No Python-repr artifacts (single quotes / spaces after colon). + assert "'k'" not in csv_str + assert "{'" not in csv_str class TestCreateFromTransformation: From 3e3d41c408ad8b64d46dc42dbf8e0b8a003a053a Mon Sep 17 00:00:00 2001 From: Petr Date: Wed, 10 Jun 2026 12:46:08 +0200 Subject: [PATCH 3/3] refactor(workspace): address review nits on the inline-results pagination - NIT-1: break the /results pagination loop as soon as the accumulated offset reaches the reported total_rows on a page boundary, instead of spending a round-trip on the empty next page (e.g. total == a multiple of the page size with --limit larger than total). - NIT-2: clarify in constants.py that QUERY_RESULTS_PAGE_SIZE (API wire constraint, 100..100000) and QUERY_RESULTS_DEFAULT_LIMIT (user-facing row cap) are distinct concepts that merely share the value 500. - test: the early break makes no wasted round-trip on a page boundary. --- src/keboola_agent_cli/constants.py | 7 +++- .../services/workspace_service.py | 6 ++++ tests/test_workspace_service.py | 36 +++++++++++++++++++ 3 files changed, 48 insertions(+), 1 deletion(-) diff --git a/src/keboola_agent_cli/constants.py b/src/keboola_agent_cli/constants.py index 62b916c7..d98c96d8 100644 --- a/src/keboola_agent_cli/constants.py +++ b/src/keboola_agent_cli/constants.py @@ -420,8 +420,13 @@ # rows, accumulated in pages of QUERY_RESULTS_PAGE_SIZE and trimmed to the limit. # The endpoint enforces 100 <= pageSize <= 100000, so QUERY_RESULTS_PAGE_SIZE is a # fixed valid page size -- it is NOT derived from --limit (a small limit would 400). +# Two DISTINCT concepts that happen to share the value 500: +# DEFAULT_LIMIT = user-facing row cap (the --limit default; freely tunable). +# PAGE_SIZE = API wire constraint -- rows per /results request, which the +# endpoint requires to be within 100..100000. Independent of +# --limit; the service pages by PAGE_SIZE and trims to --limit. QUERY_RESULTS_DEFAULT_LIMIT: int = 500 # default --limit for `workspace query` fast path -QUERY_RESULTS_PAGE_SIZE: int = 500 # rows per /results page (within the API's 100..100000) +QUERY_RESULTS_PAGE_SIZE: int = 500 # rows per /results page (API requires 100..100000) # --- Workspace Defaults --- DEFAULT_WORKSPACE_BACKEND: str = "snowflake" diff --git a/src/keboola_agent_cli/services/workspace_service.py b/src/keboola_agent_cli/services/workspace_service.py index 8c1a5ba1..3b705c4f 100644 --- a/src/keboola_agent_cli/services/workspace_service.py +++ b/src/keboola_agent_cli/services/workspace_service.py @@ -233,6 +233,12 @@ def _collect_inline_results( exhausted = True break offset += len(page_rows) + # Reached the reported total on a page boundary: stop without spending a + # round-trip on the empty next page (e.g. total == a multiple of the + # page size, limit larger than total). + if total_rows is not None and offset >= total_rows: + exhausted = True + break rows = collected[:limit] if total_rows is not None: diff --git a/tests/test_workspace_service.py b/tests/test_workspace_service.py index 3acb1569..345b5dbf 100644 --- a/tests/test_workspace_service.py +++ b/tests/test_workspace_service.py @@ -1305,6 +1305,42 @@ def test_execute_query_truncated_when_number_of_rows_missing( assert stmt["total_rows"] is None assert stmt["truncated"] is True # full last page, capped at limit -> maybe more + def test_execute_query_stops_at_total_on_page_boundary( + self, tmp_config_dir: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """When total_rows lands on a page boundary, do not spend a round-trip on + the empty next page (NIT-1).""" + monkeypatch.setattr(workspace_service_module, "QUERY_RESULTS_PAGE_SIZE", 2) + mock_client = MagicMock() + mock_client.list_dev_branches.return_value = SAMPLE_BRANCHES + mock_client.submit_query.return_value = {"id": "qj-boundary"} + mock_client.wait_for_query_job.return_value = { + "status": "completed", + "statements": [{"id": "stmt-1", "status": "completed", "numberOfRows": 2}], + } + # A full page (== page_size) that already covers numberOfRows. + mock_client.get_query_results.return_value = { + "status": "completed", + "columns": [{"name": "id"}], + "data": [[1], [2]], + "numberOfRows": 2, + } + + store = setup_single_project(tmp_config_dir) + svc = WorkspaceService( + config_store=store, + client_factory=lambda url, token: mock_client, + ) + + # limit (3) > total (2): without the early break this would make a 2nd + # call to discover the empty page. + result = svc.execute_query(alias="prod", workspace_id=42, sql="SELECT id FROM t", limit=3) + + stmt = result["statements"][0] + assert stmt["row_count"] == 2 + assert stmt["truncated"] is False + assert mock_client.get_query_results.call_count == 1 # no wasted round-trip + class TestRowsToCsv: """Tests for the module-level _rows_to_csv / _csv_cell helpers."""