From 627baf2a823930bb7f5074f935690820e6e2a716 Mon Sep 17 00:00:00 2001 From: Petr Date: Tue, 16 Jun 2026 22:27:24 +0200 Subject: [PATCH] feat(lib): typed SDK return models + py.typed marker (#428) Ship a PEP 561 `py.typed` marker and a set of typed pydantic return models for the in-process SDK facade, so downstream consumers (e.g. the FIIA Scaffold Kit) type against a versioned contract instead of undocumented `dict[str, Any]` shapes -- a contract change then surfaces at type-check time, not at runtime against a customer build. - result_models.py: JobResult, QueryResult, UploadTableResult, SyncPushResult, ConfigDetailResult. All set `extra="allow"` so backend field drift never raises, and `populate_by_name=True` + AliasChoices so `Model.model_validate(service_dict)` works on raw API or snake_case keys. - lib.py: typed facade wrappers returning the models -- `Client.run_job` (-> JobResult), `query_result` (-> QueryResult), `config_detail` (-> ConfigDetailResult), `upload_table` (-> UploadTableResult). `query()` refactored onto a shared `_run_query` (behavior unchanged, still returns list[dict]). - py.typed forced into the wheel via hatch force-include. - Export the models from the package root; document in library-workflow.md. Tests: contract tests for every model (real camelCase keys, extra-field survival, convenience props) + facade-method tests. make check green. --- .../kbagent/references/library-workflow.md | 39 ++- pyproject.toml | 5 + src/keboola_agent_cli/__init__.py | 19 +- src/keboola_agent_cli/lib.py | 227 ++++++++++++++--- src/keboola_agent_cli/py.typed | 3 + src/keboola_agent_cli/result_models.py | 228 ++++++++++++++++++ tests/test_lib.py | 158 +++++++++++- tests/test_result_models.py | 182 ++++++++++++++ 8 files changed, 826 insertions(+), 35 deletions(-) create mode 100644 src/keboola_agent_cli/py.typed create mode 100644 src/keboola_agent_cli/result_models.py create mode 100644 tests/test_result_models.py diff --git a/plugins/kbagent/skills/kbagent/references/library-workflow.md b/plugins/kbagent/skills/kbagent/references/library-workflow.md index 144f9097..f420fa93 100644 --- a/plugins/kbagent/skills/kbagent/references/library-workflow.md +++ b/plugins/kbagent/skills/kbagent/references/library-workflow.md @@ -15,6 +15,10 @@ shell operations, use the `kbagent` CLI. |--------|---------| | `Client(url, token, *, branch_id=None)` | Stateless entry point to one project; context manager | | `Client.query(workspace_id, sql, *, transactional=False, limit=500)` | Run SQL in a workspace -> `list[dict]` | +| `Client.query_result(workspace_id, sql, ...)` | Same, but typed -> `QueryResult` (columns + truncation) | +| `Client.run_job(component_id, config_id, *, wait=False, ...)` | Run a Queue job -> `JobResult` | +| `Client.config_detail(component_id, config_id, *, branch_id=None)` | One config's detail -> `ConfigDetailResult` | +| `Client.upload_table(table_id, file_path, *, incremental=False, ...)` | Import a CSV into an existing table -> `UploadTableResult` | | `Client.files.upload(source, *, name=None, tags=None, permanent=False)` | Upload a path **or** bytes -> `FileEntry` | | `Client.files.read_bytes(file_id)` | Download a file fully into memory -> `bytes` | | `Client.files.list(*, tags=None, query=None, limit=100, ...)` | List files -> `list[FileEntry]` | @@ -22,8 +26,35 @@ shell operations, use the `kbagent` CLI. | `Client.raw` | The underlying `KeboolaClient` for endpoints the facade omits | | `FileEntry` | Uniform file shape: `id, name, tags, created, size_bytes, is_permanent, raw` | -Everything exported from `keboola_agent_cli` (`Client`, `Files`, `FileEntry`) is -committed public API and follows semver. +Everything exported from `keboola_agent_cli` (`Client`, `Files`, `FileEntry`, and +the typed result models `JobResult`, `QueryResult`, `UploadTableResult`, +`SyncPushResult`, `ConfigDetailResult`) is committed public API and follows +semver. Since 0.63.0 the package ships a **`py.typed`** marker (PEP 561), so +`mypy` / `ty` / IDEs treat the SDK as typed -- a contract change surfaces at +type-check time, not at runtime. + +## Typed result models + +The high-traffic operations return pydantic models (`result_models.py`) instead +of bare dicts, so you get autocomplete and a versioned contract: + +```python +job = kbc.run_job("keboola.ex-db-snowflake", "12345", wait=True) +if job.succeeded: # -> JobResult + print(job.id, job.result) + +res = kbc.query_result(ws, 'SELECT id, name FROM t') # -> QueryResult +print(res.columns, res.row_count, res.truncated) + +cfg = kbc.config_detail("keboola.ex-http", "98765") # -> ConfigDetailResult +print(cfg.name, cfg.version, cfg.configuration) +``` + +Every model is **tolerant of extra fields** (`extra="allow"`): the named fields +are the stable surface, but anything else the API returns is preserved +(reachable via attribute access and `model_dump()`), so a new backend field +never raises. They also accept the raw API key *or* the snake_case field name, so +`JobResult.model_validate(service_dict)` works directly on a service-layer dict. ## Auth & construction @@ -100,8 +131,8 @@ kbc.files.delete(meta.id) ## Lower-level access -For endpoints the facade does not wrap (buckets, tables, jobs, branches, ...), -reach for the underlying client: +For endpoints the facade does not wrap (buckets, tables, branches, job polling +internals, ...), reach for the underlying client: ```python client = kbc.raw # a KeboolaClient diff --git a/pyproject.toml b/pyproject.toml index 83d099fa..6aa9a130 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,6 +53,11 @@ packages = ["src/keboola_agent_cli"] [tool.hatch.build.targets.wheel.force-include] "src/keboola_agent_cli/_ui_dist" = "keboola_agent_cli/_ui_dist" +# PEP 561: ship the py.typed marker so downstream type checkers treat the +# importable SDK (Client + result_models contracts, issue #428) as typed. +# Forced in explicitly rather than relying on hatchling's default file +# selection, so the marker is guaranteed to land in the wheel. +"src/keboola_agent_cli/py.typed" = "keboola_agent_cli/py.typed" [tool.hatch.build.targets.wheel.hooks.custom] path = "scripts/hatch_build.py" diff --git a/src/keboola_agent_cli/__init__.py b/src/keboola_agent_cli/__init__.py index 66680285..ed815b0a 100644 --- a/src/keboola_agent_cli/__init__.py +++ b/src/keboola_agent_cli/__init__.py @@ -4,10 +4,27 @@ from .constants import APP_NAME from .lib import Client, FileEntry, Files +from .result_models import ( + ConfigDetailResult, + JobResult, + QueryResult, + SyncPushResult, + UploadTableResult, +) try: __version__ = version(APP_NAME) except PackageNotFoundError: __version__ = "0.0.0-dev" -__all__ = ["Client", "FileEntry", "Files", "__version__"] +__all__ = [ + "Client", + "ConfigDetailResult", + "FileEntry", + "Files", + "JobResult", + "QueryResult", + "SyncPushResult", + "UploadTableResult", + "__version__", +] diff --git a/src/keboola_agent_cli/lib.py b/src/keboola_agent_cli/lib.py index 096a60a7..ae7e7747 100644 --- a/src/keboola_agent_cli/lib.py +++ b/src/keboola_agent_cli/lib.py @@ -12,12 +12,16 @@ 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] + job = kbc.run_job("keboola.ex-db-snowflake", "12345", wait=True) # JobResult 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. +file reads, a stable :class:`FileEntry`, and the typed result models in +``result_models.py`` (:class:`JobResult`, :class:`QueryResult`, +:class:`UploadTableResult`, :class:`ConfigDetailResult`) -- 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 @@ -33,8 +37,14 @@ from typing import Any from .client import KeboolaClient, _collect_inline_results -from .constants import QUERY_RESULTS_DEFAULT_LIMIT +from .constants import ( + DEFAULT_JOB_MODE, + DEFAULT_JOB_RUN_TIMEOUT, + DEFAULT_POLL_STRATEGY, + QUERY_RESULTS_DEFAULT_LIMIT, +) from .errors import ErrorCode, KeboolaApiError +from .result_models import ConfigDetailResult, JobResult, QueryResult, UploadTableResult logger = logging.getLogger(__name__) @@ -234,6 +244,54 @@ def _effective_branch_id(self) -> int: error_code=ErrorCode.NOT_FOUND, ) + def _run_query( + self, + workspace_id: int, + sql: str, + *, + transactional: bool, + limit: int, + ) -> QueryResult: + """Submit SQL, wait for completion, and collect the last result set. + + Shared by :meth:`query` (which returns just the rows) and + :meth:`query_result` (which returns the full typed shape). Mirrors the + Query Service inline-results fast path: the rows of the *last* + result-producing statement win, statements without a result set yield + nothing, and an over-``limit`` result is capped with a logged warning. + """ + 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) + + result = QueryResult() + 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] + result = QueryResult( + columns=col_names, + rows=[dict(zip(col_names, row, strict=False)) for row in inline.rows], + truncated=inline.truncated, + total_rows=inline.total_rows, + ) + if inline.truncated: + logger.warning( + "query result truncated to %d rows (warehouse has %s); " + "raise limit= to fetch more", + result.row_count, + inline.total_rows, + ) + return result + def query( self, workspace_id: int, @@ -259,6 +317,8 @@ def query( statement that produced a result set are returned (so ``USE ...; SELECT ...`` yields the SELECT). Statements without a result set yield ``[]``. + For column order and truncation metadata, use :meth:`query_result`. + Args: workspace_id: Target workspace ID. sql: One or more SQL statements. @@ -267,32 +327,143 @@ def query( 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, + return self._run_query(workspace_id, sql, transactional=transactional, limit=limit).rows + + def query_result( + self, + workspace_id: int, + sql: str, + *, + transactional: bool = False, + limit: int = QUERY_RESULTS_DEFAULT_LIMIT, + ) -> QueryResult: + """Run SQL in a workspace and return a typed :class:`QueryResult`. + + Same execution as :meth:`query`, but returns the full tabular shape -- + ``columns`` (in warehouse order), ``rows`` (list of dicts), ``truncated`` + and ``total_rows`` -- instead of just the row list. Use this when you + need the column ordering or want to detect a ``limit`` cap. The + string-typing gotcha from :meth:`query` applies to ``rows`` here too. + """ + return self._run_query(workspace_id, sql, transactional=transactional, limit=limit) + + def run_job( + self, + component_id: str, + config_id: str, + *, + config_row_ids: list[str] | None = None, + variable_values_id: str | None = None, + branch_id: int | None = None, + mode: str = DEFAULT_JOB_MODE, + wait: bool = False, + timeout: float = DEFAULT_JOB_RUN_TIMEOUT, + poll_strategy: str = DEFAULT_POLL_STRATEGY, + ) -> JobResult: + """Run a Queue API job and return a typed :class:`JobResult`. + + Creates the job, and -- when ``wait=True`` -- polls until it reaches a + terminal state (or ``timeout`` elapses). Unlike ``JobService.run_job`` + this thin facade does **not** auto-resolve linked variable values; pass + ``variable_values_id`` explicitly if the config needs a values row. + + Args: + component_id: Component to run, e.g. ``keboola.ex-db-snowflake``. + config_id: Configuration ID to run. + config_row_ids: Optional row IDs (omit to run the whole config). + variable_values_id: Optional explicit values row for linked variables. + branch_id: Dev branch to run on. Defaults to the client's branch + (``None`` = production). + mode: Queue job mode (``run`` | ``debug`` | ``forceRun``). + wait: If True, poll until the job finishes or ``timeout`` elapses. + timeout: Max seconds to wait (only used when ``wait=True``). + poll_strategy: Wait cadence, one of ``VALID_POLL_STRATEGIES``. + """ + effective_branch = branch_id if branch_id is not None else self._resolved_branch_id + job = self._client.create_job( + component_id=component_id, + config_id=config_id, + config_row_ids=config_row_ids, + mode=mode, + branch_id=effective_branch, + variable_values_id=variable_values_id, ) - job_id = str(job.get("queryJobId", job.get("id", ""))) - completed = self._client.wait_for_query_job(job_id) + job_id = str(job.get("id", "")) + if wait and job_id: + job = self._client.wait_for_queue_job( + job_id, max_wait=timeout, poll_strategy=poll_strategy + ) + return JobResult.model_validate(job) - 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 config_detail( + self, + component_id: str, + config_id: str, + *, + branch_id: int | None = None, + ) -> ConfigDetailResult: + """Fetch one configuration's detail as a typed :class:`ConfigDetailResult`. + + Args: + component_id: Owning component ID. + config_id: Configuration ID. + branch_id: Dev branch to read from. Defaults to the client's branch + (``None`` = production). + """ + effective_branch = branch_id if branch_id is not None else self._resolved_branch_id + detail = dict( + self._client.get_config_detail(component_id, config_id, branch_id=effective_branch) + ) + detail.setdefault("component_id", component_id) + if effective_branch is not None: + detail.setdefault("branch_id", effective_branch) + return ConfigDetailResult.model_validate(detail) + + def upload_table( + self, + table_id: str, + file_path: str | Path, + *, + incremental: bool = False, + delimiter: str = ",", + enclosure: str = '"', + branch_id: int | None = None, + ) -> UploadTableResult: + """Import a CSV into an **existing** Storage table -> :class:`UploadTableResult`. + + Unlike ``StorageService.upload_table`` the facade does **not** auto-create + a missing bucket/table (it has no config-dir / service context); the + target table must already exist. Use the CLI (``kbagent storage + upload-table``) for the auto-create path. + + Args: + table_id: Target table ID (must exist). + file_path: Local CSV path. + incremental: Append rows (True) or full load (False). + delimiter: CSV column delimiter. + enclosure: CSV value enclosure character. + branch_id: Dev branch to target. Defaults to the client's branch + (``None`` = production). + """ + effective_branch = branch_id if branch_id is not None else self._resolved_branch_id + file_size_bytes = Path(file_path).stat().st_size + results = self._client.upload_table( + table_id=table_id, + file_path=str(file_path), + incremental=incremental, + delimiter=delimiter, + enclosure=enclosure, + branch_id=effective_branch, + ) + return UploadTableResult.model_validate( + { + "table_id": table_id, + "incremental": incremental, + "file_size_bytes": file_size_bytes, + "imported_rows": results.get("importedRowsCount"), + "warnings": results.get("warnings", []), + } + ) def close(self) -> None: """Close the underlying HTTP client.""" diff --git a/src/keboola_agent_cli/py.typed b/src/keboola_agent_cli/py.typed new file mode 100644 index 00000000..66901d07 --- /dev/null +++ b/src/keboola_agent_cli/py.typed @@ -0,0 +1,3 @@ +# PEP 561 marker: keboola_agent_cli ships inline type annotations. +# Downstream mypy / ty / IDEs treat the importable SDK (Client, the +# result_models.* contracts) as typed. See issue #428. diff --git a/src/keboola_agent_cli/result_models.py b/src/keboola_agent_cli/result_models.py new file mode 100644 index 00000000..b138d2bc --- /dev/null +++ b/src/keboola_agent_cli/result_models.py @@ -0,0 +1,228 @@ +"""Typed return models for the in-process SDK facade (issue #428). + +These pydantic models document the **stable** return shapes of the high-traffic +service / facade operations. They are exported from the package root so a +downstream in-process consumer (a Keboola Data App, a transformation, a hosted +service, the FIIA Scaffold Kit) gets static typing (mypy / IDE autocomplete) and +a semver-versioned contract instead of coupling to undocumented +``dict[str, Any]`` shapes -- a contract change then surfaces at type-check time, +not at runtime against a customer build. + +Every model sets ``extra="allow"``: the backing Keboola APIs grow fields across +stack versions, and an SDK contract must not raise when the server returns more +than the documented subset. Extra keys are preserved (reachable via attribute +access and ``model_dump()``), so nothing is lost -- only the **named** fields are +the committed, semver-stable surface. ``populate_by_name=True`` means each model +accepts both the snake_case field name and the raw API key (declared via +``AliasChoices``), so ``Model.model_validate(service_dict)`` works directly on a +service-layer dict without renaming. +""" + +from typing import Any + +from pydantic import AliasChoices, BaseModel, ConfigDict, Field + + +class _ApiResultModel(BaseModel): + """Base for SDK result contracts: tolerant of unknown / extra API fields. + + Subclasses type only the stable subset; everything else the API returns is + kept as model extras (``extra="allow"``) rather than dropped or raised on. + """ + + model_config = ConfigDict(extra="allow", populate_by_name=True) + + +class JobResult(_ApiResultModel): + """Result of running a Queue API job. + + Returned by :meth:`keboola_agent_cli.Client.run_job` and produced by + ``JobService.run_job``. The named fields below are the committed contract; + everything else the Queue API returns (``branchId``, ``createdTime``, + ``startTime``, ``endTime``, ``durationSeconds``, ``runId``, ``url``, ...) is + preserved as model extras. + """ + + id: str = Field(default="", description="Queue job ID.") + status: str = Field( + default="", + description=( + "Job status: created | waiting | processing | success | error | warning | terminated." + ), + ) + is_finished: bool = Field( + default=False, + validation_alias=AliasChoices("isFinished", "is_finished"), + description="True once the job reached a terminal state.", + ) + component_id: str = Field( + default="", + validation_alias=AliasChoices("component", "componentId", "component_id"), + description="Component that ran.", + ) + config_id: str = Field( + default="", + validation_alias=AliasChoices("configId", "config", "config_id"), + description="Configuration that ran.", + ) + mode: str = Field(default="", description="Job mode: run | debug | forceRun.") + result: dict[str, Any] | None = Field( + default=None, + description="Component result payload (message, import stats, ...) when present.", + ) + project_alias: str = Field( + default="", + description=( + "Project alias the job ran in (CLI / service path). Empty for the " + "in-process facade, which is not config-dir aware." + ), + ) + resolved_variable_values_id: str | None = Field( + default=None, + validation_alias=AliasChoices("resolvedVariableValuesId", "resolved_variable_values_id"), + description="Values row resolved for linked variables, if any.", + ) + log_tail: list[dict[str, Any]] | None = Field( + default=None, + validation_alias=AliasChoices("logTail", "log_tail"), + description="Trailing job events surfaced on a non-success terminal state (wait mode).", + ) + + @property + def succeeded(self) -> bool: + """True iff the job finished in the ``success`` state.""" + return self.status == "success" + + @property + def failed(self) -> bool: + """True iff the job finished in a terminal failure state.""" + return self.status in {"error", "terminated", "cancelled"} + + +class QueryResult(_ApiResultModel): + """Tabular result of a workspace SQL query. + + Returned by :meth:`keboola_agent_cli.Client.query_result`. Carries the + column order and truncation metadata that the plain ``Client.query`` (which + returns ``list[dict]``) drops. Values are **not** coerced -- for Snowflake + every scalar comes back as a string (see the ``query()`` docstring / gotchas). + """ + + columns: list[str] = Field( + default_factory=list, description="Result column names, in warehouse order." + ) + rows: list[dict[str, Any]] = Field( + default_factory=list, description="Rows as dicts keyed by column name." + ) + truncated: bool = Field( + default=False, + description="True if the result was capped at the requested ``limit``.", + ) + total_rows: int | None = Field( + default=None, + description="Total rows the warehouse reported for the statement, when known.", + ) + + @property + def row_count(self) -> int: + """Number of rows actually returned (after any ``limit`` cap).""" + return len(self.rows) + + +class UploadTableResult(_ApiResultModel): + """Result of importing a CSV into a Storage table. + + Returned by :meth:`keboola_agent_cli.Client.upload_table` and produced by + ``StorageService.upload_table``. The ``auto_created_*`` flags are only set by + the service path (which can create a missing bucket/table); the in-process + facade requires the table to exist and leaves them ``False``. + """ + + table_id: str = Field(default="", description="Target table ID.") + incremental: bool = Field(default=False, description="True = rows appended; False = full load.") + imported_rows: int | None = Field( + default=None, + validation_alias=AliasChoices("imported_rows", "importedRowsCount"), + description="Rows imported, when the backend reports a count.", + ) + file_size_bytes: int | None = Field( + default=None, description="Size of the uploaded CSV on disk." + ) + warnings: list[Any] = Field( + default_factory=list, description="Import warnings surfaced by Storage." + ) + auto_created_bucket: bool = Field( + default=False, description="True if the service path created the bucket." + ) + auto_created_table: bool = Field( + default=False, description="True if the service path created the table." + ) + project_alias: str = Field( + default="", description="Project alias (service path; empty for the facade)." + ) + + +class SyncPushResult(_ApiResultModel): + """Result of a GitOps ``sync push``. + + Documents the shape of ``SyncService.push`` and is embedded in + ``CloneResult`` (issue #426). ``status`` is ``pushed`` | ``no_changes`` | + ``dry_run``; the counters and ``pushed_details`` describe what changed. + """ + + status: str = Field(default="", description="pushed | no_changes | dry_run.") + created: int = Field(default=0, description="Configs/rows created.") + updated: int = Field(default=0, description="Configs/rows updated.") + deleted: int = Field(default=0, description="Configs/rows deleted.") + errors: list[dict[str, Any]] = Field( + default_factory=list, + description="Per-change failures (change_type, component_id, config_id, message).", + ) + pushed_details: list[dict[str, Any]] = Field( + default_factory=list, description="One entry per applied change." + ) + name_drift_warnings: list[dict[str, Any]] = Field( + default_factory=list, + description="Local-vs-remote name drift warnings, when surfaced.", + ) + + @property + def ok(self) -> bool: + """True iff the push completed with no per-change errors.""" + return not self.errors + + +class ConfigDetailResult(_ApiResultModel): + """Detail of a single configuration. + + Returned by :meth:`keboola_agent_cli.Client.config_detail` and produced by + ``ConfigService.get_config_detail`` in single-config mode. ``id`` is the + configuration ID (Storage returns it under ``id``); the full Storage detail + (``created``, ``isDisabled``, ``state``, ...) is preserved as extras. + """ + + id: str = Field(default="", description="Configuration ID.") + name: str = Field(default="", description="Configuration name.") + description: str = Field(default="", description="Configuration description.") + version: int | None = Field( + default=None, + validation_alias=AliasChoices("version", "currentVersion"), + description="Current configuration version.", + ) + configuration: dict[str, Any] = Field( + default_factory=dict, description="The configuration body." + ) + rows: list[dict[str, Any]] = Field( + default_factory=list, description="Config rows, when the config has any." + ) + component_id: str = Field( + default="", + validation_alias=AliasChoices("component_id", "componentId"), + description="Owning component ID.", + ) + project_alias: str = Field( + default="", description="Project alias (service path; empty for the facade)." + ) + branch_id: int | None = Field( + default=None, description="Dev branch ID the detail was read from (None = production)." + ) diff --git a/tests/test_lib.py b/tests/test_lib.py index f149a07e..ed6b8f87 100644 --- a/tests/test_lib.py +++ b/tests/test_lib.py @@ -12,7 +12,15 @@ import pytest -from keboola_agent_cli import Client, FileEntry, Files +from keboola_agent_cli import ( + Client, + ConfigDetailResult, + FileEntry, + Files, + JobResult, + QueryResult, + UploadTableResult, +) from keboola_agent_cli.errors import KeboolaApiError # Canonical fake token (projectId-tokenId-secret); never a realistic secret. @@ -239,6 +247,143 @@ def test_delete_forwards(self, client: Client, mock_kc: MagicMock) -> None: mock_kc.delete_file.assert_called_once_with(11, branch_id=None) +class TestQueryResult: + def test_returns_typed_shape_with_columns_and_truncation( + 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": 100}] + } + mock_kc.get_query_results.return_value = { + "columns": [{"name": "id"}, {"name": "name"}], + "data": [["1", "alice"], ["2", "bob"]], + "numberOfRows": 100, + } + result = client.query_result(456, "SELECT id, name FROM big", limit=2) + assert isinstance(result, QueryResult) + assert result.columns == ["id", "name"] + assert result.rows == [{"id": "1", "name": "alice"}, {"id": "2", "name": "bob"}] + assert result.row_count == 2 + assert result.truncated is True + assert result.total_rows == 100 + + def test_query_and_query_result_share_rows(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": 1}] + } + mock_kc.get_query_results.return_value = { + "columns": [{"name": "n"}], + "data": [["7"]], + "numberOfRows": 1, + } + assert client.query_result(1, "SELECT 7 AS n").rows == client.query(1, "SELECT 7 AS n") + + +class TestRunJob: + def test_create_only_returns_job_result(self, client: Client, mock_kc: MagicMock) -> None: + mock_kc.create_job.return_value = { + "id": "job-1", + "status": "processing", + "component": "keboola.ex-db-snowflake", + "configId": "cfg-9", + } + result = client.run_job("keboola.ex-db-snowflake", "cfg-9") + assert isinstance(result, JobResult) + assert result.id == "job-1" and result.component_id == "keboola.ex-db-snowflake" + assert result.config_id == "cfg-9" + mock_kc.wait_for_queue_job.assert_not_called() + # production (no branch) by default, no variable resolution in the facade + call = mock_kc.create_job.call_args.kwargs + assert call["branch_id"] is None and call["variable_values_id"] is None + + def test_wait_polls_for_terminal_state(self, client: Client, mock_kc: MagicMock) -> None: + mock_kc.create_job.return_value = {"id": "job-2", "status": "processing"} + mock_kc.wait_for_queue_job.return_value = { + "id": "job-2", + "status": "success", + "isFinished": True, + } + result = client.run_job("c", "cfg", wait=True, timeout=30, poll_strategy="fixed") + assert result.succeeded and result.is_finished + mock_kc.wait_for_queue_job.assert_called_once_with( + "job-2", max_wait=30, poll_strategy="fixed" + ) + + def test_branch_and_row_ids_forwarded(self, mock_kc: MagicMock) -> None: + c = _make_client(mock_kc, branch_id=77) + mock_kc.create_job.return_value = {"id": "j", "status": "created"} + c.run_job("c", "cfg", config_row_ids=["r1"], variable_values_id="vv1", mode="debug") + call = mock_kc.create_job.call_args.kwargs + assert call["branch_id"] == 77 + assert call["config_row_ids"] == ["r1"] + assert call["variable_values_id"] == "vv1" + assert call["mode"] == "debug" + + def test_explicit_branch_overrides_client_branch(self, mock_kc: MagicMock) -> None: + c = _make_client(mock_kc, branch_id=77) + mock_kc.create_job.return_value = {"id": "j", "status": "created"} + c.run_job("c", "cfg", branch_id=88) + assert mock_kc.create_job.call_args.kwargs["branch_id"] == 88 + + +class TestConfigDetail: + def test_returns_typed_detail(self, client: Client, mock_kc: MagicMock) -> None: + mock_kc.get_config_detail.return_value = { + "id": "cfg-1", + "name": "My Config", + "currentVersion": 4, + "configuration": {"parameters": {"x": 1}}, + "rows": [], + } + detail = client.config_detail("keboola.ex-http", "cfg-1") + assert isinstance(detail, ConfigDetailResult) + assert detail.id == "cfg-1" and detail.version == 4 + assert detail.component_id == "keboola.ex-http" # injected from the arg + assert detail.configuration == {"parameters": {"x": 1}} + mock_kc.get_config_detail.assert_called_once_with( + "keboola.ex-http", "cfg-1", branch_id=None + ) + + def test_branch_scoped(self, mock_kc: MagicMock) -> None: + c = _make_client(mock_kc, branch_id=55) + mock_kc.get_config_detail.return_value = {"id": "cfg-2", "name": "n"} + detail = c.config_detail("keboola.ex-http", "cfg-2") + assert mock_kc.get_config_detail.call_args.kwargs["branch_id"] == 55 + assert detail.branch_id == 55 + + +class TestUploadTable: + def test_computes_size_and_maps_imported_rows( + self, client: Client, mock_kc: MagicMock, tmp_path: Path + ) -> None: + csv = tmp_path / "data.csv" + csv.write_text("a,b\n1,2\n3,4\n") + mock_kc.upload_table.return_value = {"importedRowsCount": 2, "warnings": []} + result = client.upload_table("in.c-x.t", csv, incremental=True) + assert isinstance(result, UploadTableResult) + assert result.table_id == "in.c-x.t" + assert result.incremental is True + assert result.imported_rows == 2 + assert result.file_size_bytes == csv.stat().st_size + # facade never auto-creates + assert result.auto_created_bucket is False and result.auto_created_table is False + call = mock_kc.upload_table.call_args.kwargs + assert call["table_id"] == "in.c-x.t" and call["incremental"] is True + assert call["branch_id"] is None + + def test_branch_scoped_upload(self, mock_kc: MagicMock, tmp_path: Path) -> None: + c = _make_client(mock_kc, branch_id=33) + csv = tmp_path / "d.csv" + csv.write_text("a\n1\n") + mock_kc.upload_table.return_value = {"importedRowsCount": 1} + c.upload_table("in.c-x.t", csv) + assert mock_kc.upload_table.call_args.kwargs["branch_id"] == 33 + + class TestModuleLayout: def test_pagination_helper_relocated_to_client(self) -> None: """_collect_inline_results lives in client.py; workspace_service re-exports it.""" @@ -251,4 +396,13 @@ def test_pagination_helper_relocated_to_client(self) -> None: def test_public_surface(self) -> None: import keboola_agent_cli as pkg - assert set(pkg.__all__) >= {"Client", "Files", "FileEntry"} + assert set(pkg.__all__) >= { + "Client", + "Files", + "FileEntry", + "JobResult", + "QueryResult", + "UploadTableResult", + "SyncPushResult", + "ConfigDetailResult", + } diff --git a/tests/test_result_models.py b/tests/test_result_models.py new file mode 100644 index 00000000..112bf700 --- /dev/null +++ b/tests/test_result_models.py @@ -0,0 +1,182 @@ +"""Contract tests for the typed SDK return models (result_models.py, issue #428). + +These models are the committed, semver-stable surface a downstream in-process +consumer types against. The tests pin three invariants per model: + + 1. the real API keys (camelCase, sometimes asymmetric like ``component`` vs + ``configId``) map onto the snake_case fields, + 2. unknown/extra fields survive (``extra="allow"``) so backend drift never + raises, + 3. the convenience properties compute correctly. +""" + +from keboola_agent_cli import ( + ConfigDetailResult, + JobResult, + QueryResult, + SyncPushResult, + UploadTableResult, +) +from keboola_agent_cli.result_models import _ApiResultModel + + +class TestJobResult: + def test_maps_real_queue_keys(self) -> None: + j = JobResult.model_validate( + { + "id": "98765", + "status": "success", + "isFinished": True, + "component": "keboola.ex-db-snowflake", + "configId": "12345", + "mode": "run", + "result": {"message": "Extraction finished"}, + "resolvedVariableValuesId": "row-7", + "logTail": [{"message": "done"}], + } + ) + assert j.id == "98765" + assert j.is_finished is True + # NB: the Queue job response is asymmetric -- component but configId. + assert j.component_id == "keboola.ex-db-snowflake" + assert j.config_id == "12345" + assert j.result == {"message": "Extraction finished"} + assert j.resolved_variable_values_id == "row-7" + assert j.log_tail == [{"message": "done"}] + + def test_succeeded_and_failed_properties(self) -> None: + assert JobResult(status="success").succeeded is True + assert JobResult(status="success").failed is False + for bad in ("error", "terminated", "cancelled"): + assert JobResult(status=bad).failed is True + assert JobResult(status=bad).succeeded is False + # processing is neither + assert JobResult(status="processing").failed is False + assert JobResult(status="processing").succeeded is False + + def test_extras_preserved(self) -> None: + j = JobResult.model_validate( + {"id": "1", "status": "success", "branchId": 42, "durationSeconds": 9, "url": "u"} + ) + dumped = j.model_dump() + assert dumped["branchId"] == 42 + assert dumped["durationSeconds"] == 9 + assert dumped["url"] == "u" + + def test_populate_by_field_name(self) -> None: + j = JobResult(id="x", status="error", is_finished=True, component_id="c", config_id="cc") + assert j.component_id == "c" and j.config_id == "cc" and j.is_finished is True + + def test_empty_defaults(self) -> None: + j = JobResult() + assert j.id == "" and j.status == "" and j.is_finished is False + assert j.result is None and j.log_tail is None + + +class TestQueryResult: + def test_shape_and_row_count(self) -> None: + q = QueryResult( + columns=["ID", "NAME"], + rows=[{"ID": "1", "NAME": "alice"}, {"ID": "2", "NAME": "bob"}], + truncated=True, + total_rows=999, + ) + assert q.columns == ["ID", "NAME"] + assert q.row_count == 2 + assert q.truncated is True + assert q.total_rows == 999 + + def test_defaults_empty(self) -> None: + q = QueryResult() + assert q.columns == [] and q.rows == [] and q.row_count == 0 + assert q.truncated is False and q.total_rows is None + + +class TestUploadTableResult: + def test_maps_imported_rows_count_alias(self) -> None: + up = UploadTableResult.model_validate( + {"table_id": "in.c-x.t", "importedRowsCount": 50, "warnings": ["w"]} + ) + assert up.imported_rows == 50 + assert up.table_id == "in.c-x.t" + assert up.warnings == ["w"] + assert up.auto_created_bucket is False and up.auto_created_table is False + + def test_service_shape_with_snake_case(self) -> None: + up = UploadTableResult.model_validate( + { + "project_alias": "prod", + "table_id": "in.c-x.t", + "incremental": True, + "file_size_bytes": 1024, + "imported_rows": 7, + "auto_created_bucket": True, + "auto_created_table": True, + } + ) + assert up.imported_rows == 7 and up.incremental is True + assert up.file_size_bytes == 1024 + assert up.auto_created_bucket and up.auto_created_table + assert up.project_alias == "prod" + + +class TestSyncPushResult: + def test_ok_property(self) -> None: + clean = SyncPushResult.model_validate( + {"status": "pushed", "created": 3, "updated": 1, "errors": []} + ) + assert clean.ok is True and clean.created == 3 and clean.updated == 1 + + failed = SyncPushResult.model_validate( + {"status": "pushed", "errors": [{"message": "boom"}]} + ) + assert failed.ok is False + + def test_defaults(self) -> None: + r = SyncPushResult() + assert r.created == 0 and r.updated == 0 and r.deleted == 0 + assert r.errors == [] and r.pushed_details == [] and r.name_drift_warnings == [] + assert r.ok is True + + +class TestConfigDetailResult: + def test_current_version_alias(self) -> None: + cd = ConfigDetailResult.model_validate( + { + "id": "cfg-1", + "name": "My Config", + "description": "d", + "currentVersion": 5, + "configuration": {"parameters": {"x": 1}}, + "rows": [{"id": "r1"}], + } + ) + assert cd.id == "cfg-1" and cd.version == 5 + assert cd.configuration == {"parameters": {"x": 1}} + assert cd.rows == [{"id": "r1"}] + + def test_plain_version_field(self) -> None: + cd = ConfigDetailResult.model_validate({"id": "c", "version": 2}) + assert cd.version == 2 + + def test_component_id_alias_and_extras(self) -> None: + cd = ConfigDetailResult.model_validate( + {"id": "c", "componentId": "keboola.ex-http", "isDisabled": True} + ) + assert cd.component_id == "keboola.ex-http" + # untyped Storage fields survive + assert cd.model_dump()["isDisabled"] is True + + +class TestBaseConfig: + def test_all_models_allow_extra(self) -> None: + for model in ( + JobResult, + QueryResult, + UploadTableResult, + SyncPushResult, + ConfigDetailResult, + ): + assert issubclass(model, _ApiResultModel) + assert model.model_config.get("extra") == "allow" + assert model.model_config.get("populate_by_name") is True