diff --git a/.claude-plugin/marketplace.json b/.claude-plugin/marketplace.json index 1d1a04d5..8cd1b16d 100644 --- a/.claude-plugin/marketplace.json +++ b/.claude-plugin/marketplace.json @@ -10,7 +10,7 @@ "plugins": [ { "name": "kbagent", - "version": "0.60.4", + "version": "0.61.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/README.md b/README.md index 9017a0df..11c6cc33 100644 --- a/README.md +++ b/README.md @@ -93,6 +93,24 @@ kbagent init --from-global --read-only Three protection layers (kbagent policy + filesystem chmod + Claude Code deny rules) prevent the agent from writing, deleting, or bypassing restrictions. See [Permissions Guide](docs/guide.md#permissions) for details. +## Use as a library + +Besides the CLI and `kbagent serve`, kbagent exposes a small **stateless, importable client** for in-process use -- a Keboola Data App, a transformation, or any Python service can run Query Service SQL and read/write Storage Files without spawning the CLI, running the daemon, or maintaining a config-dir. Auth is the storage token you pass in (12-factor); nothing is written to disk. + +```python +import os +from keboola_agent_cli import Client + +with Client(url=os.environ["KBC_URL"], token=os.environ["KBC_TOKEN"]) as kbc: + rows = kbc.query(workspace_id, "SELECT id, name FROM customers") # list[dict] + + meta = kbc.files.upload(b"hello", name="greeting.txt", tags=["demo"]) + data = kbc.files.read_bytes(meta.id) # bytes + files = kbc.files.list(tags=["demo"]) # list[FileEntry] +``` + +`query()` reads results inline (fast, native JSON types) and returns rows keyed by column name; `files` returns a uniform `FileEntry` shape and reads bytes straight into memory. Everything exported from `keboola_agent_cli` is committed public API (semver). For lower-level endpoints, reach for `Client.raw` (the underlying `KeboolaClient`). + ## 30-second demo ![30-second demo](docs/assets/demo-readme-main.gif) diff --git a/plugins/kbagent/.claude-plugin/plugin.json b/plugins/kbagent/.claude-plugin/plugin.json index da390e3f..c18ca528 100644 --- a/plugins/kbagent/.claude-plugin/plugin.json +++ b/plugins/kbagent/.claude-plugin/plugin.json @@ -1,6 +1,6 @@ { "name": "kbagent", - "version": "0.60.4", + "version": "0.61.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/agents/keboola-expert.md b/plugins/kbagent/agents/keboola-expert.md index 40d25b0f..427fbccf 100644 --- a/plugins/kbagent/agents/keboola-expert.md +++ b/plugins/kbagent/agents/keboola-expert.md @@ -101,6 +101,7 @@ a critical failure. | 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 "..."` (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 | +| Run Keboola SQL or read/write Storage Files from INSIDE a Python process you control (Data App, transformation, hosted service) | `from keboola_agent_cli import Client` (0.61.0+) -- stateless `Client(url, token)`; `.query(workspace_id, sql) -> list[dict]`, `.files.upload(path_or_bytes)` / `.files.read_bytes(id) -> bytes` / `.files.list() -> [FileEntry]`; no CLI subprocess, no `serve`, no config-dir | the `kbagent` CLI or `kbagent serve` REST when you are NOT already inside Python | shelling out to the `kbagent` binary from a Python process you control (import the library instead); using it for AI-driven exploration (it is fixed typed ops, not MCP tools) | | 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 | diff --git a/plugins/kbagent/skills/kbagent/references/gotchas.md b/plugins/kbagent/skills/kbagent/references/gotchas.md index 7bea2f5b..aea3d2be 100644 --- a/plugins/kbagent/skills/kbagent/references/gotchas.md +++ b/plugins/kbagent/skills/kbagent/references/gotchas.md @@ -2663,3 +2663,16 @@ cross-project flow migration is a manual dance (`sync pull` source, edit, A `sync pull` without the right flags leaves column metadata empty in the local JSON. That does NOT mean Keboola has no metadata -- always re-fetch via `kbagent storage table-detail` when deciding about types. + +### `Client` library: `query()` needs a provisioned workspace; `branch_id=None` costs a branch-list call (since v0.61.0) + +The in-process library facade (`from keboola_agent_cli import Client`, 0.61.0+) +is a thin wrapper, not a workspace manager. Two non-obvious behaviors: + +- **`query(workspace_id, sql)` does NOT create a workspace.** The `workspace_id` + must already exist (make one via `kbagent workspace create` or the Storage + API first). An unknown id surfaces the Query Service error verbatim. +- **`Client(url, token)` with no `branch_id` resolves the default branch lazily + on the first `query()`** -- one extra `list_dev_branches` API call, cached + after. Pass `branch_id=` to skip it (and to target a dev branch). Storage + Files default to the production scope when `branch_id` is unset. diff --git a/pyproject.toml b/pyproject.toml index ea54f9a2..7f8b44be 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "keboola-agent-cli" -version = "0.60.4" +version = "0.61.0" description = "AI-friendly CLI for managing Keboola projects" readme = "README.md" requires-python = ">=3.12" diff --git a/src/keboola_agent_cli/__init__.py b/src/keboola_agent_cli/__init__.py index be737303..66680285 100644 --- a/src/keboola_agent_cli/__init__.py +++ b/src/keboola_agent_cli/__init__.py @@ -3,8 +3,11 @@ from importlib.metadata import PackageNotFoundError, version from .constants import APP_NAME +from .lib import Client, FileEntry, Files try: __version__ = version(APP_NAME) except PackageNotFoundError: __version__ = "0.0.0-dev" + +__all__ = ["Client", "FileEntry", "Files", "__version__"] diff --git a/src/keboola_agent_cli/changelog.py b/src/keboola_agent_cli/changelog.py index 6d533fc6..029eedba 100644 --- a/src/keboola_agent_cli/changelog.py +++ b/src/keboola_agent_cli/changelog.py @@ -24,6 +24,22 @@ # Ordered newest-first. Each value is a list of brief one-line descriptions. CHANGELOG: dict[str, list[str]] = { + "0.61.0": [ + "New (#415): kbagent now ships a stateless, importable library facade -- " + "`from keboola_agent_cli import Client` -- so any in-process Python consumer (a Keboola " + "Data App, a transformation, a hosted service) can run Query Service SQL and read/write " + "Storage Files without a CLI subprocess, a `kbagent serve` daemon, or a config-dir. " + "`Client(url, token)` wraps the existing `KeboolaClient`; `client.query(workspace_id, sql)` " + "returns `list[dict]` rows over the fast inline `/results` path (native JSON types; " + "truncation is warned, never silently capped), and `client.files` offers " + "`upload(path_or_bytes)`, `read_bytes(file_id) -> bytes`, `list() -> list[FileEntry]` " + "(one uniform shape, read via `read_bytes` so callers never branch on a signed URL) and " + "`delete()`. The Query Service pagination helper moved from the workspace service into " + "`client.py` (re-exported, no behavior change) so the CLI and the library share one " + "implementation. Everything under `keboola_agent_cli.__all__` is committed public API. " + "Addresses the jasnost feedback points 1, 2, and 4 (point 3 -- structured query results -- " + "shipped in 0.59.0).", + ], "0.60.4": [ "Security: `kbagent serve --ui` no longer lets `GET /doctor`, `/version`, and `/changelog` " "(and any other registered endpoint) bypass bearer auth. In single-process UI mode the auth " diff --git a/src/keboola_agent_cli/client.py b/src/keboola_agent_cli/client.py index 0f207e77..ebd028ea 100644 --- a/src/keboola_agent_cli/client.py +++ b/src/keboola_agent_cli/client.py @@ -12,6 +12,7 @@ import re import time from collections.abc import Iterator +from dataclasses import dataclass from pathlib import Path from typing import Any from urllib.parse import quote @@ -47,6 +48,80 @@ logger = logging.getLogger(__name__) +@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 _collect_inline_results( + client: "KeboolaClient", + 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. + + Lives in the client layer (not a service) because it is pure Query Service + pagination over :meth:`KeboolaClient.get_query_results` -- no config, no + business logic -- so both ``WorkspaceService`` and the public library facade + (:mod:`keboola_agent_cli.lib`) can share it. + """ + collected: list[list[Any]] = [] + 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 + ) + 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: + 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: + 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 when the true count is unknown. + truncated = not exhausted and len(collected) >= limit + return InlineQueryResult( + columns=columns, + rows=rows, + total_rows=total_rows, + truncated=truncated, + ) + + def _iter_poll_intervals(strategy: str) -> Iterator[float]: """Yield sleep intervals (seconds) for Queue job polling. diff --git a/src/keboola_agent_cli/lib.py b/src/keboola_agent_cli/lib.py new file mode 100644 index 00000000..dad5d94c --- /dev/null +++ b/src/keboola_agent_cli/lib.py @@ -0,0 +1,301 @@ +"""Public, in-process library facade for Keboola (issue #415). + +A stateless importable surface so any Python consumer -- a Keboola Data App, a +transformation, a hosted service -- can use kbagent's Query Service and Storage +Files *in-process*: no CLI subprocess, no ``kbagent serve`` daemon, no +config-dir, no ``project add`` ceremony. + + from keboola_agent_cli import Client + + with Client(url=KBC_URL, token=KBC_TOKEN) as kbc: + rows = kbc.query(workspace_id, "SELECT id, name FROM t") # list[dict] + meta = kbc.files.upload(b"hello", name="greeting.txt", tags=["x"]) + data = kbc.files.read_bytes(meta.id) # bytes + metas = kbc.files.list(tags=["x"]) # list[FileEntry] + +The facade is a thin convenience wrapper over :class:`KeboolaClient` +(``client.py``); it adds the high-level shapes -- ``list[dict]`` rows, ``bytes`` +file reads, a stable :class:`FileEntry` -- that the CLI used to assemble inside +its service layer. Auth is the storage token passed at construction (12-factor: +read it from ``KBC_TOKEN`` yourself); nothing is persisted to disk. + +Everything exported here is committed public API and changes follow semver. For +lower-level access (raw Queue/Storage endpoints) reach for the underlying +:class:`KeboolaClient` via :attr:`Client.raw`. +""" + +from __future__ import annotations + +import logging +import tempfile +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from .client import KeboolaClient, _collect_inline_results +from .constants import QUERY_RESULTS_DEFAULT_LIMIT +from .errors import ErrorCode, KeboolaApiError + +logger = logging.getLogger(__name__) + +# The public ``Files.list`` method shadows the builtin ``list`` inside that +# class body, so ``list[...]`` *annotations* there would resolve to the method +# (under both runtime evaluation and ``ty``). Method *bodies* are unaffected -- +# class scope is skipped in function name resolution -- so this alias is only +# needed for annotations inside ``Files``. +_list = list + + +@dataclass(frozen=True) +class FileEntry: + """Stable, uniform shape for a Storage File across list and upload results. + + Always carries the same fields regardless of whether the underlying API + response happened to include a (sometimes-absent) signed download URL. + Read the bytes with :meth:`Files.read_bytes`, never by branching on a URL. + ``raw`` is the untouched API dict for fields the facade does not surface. + """ + + id: int + name: str + tags: list[str] + created: str | None + size_bytes: int | None + is_permanent: bool + raw: dict[str, Any] + + @classmethod + def _from_api(cls, data: dict[str, Any]) -> FileEntry: + return cls( + id=int(data["id"]), + name=data.get("name", ""), + tags=list(data.get("tags") or []), + created=data.get("created"), + size_bytes=data.get("sizeBytes"), + is_permanent=bool(data.get("isPermanent", False)), + raw=data, + ) + + +class Files: + """Storage Files operations bound to one project (and optional branch). + + Obtained via :attr:`Client.files`; not constructed directly. + """ + + def __init__(self, client: KeboolaClient, branch_id: int | None) -> None: + self._client = client + self._branch_id = branch_id + + def list( + self, + *, + tags: _list[str] | None = None, + query: str | None = None, + limit: int = 100, + offset: int = 0, + since_id: int | None = None, + ) -> _list[FileEntry]: + """List Storage Files as uniform :class:`FileEntry` records. + + Args: + tags: Filter to files carrying all of these tags (AND logic). + query: Full-text search over the file name. + limit: Max files to return. + offset: Pagination offset. + since_id: Return only files with an ID greater than this. + """ + raw = self._client.list_files( + limit=limit, + offset=offset, + tags=tags, + since_id=since_id, + query=query, + branch_id=self._branch_id, + ) + return [FileEntry._from_api(item) for item in raw] + + def upload( + self, + source: str | Path | bytes | bytearray, + *, + name: str | None = None, + tags: _list[str] | None = None, + permanent: bool = False, + ) -> FileEntry: + """Upload a local path or in-memory bytes to Storage Files. + + Args: + source: A filesystem path (``str``/``Path``) or raw ``bytes`` to + upload. When passing bytes, ``name`` is required (Storage needs + a file name and there is no path to derive it from). + name: Storage file name. Defaults to the path basename for path + sources; required for bytes sources. + tags: Tags to assign. + permanent: If True the file is not auto-expired after 15 days. + + Raises: + ValueError: If ``source`` is bytes and ``name`` is not given. + """ + if isinstance(source, (bytes, bytearray)): + if not name: + raise ValueError("name is required when uploading raw bytes") + # Reuse the battle-tested prepare + multi-cloud upload path + # (S3/GCS/Azure, issue #187 streaming) by staging the bytes in a + # temp dir rather than duplicating _upload_to_cloud. A TemporaryDirectory + # is auto-removed even if the upload raises (no leaked temp file); the + # staged file name is irrelevant -- `name` is what Storage records. + with tempfile.TemporaryDirectory() as tmpdir: + staged = Path(tmpdir) / "upload" + staged.write_bytes(bytes(source)) + info = self._client.upload_file( + file_path=str(staged), + name=name, + tags=tags, + is_permanent=permanent, + branch_id=self._branch_id, + ) + else: + info = self._client.upload_file( + file_path=str(source), + name=name, + tags=tags, + is_permanent=permanent, + branch_id=self._branch_id, + ) + return FileEntry._from_api(info) + + def read_bytes(self, file_id: int) -> bytes: + """Download a Storage File fully into memory and return its bytes. + + Hides the file-info -> signed-URL -> stream dance and transparently + handles sliced files and gzip. The whole payload is held in RAM, so use + this for reasonably sized files (results, manifests, small exports); for + multi-GB tables stream to disk via :attr:`Client.raw` instead. + """ + info = self._client.get_file_info(file_id, branch_id=self._branch_id) + with tempfile.TemporaryDirectory() as tmpdir: + out = Path(tmpdir) / "download" + if info.get("isSliced", False): + self._client.download_sliced_file(info, str(out)) + else: + url = info.get("url") + if not url: + raise KeboolaApiError( + f"Storage file {file_id} has no download URL.", + error_code=ErrorCode.API_ERROR, + ) + self._client.download_file(url, str(out)) + return out.read_bytes() + + def delete(self, file_id: int) -> None: + """Delete a Storage File.""" + self._client.delete_file(file_id, branch_id=self._branch_id) + + +class Client: + """Stateless in-process entry point to one Keboola project. + + Holds nothing but the stack URL, token, and a single :class:`KeboolaClient` + (which carries the shared retry/backoff). No config-dir, no ``project add``. + + Args: + url: Stack URL, e.g. ``https://connection.keboola.com``. + token: Storage API token for the project. + branch_id: Dev branch to scope every operation to. ``None`` (default) + targets production: Storage Files use the production scope and + :meth:`query` resolves the project's default branch on first use. + """ + + def __init__(self, url: str, token: str, *, branch_id: int | None = None) -> None: + if not url: + raise ValueError("url is required") + if not token: + raise ValueError("token is required") + self._client = KeboolaClient(stack_url=url, token=token) + self._resolved_branch_id = branch_id + self.files = Files(self._client, branch_id) + + @property + def raw(self) -> KeboolaClient: + """The underlying low-level client, for endpoints the facade omits.""" + return self._client + + def _effective_branch_id(self) -> int: + """Resolve the branch ID for query submission (caches the default).""" + if self._resolved_branch_id is not None: + return self._resolved_branch_id + for branch in self._client.list_dev_branches(): + if branch.get("isDefault", False): + self._resolved_branch_id = int(branch["id"]) + return self._resolved_branch_id + raise KeboolaApiError( + "No default branch found for this project.", + error_code=ErrorCode.NOT_FOUND, + ) + + def query( + self, + workspace_id: int, + sql: str, + *, + transactional: bool = False, + limit: int = QUERY_RESULTS_DEFAULT_LIMIT, + ) -> list[dict[str, Any]]: + """Run SQL in a workspace and return rows as a list of dicts. + + Submits the statement via Query Service, waits for completion, and reads + results inline (the fast ``/results`` path, no CSV-file materialization). + Each row is a dict keyed by the result column names exactly as the + warehouse reports them -- note Snowflake folds unquoted aliases to + UPPERCASE, so quote aliases if you want lowercase keys. Values arrive as + native JSON types (int/float/bool/None; VARIANT/STRUCT as dict/list). + + When ``sql`` contains multiple statements, the rows of the *last* + statement that produced a result set are returned (so ``USE ...; SELECT + ...`` yields the SELECT). Statements without a result set yield ``[]``. + + Args: + workspace_id: Target workspace ID. + sql: One or more SQL statements. + transactional: Wrap the statements in a transaction. + limit: Max rows to fetch (default ``QUERY_RESULTS_DEFAULT_LIMIT``). + If the warehouse has more, the result is capped and a warning is + logged -- raise ``limit`` to fetch more. + """ + branch_id = self._effective_branch_id() + job = self._client.submit_query( + branch_id=branch_id, + workspace_id=workspace_id, + statements=[sql], + transactional=transactional, + ) + job_id = str(job.get("queryJobId", job.get("id", ""))) + completed = self._client.wait_for_query_job(job_id) + + rows: list[dict[str, Any]] = [] + for stmt in completed.get("statements", []): + num_rows = stmt.get("numberOfRows", stmt.get("resultRows", 0)) + if stmt.get("status") != "completed" or not num_rows: + continue + inline = _collect_inline_results(self._client, job_id, str(stmt.get("id", "")), limit) + col_names = [col.get("name", "") for col in inline.columns] + rows = [dict(zip(col_names, row, strict=False)) for row in inline.rows] + if inline.truncated: + logger.warning( + "query result truncated to %d rows (warehouse has %s); " + "raise limit= to fetch more", + len(rows), + inline.total_rows, + ) + return rows + + def close(self) -> None: + """Close the underlying HTTP client.""" + self._client.close() + + def __enter__(self) -> Client: + return self + + def __exit__(self, *exc: object) -> None: + self.close() diff --git a/src/keboola_agent_cli/services/workspace_service.py b/src/keboola_agent_cli/services/workspace_service.py index 3b705c4f..a20412dd 100644 --- a/src/keboola_agent_cli/services/workspace_service.py +++ b/src/keboola_agent_cli/services/workspace_service.py @@ -15,10 +15,10 @@ from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa +from ..client import _collect_inline_results 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, @@ -156,16 +156,6 @@ 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 _csv_cell(value: Any) -> Any: """Coerce one `/results` JSON cell to its CSV representation. @@ -198,65 +188,6 @@ def _rows_to_csv(columns: list[dict[str, Any]], rows: list[list[Any]]) -> str: 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 - exhausted = False - 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: - 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: - 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, - total_rows=total_rows, - truncated=truncated, - ) - - class WorkspaceService(BaseService): """Business logic for managing Keboola workspaces. diff --git a/tests/test_e2e.py b/tests/test_e2e.py index bdc25ce0..c4826fd5 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -526,6 +526,9 @@ def test_full_cli_e2e(self) -> None: _step(26, "workspace query", "run SQL in workspace") self._test_workspace_query(workspace_id, table_id) + _step(26.5, "library facade", "Client.query + files round-trip, in-process") + self._test_library_facade(workspace_id, table_id) + _step(27, "workspace delete") self._test_workspace_delete(workspace_id) @@ -1949,6 +1952,50 @@ def _test_workspace_password(self, workspace_id: int) -> None: data = _json_ok(result) assert data["data"]["password"] # non-empty password + def _test_library_facade(self, workspace_id: int, table_id: str) -> None: + """Exercise the public in-process library facade against the live stack. + + Imports ``keboola_agent_cli.Client`` and runs a real query + Storage + Files round-trip with no CLI subprocess -- the in-process path the + jasnost feedback (#415) consumes. Backend-specific identifier quoting + mirrors ``_test_workspace_query``. + """ + from keboola_agent_cli import Client, FileEntry + + ws_table_name = table_id.rsplit(".", 1)[-1] + detail = self._run_ok( + "workspace", + "detail", + "--project", + self.alias, + "--workspace-id", + str(workspace_id), + )["data"] + quote = "`" if detail.get("backend") == "bigquery" else '"' + sql = f"SELECT COUNT(*) AS cnt FROM {quote}{ws_table_name}{quote}" + + with Client(url=self.url, token=self.token) as kbc: + # query() -> list[dict] keyed by column name + rows = kbc.query(workspace_id, sql) + assert isinstance(rows, list) and rows, "facade query must return rows" + assert isinstance(rows[0], dict), "facade query rows must be dicts" + assert "cnt" in {k.lower() for k in rows[0]}, f"expected cnt column, got {rows[0]}" + + # files: upload bytes -> read_bytes -> list -> delete, all in-process + facade_tag = f"{RUN_ID}-facade" + payload = b"facade-e2e-roundtrip" + meta = kbc.files.upload(payload, name=f"{RUN_ID}-facade.txt", tags=[facade_tag]) + assert isinstance(meta, FileEntry) and meta.id > 0 + self._created_file_ids.append(meta.id) + + assert kbc.files.read_bytes(meta.id) == payload, "read_bytes round-trip mismatch" + + listed = kbc.files.list(tags=[facade_tag]) + assert any(f.id == meta.id for f in listed), "uploaded file must appear in list" + + kbc.files.delete(meta.id) + self._created_file_ids.remove(meta.id) + def _test_workspace_load(self, workspace_id: int, table_id: str) -> None: """Load a table into the workspace.""" data = self._run_ok( diff --git a/tests/test_lib.py b/tests/test_lib.py new file mode 100644 index 00000000..eb59b0b6 --- /dev/null +++ b/tests/test_lib.py @@ -0,0 +1,250 @@ +"""Tests for the public in-process library facade (keboola_agent_cli.lib). + +The facade wraps a single KeboolaClient, so every test patches +``keboola_agent_cli.lib.KeboolaClient`` with a MagicMock and asserts the facade +translates between the high-level shapes (list[dict] rows, bytes, FileEntry) and +the low-level client calls. No network. +""" + +import logging +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +from keboola_agent_cli import Client, FileEntry, Files +from keboola_agent_cli.errors import KeboolaApiError + +# Canonical fake token (projectId-tokenId-secret); never a realistic secret. +FAKE_TOKEN = "901-55555-fakeTestTokenDoNotUseXXXXXXXX" +STACK_URL = "https://connection.keboola.com" + + +@pytest.fixture +def mock_kc() -> MagicMock: + return MagicMock() + + +@pytest.fixture +def client(mock_kc: MagicMock) -> Client: + """A Client whose underlying KeboolaClient is the shared mock.""" + with patch("keboola_agent_cli.lib.KeboolaClient", return_value=mock_kc): + return Client(url=STACK_URL, token=FAKE_TOKEN) + + +def _make_client(mock_kc: MagicMock, *, branch_id: int | None = None) -> Client: + with patch("keboola_agent_cli.lib.KeboolaClient", return_value=mock_kc): + return Client(url=STACK_URL, token=FAKE_TOKEN, branch_id=branch_id) + + +class TestConstruction: + def test_requires_url(self) -> None: + with pytest.raises(ValueError, match="url is required"): + Client(url="", token=FAKE_TOKEN) + + def test_requires_token(self) -> None: + with pytest.raises(ValueError, match="token is required"): + Client(url=STACK_URL, token="") + + def test_raw_exposes_underlying_client(self, client: Client, mock_kc: MagicMock) -> None: + assert client.raw is mock_kc + + def test_files_namespace_present(self, client: Client) -> None: + assert isinstance(client.files, Files) + + def test_context_manager_closes(self, mock_kc: MagicMock) -> None: + with _make_client(mock_kc) as c: + assert c is not None + mock_kc.close.assert_called_once() + + +class TestQuery: + @staticmethod + def _wire_single_select(mock_kc: MagicMock, *, num_rows: int = 2) -> None: + mock_kc.list_dev_branches.return_value = [{"id": 123, "isDefault": True}] + mock_kc.submit_query.return_value = {"queryJobId": "qj1"} + mock_kc.wait_for_query_job.return_value = { + "statements": [{"id": "s1", "status": "completed", "numberOfRows": num_rows}] + } + mock_kc.get_query_results.return_value = { + "columns": [{"name": "id"}, {"name": "name"}], + "data": [[1, "alice"], [2, "bob"]], + "numberOfRows": num_rows, + } + + def test_maps_columns_and_rows_to_dicts(self, client: Client, mock_kc: MagicMock) -> None: + self._wire_single_select(mock_kc) + rows = client.query(456, "SELECT id, name FROM t") + assert rows == [{"id": 1, "name": "alice"}, {"id": 2, "name": "bob"}] + + def test_submits_with_resolved_default_branch(self, client: Client, mock_kc: MagicMock) -> None: + self._wire_single_select(mock_kc) + client.query(456, "SELECT 1") + mock_kc.submit_query.assert_called_once_with( + branch_id=123, workspace_id=456, statements=["SELECT 1"], transactional=False + ) + + def test_explicit_branch_skips_resolution(self, mock_kc: MagicMock) -> None: + c = _make_client(mock_kc, branch_id=999) + self._wire_single_select(mock_kc) + c.query(456, "SELECT 1") + mock_kc.list_dev_branches.assert_not_called() + assert mock_kc.submit_query.call_args.kwargs["branch_id"] == 999 + + def test_no_default_branch_raises(self, client: Client, mock_kc: MagicMock) -> None: + mock_kc.list_dev_branches.return_value = [{"id": 5, "isDefault": False}] + with pytest.raises(KeboolaApiError, match="No default branch"): + client.query(1, "SELECT 1") + + def test_statement_without_result_set_returns_empty( + self, client: Client, mock_kc: MagicMock + ) -> None: + mock_kc.list_dev_branches.return_value = [{"id": 1, "isDefault": True}] + mock_kc.submit_query.return_value = {"queryJobId": "qj"} + mock_kc.wait_for_query_job.return_value = { + "statements": [{"id": "s1", "status": "completed", "numberOfRows": 0}] + } + assert client.query(1, "CREATE TABLE t (id INT)") == [] + mock_kc.get_query_results.assert_not_called() + + def test_multi_statement_returns_last_result_set( + self, client: Client, mock_kc: MagicMock + ) -> None: + mock_kc.list_dev_branches.return_value = [{"id": 1, "isDefault": True}] + mock_kc.submit_query.return_value = {"queryJobId": "qj"} + mock_kc.wait_for_query_job.return_value = { + "statements": [ + {"id": "s1", "status": "completed", "numberOfRows": 0}, + {"id": "s2", "status": "completed", "numberOfRows": 1}, + ] + } + mock_kc.get_query_results.return_value = { + "columns": [{"name": "n"}], + "data": [[7]], + "numberOfRows": 1, + } + rows = client.query(1, "USE WAREHOUSE x; SELECT 7 AS n") + assert rows == [{"n": 7}] + # Only the result-producing statement triggers a results fetch. + mock_kc.get_query_results.assert_called_once() + assert mock_kc.get_query_results.call_args.args[1] == "s2" + + def test_truncation_logs_warning( + self, client: Client, mock_kc: MagicMock, caplog: pytest.LogCaptureFixture + ) -> None: + mock_kc.list_dev_branches.return_value = [{"id": 1, "isDefault": True}] + mock_kc.submit_query.return_value = {"queryJobId": "qj"} + mock_kc.wait_for_query_job.return_value = { + "statements": [{"id": "s1", "status": "completed", "numberOfRows": 100}] + } + mock_kc.get_query_results.return_value = { + "columns": [{"name": "id"}], + "data": [[1], [2]], + "numberOfRows": 100, # warehouse has more than we keep + } + with caplog.at_level(logging.WARNING, logger="keboola_agent_cli.lib"): + rows = client.query(1, "SELECT id FROM big", limit=2) + assert len(rows) == 2 + assert "truncated" in caplog.text + + +class TestFilesList: + def test_returns_uniform_file_entries(self, client: Client, mock_kc: MagicMock) -> None: + mock_kc.list_files.return_value = [ + {"id": 1, "name": "a.csv", "tags": ["x"], "created": "2026-01-01", "sizeBytes": 9}, + {"id": 2, "name": "b.csv", "tags": [], "created": "2026-01-02", "isPermanent": True}, + ] + entries = client.files.list(tags=["x"]) + assert all(isinstance(e, FileEntry) for e in entries) + assert entries[0].id == 1 and entries[0].tags == ["x"] + assert entries[1].is_permanent is True + mock_kc.list_files.assert_called_once_with( + limit=100, offset=0, tags=["x"], since_id=None, query=None, branch_id=None + ) + + def test_branch_scoped_list(self, mock_kc: MagicMock) -> None: + c = _make_client(mock_kc, branch_id=42) + mock_kc.list_files.return_value = [] + c.files.list() + assert mock_kc.list_files.call_args.kwargs["branch_id"] == 42 + + +class TestFilesUpload: + def test_upload_from_path(self, client: Client, mock_kc: MagicMock, tmp_path: Path) -> None: + p = tmp_path / "data.csv" + p.write_text("a,b\n1,2\n") + mock_kc.upload_file.return_value = {"id": 5, "name": "data.csv", "tags": ["t"]} + entry = client.files.upload(str(p), tags=["t"], permanent=True) + assert entry.id == 5 + call = mock_kc.upload_file.call_args.kwargs + assert call["file_path"] == str(p) + assert call["name"] is None and call["is_permanent"] is True and call["tags"] == ["t"] + + def test_upload_from_bytes_stages_temp_file(self, client: Client, mock_kc: MagicMock) -> None: + captured: dict[str, object] = {} + + def fake_upload(*, file_path: str, name: str, tags, is_permanent, branch_id): + captured["path"] = file_path + captured["content"] = Path(file_path).read_bytes() + return {"id": 9, "name": name, "tags": tags or []} + + mock_kc.upload_file.side_effect = fake_upload + entry = client.files.upload(b"hello bytes", name="greeting.txt") + assert entry.id == 9 + assert captured["content"] == b"hello bytes" + # temp file is cleaned up after the upload returns + assert not Path(str(captured["path"])).exists() + + def test_upload_bytes_without_name_raises(self, client: Client) -> None: + with pytest.raises(ValueError, match="name is required"): + client.files.upload(b"x") + + +class TestFilesReadBytes: + def test_non_sliced_download(self, client: Client, mock_kc: MagicMock) -> None: + mock_kc.get_file_info.return_value = {"isSliced": False, "url": "https://signed/url"} + + def fake_download(url: str, output_path: str) -> int: + Path(output_path).write_bytes(b"file-content") + return 12 + + mock_kc.download_file.side_effect = fake_download + assert client.files.read_bytes(42) == b"file-content" + mock_kc.get_file_info.assert_called_once_with(42, branch_id=None) + + def test_sliced_download(self, client: Client, mock_kc: MagicMock) -> None: + mock_kc.get_file_info.return_value = {"isSliced": True} + + def fake_sliced(info: dict, output_path: str) -> int: + Path(output_path).write_bytes(b"sliced-content") + return 14 + + mock_kc.download_sliced_file.side_effect = fake_sliced + assert client.files.read_bytes(7) == b"sliced-content" + mock_kc.download_file.assert_not_called() + + def test_missing_url_raises(self, client: Client, mock_kc: MagicMock) -> None: + mock_kc.get_file_info.return_value = {"isSliced": False} # no url + with pytest.raises(KeboolaApiError, match="no download URL"): + client.files.read_bytes(9) + + +class TestFilesDelete: + def test_delete_forwards(self, client: Client, mock_kc: MagicMock) -> None: + client.files.delete(11) + mock_kc.delete_file.assert_called_once_with(11, branch_id=None) + + +class TestModuleLayout: + def test_pagination_helper_relocated_to_client(self) -> None: + """_collect_inline_results lives in client.py; workspace_service re-exports it.""" + import keboola_agent_cli.services.workspace_service as ws + from keboola_agent_cli.client import InlineQueryResult, _collect_inline_results + + assert ws._collect_inline_results is _collect_inline_results + assert InlineQueryResult.__module__ == "keboola_agent_cli.client" + + def test_public_surface(self) -> None: + import keboola_agent_cli as pkg + + assert set(pkg.__all__) >= {"Client", "Files", "FileEntry"} diff --git a/tests/test_workspace_service.py b/tests/test_workspace_service.py index 345b5dbf..6d5ee135 100644 --- a/tests/test_workspace_service.py +++ b/tests/test_workspace_service.py @@ -11,6 +11,7 @@ import pytest from helpers import setup_single_project, setup_two_projects +from keboola_agent_cli import client as client_module from keboola_agent_cli.config_store import ConfigStore from keboola_agent_cli.errors import ConfigError, KeboolaApiError from keboola_agent_cli.models import ProjectConfig, TokenVerifyResponse @@ -994,7 +995,7 @@ def test_execute_query_success(self, tmp_config_dir: Path) -> None: "qj-abc123", "stmt-1", offset=0, - page_size=workspace_service_module.QUERY_RESULTS_PAGE_SIZE, + page_size=client_module.QUERY_RESULTS_PAGE_SIZE, ) mock_client.export_query_results.assert_not_called() @@ -1043,7 +1044,7 @@ def test_execute_query_inline_pagination( ) -> 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) + monkeypatch.setattr(client_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"} @@ -1115,9 +1116,9 @@ def test_execute_query_small_limit_keeps_valid_page_size(self, tmp_config_dir: P "qj-small", "stmt-1", offset=0, - page_size=workspace_service_module.QUERY_RESULTS_PAGE_SIZE, + page_size=client_module.QUERY_RESULTS_PAGE_SIZE, ) - assert workspace_service_module.QUERY_RESULTS_PAGE_SIZE >= 100 + assert client_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: @@ -1278,7 +1279,7 @@ def test_execute_query_truncated_when_number_of_rows_missing( 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) + monkeypatch.setattr(client_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"} @@ -1310,7 +1311,7 @@ def test_execute_query_stops_at_total_on_page_boundary( ) -> 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) + monkeypatch.setattr(client_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"} diff --git a/uv.lock b/uv.lock index bece0334..517474cc 100644 --- a/uv.lock +++ b/uv.lock @@ -580,7 +580,7 @@ wheels = [ [[package]] name = "keboola-agent-cli" -version = "0.60.4" +version = "0.61.0" source = { editable = "." } dependencies = [ { name = "croniter" },