diff --git a/CLAUDE.md b/CLAUDE.md index c7da3356..a6b0fcde 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -124,6 +124,7 @@ tests/ test_ai_client.py # AI Service client tests test_component_service.py # Component service tests test_component_cli.py # Component CLI tests via CliRunner + test_e2e.py # E2E tests against real API (make test-e2e) test_integration.py # Integration tests (edge cases, linting) ``` @@ -199,6 +200,8 @@ All three inherit from `BaseHttpClient` (`http_base.py`) which provides shared r 15. **Pre-commit checks are mandatory.** Before every `git commit`, run `ruff check` and `ruff format --check` on changed files. A pre-commit hook (`scripts/pre-commit`, install via `make hooks`) does this automatically. **Never commit without passing lint + format.** If using sub-agents that write code, always run `make check` (or at minimum `ruff check src/ tests/ && ruff format . --check`) before committing their output. +16. **E2E test coverage**: Every new CLI command MUST have a corresponding E2E test in `tests/test_e2e.py`. Run `make test-e2e` to verify. E2E tests require `E2E_API_TOKEN` and `E2E_URL` env vars and exercise the full CLI against a real Keboola project. + ## Claude Code Plugin (Marketplace) This repo doubles as a Claude Code plugin marketplace. The plugin lives in `plugins/kbagent/` and contains a skill that teaches Claude how to use kbagent. @@ -241,6 +244,7 @@ kbagent config list [--project NAME] [--component-type TYPE] [--component-id ID] kbagent config detail --project NAME --component-id ID --config-id ID [--branch ID] kbagent config search --query PATTERN [--project NAME] [--component-type TYPE] [--ignore-case] [--regex] [--branch ID] kbagent config update --project NAME --component-id ID --config-id ID [--name N] [--description D] [--configuration JSON|@file|-] [--configuration-file PATH] [--set PATH=VALUE ...] [--merge] [--dry-run] [--branch ID] +kbagent config rename --project NAME --component-id ID --config-id ID --name "New Name" [--branch ID] [--directory DIR] kbagent job list [--project NAME] [--component-id ID] [--status STATUS] [--limit N] kbagent job detail --project NAME --job-id ID @@ -255,6 +259,7 @@ kbagent storage create-table --project NAME --bucket-id ID --name NAME --column kbagent storage upload-table --project NAME --table-id ID --file PATH [--incremental] [--branch ID] kbagent storage download-table --project NAME --table-id ID [--output FILE] [--columns COL ...] [--limit N] [--branch ID] kbagent storage delete-table --project NAME --table-id ID [--table-id ...] [--dry-run] [--yes] [--branch ID] +kbagent storage delete-column --project NAME --table-id ID --column COL [--column ...] [--dry-run] [--yes] [--branch ID] kbagent storage delete-bucket --project NAME --bucket-id ID [--bucket-id ...] [--force] [--dry-run] [--yes] [--branch ID] kbagent storage files --project NAME [--tag TAG ...] [--limit N] [--offset N] [--query Q] [--branch ID] kbagent storage file-upload --project NAME --file PATH [--name NAME] [--tag TAG ...] [--permanent] [--branch ID] @@ -302,6 +307,11 @@ kbagent config new --component-id ID [--name NAME] [--project NAME] [--output-di kbagent encrypt values --project ALIAS --component-id ID --input JSON|@file|- [--output-file PATH] +kbagent kai ping [--project NAME] +kbagent kai ask --message "question" [--project NAME] +kbagent kai chat --message "msg" [--chat-id ID] [--project NAME] +kbagent kai history [--project NAME] [--limit N] + kbagent context kbagent init [--from-global] kbagent doctor [--fix] diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 87ac3bce..be051647 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -188,22 +188,29 @@ When adding a new command (e.g., `kbagent storage create-foo`), you must update - [ ] **Client method** in `client.py` (or `manage_client.py`) -- HTTP layer - [ ] **Service method** in `services/` -- business logic, validation, orchestration - [ ] **Command function** in `commands/` -- Typer options, formatter, error handling -- [ ] **Hint definition** in `hints/definitions/` -- register a `CommandHint` for `--hint` code generation (see existing files for pattern) -- [ ] **Hint short-circuit** in the command function -- add `if should_hint(ctx): emit_hint(...)` before service call +- [ ] **`--hint` support** -- every command must support `--hint client` and `--hint service` code generation: + - [ ] **Hint definition** in `hints/definitions/` -- register a `CommandHint` with `ClientCall` + `ServiceCall` (see existing files for pattern) + - [ ] **Hint short-circuit** in the command function -- add `if should_hint(ctx): emit_hint(...)` **before** the service call + - [ ] **Verify** both modes produce valid Python: `kbagent --hint client ...` and `kbagent --hint service ...` - [ ] **Permission registration** in `permissions.py` (`OPERATION_REGISTRY` dict) - [ ] **Service wiring** in `cli.py` if adding a new service class ### Documentation changes (mandatory!) -- [ ] **`kbagent context`** -- update `AGENT_CONTEXT` string in `commands/context.py` (this is the primary reference for AI agents) +- [ ] **`kbagent context`** -- update `AGENT_CONTEXT` string in `commands/context.py` (this is the primary reference for AI agents; if it's missing there, AI agents won't know the command exists) - [ ] **SKILL.md** -- run `make skill-gen` to regenerate the decision table (CI has a freshness check that will fail if the generated output doesn't match). **Do not edit SKILL.md by hand** -- the table is auto-generated from CLI command metadata - [ ] **CLAUDE.md** -- add command signature to the `## All CLI Commands` section +- [ ] **Plugin references** -- update `plugins/kbagent/skills/kbagent/references/`: + - [ ] **`commands-reference.md`** -- add the new command to the appropriate section (this is a hand-maintained file, NOT auto-generated) + - [ ] **New reference file** -- if the command introduces a new workflow or topic area (e.g. a new subcommand group), create a dedicated `-workflow.md` in the references directory. Existing examples: `workspace-workflow.md`, `branch-workflow.md`, `sync-workflow.md`, `storage-files-workflow.md` + - [ ] **`gotchas.md`** -- if the command has non-obvious behavior, response format quirks, or common mistakes, document them here - [ ] **`--help` text** -- Typer docstring and option help strings should be clear and complete ### Tests (mandatory!) - [ ] **Service-layer tests** -- mock the client, test business logic, edge cases, error propagation - [ ] **CLI-layer tests** -- use `CliRunner`, test JSON output, error exit codes +- [ ] **E2E tests** -- add a test in `tests/test_e2e.py` that exercises the command against a real Keboola project (requires `E2E_API_TOKEN` + `E2E_URL`). Run `make test-e2e` to verify. Every CLI command must have E2E coverage - [ ] **Run `make check`** before committing (lint + format + full test suite) ### UX considerations diff --git a/Makefile b/Makefile index ac2efe87..188af848 100644 --- a/Makefile +++ b/Makefile @@ -1,9 +1,9 @@ .DEFAULT_GOAL := help -.PHONY: help install install-mcp sync test test-unit test-integration test-file lint lint-fix format format-check skill-check skill-gen version-sync version-check changelog changelog-check check clean hooks +.PHONY: help install install-mcp sync test test-unit test-integration test-e2e test-file lint lint-fix format format-check skill-check skill-gen version-sync version-check changelog changelog-check check clean hooks help: ## Show this help message - @grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | awk 'BEGIN {FS = ":.*?## "}; {printf " \033[36m%-18s\033[0m %s\n", $$1, $$2}' + @grep -E '^[a-zA-Z0-9_-]+:.*?## .*$$' $(MAKEFILE_LIST) | awk 'BEGIN {FS = ":.*?## "}; {printf " \033[36m%-18s\033[0m %s\n", $$1, $$2}' install: ## Install in development mode (editable) uv pip install -e ".[dev]" @@ -14,15 +14,18 @@ install-mcp: ## Install Keboola MCP server (required for 'tool' commands) sync: ## Sync dependencies from lockfile uv sync -test: ## Run all tests - uv run pytest tests/ -v +test: ## Run all tests (excluding e2e — use test-e2e separately) + uv run pytest tests/ -v -m "not e2e" -test-unit: ## Run unit tests only (exclude integration) - uv run pytest tests/ -v -m "not integration" +test-unit: ## Run unit tests only (exclude integration and e2e) + uv run pytest tests/ -v -m "not integration and not e2e" test-integration: ## Run integration tests only uv run pytest tests/ -v -m integration +test-e2e: ## Run E2E tests (E2E_API_TOKEN and E2E_URL required) + uv run pytest tests/test_e2e.py -v -s --tb=long + test-file: ## Run a specific test file (FILE=tests/test_cli.py) uv run pytest $(FILE) -v diff --git a/README.md b/README.md index ed3d2221..43e4d821 100644 --- a/README.md +++ b/README.md @@ -16,7 +16,7 @@ Auto-updates on every launch. Run `kbagent changelog` to see what changed. This CLI is built AI-first. Every command outputs structured JSON (`--json`), errors include machine-readable codes, and the permission firewall enforces safety at the code level -- not via prompt instructions. -**Claude Code plugin** (agent learns all 74 commands automatically): +**Claude Code plugin** (agent learns all 80 commands automatically): ``` /plugin marketplace add padak/keboola_agent_cli @@ -77,6 +77,7 @@ kbagent workspace query --project prod --workspace-id WS_ID \ | **MCP tools** | Call `keboola-mcp-server` tools with auto-expand, multi-project fan-out, branch propagation, schema validation. | | **Workspaces** | Create Snowflake/BQ workspace, load tables, run SQL. Create from transformation config for instant debugging. | | **Sharing & lineage** | Cross-project data lineage via bucket sharing. Share/link/unlink with org/project/user access control. | +| **Kai (AI Assistant)** | Ask Keboola's built-in AI questions about your project. One-shot or chat sessions with full MCP context. | | **Encryption** | Encrypt secrets (`#password`, `#api_token`) via Keboola Encryption API. Works with sync push and MCP. | | **Permissions** | Firewall for AI agents: read-only, deny-writes, deny-destructive. Code-level enforcement, not prompt tricks. | | **Auto-update** | Self-updates on startup. "What's new" after each update. Full changelog via `kbagent changelog`. | @@ -110,10 +111,10 @@ Full command reference with flags: [SKILL.md](plugins/kbagent/skills/kbagent/SKI kbagent project add | list | remove | edit | status | refresh kbagent org setup kbagent component list | detail -kbagent config list | detail | search | update | delete | new +kbagent config list | detail | search | update | rename | delete | new kbagent job list | detail | run kbagent storage buckets | bucket-detail | create-bucket | delete-bucket - tables | table-detail | create-table | upload-table | download-table | delete-table + tables | table-detail | create-table | upload-table | download-table | delete-table | delete-column files | file-detail | file-upload | file-download | file-tag | file-delete load-file | unload-table kbagent sharing list | share | unshare | link | unlink @@ -122,6 +123,7 @@ kbagent branch list | create | use | reset | delete | merge kbagent workspace create | list | detail | delete | password | load | query | from-transformation kbagent tool list | call kbagent sync init | pull | status | diff | push | branch-link | branch-unlink | branch-status +kbagent kai ping | ask | chat | history kbagent encrypt values kbagent permissions list | show | set | reset | check kbagent init | context | doctor | version | update | changelog diff --git a/docs/e2e-scenarios.md b/docs/e2e-scenarios.md new file mode 100644 index 00000000..e54d794a --- /dev/null +++ b/docs/e2e-scenarios.md @@ -0,0 +1,223 @@ +# E2E Test Scenarios + +End-to-end tests that exercise the full CLI against a real Keboola project. + +## Running + +```bash +# All E2E tests (~2.5 min) +E2E_API_TOKEN=xxx E2E_URL=connection.keboola.com make test-e2e + +# Only the main scenario (36 steps) +E2E_API_TOKEN=xxx E2E_URL=connection.keboola.com \ + uv run pytest tests/test_e2e.py::TestFullE2E -v -s + +# Without credentials -- all tests are skipped automatically +make test-e2e +``` + +## Test classes + +| Class | Tests | What it covers | +|-------|------:|----------------| +| `TestFullE2E` | 1 (41 steps) | Progressive scenario building state from empty project | +| `TestE2EErrorHandling` | 6 | Invalid tokens, nonexistent resources, correct exit codes | +| `TestE2EJsonConsistency` | 2 | All read commands return valid JSON; token never leaks | +| `TestE2ESyncWorkflow` | 1 (5 steps) | Sync init/pull/status/diff/push in a temp git repo | +| `TestE2EToolCommands` | 2 | MCP tool list + tool call (skipped if no MCP server) | + +--- + +## TestFullE2E -- Main scenario (36 steps) + +All steps run sequentially. Each step builds on the state created by previous steps. +Resources are prefixed with `e2e-{timestamp}` and cleaned up via yield fixture even on failure. + +### Phase 1: Setup + +| Step | Command | What is verified | +|-----:|---------|------------------| +| 1 | `version`, `changelog`, `context` | Offline commands work, version contains dot, context mentions kbagent | +| 2 | `init` | Creates `.kbagent/` directory, returns `created: true` | +| 3 | `project add` | Registers project, returns alias/name/id, token is masked | +| 4 | `project list`, `project status` | Project appears in list, status is `ok` with response time | +| 5 | `doctor` | Health check passes (`summary.healthy: true`) | + +### Phase 2: Read empty project + +| Step | Command | What is verified | +|-----:|---------|------------------| +| 6 | `config list`, `storage buckets`, `job list` | Empty lists with correct JSON structure, no errors | + +### Phase 3: Storage CRUD + +| Step | Command | What is verified | +|-----:|---------|------------------| +| 7 | `storage create-bucket` | Bucket ID starts with `in.c-`, tracked for cleanup | +| 8 | `storage buckets`, `storage bucket-detail` | Bucket appears in listing, detail returns correct ID | +| 9 | `storage create-table` | Table created with typed columns (INTEGER, STRING) and primary key | +| 10 | `storage upload-table` | 5-row CSV uploaded, `imported_rows: 5` | +| 11 | `storage upload-table --incremental` | 3 more rows appended, download verifies 8 total rows | +| 12 | `storage tables`, `storage table-detail` | Table in listing, column details match (id, name, value) | +| 13 | `storage download-table` | Full download: 8 rows, correct IDs. With `--columns`/`--limit`: subset verified | +| 14 | `storage unload-table --download` | Table exported to file storage, file_id > 0, file downloaded | +| 15 | `storage load-file` | CSV uploaded as file, then loaded into table via `load-file` | + +### Phase 4: Config operations + +| Step | Command | What is verified | +|-----:|---------|------------------| +| 16 | `config list`, `config detail`, `config search`, `config search --ignore-case` | Config found by name, detail has correct parameters, search matches by pattern | +| 17 | `config update --set`, `--dry-run`, `--name/--description`, `--configuration` | Nested key set preserves siblings, dry-run shows diff, full replace removes old keys | +| 18 | `config update --merge` | Partial JSON deep-merged, existing keys preserved alongside new ones | +| 19 | `config new --component-id keboola.ex-http` | Scaffold generated with `_config.yml` file | + +### Phase 5: Component discovery + +| Step | Command | What is verified | +|-----:|---------|------------------| +| 20 | `component list`, `component list --type extractor`, `component detail` | Components listed (after config exists), type filter works, detail has schema info | + +### Phase 6: Workspace lifecycle + +| Step | Command | What is verified | +|-----:|---------|------------------| +| 21 | `workspace create` | Returns workspace_id, host, schema, user, password | +| 22 | `workspace list` | Workspace appears in project listing | +| 23 | `workspace detail` | Returns backend, host, schema, user | +| 24 | `workspace password` | New password returned (non-empty) | +| 25 | `workspace load --tables TABLE_ID` | Test table loaded into workspace | +| 26 | `workspace query --sql "SELECT COUNT(*)"` | SQL executed successfully | +| 27 | `workspace delete` | Workspace removed | + +Steps 21-27 are wrapped in try/except -- if workspace API is unavailable on the stack, they are skipped gracefully. + +### Phase 7: Transformation job run (Snowflake SQL) + +| Step | Command | What is verified | +|-----:|---------|------------------| +| 28 | `storage create-bucket` (out stage) + API `create_config` | Output bucket created, Snowflake transformation config created with SQL: `SELECT id, name, CAST(value AS INT) AS value, CAST(value AS INT) * 2 AS doubled_value` | +| 29 | `job run --wait --timeout 300` | Transformation executes, job status is `success` | +| 30 | `job detail --job-id ID` | Completed job detail: `status=success`, `isFinished=true`, component is `keboola.snowflake-transformation` | +| 31 | `storage download-table` (output table) | Output downloaded, 9 rows, every row has `doubled_value == value * 2` | +| 32 | `config delete` + `storage delete-bucket --force` | Transformation config and output bucket cleaned up | + +### Phase 8: File operations + +| Step | Command | What is verified | +|-----:|---------|------------------| +| 33 | `file-upload`, `files`, `file-detail`, `file-download`, `file-tag --add/--remove`, `file-delete --dry-run`, `file-delete --yes` | Full lifecycle: upload with tags, list by tag, detail shows tags, download content matches, tag add/remove verified, dry-run shows would_delete, actual delete confirmed | + +### Phase 9: Encrypt + +| Step | Command | What is verified | +|-----:|---------|------------------| +| 34 | `encrypt values` | Input `#password`/`#api_key` encrypted to `KBC::ProjectSecure::...` format | + +### Phase 10: Branch lifecycle + +| Step | Command | What is verified | +|-----:|---------|------------------| +| 35 | `branch list`, `branch create`, `branch use`, `branch reset`, `branch merge`, `branch delete` | Main branch exists, dev branch created (auto-activates), use/reset toggle active branch in project status, merge returns URL, branch deleted and gone from list | + +### Phase 11: Permissions + +| Step | Command | What is verified | +|-----:|---------|------------------| +| 36 | `permissions list`, `permissions show`, `permissions check branch.delete` | List returns operations array, show returns policy status, check returns `allowed: true` | + +### Phase 12: Sharing and lineage + +| Step | Command | What is verified | +|-----:|---------|------------------| +| 37 | `sharing list`, `lineage show` | Both return valid responses (may be empty on single project) | + +### Phase 12.5: Kai (Keboola AI Assistant) + +| Step | Command | What is verified | +|-----:|---------|------------------| +| 38 | `kai ping` | Server health, timestamp, MCP status. Gracefully skips all kai tests if `agent-chat` feature not enabled | +| 38 | `kai ask -m "..."` | One-shot question, verify response text + chat_id. Skips if auth fails (token type) | +| 38 | `kai history --limit 5` | At least 1 chat after asking | + +Steps are wrapped in graceful skip logic — if Kai is not available (feature flag or auth), remaining kai tests are skipped without failing. + +### Phase 13: Job commands + +| Step | Command | What is verified | +|-----:|---------|------------------| +| 38 | `job list`, `job list --component-id`, `job detail` | List structure correct, component filter works. If jobs exist from uploads: detail returns full job data with status field | + +### Phase 14: Cleanup via CLI + +| Step | Command | What is verified | +|-----:|---------|------------------| +| 39 | `config delete` | Config removed, confirmed by config_id in response | +| 40 | `storage delete-table --dry-run`, `--yes`, `storage delete-bucket --dry-run`, `--yes` | Dry-run shows would_delete, actual delete confirmed | +| 41 | `project edit`, `project remove` | Edit preserves alias, remove confirmed, project gone from list | + +--- + +## TestE2EErrorHandling + +| Test | Command | Expected | +|------|---------|----------| +| `test_add_with_invalid_token` | `project add --token 000-invalid` | Exit code 3, `INVALID_TOKEN` | +| `test_status_of_nonexistent_project` | `project status --project nonexistent` | Exit code 5 | +| `test_remove_nonexistent_project` | `project remove --project nonexistent` | Exit code 5 | +| `test_config_detail_nonexistent` | `config detail --config-id 999999999` | Exit code != 0 | +| `test_download_nonexistent_table` | `download-table --table-id in.c-nonexistent.nonexistent` | Exit code != 0 | +| `test_delete_nonexistent_bucket` | `delete-bucket --bucket-id in.c-nonexistent-bucket-xyz` | Exit code != 0 | + +--- + +## TestE2EJsonConsistency + +| Test | What is verified | +|------|------------------| +| `test_all_read_commands_return_valid_json` | `project list`, `project status`, `config list`, `storage buckets`, `job list`, `component list`, `branch list`, `sharing list`, `lineage show`, `doctor`, `permissions list`, `permissions show` -- all return parseable JSON with `status` key | +| `test_token_never_appears_in_any_output` | Full API token never appears in output of `project list`, `project status`, `doctor` | + +--- + +## TestE2ESyncWorkflow + +Runs in a temporary git repository (`git init` + initial commit). + +| Step | Command | What is verified | +|-----:|---------|------------------| +| 1 | `sync init --project ALIAS --directory DIR` | Returns project_alias in response | +| 2 | `sync pull --project ALIAS --directory DIR` | Exit code 0, files pulled | +| 3 | `sync status --directory DIR` | Exit code 0, returns status structure | +| 4 | `sync diff --project ALIAS --directory DIR` | Exit code 0, returns diff structure | +| 5 | `sync push --project ALIAS --directory DIR --dry-run` | Exit code 0, dry-run shows what would be pushed | + +--- + +## TestE2EToolCommands + +Skipped if `keboola-mcp-server` is not installed. + +| Test | Command | What is verified | +|------|---------|------------------| +| `test_tool_list` | `tool list --project ALIAS` | Exit code 0, tools returned | +| `test_tool_call_get_buckets` | `tool call get_buckets --project ALIAS` | Exit code 0, bucket data returned | + +--- + +## Commands NOT covered by E2E (with reasons) + +| Command | Reason | +|---------|--------| +| `project refresh` | Requires Manage API token (`KBC_MANAGE_API_TOKEN`) | +| `org setup` | Requires Manage API token + destructive (registers projects in org) | +| `sharing share/unshare` | Requires org-level permissions or second project | +| `sharing link/unlink` | Requires shared bucket from another project | +| `permissions set/reset` | Interactive random-code confirmation blocks automated testing | +| `workspace from-transformation` | Requires existing transformation config with input mappings | +| `workspace query --file` | Equivalent to `--sql`; only the input source differs | +| `update` | Would actually update the installed package via PyPI | +| `repl` | Interactive REPL, not testable via CliRunner | +| `doctor --fix` | Installs MCP server binary; side effect not suitable for E2E | +| `init --from-global` | Requires global config with projects; tested via unit tests | +| `init --read-only` | Creates Claude Code permission rules; tested via unit tests | diff --git a/plugins/kbagent/.claude-plugin/plugin.json b/plugins/kbagent/.claude-plugin/plugin.json index 9f8b410d..aa3f643d 100644 --- a/plugins/kbagent/.claude-plugin/plugin.json +++ b/plugins/kbagent/.claude-plugin/plugin.json @@ -1,6 +1,6 @@ { "name": "kbagent", - "version": "0.18.6", + "version": "0.19.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/SKILL.md b/plugins/kbagent/skills/kbagent/SKILL.md index 1fa1fbd1..204a79c4 100644 --- a/plugins/kbagent/skills/kbagent/SKILL.md +++ b/plugins/kbagent/skills/kbagent/SKILL.md @@ -68,6 +68,7 @@ If kbagent is not installed or you need the full standalone reference, run `kbag | Show detailed information about a specific configuration | `kbagent config detail --project PROJECT --component-id COMPONENT-ID --config-id CONFIG-ID` | | Search through configuration bodies for a string or pattern | `kbagent config search --query QUERY` | | Update a configuration's metadata and/or content | `kbagent config update --project PROJECT --component-id COMPONENT-ID --config-id CONFIG-ID` | +| Rename a configuration (update name via API + rename local sync directory) | `kbagent config rename --project PROJECT --component-id COMPONENT-ID --config-id CONFIG-ID --name NAME` | | Delete a configuration from a project | `kbagent config delete --project PROJECT --component-id COMPONENT-ID --config-id CONFIG-ID` | | Generate boilerplate configuration files for a Keboola component | `kbagent config new --component-id COMPONENT-ID` | | List jobs from connected projects | `kbagent job list` | @@ -82,6 +83,7 @@ If kbagent is not installed or you need the full standalone reference, run `kbag | Upload a CSV file into a storage table | `kbagent storage upload-table --project PROJECT --table-id TABLE-ID --file FILE` | | Export a storage table to a local CSV file | `kbagent storage download-table --project PROJECT --table-id TABLE-ID` | | Delete one or more storage tables | `kbagent storage delete-table --project PROJECT --table-id TABLE-ID` | +| Delete one or more columns from a storage table | `kbagent storage delete-column --project PROJECT --table-id TABLE-ID --column COLUMN` | | Delete one or more storage buckets | `kbagent storage delete-bucket --project PROJECT --bucket-id BUCKET-ID` | | List Storage Files with optional tag filtering | `kbagent storage files --project PROJECT` | | Show Storage File metadata (without downloading) | `kbagent storage file-detail --project PROJECT --file-id FILE-ID` | @@ -97,6 +99,10 @@ If kbagent is not installed or you need the full standalone reference, run `kbag | Link a shared bucket into a project | `kbagent sharing link --project PROJECT --source-project-id SOURCE-PROJECT-ID --bucket-id BUCKET-ID` | | Remove a linked bucket from a project | `kbagent sharing unlink --project PROJECT --bucket-id BUCKET-ID` | | Show cross-project data lineage via bucket sharing | `kbagent lineage show` | +| Check Kai server health and MCP connection status | `kbagent kai ping` | +| Ask Kai a one-shot question and get the full response | `kbagent kai ask --message MESSAGE` | +| Send a message to Kai in a chat session | `kbagent kai chat --message MESSAGE` | +| List recent Kai chat sessions | `kbagent kai history` | | List development branches from connected projects | `kbagent branch list` | | Create a new development branch and auto-activate it | `kbagent branch create --project PROJECT --name NAME` | | Set an existing development branch as active | `kbagent branch use --project PROJECT --branch BRANCH` | diff --git a/plugins/kbagent/skills/kbagent/references/commands-reference.md b/plugins/kbagent/skills/kbagent/references/commands-reference.md index 04c2f6d7..ecda4ddc 100644 --- a/plugins/kbagent/skills/kbagent/references/commands-reference.md +++ b/plugins/kbagent/skills/kbagent/references/commands-reference.md @@ -30,6 +30,7 @@ All commands support `--json` for structured output. Multi-project flags (`--pro - `config detail --project NAME --component-id ID --config-id ID [--branch ID]` -- full config with parameters and rows (branch-aware) - `config search --query PATTERN [--project NAME] [-i] [-r] [--branch ID]` -- search config bodies for string/regex (branch-aware) - `config update --project NAME --component-id ID --config-id ID [--name N] [--description D] [--configuration JSON|@file|-] [--configuration-file PATH] [--set PATH=VALUE ...] [--merge] [--dry-run] [--branch ID]` -- update metadata and/or configuration content. `--set` targets a nested key (e.g. `parameters.db.host=new-host`). `--merge` deep-merges into existing config (preserves sibling keys). `--dry-run` previews changes without applying. Paths are relative to the configuration root (unlike MCP's `update_config` which uses paths relative to `parameters`) +- `config rename --project NAME --component-id ID --config-id ID --name "New Name" [--branch ID] [--directory DIR]` -- rename a configuration (API update + local sync directory rename with git mv support) - `config delete --project NAME --component-id ID --config-id ID [--branch ID]` -- delete a configuration - `config new --component-id ID [--project NAME] [--name NAME] [--output-dir DIR]` -- scaffold new config from component schema @@ -48,6 +49,7 @@ All commands support `--json` for structured output. Multi-project flags (`--pro - `storage upload-table --project NAME --table-id ID --file PATH [--incremental] [--branch ID]` -- upload CSV (branch-aware) - `storage download-table --project NAME --table-id ID [--output FILE] [--columns COL ...] [--limit N] [--branch ID]` -- export table to CSV (branch-aware) - `storage delete-table --project NAME --table-id ID [--table-id ...] [--dry-run] [--yes] [--branch ID]` -- delete tables (branch-aware) +- `storage delete-column --project NAME --table-id ID --column COL [--column ...] [--dry-run] [--yes] [--branch ID]` -- delete columns from a table (branch-aware) - `storage delete-bucket --project NAME --bucket-id ID [--bucket-id ...] [--force] [--dry-run] [--yes] [--branch ID]` -- delete buckets (branch-aware) ## Data Lineage @@ -75,6 +77,12 @@ All commands support `--json` for structured output. Multi-project flags (`--pro - `tool list [--project NAME] [--branch ID]` -- list available MCP tools (multi_project annotation) - `tool call TOOL_NAME [--project NAME] [--input JSON|@file|-] [--branch ID]` -- call MCP tool (read = all projects, write = single). `--input` accepts inline JSON, `@file.json`, or `-` (stdin) +## Kai (Keboola AI Assistant) +- `kai ping [--project NAME]` -- check Kai server health and MCP connection status +- `kai ask --message "question" [--project NAME]` -- one-shot question to Kai, collects full response +- `kai chat --message "msg" [--chat-id ID] [--project NAME]` -- send message in a chat session, returns chat_id for continuation +- `kai history [--project NAME] [--limit N]` -- list recent Kai chat sessions (default limit: 10) + ## Sync (GitOps) - `sync init --project ALIAS [--directory DIR] [--git-branching]` -- initialize sync working directory - `sync pull --project ALIAS [--all-projects] [--force] [--dry-run] [--with-samples] [--no-storage] [--no-jobs] [--job-limit N]` -- download configs to local files. For large projects (>100 configs), automatically fetches jobs per-config when the grouped API limit is insufficient diff --git a/plugins/kbagent/skills/kbagent/references/kai-workflow.md b/plugins/kbagent/skills/kbagent/references/kai-workflow.md new file mode 100644 index 00000000..3fc53a32 --- /dev/null +++ b/plugins/kbagent/skills/kbagent/references/kai-workflow.md @@ -0,0 +1,72 @@ +# Kai (Keboola AI Assistant) Workflow + +Kai is Keboola's cloud AI assistant with MCP access to project data. +kbagent bridges Claude Code (local) to Kai (cloud) for Keboola-specific questions. + +> **BETA**: Kai commands require a project with the `agent-chat` feature enabled. +> Token authentication requirements are being finalized. + +## When to use Kai vs local tools + +| Situation | Use | +|-----------|-----| +| Need project-specific context (tables, configs, lineage) | `kbagent kai ask` | +| Simple data listing (buckets, tables, configs) | `kbagent config list`, `kbagent storage tables` | +| Need Keboola domain knowledge (component behavior, best practices) | `kbagent kai ask` | +| Need to modify data (upload, create, delete) | Direct CLI commands | + +## Quick start + +```bash +# Check if Kai is available +kbagent kai ping --project my-project + +# Ask a question about the project +kbagent kai ask --project my-project -m "What tables do I have?" + +# Multi-turn conversation +kbagent kai chat --project my-project -m "Help me debug my pipeline" +# Note the chat_id in the response, then continue: +kbagent kai chat --project my-project --chat-id CHAT_ID -m "What about the error in step 3?" + +# View chat history +kbagent kai history --project my-project --limit 10 +``` + +## Feature detection + +Kai requires the `agent-chat` feature flag on the project. If not enabled, +kai commands return error code `KAI_NOT_ENABLED` with a clear message. + +Check via: `kbagent --json kai ping --project ALIAS` — exit code 0 means Kai is available. + +## JSON output + +All kai commands support `--json` for structured output: + +```bash +# Ping +kbagent --json kai ping --project my-project +# {"status": "ok", "data": {"timestamp": "...", "mcp_status": "ok", ...}} + +# Ask +kbagent --json kai ask --project my-project -m "How many tables?" +# {"status": "ok", "data": {"chat_id": "uuid", "response": "You have 19 tables."}} + +# History +kbagent --json kai history --project my-project +# {"status": "ok", "data": {"chats": [...], "has_more": false}} +``` + +## Common patterns for Claude Code + +```bash +# Use kai ask as a Keboola knowledge oracle +kbagent --json kai ask --project prod -m "Is it safe to drop bucket in.c-legacy?" + +# Get project overview for onboarding +kbagent --json kai ask --project prod -m "Describe the data flow in this project" + +# Debug a failed job +kbagent --json kai ask --project prod -m "Why did job 12345 fail?" +``` diff --git a/pyproject.toml b/pyproject.toml index a0bf30cf..6706476b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "keboola-agent-cli" -version = "0.18.6" +version = "0.19.0" description = "AI-friendly CLI for managing Keboola projects" readme = "README.md" requires-python = ">=3.12" @@ -19,6 +19,7 @@ dependencies = [ "pyyaml>=6", "packaging>=23", "prompt-toolkit>=3.0", + "kai-client>=0.11.0", ] [project.scripts] @@ -36,6 +37,7 @@ testpaths = ["tests"] pythonpath = ["src", "tests"] markers = [ "integration: marks tests as integration tests requiring real API credentials (deselect with '-m \"not integration\"')", + "e2e: marks tests as end-to-end tests requiring real API credentials (deselect with '-m \"not e2e\"')", ] [tool.ruff] diff --git a/src/keboola_agent_cli/changelog.py b/src/keboola_agent_cli/changelog.py index e04339c5..6824636f 100644 --- a/src/keboola_agent_cli/changelog.py +++ b/src/keboola_agent_cli/changelog.py @@ -8,6 +8,15 @@ # Ordered newest-first. Each value is a list of brief one-line descriptions. CHANGELOG: dict[str, list[str]] = { + "0.19.0": [ + "New: Kai (Keboola AI Assistant) -- kai ping, ask, chat, history (BETA) (#164)", + "New: config rename -- rename via API + auto-rename local sync directory (#160)", + "New: sync pull auto-rename -- detects remote name changes and renames local dirs (#160)", + "New: sync push warning -- alerts when local dir names drift from config names (#160)", + "New: storage delete-column -- remove columns from tables with --dry-run (#159)", + "Fix: branch-scoped file operations (get_file_info, delete, tag, untag) (#161)", + "Test: comprehensive E2E test suite covering all CLI commands (#158)", + ], "0.18.6": [ "New: config update --set PATH=VALUE -- set nested config keys without losing siblings (#156)", "New: config update --merge -- deep-merge partial JSON into existing configuration (#156)", diff --git a/src/keboola_agent_cli/cli.py b/src/keboola_agent_cli/cli.py index 384b9dd2..de8a1844 100644 --- a/src/keboola_agent_cli/cli.py +++ b/src/keboola_agent_cli/cli.py @@ -15,6 +15,7 @@ from .commands.encrypt import encrypt_app from .commands.init import init_command from .commands.job import job_app +from .commands.kai import kai_app from .commands.lineage import lineage_app from .commands.org import org_app from .commands.permissions import permissions_app @@ -37,6 +38,7 @@ from .services.doctor_service import DoctorService from .services.encrypt_service import EncryptService from .services.job_service import JobService +from .services.kai_service import KaiService from .services.lineage_service import LineageService from .services.mcp_service import McpService from .services.org_service import OrgService @@ -77,6 +79,7 @@ app.add_typer(storage_app, name="storage", rich_help_panel=_BROWSE) app.add_typer(sharing_app, name="sharing", rich_help_panel=_BROWSE) app.add_typer(lineage_app, name="lineage", rich_help_panel=_BROWSE) +app.add_typer(kai_app, name="kai", rich_help_panel=_BROWSE) # -- Development -- _DEV = "Development" @@ -186,6 +189,7 @@ def main( sync_service = SyncService(config_store=config_store) encrypt_service = EncryptService(config_store=config_store) workspace_service = WorkspaceService(config_store=config_store) + kai_service = KaiService(config_store=config_store) doctor_service = DoctorService(config_store=config_store, mcp_service=mcp_service) version_service = VersionService() @@ -224,6 +228,7 @@ def main( ctx.obj["sync_service"] = sync_service ctx.obj["encrypt_service"] = encrypt_service ctx.obj["workspace_service"] = workspace_service + ctx.obj["kai_service"] = kai_service ctx.obj["doctor_service"] = doctor_service ctx.obj["version_service"] = version_service diff --git a/src/keboola_agent_cli/client.py b/src/keboola_agent_cli/client.py index cab4cee5..fe6aedb9 100644 --- a/src/keboola_agent_cli/client.py +++ b/src/keboola_agent_cli/client.py @@ -189,6 +189,7 @@ def verify_token(self) -> TokenVerifyResponse: project_name=owner.get("name", ""), owner_name=owner.get("name", ""), default_backend=owner.get("defaultBackend", "snowflake"), + features=owner.get("features", []), ) def list_components( @@ -1128,6 +1129,21 @@ def delete_table(self, table_id: str, branch_id: int | None = None) -> dict[str, response = self._request("DELETE", f"{prefix}/tables/{safe_id}", params={"async": "true"}) return self._wait_for_storage_job(response.json()) + def delete_column(self, table_id: str, column_name: str, branch_id: int | None = None) -> None: + """Delete a column from a storage table. + + This is a synchronous operation (no async job). + + Args: + table_id: Full table ID (e.g. "in.c-bucket.table"). + column_name: Name of the column to delete. + branch_id: If set, target a specific dev branch. + """ + prefix = f"/v2/storage/branch/{branch_id}" if branch_id else "/v2/storage" + safe_table_id = quote(table_id, safe="") + safe_column = quote(column_name, safe="") + self._request("DELETE", f"{prefix}/tables/{safe_table_id}/columns/{safe_column}") + def list_tables_with_metadata(self) -> list[dict[str, Any]]: """List all storage tables with columns and metadata. @@ -1196,18 +1212,20 @@ def export_table_async( ) return self._wait_for_storage_job(response.json(), max_wait=EXPORT_JOB_MAX_WAIT) - def get_file_info(self, file_id: int) -> dict[str, Any]: + def get_file_info(self, file_id: int, branch_id: int | None = None) -> dict[str, Any]: """Get file metadata including download URL. Args: file_id: Storage file ID (from export job results). + branch_id: If set, query file from a specific dev branch scope. Returns: File resource dict with 'url', 'isSliced', 'sizeBytes', etc. """ + prefix = f"/v2/storage/branch/{branch_id}" if branch_id else "/v2/storage" response = self._request( "GET", - f"/v2/storage/files/{file_id}", + f"{prefix}/files/{file_id}", params={"federationToken": "1"}, ) return response.json() @@ -1291,32 +1309,38 @@ def upload_file( "created": upload_info.get("created"), } - def delete_file(self, file_id: int) -> None: + def delete_file(self, file_id: int, branch_id: int | None = None) -> None: """Delete a Storage File. Args: file_id: Storage file ID. + branch_id: If set, target a file in a specific dev branch scope. """ - self._request("DELETE", f"/v2/storage/files/{file_id}") + prefix = f"/v2/storage/branch/{branch_id}" if branch_id else "/v2/storage" + self._request("DELETE", f"{prefix}/files/{file_id}") - def tag_file(self, file_id: int, tag: str) -> None: + def tag_file(self, file_id: int, tag: str, branch_id: int | None = None) -> None: """Add a tag to a Storage File. Args: file_id: Storage file ID. tag: Tag string to add. + branch_id: If set, target a file in a specific dev branch scope. """ - self._request("POST", f"/v2/storage/files/{file_id}/tags", data={"tag": tag}) + prefix = f"/v2/storage/branch/{branch_id}" if branch_id else "/v2/storage" + self._request("POST", f"{prefix}/files/{file_id}/tags", data={"tag": tag}) - def untag_file(self, file_id: int, tag: str) -> None: + def untag_file(self, file_id: int, tag: str, branch_id: int | None = None) -> None: """Remove a tag from a Storage File. Args: file_id: Storage file ID. tag: Tag string to remove. + branch_id: If set, target a file in a specific dev branch scope. """ + prefix = f"/v2/storage/branch/{branch_id}" if branch_id else "/v2/storage" safe_tag = quote(tag, safe="") - self._request("DELETE", f"/v2/storage/files/{file_id}/tags/{safe_tag}") + self._request("DELETE", f"{prefix}/files/{file_id}/tags/{safe_tag}") def download_sliced_file(self, file_detail: dict[str, Any], output_path: str) -> int: """Download a sliced file by fetching manifest and concatenating slices. diff --git a/src/keboola_agent_cli/commands/config.py b/src/keboola_agent_cli/commands/config.py index 70dcea67..c9cbaf63 100644 --- a/src/keboola_agent_cli/commands/config.py +++ b/src/keboola_agent_cli/commands/config.py @@ -539,6 +539,117 @@ def config_update( ) +@config_app.command("rename") +def config_rename( + ctx: typer.Context, + project: str = typer.Option( + ..., + "--project", + help="Project alias", + ), + component_id: str = typer.Option( + ..., + "--component-id", + help="Component ID (e.g. keboola.python-transformation-v2)", + ), + config_id: str = typer.Option( + ..., + "--config-id", + help="Configuration ID to rename", + ), + name: str = typer.Option( + ..., + "--name", + help="New name for the configuration", + ), + branch: int | None = typer.Option( + None, + "--branch", + help="Rename in a specific dev branch ID (defaults to active branch)", + ), + directory: Path | None = typer.Option( + None, + "--directory", + "-d", + help="Sync working directory (auto-detects .keboola/manifest.json in CWD if omitted)", + ), +) -> None: + """Rename a configuration (update name via API + rename local sync directory). + + Updates the configuration name in the Keboola project. If a local sync + directory is detected (either via --directory or the current working + directory), the local folder is renamed and the manifest is updated + to match. + + \b + Examples: + # Simple rename + kbagent config rename --project prod --component-id kds-team.app-custom-python \\ + --config-id abc123 --name "Stripe Extractor" + + # Rename with explicit sync directory + kbagent config rename --project prod --component-id kds-team.app-custom-python \\ + --config-id abc123 --name "Stripe Extractor" --directory ./my-project + """ + if should_hint(ctx): + emit_hint( + ctx, + "config.rename", + project=project, + component_id=component_id, + config_id=config_id, + name=name, + branch=branch, + ) + + formatter = get_formatter(ctx) + service = get_service(ctx, "config_service") + + # Auto-detect sync directory from CWD if not specified + effective_directory = directory + if effective_directory is None: + cwd = Path.cwd() + if (cwd / KEBOOLA_DIR_NAME / MANIFEST_FILENAME).exists(): + effective_directory = cwd + + try: + result = service.rename_config( + alias=project, + component_id=component_id, + config_id=config_id, + name=name, + branch_id=branch, + directory=effective_directory, + ) + except ConfigError as exc: + formatter.error(message=exc.message, error_code="CONFIG_ERROR") + raise typer.Exit(code=5) from None + except KeboolaApiError as exc: + formatter.error( + message=exc.message, + error_code=exc.error_code, + retryable=exc.retryable, + ) + raise typer.Exit(code=map_error_to_exit_code(exc)) from None + + if formatter.json_mode: + formatter.output(result) + else: + branch_info = "" + if result.get("branch_id"): + branch_info = f" (branch {result['branch_id']})" + formatter.success( + f'Renamed "{result["old_name"]}" -> "{result["new_name"]}"' + f" ({component_id}/{config_id}){branch_info}" + ) + sync_info = result.get("sync") + if sync_info: + formatter.console.print( + f" Sync: {sync_info['old_path']}/ -> {sync_info['new_path']}/" + f" ({sync_info['method']})" + ) + + @config_app.command("delete") def config_delete( ctx: typer.Context, diff --git a/src/keboola_agent_cli/commands/context.py b/src/keboola_agent_cli/commands/context.py index 53b1175a..0dbf274e 100644 --- a/src/keboola_agent_cli/commands/context.py +++ b/src/keboola_agent_cli/commands/context.py @@ -96,6 +96,11 @@ existing config (preserves sibling keys). --dry-run previews changes. Paths are always relative to the configuration root. + kbagent config rename --project NAME --component-id ID --config-id ID --name "New Name" [--branch ID] [--directory DIR] + Rename a configuration. Updates name via API. If a local sync directory + exists (.keboola/manifest.json), renames the directory and updates the + manifest path. Uses git mv when inside a git repo for cleaner history. + kbagent config delete --project NAME --component-id ID --config-id ID [--branch ID] Delete a configuration. Branch-aware. @@ -151,6 +156,9 @@ kbagent storage delete-table --project NAME --table-id ID [--table-id ...] [--dry-run] [--yes] [--branch ID] Delete one or more tables. Batch: repeat --table-id. --dry-run to preview. Branch-aware. + kbagent storage delete-column --project NAME --table-id ID --column COL [--column ...] [--dry-run] [--yes] [--branch ID] + Delete one or more columns from a table. Batch: repeat --column. --dry-run to preview. Branch-aware. + kbagent storage delete-bucket --project NAME --bucket-id ID [--bucket-id ...] [--force] [--dry-run] [--yes] [--branch ID] Delete one or more buckets. --force cascade-deletes tables. Linked/shared buckets protected. Branch-aware. @@ -273,6 +281,7 @@ Download configs as local files. Idempotent, protects local modifications. --job-limit controls max recent jobs per config (default 5). For large projects, automatically falls back to per-config job fetching to ensure all configs get job history. + Auto-detects renamed configs and renames local directories to match (uses git mv in git repos). kbagent sync status [--directory DIR] Show local changes since last pull (SHA256-based). @@ -312,6 +321,24 @@ --input accepts: inline JSON, @file.json (from file), or - (from stdin). --branch is a CLI flag (NOT a tool input param). Do not pass branch_id in --input. +### Kai -- Keboola AI Assistant (BETA) + + kbagent kai ping [--project NAME] + Check Kai server health and MCP connection status. + Fails with KAI_NOT_ENABLED if the project lacks the 'agent-chat' feature. + + kbagent kai ask --message "question" [--project NAME] + One-shot question to Kai. Collects full response. Use --json for structured output. + Kai has MCP access to project data -- use for Keboola-specific questions + (e.g. "What tables do I have?", "Is it safe to drop bucket X?"). + + kbagent kai chat --message "msg" [--chat-id ID] [--project NAME] + Send message in a chat session. Use --chat-id to continue a conversation. + Without --chat-id starts a new chat. Returns chat_id for continuation. + + kbagent kai history [--project NAME] [--limit N] + List recent Kai chat sessions. Default limit: 10. + ### Utility Commands kbagent init [--from-global] diff --git a/src/keboola_agent_cli/commands/kai.py b/src/keboola_agent_cli/commands/kai.py new file mode 100644 index 00000000..c7a07069 --- /dev/null +++ b/src/keboola_agent_cli/commands/kai.py @@ -0,0 +1,207 @@ +"""CLI commands for Kai (Keboola AI Assistant) integration. + +Bridges Claude Code (local) to Kai (cloud) via kbagent CLI. +Kai has MCP access to project data and can answer Keboola-specific questions. +""" + +import typer + +from ..errors import ConfigError, KeboolaApiError +from ._helpers import ( + check_cli_permission, + emit_hint, + get_formatter, + get_service, + map_error_to_exit_code, + should_hint, +) + +kai_app = typer.Typer(help="(BETA) Keboola AI Assistant (Kai) — ask questions about your project") + + +@kai_app.callback(invoke_without_command=True) +def _kai_permission_check(ctx: typer.Context) -> None: + check_cli_permission(ctx, "kai") + + +@kai_app.command("ping") +def kai_ping( + ctx: typer.Context, + project: str | None = typer.Option( + None, + "--project", + help="Project alias (uses default if omitted).", + ), +) -> None: + """Check Kai server health and MCP connection status.""" + if should_hint(ctx): + emit_hint(ctx, "kai.ping", project=project) + + formatter = get_formatter(ctx) + service = get_service(ctx, "kai_service") + + try: + alias = service.resolve_alias(project) + result = service.ping(alias) + except ConfigError as exc: + formatter.error(message=exc.message, error_code="CONFIG_ERROR") + raise typer.Exit(code=5) from None + except KeboolaApiError as exc: + formatter.error(message=exc.message, error_code=exc.error_code) + raise typer.Exit(code=map_error_to_exit_code(exc)) from None + + def _human(console, data): + console.print(f"[bold green]Kai is alive[/bold green] ({data['project_alias']})") + console.print(f" Timestamp: {data['timestamp']}") + console.print(f" App: {data['app_name']} {data['app_version']}") + console.print(f" Server: {data['server_version']}") + console.print(f" MCP connection: {data['mcp_status']}") + + formatter.output(result, _human) + + +@kai_app.command("ask") +def kai_ask( + ctx: typer.Context, + message: str = typer.Option( + ..., + "--message", + "-m", + help="Question to ask Kai about your project.", + ), + project: str | None = typer.Option( + None, + "--project", + help="Project alias (uses default if omitted).", + ), +) -> None: + """Ask Kai a one-shot question and get the full response. + + Kai has access to your project's data, configurations, and lineage + via MCP tools. Use this for Keboola-specific questions that require + project context. + """ + if should_hint(ctx): + emit_hint(ctx, "kai.ask", project=project, message=message) + + formatter = get_formatter(ctx) + service = get_service(ctx, "kai_service") + + try: + alias = service.resolve_alias(project) + result = service.ask(alias, message) + except ConfigError as exc: + formatter.error(message=exc.message, error_code="CONFIG_ERROR") + raise typer.Exit(code=5) from None + except KeboolaApiError as exc: + formatter.error(message=exc.message, error_code=exc.error_code) + raise typer.Exit(code=map_error_to_exit_code(exc)) from None + + def _human(console, data): + console.print(data["response"]) + + formatter.output(result, _human) + + +@kai_app.command("chat") +def kai_chat( + ctx: typer.Context, + message: str = typer.Option( + ..., + "--message", + "-m", + help="Message to send to Kai.", + ), + chat_id: str | None = typer.Option( + None, + "--chat-id", + help="Continue an existing chat session.", + ), + project: str | None = typer.Option( + None, + "--project", + help="Project alias (uses default if omitted).", + ), +) -> None: + """Send a message to Kai in a chat session. + + Use --chat-id to continue a previous conversation. + Without --chat-id, starts a new chat. + """ + if should_hint(ctx): + emit_hint(ctx, "kai.chat", project=project, message=message, chat_id=chat_id) + + formatter = get_formatter(ctx) + service = get_service(ctx, "kai_service") + + try: + alias = service.resolve_alias(project) + result = service.chat_message(alias, message, chat_id=chat_id) + except ConfigError as exc: + formatter.error(message=exc.message, error_code="CONFIG_ERROR") + raise typer.Exit(code=5) from None + except KeboolaApiError as exc: + formatter.error(message=exc.message, error_code=exc.error_code) + raise typer.Exit(code=map_error_to_exit_code(exc)) from None + + def _human(console, data): + console.print(data["response"]) + console.print(f"\n[dim]Chat ID: {data['chat_id']}[/dim]") + + formatter.output(result, _human) + + +@kai_app.command("history") +def kai_history( + ctx: typer.Context, + project: str | None = typer.Option( + None, + "--project", + help="Project alias (uses default if omitted).", + ), + limit: int = typer.Option( + 10, + "--limit", + "-n", + help="Maximum number of chats to return.", + ), +) -> None: + """List recent Kai chat sessions.""" + if should_hint(ctx): + emit_hint(ctx, "kai.history", project=project, limit=limit) + + formatter = get_formatter(ctx) + service = get_service(ctx, "kai_service") + + try: + alias = service.resolve_alias(project) + result = service.get_history(alias, limit=limit) + except ConfigError as exc: + formatter.error(message=exc.message, error_code="CONFIG_ERROR") + raise typer.Exit(code=5) from None + except KeboolaApiError as exc: + formatter.error(message=exc.message, error_code=exc.error_code) + raise typer.Exit(code=map_error_to_exit_code(exc)) from None + + def _human(console, data): + chats = data["chats"] + if not chats: + console.print("[dim]No chat history.[/dim]") + return + from rich.table import Table + + table = Table(title=f"Kai Chat History ({data['project_alias']})") + table.add_column("Chat ID", style="cyan", no_wrap=True) + table.add_column("Title") + table.add_column("Created", style="dim") + for chat in chats: + table.add_row( + chat["id"][:12] + "...", + chat["title"], + chat["created_at"] or "", + ) + console.print(table) + if data["has_more"]: + console.print("[dim]More chats available. Use --limit to see more.[/dim]") + + formatter.output(result, _human) diff --git a/src/keboola_agent_cli/commands/storage.py b/src/keboola_agent_cli/commands/storage.py index 80a19793..1fb12095 100644 --- a/src/keboola_agent_cli/commands/storage.py +++ b/src/keboola_agent_cli/commands/storage.py @@ -835,6 +835,118 @@ def storage_delete_table( raise typer.Exit(code=1) +@storage_app.command("delete-column", rich_help_panel=_TABLES) +def storage_delete_column( + ctx: typer.Context, + project: str = typer.Option( + ..., + "--project", + help="Project alias", + ), + table_id: str = typer.Option( + ..., + "--table-id", + help="Table ID containing the column(s) (e.g. 'in.c-bucket.table')", + ), + column: list[str] = typer.Option( + ..., + "--column", + help="Column name to delete. Can be repeated.", + ), + dry_run: bool = typer.Option( + False, + "--dry-run", + help="Show what would be deleted without executing", + ), + yes: bool = typer.Option( + False, + "--yes", + "-y", + help="Skip confirmation prompt", + ), + branch: int | None = typer.Option( + None, + "--branch", + help="Dev branch ID (defaults to active branch if set via 'branch use')", + ), +) -> None: + """Delete one or more columns from a storage table. + + Supports batch deletion with multiple --column flags. + """ + if should_hint(ctx): + emit_hint( + ctx, + "storage.delete-column", + project=project, + table_id=table_id, + column=column, + dry_run=dry_run, + branch=branch, + ) + + formatter = get_formatter(ctx) + service = get_service(ctx, "storage_service") + config_store: ConfigStore = ctx.obj["config_store"] + _, effective_branch = resolve_branch(config_store, formatter, project, branch) + + if dry_run: + try: + result = service.delete_columns( + alias=project, + table_id=table_id, + columns=column, + dry_run=True, + branch_id=effective_branch, + ) + except ConfigError as exc: + formatter.error(message=exc.message, error_code="CONFIG_ERROR") + raise typer.Exit(code=5) from None + + if formatter.json_mode: + formatter.output(result) + else: + for col in result.get("would_delete", []): + formatter.console.print( + f"[bold blue]Would delete:[/bold blue] {col} from {table_id}" + ) + return + + if ( + not yes + and not formatter.json_mode + and not typer.confirm( + f"Delete {len(column)} column(s) from table '{table_id}' in project '{project}'?" + ) + ): + formatter.console.print("Aborted.") + raise typer.Exit(code=0) + + try: + result = service.delete_columns( + alias=project, + table_id=table_id, + columns=column, + branch_id=effective_branch, + ) + except ConfigError as exc: + formatter.error(message=exc.message, error_code="CONFIG_ERROR") + raise typer.Exit(code=5) from None + + if formatter.json_mode: + formatter.output(result) + else: + for col in result["deleted"]: + formatter.console.print(f"[bold green]Deleted:[/bold green] {col} from {table_id}") + for f_item in result["failed"]: + formatter.console.print( + f"[bold red]Failed:[/bold red] {f_item['column']}: {f_item['error']}" + ) + + if result["failed"]: + raise typer.Exit(code=1) + + @storage_app.command("delete-bucket", rich_help_panel=_BUCKETS) def storage_delete_bucket( ctx: typer.Context, diff --git a/src/keboola_agent_cli/commands/sync.py b/src/keboola_agent_cli/commands/sync.py index 0f00946d..dc560ec5 100644 --- a/src/keboola_agent_cli/commands/sync.py +++ b/src/keboola_agent_cli/commands/sync.py @@ -134,9 +134,10 @@ def _format_pull_result(formatter: Any, result: dict) -> None: new_cfgs = [d for d in details if d["action"] == "new"] updated_cfgs = [d for d in details if d["action"] == "updated"] removed_cfgs = [d for d in details if d["action"] == "removed"] + renamed_cfgs = [d for d in details if d["action"] == "renamed"] skipped_cfgs = [d for d in details if d["action"] == "skipped"] - has_changes = bool(new_cfgs or updated_cfgs or removed_cfgs) + has_changes = bool(new_cfgs or updated_cfgs or removed_cfgs or renamed_cfgs) storage = result.get("storage", {}) jobs_written = result.get("jobs_written", 0) @@ -168,6 +169,12 @@ def _format_pull_result(formatter: Any, result: dict) -> None: if jobs_written: formatter.console.print(f" Jobs: {jobs_written} configs with job history") + if renamed_cfgs: + formatter.console.print(f" [magenta]Renamed ({len(renamed_cfgs)}):[/magenta]") + for d in renamed_cfgs: + formatter.console.print( + f" > {d.get('old_path', '?')} -> {d['component_id']}/{d['config_name']}" + ) if new_cfgs: formatter.console.print(f" [green]New ({len(new_cfgs)}):[/green]") for d in new_cfgs: @@ -241,6 +248,19 @@ def _format_push_result(formatter: Any, result: dict) -> None: f"{result.get('updated', 0)} updated, " f"{result.get('deleted', 0)} deleted" ) + # Show name drift warnings + drift_warnings = result.get("name_drift_warnings", []) + if drift_warnings: + formatter.console.print( + f"\n [yellow]Warning: {len(drift_warnings)} config(s) have " + f"local directory names that don't match their config name:[/yellow]" + ) + for w in drift_warnings: + formatter.console.print( + f" '{w['local_dirname']}' should be " + f"'{w['expected_dirname']}' (config: {w['config_name']})" + ) + formatter.console.print(" Run 'kbagent config rename' or 'kbagent sync pull' to fix.") def _pull_one_liner(result: dict) -> str: @@ -249,10 +269,13 @@ def _pull_one_liner(result: dict) -> str: new_n = sum(1 for d in details if d["action"] == "new") upd_n = sum(1 for d in details if d["action"] == "updated") rem_n = sum(1 for d in details if d["action"] == "removed") + ren_n = sum(1 for d in details if d["action"] == "renamed") skip_n = sum(1 for d in details if d["action"] == "skipped") - if not new_n and not upd_n and not rem_n and not skip_n: + if not new_n and not upd_n and not rem_n and not ren_n and not skip_n: return "[green]up to date[/green]" parts = [] + if ren_n: + parts.append(f"[magenta]>{ren_n} renamed[/magenta]") if new_n: parts.append(f"[green]+{new_n} new[/green]") if upd_n: diff --git a/src/keboola_agent_cli/constants.py b/src/keboola_agent_cli/constants.py index bd6933f8..7108d381 100644 --- a/src/keboola_agent_cli/constants.py +++ b/src/keboola_agent_cli/constants.py @@ -106,6 +106,11 @@ # --- AI Service --- AI_SERVICE_TIMEOUT: httpx.Timeout = httpx.Timeout(connect=5.0, read=15.0, write=5.0, pool=5.0) + +# --- Kai (Keboola AI Assistant) --- +KAI_FEATURE_FLAG: str = "agent-chat" +KAI_REQUEST_TIMEOUT: float = 300.0 # 5 min for non-streaming requests +KAI_STREAM_TIMEOUT: float = 600.0 # 10 min for SSE streaming responses SECRET_PLACEHOLDER: str = "" # --- Job Run --- diff --git a/src/keboola_agent_cli/hints/definitions/__init__.py b/src/keboola_agent_cli/hints/definitions/__init__.py index 189587d1..ac464e24 100644 --- a/src/keboola_agent_cli/hints/definitions/__init__.py +++ b/src/keboola_agent_cli/hints/definitions/__init__.py @@ -6,6 +6,7 @@ config, # noqa: F401 encrypt, # noqa: F401 job, # noqa: F401 + kai, # noqa: F401 lineage, # noqa: F401 org, # noqa: F401 sharing, # noqa: F401 diff --git a/src/keboola_agent_cli/hints/definitions/config.py b/src/keboola_agent_cli/hints/definitions/config.py index ed1ebe49..b02a713f 100644 --- a/src/keboola_agent_cli/hints/definitions/config.py +++ b/src/keboola_agent_cli/hints/definitions/config.py @@ -1,4 +1,4 @@ -"""Hint definitions for config commands (list, detail, search).""" +"""Hint definitions for config commands (list, detail, search, rename).""" from .. import HintRegistry from ..models import ClientCall, CommandHint, HintStep, ServiceCall @@ -118,3 +118,45 @@ ], ) ) + +# ── config rename ───────────────────────────────────────────────── + +HintRegistry.register( + CommandHint( + cli_command="config.rename", + description="Rename a configuration (update name via API + local sync dir)", + steps=[ + HintStep( + comment="Rename configuration via API", + client=ClientCall( + method="update_config", + args={ + "component_id": "{component_id}", + "config_id": "{config_id}", + "name": "{name}", + "branch_id": "{branch}", + }, + result_var="result", + result_hint="dict", + ), + service=ServiceCall( + service_class="ConfigService", + service_module="config_service", + method="rename_config", + args={ + "alias": "{project}", + "component_id": "{component_id}", + "config_id": "{config_id}", + "name": "{name}", + "branch_id": "{branch}", + }, + ), + ), + ], + notes=[ + "Only the name is updated; configuration content is unchanged.", + "If a local sync directory exists, the folder is renamed and " + "manifest.json is updated automatically.", + ], + ) +) diff --git a/src/keboola_agent_cli/hints/definitions/kai.py b/src/keboola_agent_cli/hints/definitions/kai.py new file mode 100644 index 00000000..72514a7c --- /dev/null +++ b/src/keboola_agent_cli/hints/definitions/kai.py @@ -0,0 +1,141 @@ +"""Hint definitions for Kai (Keboola AI Assistant) commands (ping, ask, chat, history).""" + +from .. import HintRegistry +from ..models import ClientCall, CommandHint, HintStep, ServiceCall + +# ── kai ping ────────────────────────────────────────────────────── + +HintRegistry.register( + CommandHint( + cli_command="kai.ping", + description="Check Kai server health and MCP connection status", + steps=[ + HintStep( + comment="Verify token and check Kai feature flag", + client=ClientCall( + method="verify_token", + args={}, + result_var="token_info", + result_hint="TokenInfo", + ), + service=ServiceCall( + service_class="KaiService", + service_module="kai_service", + method="ping", + args={"alias": "{project}"}, + ), + ), + ], + notes=[ + "Kai commands use KaiClient from the 'kai_client' package, not KeboolaClient.", + "KaiClient.from_storage_api() auto-discovers the Kai API URL from the stack URL.", + "The service checks the 'agent-chat' feature flag before calling Kai.", + "Client hint shows verify_token (feature detection); actual ping uses KaiClient.", + ], + ) +) + +# ── kai ask ─────────────────────────────────────────────────────── + +HintRegistry.register( + CommandHint( + cli_command="kai.ask", + description="Ask Kai a one-shot question about your project", + steps=[ + HintStep( + comment="Send a one-shot question to Kai and collect the full response", + client=ClientCall( + method="verify_token", + args={}, + result_var="token_info", + result_hint="TokenInfo", + ), + service=ServiceCall( + service_class="KaiService", + service_module="kai_service", + method="ask", + args={ + "alias": "{project}", + "message": "{message}", + }, + ), + ), + ], + notes=[ + "Kai commands use KaiClient from the 'kai_client' package, not KeboolaClient.", + "KaiClient.chat(message) sends a question and returns (chat_id, response_text).", + "Service returns {'project_alias', 'chat_id', 'response'}.", + ], + ) +) + +# ── kai chat ────────────────────────────────────────────────────── + +HintRegistry.register( + CommandHint( + cli_command="kai.chat", + description="Send a message in a Kai chat session (new or continued)", + steps=[ + HintStep( + comment="Send a chat message to Kai, optionally continuing an existing session", + client=ClientCall( + method="verify_token", + args={}, + result_var="token_info", + result_hint="TokenInfo", + ), + service=ServiceCall( + service_class="KaiService", + service_module="kai_service", + method="chat_message", + args={ + "alias": "{project}", + "message": "{message}", + "chat_id": "{chat_id}", + }, + ), + ), + ], + notes=[ + "Kai commands use KaiClient from the 'kai_client' package, not KeboolaClient.", + "Without --chat-id, starts a new chat session.", + "With --chat-id, continues an existing conversation.", + "KaiClient.send_message() returns an async stream of events; service collects text events.", + "Service returns {'project_alias', 'chat_id', 'response'}.", + ], + ) +) + +# ── kai history ─────────────────────────────────────────────────── + +HintRegistry.register( + CommandHint( + cli_command="kai.history", + description="List recent Kai chat sessions", + steps=[ + HintStep( + comment="Get chat history for the current user", + client=ClientCall( + method="verify_token", + args={}, + result_var="token_info", + result_hint="TokenInfo", + ), + service=ServiceCall( + service_class="KaiService", + service_module="kai_service", + method="get_history", + args={ + "alias": "{project}", + "limit": "{limit}", + }, + ), + ), + ], + notes=[ + "Kai commands use KaiClient from the 'kai_client' package, not KeboolaClient.", + "Service returns {'project_alias', 'chats': [...], 'has_more': bool}.", + "Each chat has: id, title, created_at, visibility.", + ], + ) +) diff --git a/src/keboola_agent_cli/hints/definitions/storage.py b/src/keboola_agent_cli/hints/definitions/storage.py index 6381295d..7c7be5c0 100644 --- a/src/keboola_agent_cli/hints/definitions/storage.py +++ b/src/keboola_agent_cli/hints/definitions/storage.py @@ -350,6 +350,45 @@ ) ) +# ── storage delete-column ───────────────────────────────────────── + +HintRegistry.register( + CommandHint( + cli_command="storage.delete-column", + description="Delete one or more columns from a table", + steps=[ + HintStep( + comment="Delete column(s)", + client=ClientCall( + method="delete_column", + args={ + "table_id": "{table_id}", + "column_name": "{column}", + "branch_id": "{branch}", + }, + result_var=None, + ), + service=ServiceCall( + service_class="StorageService", + service_module="storage_service", + method="delete_columns", + args={ + "alias": "{project}", + "table_id": "{table_id}", + "columns": "{column}", + "dry_run": "{dry_run}", + "branch_id": "{branch}", + }, + ), + ), + ], + notes=[ + "Client layer deletes one column at a time. Loop for batch.", + "Synchronous API — no async job polling needed.", + ], + ) +) + # ── storage files ──────��──────────────────────────────────���──────── HintRegistry.register( diff --git a/src/keboola_agent_cli/models.py b/src/keboola_agent_cli/models.py index 90fac7c6..76475e31 100644 --- a/src/keboola_agent_cli/models.py +++ b/src/keboola_agent_cli/models.py @@ -100,6 +100,10 @@ class TokenVerifyResponse(BaseModel): default="snowflake", description="Project default backend (snowflake, bigquery, etc.)", ) + features: list[str] = Field( + default_factory=list, + description="Project feature flags (e.g. agent-chat, storage-types)", + ) class ComponentDetail(BaseModel): diff --git a/src/keboola_agent_cli/permissions.py b/src/keboola_agent_cli/permissions.py index f4a908a6..74517335 100644 --- a/src/keboola_agent_cli/permissions.py +++ b/src/keboola_agent_cli/permissions.py @@ -25,6 +25,7 @@ "config.detail": "read", "config.search": "read", "config.update": "write", + "config.rename": "write", "config.delete": "destructive", "config.new": "write", # Job history @@ -60,6 +61,11 @@ # MCP tools "tool.list": "read", "tool.call": "write", + # Kai (Keboola AI Assistant) + "kai.ping": "read", + "kai.ask": "read", + "kai.chat": "write", + "kai.history": "read", # Component discovery "component.list": "read", "component.detail": "read", @@ -83,6 +89,7 @@ "storage.unload-table": "read", # Storage destructive "storage.delete-table": "destructive", + "storage.delete-column": "destructive", "storage.delete-bucket": "destructive", "storage.file-delete": "destructive", # Encryption diff --git a/src/keboola_agent_cli/services/config_service.py b/src/keboola_agent_cli/services/config_service.py index 2fbd12fd..6ab01835 100644 --- a/src/keboola_agent_cli/services/config_service.py +++ b/src/keboola_agent_cli/services/config_service.py @@ -5,14 +5,22 @@ """ import json +import logging import re +import shutil +import subprocess +from pathlib import Path from typing import Any from ..errors import KeboolaApiError from ..json_utils import compute_diff, deep_merge, set_nested_value from ..models import ProjectConfig +from ..sync.manifest import Manifest, load_manifest, save_manifest +from ..sync.naming import sanitize_name from .base import BaseService +logger = logging.getLogger(__name__) + def _find_matches_in_json( obj: Any, @@ -437,6 +445,203 @@ def delete_config( "branch_id": effective_branch_id, } + def rename_config( + self, + alias: str, + component_id: str, + config_id: str, + name: str, + branch_id: int | None = None, + directory: Path | None = None, + ) -> dict[str, Any]: + """Rename a configuration (update name via API + rename local sync dir). + + Args: + alias: Project alias. + component_id: The component ID. + config_id: The configuration ID to rename. + name: The new configuration name. + branch_id: If set, rename in a specific dev branch. + If None, uses the project's active branch (if any). + directory: Optional sync working directory. If a manifest exists + here and tracks this config, the local directory is + renamed and the manifest path is updated. + + Returns: + Dict with old name, new name, and optional sync rename details. + + Raises: + ConfigError: If the alias is not found. + KeboolaApiError: If the API call fails. + """ + projects = self.resolve_projects([alias]) + project = projects[alias] + effective_branch_id = branch_id or project.active_branch_id + + client = self._client_factory(project.stack_url, project.token) + try: + # Fetch current state to get old name + current = client.get_config_detail( + component_id, config_id, branch_id=effective_branch_id + ) + old_name = current.get("name", "") + + # Update name via API + client.update_config( + component_id=component_id, + config_id=config_id, + name=name, + change_description=f"Renamed via kbagent config rename: {old_name} -> {name}", + branch_id=effective_branch_id, + ) + finally: + client.close() + + result: dict[str, Any] = { + "status": "renamed", + "project_alias": alias, + "component_id": component_id, + "config_id": config_id, + "old_name": old_name, + "new_name": name, + "branch_id": effective_branch_id, + } + + # Attempt local sync directory rename if applicable + sync_result = self._rename_sync_directory( + directory=directory, + component_id=component_id, + config_id=config_id, + new_name=name, + ) + if sync_result: + result["sync"] = sync_result + + return result + + def _rename_sync_directory( + self, + directory: Path | None, + component_id: str, + config_id: str, + new_name: str, + ) -> dict[str, str] | None: + """Rename the local sync directory for a config if a manifest tracks it. + + Returns a dict with old_path/new_path on success, or None if no + sync directory was found or rename was not needed. + """ + if directory is None: + return None + + from ..constants import KEBOOLA_DIR_NAME, MANIFEST_FILENAME + + manifest_path = directory / KEBOOLA_DIR_NAME / MANIFEST_FILENAME + if not manifest_path.exists(): + return None + + try: + manifest = load_manifest(directory) + except (FileNotFoundError, ValueError): + return None + + # Find the config entry in the manifest + target_cfg = None + for cfg in manifest.configurations: + if cfg.component_id == component_id and cfg.id == config_id: + target_cfg = cfg + break + + if target_cfg is None: + return None + + # Compute new path using the naming template + old_path = target_cfg.path + old_basename = old_path.rsplit("/", 1)[-1] if "/" in old_path else old_path + new_basename = sanitize_name(new_name) + + if old_basename == new_basename: + return None # No rename needed + + # Build new path: replace only the last segment (config name) + if "/" in old_path: + parent = old_path.rsplit("/", 1)[0] + new_path = f"{parent}/{new_basename}" + else: + new_path = new_basename + + # Collision detection: if target already exists, append numeric suffix + branch_dir = self._find_sync_branch_dir(manifest, directory) + if branch_dir is None: + return None + + target_dir = branch_dir / new_path + if target_dir.exists(): + counter = 2 + while (branch_dir / f"{new_path}-{counter}").exists(): + counter += 1 + new_path = f"{new_path}-{counter}" + target_dir = branch_dir / new_path + + # Perform the rename + source_dir = branch_dir / old_path + if not source_dir.exists(): + # Directory doesn't exist locally, just update manifest + target_cfg.path = new_path + target_cfg.metadata.pop("pull_hash", None) + target_cfg.metadata.pop("pull_config_hash", None) + save_manifest(directory, manifest) + return {"old_path": old_path, "new_path": new_path, "method": "manifest_only"} + + # Try git mv first for cleaner history, fall back to shutil.move + method = self._move_directory(source_dir, target_dir) + + # Update manifest + target_cfg.path = new_path + target_cfg.metadata.pop("pull_hash", None) + target_cfg.metadata.pop("pull_config_hash", None) + save_manifest(directory, manifest) + + # Clean up empty parent directories + parent_dir = source_dir.parent + while parent_dir != branch_dir and parent_dir.exists(): + if not any(parent_dir.iterdir()): + parent_dir.rmdir() + parent_dir = parent_dir.parent + else: + break + + return {"old_path": old_path, "new_path": new_path, "method": method} + + @staticmethod + def _move_directory(source: Path, target: Path) -> str: + """Move a directory, using git mv if in a git repo, else shutil.move.""" + target.parent.mkdir(parents=True, exist_ok=True) + try: + result = subprocess.run( + ["git", "mv", str(source), str(target)], + cwd=source.parent, + capture_output=True, + text=True, + check=False, + ) + if result.returncode == 0: + return "git_mv" + except FileNotFoundError: + pass # git not installed + shutil.move(str(source), str(target)) + return "shutil_move" + + @staticmethod + def _find_sync_branch_dir(manifest: Manifest, project_root: Path) -> Path | None: + """Find the branch directory within a sync project root.""" + if not manifest.branches: + return None + # Use the first branch (typically "main") + branch_path = manifest.branches[0].path + branch_dir = project_root / branch_path + return branch_dir if branch_dir.exists() else None + def search_configs( self, query: str, diff --git a/src/keboola_agent_cli/services/kai_service.py b/src/keboola_agent_cli/services/kai_service.py new file mode 100644 index 00000000..15f4b6c7 --- /dev/null +++ b/src/keboola_agent_cli/services/kai_service.py @@ -0,0 +1,242 @@ +"""Kai (Keboola AI Assistant) service — bridge between CLI and cloud Kai API. + +Provides sync wrappers around the async kai-client library, with feature +detection (agent-chat flag) and project resolution via BaseService. +""" + +import asyncio +import logging +from typing import Any + +from kai_client import KaiClient, KaiError + +from ..constants import KAI_FEATURE_FLAG, KAI_REQUEST_TIMEOUT, KAI_STREAM_TIMEOUT +from ..errors import ConfigError, KeboolaApiError +from .base import BaseService + +logger = logging.getLogger(__name__) + + +class KaiService(BaseService): + """Business logic for Kai AI Assistant integration. + + All public methods are synchronous — they wrap the async KaiClient + via asyncio.run() so Typer commands can call them directly. + """ + + # ------------------------------------------------------------------ + # Project resolution + # ------------------------------------------------------------------ + + def resolve_alias(self, alias: str | None) -> str: + """Resolve a project alias, falling back to the default project. + + Args: + alias: Explicit alias, or None for default. + + Returns: + Resolved alias string. + + Raises: + ConfigError: If no projects configured or alias not found. + """ + if alias: + # Validate it exists + self.resolve_projects([alias]) + return alias + # Fall back to default (first project) + projects = self.resolve_projects() + if not projects: + raise ConfigError("No projects configured. Run 'kbagent project add' first.") + return next(iter(projects)) + + # ------------------------------------------------------------------ + # Feature detection + # ------------------------------------------------------------------ + + def _check_kai_enabled(self, alias: str) -> None: + """Raise KeboolaApiError if Kai is not enabled for the project. + + Calls verify_token to check owner.features for the agent-chat flag. + """ + projects = self.resolve_projects([alias]) + project = projects[alias] + client = self._client_factory(project.stack_url, project.token) + try: + token_info = client.verify_token() + finally: + client.close() + + if KAI_FEATURE_FLAG not in token_info.features: + raise KeboolaApiError( + message=( + f"Kai is not enabled for project '{alias}'. " + "Enable the 'AI Agent Chat' feature in project settings." + ), + status_code=0, + error_code="KAI_NOT_ENABLED", + ) + + # ------------------------------------------------------------------ + # Async helpers + # ------------------------------------------------------------------ + + async def _create_kai_client(self, alias: str) -> KaiClient: + """Create a KaiClient with auto-discovered URL for the given project.""" + projects = self.resolve_projects([alias]) + project = projects[alias] + return await KaiClient.from_storage_api( + storage_api_token=project.token, + storage_api_url=project.stack_url, + timeout=KAI_REQUEST_TIMEOUT, + stream_timeout=KAI_STREAM_TIMEOUT, + ) + + # ------------------------------------------------------------------ + # Public methods (sync wrappers) + # ------------------------------------------------------------------ + + def ping(self, alias: str) -> dict[str, Any]: + """Check Kai server health for a project. + + Returns: + Dict with timestamp and server info. + """ + self._check_kai_enabled(alias) + + async def _ping() -> dict[str, Any]: + client = await self._create_kai_client(alias) + async with client: + ping_resp = await client.ping() + info_resp = await client.info() + + return { + "project_alias": alias, + "timestamp": ping_resp.timestamp.isoformat(), + "app_name": info_resp.app_name, + "app_version": info_resp.app_version, + "server_version": info_resp.server_version, + "mcp_status": ( + info_resp.connected_mcp.get("status", "unknown") + if isinstance(info_resp.connected_mcp, dict) + else "unknown" + ), + } + + try: + return asyncio.run(_ping()) + except KaiError as exc: + raise KeboolaApiError( + message=f"Kai ping failed: {exc.message}", + status_code=0, + error_code="KAI_ERROR", + ) from exc + + def ask(self, alias: str, message: str) -> dict[str, Any]: + """Send a one-shot question to Kai and collect the full text response. + + Args: + alias: Project alias. + message: The question to ask. + + Returns: + Dict with chat_id and response text. + """ + self._check_kai_enabled(alias) + + async def _ask() -> dict[str, Any]: + client = await self._create_kai_client(alias) + async with client: + chat_id, response_text = await client.chat(message) + + return { + "project_alias": alias, + "chat_id": chat_id, + "response": response_text, + } + + try: + return asyncio.run(_ask()) + except KaiError as exc: + raise KeboolaApiError( + message=f"Kai ask failed: {exc.message}", + status_code=0, + error_code="KAI_ERROR", + ) from exc + + def chat_message(self, alias: str, message: str, chat_id: str | None = None) -> dict[str, Any]: + """Send a message in a chat session and collect the response. + + Args: + alias: Project alias. + message: The message to send. + chat_id: Optional existing chat ID to continue. + + Returns: + Dict with chat_id and response text. + """ + self._check_kai_enabled(alias) + + async def _chat() -> dict[str, Any]: + client = await self._create_kai_client(alias) + async with client: + cid = chat_id or client.new_chat_id() + response_parts: list[str] = [] + async for event in client.send_message(cid, message): + if event.type == "text": + response_parts.append(event.text) # type: ignore[attr-defined] + + return { + "project_alias": alias, + "chat_id": cid, + "response": "".join(response_parts), + } + + try: + return asyncio.run(_chat()) + except KaiError as exc: + raise KeboolaApiError( + message=f"Kai chat failed: {exc.message}", + status_code=0, + error_code="KAI_ERROR", + ) from exc + + def get_history(self, alias: str, limit: int = 10) -> dict[str, Any]: + """Get chat history for the current user. + + Args: + alias: Project alias. + limit: Max number of chats to return. + + Returns: + Dict with list of chat summaries. + """ + self._check_kai_enabled(alias) + + async def _history() -> dict[str, Any]: + client = await self._create_kai_client(alias) + async with client: + history = await client.get_history(limit=limit) + + return { + "project_alias": alias, + "chats": [ + { + "id": chat.id, + "title": chat.title or "(untitled)", + "created_at": chat.created_at.isoformat() if chat.created_at else None, + "visibility": chat.visibility, + } + for chat in history.chats + ], + "has_more": history.has_more, + } + + try: + return asyncio.run(_history()) + except KaiError as exc: + raise KeboolaApiError( + message=f"Kai history failed: {exc.message}", + status_code=0, + error_code="KAI_ERROR", + ) from exc diff --git a/src/keboola_agent_cli/services/storage_service.py b/src/keboola_agent_cli/services/storage_service.py index babecb26..27c91787 100644 --- a/src/keboola_agent_cli/services/storage_service.py +++ b/src/keboola_agent_cli/services/storage_service.py @@ -600,8 +600,8 @@ def download_table( retryable=False, ) - # Step 3: Get download URL - file_detail = client.get_file_info(file_id) + # Step 3: Get download URL (branch-scoped if exporting from dev branch) + file_detail = client.get_file_info(file_id, branch_id=branch_id) download_url = file_detail.get("url") if not download_url: raise KeboolaApiError( @@ -695,6 +695,67 @@ def delete_tables( "project_alias": alias, } + def delete_columns( + self, + alias: str, + table_id: str, + columns: list[str], + dry_run: bool = False, + branch_id: int | None = None, + ) -> dict[str, Any]: + """Delete one or more columns from a storage table. + + Batch-tolerant: accumulates errors per column, one failure does not + stop other deletes. + + Args: + alias: Project alias. + table_id: Full table ID (e.g. "in.c-bucket.table"). + columns: List of column names to delete. + dry_run: If True, only report what would be deleted. + branch_id: If set, target a specific dev branch. + + Returns: + Dict with 'deleted', 'failed', 'dry_run', 'project_alias', + 'table_id', and optionally 'would_delete'. + """ + from ..errors import KeboolaApiError + + projects = self.resolve_projects([alias]) + project = projects[alias] + + if dry_run: + return { + "deleted": [], + "failed": [], + "would_delete": list(columns), + "dry_run": True, + "project_alias": alias, + "table_id": table_id, + } + + deleted: list[str] = [] + failed: list[dict[str, str]] = [] + + client = self._client_factory(project.stack_url, project.token) + try: + for col in columns: + try: + client.delete_column(table_id, col, branch_id=branch_id) + deleted.append(col) + except KeboolaApiError as exc: + failed.append({"column": col, "error": exc.message}) + finally: + client.close() + + return { + "deleted": deleted, + "failed": failed, + "dry_run": False, + "project_alias": alias, + "table_id": table_id, + } + def delete_buckets( self, alias: str, @@ -1193,12 +1254,12 @@ def unload_table_to_file( retryable=False, ) - # Step 3: Tag the exported file + # Step 3: Tag the exported file (branch-scoped if on dev branch) for tag in tags or []: - client.tag_file(file_id, tag) + client.tag_file(file_id, tag, branch_id=branch_id) - # Step 4: Get full file detail - file_detail = client.get_file_info(file_id) + # Step 4: Get full file detail (branch-scoped if on dev branch) + file_detail = client.get_file_info(file_id, branch_id=branch_id) result: dict[str, Any] = { "project_alias": alias, diff --git a/src/keboola_agent_cli/services/sync_service.py b/src/keboola_agent_cli/services/sync_service.py index 83951e25..b5725602 100644 --- a/src/keboola_agent_cli/services/sync_service.py +++ b/src/keboola_agent_cli/services/sync_service.py @@ -9,6 +9,7 @@ import json import logging import shutil +import subprocess import threading from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path @@ -318,6 +319,40 @@ def pull( is_new = lookup_key not in existing_keys if lookup_key in existing_paths: rel_path = existing_paths[lookup_key] + + # Auto-rename: detect when remote name changed + expected_path = config_path( + manifest.naming.config, + component_type, + component_id, + config_name, + ) + if rel_path != expected_path and not dry_run: + rename_target = expected_path + # Collision: if target already used, add suffix + if rename_target in used_paths: + suffix = config_id[:8] if len(config_id) > 8 else config_id + rename_target = f"{rename_target}-{suffix}" + + old_dir = branch_dir / rel_path + new_dir = branch_dir / rename_target + if old_dir.exists() and not new_dir.exists(): + self._rename_directory(old_dir, new_dir) + pull_details.append( + { + "action": "renamed", + "component_id": component_id, + "config_name": config_name, + "path": rename_target, + "old_path": rel_path, + } + ) + rel_path = rename_target + logger.info( + "Renamed config dir: %s -> %s", + rel_path, + rename_target, + ) else: # Generate new filesystem path with collision detection rel_path = config_path( @@ -889,6 +924,9 @@ def push( branch_id = self._resolve_branch_id(project, manifest, project_root) + # Detect name drift: local dir name doesn't match config name + name_drift_warnings = self._detect_name_drift(manifest, project_root) + client = self._client_factory(project.stack_url, project.token) created = 0 updated = 0 @@ -1027,7 +1065,7 @@ def push( if manifest_dirty: save_manifest(project_root, manifest) - return { + result_data: dict[str, Any] = { "status": "pushed", "created": created, "updated": updated, @@ -1035,6 +1073,9 @@ def push( "errors": errors, "pushed_details": pushed_details, } + if name_drift_warnings: + result_data["name_drift_warnings"] = name_drift_warnings + return result_data @staticmethod def _encrypt_secrets_in_config( @@ -2047,6 +2088,65 @@ def _file_hash(self, file_path: Path) -> str: content = file_path.read_bytes() return hashlib.sha256(content).hexdigest() + def _detect_name_drift(self, manifest: Manifest, project_root: Path) -> list[dict[str, str]]: + """Detect configs where local dir name doesn't match the config name. + + Reads each tracked config's _config.yml to get the current name, + then compares sanitize_name(name) against the directory basename. + + Returns a list of warning dicts with component_id, config_id, + local_dirname, and expected_dirname. + """ + warnings: list[dict[str, str]] = [] + for cfg in manifest.configurations: + path = cfg.path + dirname = path.rsplit("/", 1)[-1] if "/" in path else path + + # Find branch dir and read _config.yml + branch_path = self._find_branch_path(manifest, cfg.branch_id) + config_dir = project_root / branch_path / path + local_data = self._read_config_file(config_dir) + if local_data is None: + continue + + config_name = local_data.get("name", "") + if not config_name: + continue + + expected_dirname = sanitize_name(config_name) + if dirname != expected_dirname: + warnings.append( + { + "component_id": cfg.component_id, + "config_id": cfg.id, + "local_dirname": dirname, + "expected_dirname": expected_dirname, + "config_name": config_name, + } + ) + return warnings + + def _rename_directory(self, source: Path, target: Path) -> str: + """Rename a directory, using git mv if in a git repo, else shutil.move. + + Returns 'git_mv' or 'shutil_move' indicating which method was used. + """ + target.parent.mkdir(parents=True, exist_ok=True) + try: + result = subprocess.run( + ["git", "mv", str(source), str(target)], + cwd=source.parent, + capture_output=True, + text=True, + check=False, + ) + if result.returncode == 0: + return "git_mv" + except FileNotFoundError: + pass # git not installed + shutil.move(str(source), str(target)) + return "shutil_move" + def _read_config_file(self, config_dir: Path) -> dict[str, Any] | None: """Read and parse a ``_config.yml`` file, returning None if missing.""" config_file = config_dir / CONFIG_FILENAME diff --git a/tests/test_config_rename.py b/tests/test_config_rename.py new file mode 100644 index 00000000..702d9cc8 --- /dev/null +++ b/tests/test_config_rename.py @@ -0,0 +1,325 @@ +"""Tests for config rename feature (API rename + local sync directory rename). + +Covers ConfigService.rename_config and ConfigService._rename_sync_directory. +""" + +import json +from pathlib import Path +from unittest.mock import MagicMock + +import pytest + +from helpers import setup_single_project +from keboola_agent_cli.errors import KeboolaApiError +from keboola_agent_cli.services.config_service import ConfigService + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +SAMPLE_CONFIG_DETAIL = { + "id": "cfg-001", + "name": "Old Name", + "description": "A test configuration", + "configuration": {"parameters": {"key": "value"}}, +} + + +def _make_service( + tmp_config_dir: Path, +) -> tuple[ConfigService, MagicMock]: + """Create a ConfigService with a mock client for rename tests.""" + store = setup_single_project(tmp_config_dir) + mock_client = MagicMock() + mock_client.get_config_detail.return_value = SAMPLE_CONFIG_DETAIL + mock_client.update_config.return_value = { + "id": "cfg-001", + "name": "New Name", + "componentId": "keboola.ex-http", + } + service = ConfigService( + config_store=store, + client_factory=lambda url, token: mock_client, + ) + return service, mock_client + + +def _create_sync_directory(tmp_path: Path) -> Path: + """Create a realistic Keboola CLI sync directory structure. + + Returns the project root directory (parent of .keboola/). + """ + keboola_dir = tmp_path / ".keboola" + keboola_dir.mkdir(parents=True) + + manifest_data = { + "version": 2, + "project": {"id": 258, "apiHost": "connection.keboola.com"}, + "allowTargetEnv": True, + "gitBranching": {"enabled": False, "defaultBranch": "main"}, + "sortBy": "id", + "naming": { + "branch": "{branch_name}", + "config": "{component_type}/{component_id}/{config_name}", + "configRow": "rows/{config_row_name}", + }, + "allowedBranches": [], + "ignoredComponents": [], + "branches": [{"id": 12345, "path": "main", "metadata": {}}], + "configurations": [ + { + "branchId": 12345, + "componentId": "keboola.ex-http", + "id": "cfg-001", + "path": "extractor/keboola.ex-http/old-name", + "metadata": {"pull_hash": "abc", "pull_config_hash": "def"}, + "rows": [], + } + ], + } + (keboola_dir / "manifest.json").write_text(json.dumps(manifest_data)) + + # Create the actual config directory with a file inside + config_dir = tmp_path / "main" / "extractor" / "keboola.ex-http" / "old-name" + config_dir.mkdir(parents=True) + (config_dir / "_config.yml").write_text("name: Old Name\n") + + return tmp_path + + +# --------------------------------------------------------------------------- +# rename_config (API-level) tests +# --------------------------------------------------------------------------- + + +class TestRenameConfigApi: + """Tests for ConfigService.rename_config API interaction.""" + + def test_rename_config_basic(self, tmp_config_dir: Path) -> None: + """Rename via API returns old_name and new_name in result.""" + service, client = _make_service(tmp_config_dir) + + result = service.rename_config( + alias="prod", + component_id="keboola.ex-http", + config_id="cfg-001", + name="New Name", + ) + + assert result["status"] == "renamed" + assert result["old_name"] == "Old Name" + assert result["new_name"] == "New Name" + assert result["project_alias"] == "prod" + assert result["component_id"] == "keboola.ex-http" + assert result["config_id"] == "cfg-001" + + client.get_config_detail.assert_called_once_with( + "keboola.ex-http", "cfg-001", branch_id=None + ) + client.update_config.assert_called_once() + call_kwargs = client.update_config.call_args.kwargs + assert call_kwargs["name"] == "New Name" + assert call_kwargs["component_id"] == "keboola.ex-http" + assert call_kwargs["config_id"] == "cfg-001" + + def test_rename_config_with_branch(self, tmp_config_dir: Path) -> None: + """Rename with branch_id passes it through to API calls.""" + service, client = _make_service(tmp_config_dir) + + result = service.rename_config( + alias="prod", + component_id="keboola.ex-http", + config_id="cfg-001", + name="New Name", + branch_id=9999, + ) + + assert result["branch_id"] == 9999 + + client.get_config_detail.assert_called_once_with( + "keboola.ex-http", "cfg-001", branch_id=9999 + ) + call_kwargs = client.update_config.call_args.kwargs + assert call_kwargs["branch_id"] == 9999 + + def test_rename_config_api_error(self, tmp_config_dir: Path) -> None: + """KeboolaApiError from the client propagates to the caller.""" + service, client = _make_service(tmp_config_dir) + client.get_config_detail.side_effect = KeboolaApiError( + status_code=404, + message="Configuration not found", + ) + + with pytest.raises(KeboolaApiError, match="Configuration not found"): + service.rename_config( + alias="prod", + component_id="keboola.ex-http", + config_id="cfg-001", + name="New Name", + ) + + +# --------------------------------------------------------------------------- +# _rename_sync_directory tests +# --------------------------------------------------------------------------- + + +class TestRenameSyncDirectory: + """Tests for ConfigService._rename_sync_directory.""" + + def test_rename_sync_directory_no_directory(self, tmp_config_dir: Path) -> None: + """Returns None when directory is None.""" + service, _ = _make_service(tmp_config_dir) + + result = service._rename_sync_directory( + directory=None, + component_id="keboola.ex-http", + config_id="cfg-001", + new_name="New Name", + ) + + assert result is None + + def test_rename_sync_directory_no_manifest(self, tmp_config_dir: Path, tmp_path: Path) -> None: + """Returns None when no manifest.json exists in the directory.""" + service, _ = _make_service(tmp_config_dir) + empty_dir = tmp_path / "no-manifest" + empty_dir.mkdir() + + result = service._rename_sync_directory( + directory=empty_dir, + component_id="keboola.ex-http", + config_id="cfg-001", + new_name="New Name", + ) + + assert result is None + + def test_rename_sync_directory_config_not_tracked( + self, tmp_config_dir: Path, tmp_path: Path + ) -> None: + """Returns None when the config is not tracked in the manifest.""" + service, _ = _make_service(tmp_config_dir) + project_root = _create_sync_directory(tmp_path / "sync") + + result = service._rename_sync_directory( + directory=project_root, + component_id="keboola.ex-http", + config_id="cfg-999", # Not in manifest + new_name="New Name", + ) + + assert result is None + + def test_rename_sync_directory_no_change_needed( + self, tmp_config_dir: Path, tmp_path: Path + ) -> None: + """Returns None when old name matches new name after sanitization.""" + service, _ = _make_service(tmp_config_dir) + project_root = _create_sync_directory(tmp_path / "sync") + + # "old-name" sanitized stays "old-name" + result = service._rename_sync_directory( + directory=project_root, + component_id="keboola.ex-http", + config_id="cfg-001", + new_name="Old Name", # sanitize_name("Old Name") == "old-name" + ) + + assert result is None + + def test_rename_sync_directory_renames_dir(self, tmp_config_dir: Path, tmp_path: Path) -> None: + """Full rename: moves files, updates manifest path, clears hashes.""" + service, _ = _make_service(tmp_config_dir) + project_root = _create_sync_directory(tmp_path / "sync") + + result = service._rename_sync_directory( + directory=project_root, + component_id="keboola.ex-http", + config_id="cfg-001", + new_name="New Name", + ) + + assert result is not None + assert result["old_path"] == "extractor/keboola.ex-http/old-name" + assert result["new_path"] == "extractor/keboola.ex-http/new-name" + assert result["method"] in ("git_mv", "shutil_move") + + # Verify old directory is gone and new one exists + old_dir = project_root / "main" / "extractor" / "keboola.ex-http" / "old-name" + new_dir = project_root / "main" / "extractor" / "keboola.ex-http" / "new-name" + assert not old_dir.exists() + assert new_dir.exists() + assert (new_dir / "_config.yml").read_text() == "name: Old Name\n" + + # Verify manifest was updated + manifest_path = project_root / ".keboola" / "manifest.json" + manifest_data = json.loads(manifest_path.read_text()) + cfg_entry = manifest_data["configurations"][0] + assert cfg_entry["path"] == "extractor/keboola.ex-http/new-name" + # Pull hashes should be cleared + assert "pull_hash" not in cfg_entry.get("metadata", {}) + assert "pull_config_hash" not in cfg_entry.get("metadata", {}) + + def test_rename_sync_directory_collision(self, tmp_config_dir: Path, tmp_path: Path) -> None: + """When target dir already exists, appends a numeric suffix.""" + service, _ = _make_service(tmp_config_dir) + project_root = _create_sync_directory(tmp_path / "sync") + + # Pre-create the target directory to cause a collision + collision_dir = project_root / "main" / "extractor" / "keboola.ex-http" / "new-name" + collision_dir.mkdir(parents=True) + (collision_dir / "_config.yml").write_text("name: Existing\n") + + result = service._rename_sync_directory( + directory=project_root, + component_id="keboola.ex-http", + config_id="cfg-001", + new_name="New Name", + ) + + assert result is not None + # Should get a numeric suffix due to collision + assert result["new_path"] == "extractor/keboola.ex-http/new-name-2" + + # Original collision dir is untouched + assert collision_dir.exists() + assert (collision_dir / "_config.yml").read_text() == "name: Existing\n" + + # New suffixed dir exists with moved content + suffixed_dir = project_root / "main" / "extractor" / "keboola.ex-http" / "new-name-2" + assert suffixed_dir.exists() + assert (suffixed_dir / "_config.yml").read_text() == "name: Old Name\n" + + def test_rename_sync_directory_source_missing( + self, tmp_config_dir: Path, tmp_path: Path + ) -> None: + """When source dir doesn't exist, updates manifest only.""" + service, _ = _make_service(tmp_config_dir) + project_root = _create_sync_directory(tmp_path / "sync") + + # Remove the source directory (simulates not-yet-pulled state) + source_dir = project_root / "main" / "extractor" / "keboola.ex-http" / "old-name" + import shutil + + shutil.rmtree(source_dir) + + result = service._rename_sync_directory( + directory=project_root, + component_id="keboola.ex-http", + config_id="cfg-001", + new_name="New Name", + ) + + assert result is not None + assert result["old_path"] == "extractor/keboola.ex-http/old-name" + assert result["new_path"] == "extractor/keboola.ex-http/new-name" + assert result["method"] == "manifest_only" + + # Verify manifest was updated + manifest_path = project_root / ".keboola" / "manifest.json" + manifest_data = json.loads(manifest_path.read_text()) + cfg_entry = manifest_data["configurations"][0] + assert cfg_entry["path"] == "extractor/keboola.ex-http/new-name" + assert "pull_hash" not in cfg_entry.get("metadata", {}) + assert "pull_config_hash" not in cfg_entry.get("metadata", {}) diff --git a/tests/test_config_rename_cli.py b/tests/test_config_rename_cli.py new file mode 100644 index 00000000..c13b5360 --- /dev/null +++ b/tests/test_config_rename_cli.py @@ -0,0 +1,277 @@ +"""Tests for config rename CLI command via CliRunner. + +Tests the `kbagent config rename` subcommand: JSON output, human-readable +output, sync directory info, API error handling, and help text. +""" + +import json +from pathlib import Path +from unittest.mock import MagicMock + +import pytest +from typer.testing import CliRunner + +from helpers import setup_single_project +from keboola_agent_cli.cli import app +from keboola_agent_cli.errors import KeboolaApiError +from keboola_agent_cli.services.config_service import ConfigService + +runner = CliRunner() + + +class TestConfigRenameCli: + """Tests for `kbagent config rename` command.""" + + def test_config_rename_json_output(self, tmp_config_dir: Path) -> None: + """config rename --json returns structured JSON with rename details.""" + store = setup_single_project(tmp_config_dir) + + mock_client = MagicMock() + mock_client.get_config_detail.return_value = { + "id": "cfg-001", + "name": "Old Name", + "componentId": "keboola.ex-http", + } + mock_client.update_config.return_value = { + "id": "cfg-001", + "name": "New Name", + "componentId": "keboola.ex-http", + } + + with pytest.MonkeyPatch.context() as mp: + mp.setattr( + "keboola_agent_cli.commands.config.get_service", + lambda ctx, name: ConfigService( + config_store=store, + client_factory=lambda url, token: mock_client, + ), + ) + + result = runner.invoke( + app, + [ + "--json", + "--config-dir", + str(tmp_config_dir), + "config", + "rename", + "--project", + "prod", + "--component-id", + "keboola.ex-http", + "--config-id", + "cfg-001", + "--name", + "New Name", + ], + ) + + assert result.exit_code == 0, f"Exit code {result.exit_code}: {result.output}" + output = json.loads(result.output) + assert output["status"] == "ok" + assert output["data"]["status"] == "renamed" + assert output["data"]["old_name"] == "Old Name" + assert output["data"]["new_name"] == "New Name" + assert output["data"]["component_id"] == "keboola.ex-http" + assert output["data"]["config_id"] == "cfg-001" + assert output["data"]["project_alias"] == "prod" + + def test_config_rename_human_output(self, tmp_config_dir: Path) -> None: + """config rename in human mode outputs success message with rename info.""" + store = setup_single_project(tmp_config_dir) + + mock_client = MagicMock() + mock_client.get_config_detail.return_value = { + "id": "cfg-001", + "name": "Old Name", + "componentId": "keboola.ex-http", + } + mock_client.update_config.return_value = { + "id": "cfg-001", + "name": "New Name", + "componentId": "keboola.ex-http", + } + + with pytest.MonkeyPatch.context() as mp: + mp.setattr( + "keboola_agent_cli.commands.config.get_service", + lambda ctx, name: ConfigService( + config_store=store, + client_factory=lambda url, token: mock_client, + ), + ) + + result = runner.invoke( + app, + [ + "--config-dir", + str(tmp_config_dir), + "config", + "rename", + "--project", + "prod", + "--component-id", + "keboola.ex-http", + "--config-id", + "cfg-001", + "--name", + "New Name", + ], + ) + + assert result.exit_code == 0, f"Exit code {result.exit_code}: {result.output}" + assert "Renamed" in result.output + assert "Old Name" in result.output + assert "New Name" in result.output + assert "keboola.ex-http" in result.output + + def test_config_rename_with_sync_info(self, tmp_config_dir: Path, tmp_path: Path) -> None: + """config rename with sync directory shows sync rename details in human output.""" + store = setup_single_project(tmp_config_dir) + + mock_client = MagicMock() + mock_client.get_config_detail.return_value = { + "id": "cfg-001", + "name": "Old Name", + "componentId": "keboola.ex-http", + } + mock_client.update_config.return_value = { + "id": "cfg-001", + "name": "New Name", + "componentId": "keboola.ex-http", + } + + # Set up a mock sync directory with manifest + sync_dir = tmp_path / "sync_project" + sync_dir.mkdir() + keboola_dir = sync_dir / ".keboola" + keboola_dir.mkdir() + manifest = { + "version": 2, + "project": {"id": 258, "apiHost": "connection.keboola.com"}, + "allowTargetEnv": True, + "gitBranching": {"enabled": False, "defaultBranch": "main"}, + "sortBy": "id", + "naming": { + "branch": "{branch_name}", + "config": "{component_type}/{component_id}/{config_name}", + "configRow": "rows/{config_row_name}", + "schedulerConfig": "schedules/{config_name}", + "sharedCodeConfig": "_shared/{target_component_id}", + "sharedCodeConfigRow": "codes/{config_row_name}", + "variablesConfig": "variables", + "variablesValuesRow": "values/{config_row_name}", + "dataAppConfig": "app/{component_id}/{config_name}", + }, + "allowedBranches": [], + "ignoredComponents": [], + "branches": [{"id": 12345, "path": "main", "metadata": {}}], + "configurations": [ + { + "branchId": 12345, + "componentId": "keboola.ex-http", + "id": "cfg-001", + "path": "extractor/keboola.ex-http/old-name", + "metadata": {}, + "rows": [], + } + ], + } + (keboola_dir / "manifest.json").write_text(json.dumps(manifest), encoding="utf-8") + + # Create the old config directory on disk + old_config_dir = sync_dir / "main" / "extractor" / "keboola.ex-http" / "old-name" + old_config_dir.mkdir(parents=True) + (old_config_dir / "_config.yml").write_text("name: Old Name\n", encoding="utf-8") + + with pytest.MonkeyPatch.context() as mp: + mp.setattr( + "keboola_agent_cli.commands.config.get_service", + lambda ctx, name: ConfigService( + config_store=store, + client_factory=lambda url, token: mock_client, + ), + ) + + result = runner.invoke( + app, + [ + "--config-dir", + str(tmp_config_dir), + "config", + "rename", + "--project", + "prod", + "--component-id", + "keboola.ex-http", + "--config-id", + "cfg-001", + "--name", + "New Name", + "--directory", + str(sync_dir), + ], + ) + + assert result.exit_code == 0, f"Exit code {result.exit_code}: {result.output}" + assert "Renamed" in result.output + # The sync info line shows old and new paths + assert "Sync:" in result.output + + def test_config_rename_api_error(self, tmp_config_dir: Path) -> None: + """config rename with API error returns appropriate exit code.""" + store = setup_single_project(tmp_config_dir) + + mock_client = MagicMock() + mock_client.get_config_detail.side_effect = KeboolaApiError( + message="Configuration 'cfg-999' not found", + status_code=404, + error_code="NOT_FOUND", + retryable=False, + ) + + with pytest.MonkeyPatch.context() as mp: + mp.setattr( + "keboola_agent_cli.commands.config.get_service", + lambda ctx, name: ConfigService( + config_store=store, + client_factory=lambda url, token: mock_client, + ), + ) + + result = runner.invoke( + app, + [ + "--json", + "--config-dir", + str(tmp_config_dir), + "config", + "rename", + "--project", + "prod", + "--component-id", + "keboola.ex-http", + "--config-id", + "cfg-999", + "--name", + "Whatever", + ], + ) + + # NOT_FOUND maps to exit code 1 (general error) + assert result.exit_code == 1 + output = json.loads(result.output) + assert output["status"] == "error" + assert "NOT_FOUND" in output["error"]["code"] + + def test_config_rename_help(self) -> None: + """config rename --help shows usage information.""" + result = runner.invoke(app, ["config", "rename", "--help"]) + + assert result.exit_code == 0 + assert "Rename a configuration" in result.output + assert "--project" in result.output + assert "--component-id" in result.output + assert "--config-id" in result.output + assert "--name" in result.output + assert "--directory" in result.output diff --git a/tests/test_e2e.py b/tests/test_e2e.py new file mode 100644 index 00000000..d961153d --- /dev/null +++ b/tests/test_e2e.py @@ -0,0 +1,2453 @@ +"""Comprehensive end-to-end tests for Keboola Agent CLI. + +Exercises the FULL CLI surface against a real (empty) Keboola project: + - Project CRUD (add / list / status / edit / remove) + - Storage CRUD (create-bucket / create-table / upload / download / delete) + - Config operations (list / detail / search / update --set / update --merge / delete) + - File operations (upload / list / detail / download / tag / delete) + - Branch lifecycle (list / create / use / reset / merge / delete) + - Workspace lifecycle (create / list / detail / password / load / query / delete) + - Component discovery (list / detail / config new scaffold) + - Job commands (list / detail with filters) + - Encrypt (values) + - Permissions (list / show / check) + - Sync workflow (init / pull / status / diff / push --dry-run) + - Tool commands (list / call) -- requires keboola-mcp-server + - Lineage, sharing, doctor, context, version, changelog, init + +All resources are prefixed with 'e2e-{run_id}' and cleaned up even on failure. + +Requires environment variables: + - E2E_API_TOKEN: Storage API token + - E2E_URL: Stack URL (e.g. connection.keboola.com) + +Run: + E2E_API_TOKEN=xxx E2E_URL=connection.keboola.com \ + uv run pytest tests/test_e2e.py -v -s --tb=long +""" + +from __future__ import annotations + +import csv +import json +import os +import shutil +import subprocess +import time +from pathlib import Path +from typing import Any +from unittest.mock import patch + +import pytest +from typer.testing import CliRunner + +from keboola_agent_cli.cli import app +from keboola_agent_cli.client import KeboolaClient +from keboola_agent_cli.config_store import ConfigStore + +# --------------------------------------------------------------------------- +# Environment & skip logic +# --------------------------------------------------------------------------- + +ENV_TOKEN = "E2E_API_TOKEN" +ENV_URL = "E2E_URL" + +HAS_CREDENTIALS = os.environ.get(ENV_TOKEN) is not None + +skip_without_credentials = pytest.mark.skipif( + not HAS_CREDENTIALS, + reason=f"E2E tests require {ENV_TOKEN} environment variable", +) + +runner = CliRunner() + +# --------------------------------------------------------------------------- +# Unique run identifier (avoids collisions between concurrent runs) +# --------------------------------------------------------------------------- + +RUN_ID = f"e2e-{int(time.time())}" + +# Component used for creating test configurations (always exists in Keboola) +TEST_COMPONENT_ID = "keboola.ex-db-snowflake" + +# --------------------------------------------------------------------------- +# Output formatting constants +# --------------------------------------------------------------------------- + +# ANSI colors for terminal output +_DIM = "\033[2m" +_CYAN = "\033[36m" +_GREEN = "\033[32m" +_RED = "\033[31m" +_YELLOW = "\033[33m" +_RESET = "\033[0m" +_BOLD = "\033[1m" + +# Maximum length for JSON response preview +_MAX_RESPONSE_LEN = 300 + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _mask_token(text: str) -> str: + """Replace any occurrence of the real token in text with a placeholder.""" + token = os.environ.get(ENV_TOKEN, "") + if token and token in text: + return text.replace(token, "***TOKEN***") + return text + + +def _format_cmd(args: list[str]) -> str: + """Format CLI args into a readable command string, masking the token.""" + cmd = "kbagent " + " ".join(args) + return _mask_token(cmd) + + +def _summarize_json(output: str, max_len: int = _MAX_RESPONSE_LEN) -> str: + """Pretty-print JSON output, truncated if too long.""" + try: + data = json.loads(output) + pretty = json.dumps(data, indent=2, ensure_ascii=False) + pretty = _mask_token(pretty) + if len(pretty) > max_len: + return pretty[:max_len] + f"\n ... ({len(pretty)} chars total)" + return pretty + except (json.JSONDecodeError, TypeError): + text = _mask_token(output.strip()) + if len(text) > max_len: + return text[:max_len] + f"... ({len(text)} chars total)" + return text + + +def _invoke(config_dir: Path, args: list[str], catch: bool = True) -> Any: + """Invoke the CLI with a custom config store backed by *config_dir*. + + Prints the command and a response summary for visibility. + """ + print(f"\n {_CYAN}$ {_format_cmd(args)}{_RESET}") + + with patch("keboola_agent_cli.cli.ConfigStore") as mock_store_cls: + mock_store_cls.return_value = ConfigStore(config_dir=config_dir) + result = runner.invoke(app, args, catch_exceptions=catch) + + # Print result summary + if result.exit_code == 0: + status_icon = f"{_GREEN}OK{_RESET}" + else: + status_icon = f"{_RED}EXIT {result.exit_code}{_RESET}" + + print(f" {_DIM}-> {status_icon} {_DIM}({len(result.output)} bytes){_RESET}") + + # Print abbreviated response + summary = _summarize_json(result.output) + for line in summary.split("\n"): + print(f" {_DIM} {line}{_RESET}") + + return result + + +def _json(result) -> dict[str, Any]: + """Parse CLI result output as JSON, with a clear error if parsing fails.""" + assert result.exit_code == 0, f"Command failed (exit {result.exit_code}):\n{result.output}" + try: + return json.loads(result.output) + except json.JSONDecodeError: + pytest.fail(f"Output is not valid JSON:\n{result.output}") + + +def _json_ok(result) -> dict[str, Any]: + """Parse CLI result as JSON and assert status == 'ok'.""" + data = _json(result) + assert data.get("status") == "ok", f"Expected status=ok, got: {data}" + return data + + +def _step(num: int, title: str, detail: str = "") -> None: + """Print a visible step marker for -s output.""" + suffix = f" — {detail}" if detail else "" + print(f"\n{_BOLD}{'=' * 60}") + print(f" STEP {num}: {title}{suffix}") + print(f"{'=' * 60}{_RESET}") + + +def _create_test_csv(path: Path, rows: int = 5) -> Path: + """Create a small CSV file for upload testing.""" + csv_path = path / f"{RUN_ID}_data.csv" + with open(csv_path, "w", newline="") as f: + writer = csv.writer(f) + writer.writerow(["id", "name", "value"]) + for i in range(1, rows + 1): + writer.writerow([i, f"item_{i}", i * 10]) + return csv_path + + +def _create_incremental_csv(path: Path, start: int = 6, rows: int = 3) -> Path: + """Create a CSV file for incremental upload testing.""" + csv_path = path / f"{RUN_ID}_incr_data.csv" + with open(csv_path, "w", newline="") as f: + writer = csv.writer(f) + writer.writerow(["id", "name", "value"]) + for i in range(start, start + rows): + writer.writerow([i, f"item_{i}", i * 10]) + return csv_path + + +def _create_test_file(path: Path, content: str = "hello e2e") -> Path: + """Create a small text file for file-upload testing.""" + file_path = path / f"{RUN_ID}_file.txt" + file_path.write_text(content) + return file_path + + +def _check_mcp_module() -> bool: + """Check if keboola-mcp-server is available as a Python module.""" + try: + result = subprocess.run( + ["python", "-m", "keboola_mcp_server", "--help"], + capture_output=True, + timeout=10, + ) + return result.returncode == 0 + except Exception: + return False + + +# MCP server availability +HAS_MCP_SERVER = shutil.which("keboola_mcp_server") is not None or _check_mcp_module() + +skip_without_mcp = pytest.mark.skipif( + not HAS_MCP_SERVER, + reason="Tool tests require keboola-mcp-server", +) + + +def _git(cwd: Path, *args: str) -> str: + """Run a git command and return stdout.""" + result = subprocess.run( + ["git", *args], + cwd=cwd, + capture_output=True, + text=True, + check=True, + ) + return result.stdout.strip() + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@skip_without_credentials +@pytest.mark.e2e +class TestFullE2E: + """Comprehensive end-to-end test exercising the entire CLI.""" + + @pytest.fixture(autouse=True) + def setup(self, tmp_path: Path) -> None: + """Prepare credentials, directories, and API client for cleanup.""" + self.token = os.environ[ENV_TOKEN] + raw_url = os.environ.get(ENV_URL, "connection.keboola.com") + self.url = raw_url if raw_url.startswith("https://") else f"https://{raw_url}" + self.alias = f"{RUN_ID}-proj" + + # Working directories + self.work_dir = tmp_path / f"kbagent_{RUN_ID}" + self.work_dir.mkdir() + self.config_dir = self.work_dir / "config" + self.config_dir.mkdir() + self.data_dir = self.work_dir / "data" + self.data_dir.mkdir() + + # Direct API client for setup / cleanup helpers + self.api = KeboolaClient(self.url, self.token) + + # Track resources for cleanup + self._created_buckets: list[str] = [] + self._created_branches: list[int] = [] + self._created_config_ids: list[tuple[str, str]] = [] # (component_id, config_id) + self._created_file_ids: list[int] = [] + self._created_workspace_ids: list[int] = [] + + @pytest.fixture(autouse=True) + def cleanup(self) -> Any: + """Guarantee cleanup of ALL created resources, even on test failure.""" + yield + print("\n--- CLEANUP ---") + # Delete workspaces + for ws_id in self._created_workspace_ids: + try: + self.api.delete_workspace(ws_id) + print(f" Deleted workspace {ws_id}") + except Exception as exc: + print(f" WARN: failed to delete workspace {ws_id}: {exc}") + + # Delete configs created via API + for comp_id, cfg_id in self._created_config_ids: + try: + self.api.delete_config(comp_id, cfg_id) + print(f" Deleted config {comp_id}/{cfg_id}") + except Exception as exc: + print(f" WARN: failed to delete config {comp_id}/{cfg_id}: {exc}") + + # Delete branches + for branch_id in self._created_branches: + try: + self.api.delete_dev_branch(branch_id) + print(f" Deleted branch {branch_id}") + except Exception as exc: + print(f" WARN: failed to delete branch {branch_id}: {exc}") + + # Delete buckets (force to cascade-delete tables) + for bucket_id in self._created_buckets: + try: + self.api.delete_bucket(bucket_id, force=True) + print(f" Deleted bucket {bucket_id}") + except Exception as exc: + print(f" WARN: failed to delete bucket {bucket_id}: {exc}") + + # Delete uploaded files + for file_id in self._created_file_ids: + try: + self.api.delete_file(file_id) + print(f" Deleted file {file_id}") + except Exception as exc: + print(f" WARN: failed to delete file {file_id}: {exc}") + + # ------------------------------------------------------------------ + # Invoke shorthand + # ------------------------------------------------------------------ + def _run(self, *args: str) -> Any: + return _invoke(self.config_dir, ["--json", *args]) + + def _run_ok(self, *args: str) -> dict[str, Any]: + return _json_ok(self._run(*args)) + + def _run_json(self, *args: str) -> dict[str, Any]: + return _json(self._run(*args)) + + def _run_raw(self, *args: str) -> Any: + """Invoke without --json (for human-readable output testing).""" + return _invoke(self.config_dir, list(args)) + + # ================================================================== + # THE BIG TEST + # ================================================================== + + def test_full_cli_e2e(self) -> None: + """Progressive scenario testing every CLI command group.""" + + # ============================================================== + # PHASE 1: Setup -- offline commands + project registration + # ============================================================== + + _step(1, "version / changelog / context", "offline commands") + self._test_offline_commands() + + _step(2, "init", "create local workspace in sub-dir") + self._test_init() + + _step(3, "project add", "register project") + self._test_project_add() + + _step(4, "project list + status", "verify connectivity") + self._test_project_list_and_status() + + _step(5, "doctor", "health check") + self._test_doctor() + + # ============================================================== + # PHASE 2: Read empty project + # ============================================================== + + _step(6, "read empty project", "config list / storage buckets / job list") + self._test_empty_reads() + + # ============================================================== + # PHASE 3: Storage CRUD + # ============================================================== + + _step(7, "storage create-bucket") + bucket_id = self._test_create_bucket() + + _step(8, "storage buckets + bucket-detail", "verify bucket exists") + self._test_bucket_listing(bucket_id) + + _step(9, "storage create-table") + table_id = self._test_create_table(bucket_id) + + _step(10, "storage upload-table", "upload CSV data") + self._test_upload_table(table_id) + + _step( + 11, + "storage upload-table --incremental", + "append rows + verify total", + ) + self._test_upload_incremental(table_id) + + _step(12, "storage tables + table-detail") + self._test_table_listing(bucket_id, table_id) + + _step(13, "storage download-table", "data round-trip verification") + self._test_download_table(table_id) + + _step(14, "storage unload-table", "export to file storage") + self._test_unload_table(table_id) + + _step(15, "storage load-file", "upload CSV as file then load into table") + self._test_load_file(table_id) + + # ============================================================== + # PHASE 4: Config operations (create via API, test via CLI) + # ============================================================== + + _step(16, "config create (via API) + CLI list / detail / search") + config_id = self._test_config_operations() + + _step(17, "config update --set / --dry-run / --name / --configuration") + self._test_config_update(config_id) + + _step(18, "config update --merge", "partial merge without losing keys") + self._test_config_merge(config_id) + + _step("18b", "config rename", "rename config via API") + self._test_config_rename(config_id) + + _step(19, "config new scaffold", "generate boilerplate for component") + self._test_config_new_scaffold() + + # ============================================================== + # PHASE 5: Component commands + # ============================================================== + + _step(20, "component list + detail", "discover components") + self._test_component_commands() + + # ============================================================== + # PHASE 6: Workspace lifecycle + # ============================================================== + + _step(21, "workspace create") + workspace_id = self._test_workspace_create() + + if workspace_id is not None: + _step(22, "workspace list") + self._test_workspace_list(workspace_id) + + _step(23, "workspace detail") + self._test_workspace_detail(workspace_id) + + _step(24, "workspace password") + self._test_workspace_password(workspace_id) + + _step(25, "workspace load", "load test table into workspace") + self._test_workspace_load(workspace_id, table_id) + + _step(26, "workspace query", "run SQL in workspace") + self._test_workspace_query(workspace_id, table_id) + + _step(27, "workspace delete") + self._test_workspace_delete(workspace_id) + + # ============================================================== + # PHASE 7: Transformation job run (Snowflake SQL) + # ============================================================== + + _step(28, "transformation setup", "create output bucket + SQL config") + out_bucket_id, transform_config_id, out_table_id = self._test_transformation_setup(table_id) + + _step(29, "job run --wait", "execute Snowflake transformation") + job_id = self._test_job_run(transform_config_id) + + _step(30, "job detail", "verify completed job") + self._test_job_detail(job_id) + + _step(31, "download transformation output", "verify transformed data") + self._test_transformation_output(out_table_id) + + _step(32, "transformation cleanup") + self._test_transformation_cleanup(out_bucket_id, transform_config_id) + + # ============================================================== + # PHASE 8: File operations + # ============================================================== + + _step(33, "file upload / list / detail / download / tag / delete") + self._test_file_operations() + + # ============================================================== + # PHASE 9: Encrypt + # ============================================================== + + _step(34, "encrypt values") + self._test_encrypt(config_id) + + # ============================================================== + # PHASE 10: Branch lifecycle (expanded with merge) + # ============================================================== + + _step(35, "branch lifecycle", "list / create / use / reset / merge / delete") + self._test_branch_lifecycle() + + # ============================================================== + # PHASE 11: Permissions + # ============================================================== + + _step(36, "permissions list / show / check", "permission system") + self._test_permissions() + + # ============================================================== + # PHASE 12: Sharing & Lineage + # ============================================================== + + _step(37, "sharing list / lineage show", "read-only checks") + self._test_sharing_and_lineage() + + # ============================================================== + # PHASE 12.5: Kai (Keboola AI Assistant) + # ============================================================== + + _step(38, "kai ping / ask / history", "Keboola AI Assistant") + self._test_kai_commands() + + # ============================================================== + # PHASE 13: Job commands (expanded) + # ============================================================== + + _step(39, "job list + detail", "verify job listing structure") + self._test_job_commands() + + # ============================================================== + # PHASE 14: Storage column delete + # ============================================================== + + _step(40, "storage delete-column", "dry-run + actual delete + verify") + self._test_delete_column(table_id) + + # ============================================================== + # PHASE 15: Cleanup + # ============================================================== + + _step(41, "config delete", "cleanup config via CLI") + self._test_config_delete(config_id) + + _step(42, "storage delete-table + delete-bucket", "CLI-driven cleanup") + self._test_storage_cleanup(bucket_id, table_id) + + _step(43, "project edit + remove", "final cleanup") + self._test_project_edit_and_remove() + + print("\n" + "=" * 60) + print(" ALL E2E STEPS PASSED") + print("=" * 60) + + # ================================================================== + # Step implementations + # ================================================================== + + def _test_offline_commands(self) -> None: + """Test version, changelog, context -- no project needed.""" + # version (not JSON, just prints version string) + result = self._run_raw("version") + assert result.exit_code == 0 + assert "." in result.output # should contain a version like "0.18.x" + + # changelog + result = self._run("changelog") + assert result.exit_code == 0 + + # context + result = self._run_raw("context") + assert result.exit_code == 0 + assert "kbagent" in result.output + + def _test_init(self) -> None: + """Test init command -- creates .kbagent/ in a sub-directory.""" + init_dir = self.work_dir / "init_test" + init_dir.mkdir() + + # Use a separate config_dir for init (it creates its own workspace) + init_config_dir = init_dir / "config_for_init" + init_config_dir.mkdir() + + # Run init from the init_dir by invoking with cwd override + # The init command uses Path.cwd(), so we patch it + with patch("keboola_agent_cli.commands.init.Path.cwd", return_value=init_dir): + result = _invoke( + init_config_dir, + ["--json", "init"], + ) + data = _json_ok(result) + assert data["data"]["created"] is True + assert "path" in data["data"] + + def _test_project_add(self) -> None: + """Add a project and verify the response.""" + data = self._run_ok( + "project", + "add", + "--project", + self.alias, + "--url", + self.url, + "--token", + self.token, + ) + proj = data["data"] + assert proj["alias"] == self.alias + assert proj["project_name"] # non-empty + assert proj["project_id"] > 0 + # Token must be masked + assert self.token not in json.dumps(data) + + def _test_project_list_and_status(self) -> None: + """Verify project appears in list and status is ok.""" + # list + data = self._run_ok("project", "list") + aliases = [p["alias"] for p in data["data"]] + assert self.alias in aliases + + # status + data = self._run_ok("project", "status", "--project", self.alias) + status_entry = data["data"][0] + assert status_entry["alias"] == self.alias + assert status_entry["status"] == "ok" + assert status_entry["response_time_ms"] >= 0 + + def _test_doctor(self) -> None: + """Run doctor health check.""" + data = self._run_ok("doctor") + assert data["data"]["summary"]["healthy"] is True + + def _test_empty_reads(self) -> None: + """Read operations on a fresh project should return empty lists.""" + # config list + data = self._run_ok("config", "list", "--project", self.alias) + assert data["data"]["errors"] == [] + # configs may or may not be empty (some projects have default configs) + + # storage buckets -- filter only our prefix later + data = self._run_ok("storage", "buckets", "--project", self.alias) + # Just check structure + assert "buckets" in data["data"] + assert "errors" in data["data"] + + # job list + data = self._run_ok("job", "list", "--project", self.alias, "--limit", "5") + assert "jobs" in data["data"] + assert data["data"]["errors"] == [] + + def _test_create_bucket(self) -> str: + """Create a test bucket and return its ID.""" + bucket_name = RUN_ID.replace("-", "_") + data = self._run_ok( + "storage", + "create-bucket", + "--project", + self.alias, + "--stage", + "in", + "--name", + bucket_name, + "--description", + "E2E test bucket", + ) + bucket_id = data["data"]["id"] + assert bucket_id.startswith("in.c-") + self._created_buckets.append(bucket_id) + return bucket_id + + def _test_bucket_listing(self, bucket_id: str) -> None: + """Verify bucket appears in listings.""" + # buckets + data = self._run_ok("storage", "buckets", "--project", self.alias) + bucket_ids = [b["id"] for b in data["data"]["buckets"]] + assert bucket_id in bucket_ids + + # bucket-detail + data = self._run_ok( + "storage", + "bucket-detail", + "--project", + self.alias, + "--bucket-id", + bucket_id, + ) + assert data["data"]["bucket_id"] == bucket_id + + def _test_create_table(self, bucket_id: str) -> str: + """Create a typed table in the bucket.""" + table_name = f"{RUN_ID.replace('-', '_')}_data" + data = self._run_ok( + "storage", + "create-table", + "--project", + self.alias, + "--bucket-id", + bucket_id, + "--name", + table_name, + "--column", + "id:INTEGER", + "--column", + "name:STRING", + "--column", + "value:INTEGER", + "--primary-key", + "id", + ) + table_id = data["data"]["table_id"] + assert table_id + return table_id + + def _test_upload_table(self, table_id: str) -> None: + """Upload CSV data to the table.""" + csv_path = _create_test_csv(self.data_dir, rows=5) + data = self._run_ok( + "storage", + "upload-table", + "--project", + self.alias, + "--table-id", + table_id, + "--file", + str(csv_path), + ) + assert data["data"]["table_id"] == table_id + + def _test_upload_incremental(self, table_id: str) -> None: + """Upload additional rows incrementally and verify total count.""" + csv_path = _create_incremental_csv(self.data_dir, start=6, rows=3) + data = self._run_ok( + "storage", + "upload-table", + "--project", + self.alias, + "--table-id", + table_id, + "--file", + str(csv_path), + "--incremental", + ) + assert data["data"]["table_id"] == table_id + + # Download and verify total rows (5 original + 3 incremental = 8) + output_path = self.data_dir / "incr_verify.csv" + self._run_ok( + "storage", + "download-table", + "--project", + self.alias, + "--table-id", + table_id, + "--output", + str(output_path), + ) + assert output_path.exists() + with open(output_path) as f: + reader = csv.DictReader(f) + rows = list(reader) + assert len(rows) == 8, f"Expected 8 rows after incremental upload, got {len(rows)}" + + def _test_table_listing(self, bucket_id: str, table_id: str) -> None: + """Verify table appears in listings and detail is correct.""" + # tables + data = self._run_ok( + "storage", + "tables", + "--project", + self.alias, + "--bucket-id", + bucket_id, + ) + table_ids = [t["id"] for t in data["data"]["tables"]] + assert table_id in table_ids + + # table-detail + data = self._run_ok( + "storage", + "table-detail", + "--project", + self.alias, + "--table-id", + table_id, + ) + detail = data["data"] + assert detail["table_id"] == table_id + col_names = [c["name"] for c in detail["column_details"]] + assert "id" in col_names + assert "name" in col_names + assert "value" in col_names + + def _test_download_table(self, table_id: str) -> None: + """Download table data and verify round-trip integrity.""" + output_path = self.data_dir / "downloaded.csv" + self._run_ok( + "storage", + "download-table", + "--project", + self.alias, + "--table-id", + table_id, + "--output", + str(output_path), + ) + assert output_path.exists() + + # Verify content (8 rows after incremental upload) + with open(output_path) as f: + reader = csv.DictReader(f) + rows = list(reader) + assert len(rows) == 8 + + # Test with --columns and --limit + limited_path = self.data_dir / "limited.csv" + self._run_ok( + "storage", + "download-table", + "--project", + self.alias, + "--table-id", + table_id, + "--output", + str(limited_path), + "--columns", + "id", + "--columns", + "name", + "--limit", + "2", + ) + assert limited_path.exists() + with open(limited_path) as f: + reader = csv.DictReader(f) + limited_rows = list(reader) + assert len(limited_rows) == 2 + # Only selected columns + assert set(limited_rows[0].keys()) == {"id", "name"} + + def _test_unload_table(self, table_id: str) -> None: + """Unload a table to file storage and optionally download.""" + unload_path = self.data_dir / "unloaded.csv" + data = self._run_ok( + "storage", + "unload-table", + "--project", + self.alias, + "--table-id", + table_id, + "--download", + "--output", + str(unload_path), + ) + result_data = data["data"] + assert result_data["table_id"] == table_id + assert result_data["file_id"] > 0 + assert unload_path.exists() + + def _test_load_file(self, table_id: str) -> None: + """Upload a CSV as a file, then load it into a table via load-file.""" + # Create a CSV file to upload + csv_path = self.data_dir / f"{RUN_ID}_loadfile.csv" + with open(csv_path, "w", newline="") as f: + writer = csv.writer(f) + writer.writerow(["id", "name", "value"]) + writer.writerow([100, "loadfile_item", 999]) + + # Upload as a Storage file + data = self._run_ok( + "storage", + "file-upload", + "--project", + self.alias, + "--file", + str(csv_path), + "--tag", + f"e2e-loadfile-{RUN_ID}", + ) + file_id = data["data"]["id"] + self._created_file_ids.append(file_id) + + # Load file into existing table + data = self._run_ok( + "storage", + "load-file", + "--project", + self.alias, + "--file-id", + str(file_id), + "--table-id", + table_id, + "--incremental", + ) + assert data["status"] == "ok" + + # Clean up the uploaded file + self._run_ok( + "storage", + "file-delete", + "--project", + self.alias, + "--file-id", + str(file_id), + "--yes", + ) + self._created_file_ids.remove(file_id) + + def _test_config_operations(self) -> str: + """Create a config via API, then test CLI read operations.""" + # Create a test configuration via API (CLI has no config create) + config_body = self.api.create_config( + component_id=TEST_COMPONENT_ID, + name=f"{RUN_ID} Test Config", + configuration={ + "parameters": { + "db": { + "host": "test.example.com", + "port": 443, + "database": "test_db", + } + } + }, + description="E2E test configuration", + ) + config_id = str(config_body["id"]) + self._created_config_ids.append((TEST_COMPONENT_ID, config_id)) + + # config list -- should find our config + data = self._run_ok("config", "list", "--project", self.alias) + config_names = [c["config_name"] for c in data["data"]["configs"]] + assert f"{RUN_ID} Test Config" in config_names + + # config list with --component-id filter + data = self._run_ok( + "config", + "list", + "--project", + self.alias, + "--component-id", + TEST_COMPONENT_ID, + ) + our_configs = [c for c in data["data"]["configs"] if c["config_id"] == config_id] + assert len(our_configs) == 1 + + # config detail + data = self._run_ok( + "config", + "detail", + "--project", + self.alias, + "--component-id", + TEST_COMPONENT_ID, + "--config-id", + config_id, + ) + detail = data["data"] + assert detail["name"] == f"{RUN_ID} Test Config" + assert detail["configuration"]["parameters"]["db"]["host"] == "test.example.com" + + # config search + data = self._run_ok( + "config", + "search", + "--project", + self.alias, + "-q", + RUN_ID, + ) + matches = data["data"]["matches"] + assert len(matches) >= 1 + matched_ids = [r["config_id"] for r in matches] + assert config_id in matched_ids + + # config search with --ignore-case + data = self._run_ok( + "config", + "search", + "--project", + self.alias, + "-q", + RUN_ID.upper(), + "--ignore-case", + ) + assert len(data["data"]["matches"]) >= 1 + + return config_id + + def _test_config_update(self, config_id: str) -> None: + """Test config update with --set, --dry-run, --name, --configuration.""" + # --dry-run first + data = self._run_ok( + "config", + "update", + "--project", + self.alias, + "--component-id", + TEST_COMPONENT_ID, + "--config-id", + config_id, + "--set", + "parameters.db.host=updated.example.com", + "--dry-run", + ) + dry_data = data["data"] + assert dry_data["dry_run"] is True + + # Apply --set + data = self._run_ok( + "config", + "update", + "--project", + self.alias, + "--component-id", + TEST_COMPONENT_ID, + "--config-id", + config_id, + "--set", + "parameters.db.host=updated.example.com", + ) + + # Verify the change via config detail + data = self._run_ok( + "config", + "detail", + "--project", + self.alias, + "--component-id", + TEST_COMPONENT_ID, + "--config-id", + config_id, + ) + assert data["data"]["configuration"]["parameters"]["db"]["host"] == "updated.example.com" + # Other fields should be preserved + assert data["data"]["configuration"]["parameters"]["db"]["port"] == 443 + + # --set a new nested key + data = self._run_ok( + "config", + "update", + "--project", + self.alias, + "--component-id", + TEST_COMPONENT_ID, + "--config-id", + config_id, + "--set", + "parameters.db.schema=public", + ) + + # Verify new key exists alongside existing ones + data = self._run_ok( + "config", + "detail", + "--project", + self.alias, + "--component-id", + TEST_COMPONENT_ID, + "--config-id", + config_id, + ) + db_config = data["data"]["configuration"]["parameters"]["db"] + assert db_config["schema"] == "public" + assert db_config["host"] == "updated.example.com" + + # Update name and description + data = self._run_ok( + "config", + "update", + "--project", + self.alias, + "--component-id", + TEST_COMPONENT_ID, + "--config-id", + config_id, + "--name", + f"{RUN_ID} Updated Config", + "--description", + "Updated by E2E test", + ) + + # Verify metadata update + data = self._run_ok( + "config", + "detail", + "--project", + self.alias, + "--component-id", + TEST_COMPONENT_ID, + "--config-id", + config_id, + ) + assert data["data"]["name"] == f"{RUN_ID} Updated Config" + assert data["data"]["description"] == "Updated by E2E test" + + # Full configuration replace via --configuration + full_config = json.dumps( + { + "parameters": { + "db": { + "host": "final.example.com", + "port": 5439, + "database": "final_db", + } + } + } + ) + data = self._run_ok( + "config", + "update", + "--project", + self.alias, + "--component-id", + TEST_COMPONENT_ID, + "--config-id", + config_id, + "--configuration", + full_config, + ) + + # Verify full replace (schema key should be gone) + data = self._run_ok( + "config", + "detail", + "--project", + self.alias, + "--component-id", + TEST_COMPONENT_ID, + "--config-id", + config_id, + ) + db_config = data["data"]["configuration"]["parameters"]["db"] + assert db_config["host"] == "final.example.com" + assert db_config["port"] == 5439 + assert "schema" not in db_config + + def _test_config_merge(self, config_id: str) -> None: + """Test config update --merge: partial merge without losing existing keys.""" + # Current state: host=final.example.com, port=5439, database=final_db + # Merge in a new key (timeout) without losing existing ones + merge_json = json.dumps({"parameters": {"db": {"timeout": 30}}}) + data = self._run_ok( + "config", + "update", + "--project", + self.alias, + "--component-id", + TEST_COMPONENT_ID, + "--config-id", + config_id, + "--configuration", + merge_json, + "--merge", + ) + assert data["status"] == "ok" + + # Verify merge: timeout added, existing keys preserved + data = self._run_ok( + "config", + "detail", + "--project", + self.alias, + "--component-id", + TEST_COMPONENT_ID, + "--config-id", + config_id, + ) + db_config = data["data"]["configuration"]["parameters"]["db"] + assert db_config["timeout"] == 30, "Merged key 'timeout' should be present" + assert db_config["host"] == "final.example.com", "Existing 'host' preserved" + assert db_config["port"] == 5439, "Existing 'port' preserved" + assert db_config["database"] == "final_db", "Existing 'database' preserved" + + def _test_config_rename(self, config_id: str) -> None: + """Test config rename: rename a config via API and verify.""" + # Rename the config + data = self._run_ok( + "config", + "rename", + "--project", + self.alias, + "--component-id", + TEST_COMPONENT_ID, + "--config-id", + config_id, + "--name", + "E2E Renamed Config", + ) + result = data["data"] + assert result["status"] == "renamed" + assert result["new_name"] == "E2E Renamed Config" + assert result["old_name"] # should have the old name + assert result["component_id"] == TEST_COMPONENT_ID + assert result["config_id"] == config_id + + # Verify via config detail that the name actually changed + data = self._run_ok( + "config", + "detail", + "--project", + self.alias, + "--component-id", + TEST_COMPONENT_ID, + "--config-id", + config_id, + ) + assert data["data"]["name"] == "E2E Renamed Config" + + # Rename back so subsequent tests are not affected + self._run_ok( + "config", + "rename", + "--project", + self.alias, + "--component-id", + TEST_COMPONENT_ID, + "--config-id", + config_id, + "--name", + "E2E Test Config", + ) + + def _test_config_new_scaffold(self) -> None: + """Test config new -- generate scaffold for a component.""" + scaffold_dir = self.data_dir / "scaffold" + scaffold_dir.mkdir() + + data = self._run_ok( + "config", + "new", + "--component-id", + "keboola.ex-http", + "--project", + self.alias, + "--output-dir", + str(scaffold_dir), + ) + result = data["data"] + assert "files_written" in result or "directory" in result + + def _test_component_commands(self) -> None: + """List components and get detail for one. + + NOTE: component list only returns components that have at least one + configuration in the project. This test runs AFTER config creation. + """ + # component list -- now that we have a keboola.ex-db-snowflake config + data = self._run_ok("component", "list", "--project", self.alias) + components = data["data"]["components"] + assert len(components) > 0, "Expected at least one component after config creation" + comp_ids = [c["component_id"] for c in components] + assert TEST_COMPONENT_ID in comp_ids + + # component list with --type filter + data = self._run_ok( + "component", + "list", + "--project", + self.alias, + "--type", + "extractor", + ) + for c in data["data"]["components"]: + assert c["component_type"] == "extractor" + + # component detail (uses AI Service) + data = self._run_ok( + "component", + "detail", + "--component-id", + TEST_COMPONENT_ID, + "--project", + self.alias, + ) + detail = data["data"] + assert detail["component_id"] == TEST_COMPONENT_ID + assert detail["component_type"] == "extractor" + + def _test_workspace_create(self) -> int | None: + """Create a workspace, return its ID or None if unsupported.""" + result = self._run( + "workspace", + "create", + "--project", + self.alias, + ) + if result.exit_code != 0: + print( + f" {_YELLOW}WARN: workspace create failed " + f"(exit {result.exit_code}), skipping workspace tests{_RESET}" + ) + return None + + data = _json_ok(result) + ws_data = data["data"] + workspace_id = ws_data["workspace_id"] + assert workspace_id > 0 + self._created_workspace_ids.append(workspace_id) + return workspace_id + + def _test_workspace_list(self, workspace_id: int) -> None: + """Verify workspace appears in the list.""" + data = self._run_ok("workspace", "list", "--project", self.alias) + ws_ids = [w["id"] for w in data["data"]["workspaces"]] + assert workspace_id in ws_ids + + def _test_workspace_detail(self, workspace_id: int) -> None: + """Get workspace detail and verify structure.""" + data = self._run_ok( + "workspace", + "detail", + "--project", + self.alias, + "--workspace-id", + str(workspace_id), + ) + detail = data["data"] + assert detail["workspace_id"] == workspace_id + + def _test_workspace_password(self, workspace_id: int) -> None: + """Reset workspace password and verify a new password is returned.""" + data = self._run_ok( + "workspace", + "password", + "--project", + self.alias, + "--workspace-id", + str(workspace_id), + ) + assert data["data"]["password"] # non-empty password + + def _test_workspace_load(self, workspace_id: int, table_id: str) -> None: + """Load a table into the workspace.""" + data = self._run_ok( + "workspace", + "load", + "--project", + self.alias, + "--workspace-id", + str(workspace_id), + "--tables", + table_id, + ) + assert data["status"] == "ok" + + def _test_workspace_query(self, workspace_id: int, table_id: str) -> None: + """Run a SQL query in the workspace and verify result.""" + # Table name in workspace is the last segment of table_id + ws_table_name = table_id.rsplit(".", 1)[-1] + sql = f'SELECT COUNT(*) AS cnt FROM "{ws_table_name}"' + data = self._run_ok( + "workspace", + "query", + "--project", + self.alias, + "--workspace-id", + str(workspace_id), + "--sql", + sql, + ) + assert data["status"] == "ok" + + def _test_workspace_delete(self, workspace_id: int) -> None: + """Delete the workspace.""" + data = self._run_ok( + "workspace", + "delete", + "--project", + self.alias, + "--workspace-id", + str(workspace_id), + ) + assert data["status"] == "ok" + self._created_workspace_ids.remove(workspace_id) + + # ------------------------------------------------------------------ + # Transformation job run + # ------------------------------------------------------------------ + + def _test_transformation_setup(self, input_table_id: str) -> tuple[str, str, str]: + """Create output bucket + Snowflake transformation config. + + Returns (out_bucket_id, transform_config_id, out_table_id). + """ + # Create output bucket for transformation results + out_bucket_name = f"{RUN_ID.replace('-', '_')}_out" + data = self._run_ok( + "storage", + "create-bucket", + "--project", + self.alias, + "--stage", + "out", + "--name", + out_bucket_name, + "--description", + "E2E transformation output", + ) + out_bucket_id = data["data"]["id"] + assert out_bucket_id.startswith("out.c-") + self._created_buckets.append(out_bucket_id) + + # Derive workspace table name (last segment of table_id) + ws_input_name = input_table_id.rsplit(".", 1)[-1] + out_table_id = f"{out_bucket_id}.{RUN_ID.replace('-', '_')}_result" + + # Create Snowflake transformation config via API + transform_config = { + "parameters": { + "blocks": [ + { + "name": "E2E Block", + "codes": [ + { + "name": "Transform", + "script": [ + ( + f'CREATE TABLE "{RUN_ID.replace("-", "_")}_result"' + f" AS SELECT" + f' "id",' + f' "name",' + f' CAST("value" AS INTEGER) AS "value",' + f' CAST("value" AS INTEGER) * 2' + f' AS "doubled_value"' + f' FROM "{ws_input_name}"' + ) + ], + } + ], + } + ] + }, + "storage": { + "input": { + "tables": [ + { + "source": input_table_id, + "destination": ws_input_name, + } + ] + }, + "output": { + "tables": [ + { + "source": f"{RUN_ID.replace('-', '_')}_result", + "destination": out_table_id, + } + ] + }, + }, + } + + config_body = self.api.create_config( + component_id="keboola.snowflake-transformation", + name=f"{RUN_ID} SQL Transform", + configuration=transform_config, + description="E2E: doubles the value column", + ) + transform_config_id = str(config_body["id"]) + self._created_config_ids.append(("keboola.snowflake-transformation", transform_config_id)) + + return out_bucket_id, transform_config_id, out_table_id + + def _test_job_run(self, transform_config_id: str) -> str: + """Run the transformation job with --wait and return the job ID.""" + data = self._run_ok( + "job", + "run", + "--project", + self.alias, + "--component-id", + "keboola.snowflake-transformation", + "--config-id", + transform_config_id, + "--wait", + "--timeout", + "300", + ) + job_data = data["data"] + assert job_data["status"] == "success", ( + f"Job failed with status={job_data['status']}: " + f"{job_data.get('result', {}).get('message', 'no message')}" + ) + job_id = str(job_data["id"]) + assert job_id + return job_id + + def _test_job_detail(self, job_id: str) -> None: + """Verify job detail for the completed transformation.""" + data = self._run_ok( + "job", + "detail", + "--project", + self.alias, + "--job-id", + job_id, + ) + detail = data["data"] + assert detail["status"] == "success" + assert detail["isFinished"] is True + assert "keboola.snowflake-transformation" in str( + detail.get("component", detail.get("operationName", "")) + ) + + def _test_transformation_output(self, out_table_id: str) -> None: + """Download the transformation output and verify doubled values.""" + output_path = self.data_dir / "transform_output.csv" + self._run_ok( + "storage", + "download-table", + "--project", + self.alias, + "--table-id", + out_table_id, + "--output", + str(output_path), + ) + assert output_path.exists() + + with open(output_path) as f: + reader = csv.DictReader(f) + rows = list(reader) + + # 5 original + 3 incremental + 1 from load-file = 9 rows + assert len(rows) >= 8, f"Expected at least 8 rows, got {len(rows)}" + + # Verify transformation: doubled_value == value * 2 + for row in rows: + value = int(row["value"]) + doubled = int(row["doubled_value"]) + assert doubled == value * 2, ( + f"Row id={row['id']}: value={value}, " + f"expected doubled_value={value * 2}, got {doubled}" + ) + + def _test_transformation_cleanup(self, out_bucket_id: str, transform_config_id: str) -> None: + """Clean up transformation resources via CLI.""" + # Delete transformation config + self._run_ok( + "config", + "delete", + "--project", + self.alias, + "--component-id", + "keboola.snowflake-transformation", + "--config-id", + transform_config_id, + ) + self._created_config_ids.remove(("keboola.snowflake-transformation", transform_config_id)) + + # Delete output bucket (--force to cascade delete output table) + self._run_ok( + "storage", + "delete-bucket", + "--project", + self.alias, + "--bucket-id", + out_bucket_id, + "--force", + "--yes", + ) + self._created_buckets.remove(out_bucket_id) + + def _test_file_operations(self) -> None: + """Test the full file lifecycle: upload, list, detail, download, tag, delete.""" + # Create a test file + test_file = _create_test_file(self.data_dir, content=f"E2E test data {RUN_ID}") + + # file-upload + data = self._run_ok( + "storage", + "file-upload", + "--project", + self.alias, + "--file", + str(test_file), + "--tag", + f"e2e-{RUN_ID}", + "--tag", + "test", + ) + file_id = data["data"]["id"] + self._created_file_ids.append(file_id) + assert file_id > 0 + + # files (list) + data = self._run_ok( + "storage", + "files", + "--project", + self.alias, + "--tag", + f"e2e-{RUN_ID}", + ) + file_ids = [f["id"] for f in data["data"]["files"]] + assert file_id in file_ids + + # file-detail + data = self._run_ok( + "storage", + "file-detail", + "--project", + self.alias, + "--file-id", + str(file_id), + ) + assert data["data"]["id"] == file_id + assert f"e2e-{RUN_ID}" in data["data"]["tags"] + + # file-download + download_path = self.data_dir / "downloaded_file.txt" + data = self._run_ok( + "storage", + "file-download", + "--project", + self.alias, + "--file-id", + str(file_id), + "--output", + str(download_path), + ) + assert download_path.exists() + downloaded_content = download_path.read_text() + assert RUN_ID in downloaded_content + + # file-tag: add a tag + data = self._run_ok( + "storage", + "file-tag", + "--project", + self.alias, + "--file-id", + str(file_id), + "--add", + "extra-tag", + ) + + # Verify tag was added + data = self._run_ok( + "storage", + "file-detail", + "--project", + self.alias, + "--file-id", + str(file_id), + ) + assert "extra-tag" in data["data"]["tags"] + + # file-tag: remove a tag + data = self._run_ok( + "storage", + "file-tag", + "--project", + self.alias, + "--file-id", + str(file_id), + "--remove", + "extra-tag", + ) + + # file-delete (with --dry-run first) + data = self._run_ok( + "storage", + "file-delete", + "--project", + self.alias, + "--file-id", + str(file_id), + "--dry-run", + ) + assert file_id in data["data"]["would_delete"] + + # Actual delete + data = self._run_ok( + "storage", + "file-delete", + "--project", + self.alias, + "--file-id", + str(file_id), + "--yes", + ) + assert file_id in data["data"]["deleted"] + # Remove from cleanup list since we already deleted it + self._created_file_ids.remove(file_id) + + def _test_encrypt(self, config_id: str) -> None: + """Test encrypting values.""" + input_json = json.dumps({"#password": "secret123", "#api_key": "key456"}) + data = self._run_ok( + "encrypt", + "values", + "--project", + self.alias, + "--component-id", + TEST_COMPONENT_ID, + "--input", + input_json, + ) + encrypted = data["data"] + # Encrypted values should start with KBC::ProjectSecure:: or similar + assert "#password" in encrypted + assert "#api_key" in encrypted + assert encrypted["#password"] != "secret123" # must be encrypted + assert encrypted["#api_key"] != "key456" + assert encrypted["#password"].startswith("KBC::") + + def _test_branch_lifecycle(self) -> None: + """Test branch create, list, use, reset, merge (or delete).""" + # branch list -- should only have main + data = self._run_ok("branch", "list", "--project", self.alias) + branches = data["data"]["branches"] + # Main branch always exists + assert len(branches) >= 1 + + # branch create + branch_name = f"{RUN_ID}-test-branch" + data = self._run_ok( + "branch", + "create", + "--project", + self.alias, + "--name", + branch_name, + "--description", + "E2E test branch", + ) + branch_data = data["data"] + branch_id = branch_data["branch_id"] + assert branch_id > 0 + assert branch_data["branch_name"] == branch_name + assert branch_data["activated"] is True + self._created_branches.append(branch_id) + # Branch create auto-activates -- reset so further tests use main + self._run_ok("branch", "reset", "--project", self.alias) + + # branch list -- should now include our branch + data = self._run_ok("branch", "list", "--project", self.alias) + branch_names = [b["name"] for b in data["data"]["branches"]] + assert branch_name in branch_names + + # branch use -- activate the dev branch + data = self._run_ok( + "branch", + "use", + "--project", + self.alias, + "--branch", + str(branch_id), + ) + + # Verify: project status should show active branch + data = self._run_ok("project", "status", "--project", self.alias) + status = data["data"][0] + assert status["active_branch_id"] == branch_id + + # Storage commands should work in branch context + data = self._run_ok("storage", "buckets", "--project", self.alias) + assert data["data"]["errors"] == [] + + # branch reset -- deactivate the dev branch + data = self._run_ok("branch", "reset", "--project", self.alias) + + # Verify: project status should show no active branch + data = self._run_ok("project", "status", "--project", self.alias) + status = data["data"][0] + assert status["active_branch_id"] is None + + # Try branch merge + merge_result = self._run( + "branch", + "merge", + "--project", + self.alias, + "--branch", + str(branch_id), + ) + # branch merge returns a URL for UI-based merge; it doesn't + # auto-merge via API. We verify the command succeeds, then delete. + if merge_result.exit_code == 0: + merge_data = json.loads(merge_result.output) + assert merge_data["status"] == "ok" + # The response contains a URL to the branch overview + assert "url" in merge_data["data"] or "message" in merge_data["data"] + + # Clean up: delete the branch + self._run_ok( + "branch", + "delete", + "--project", + self.alias, + "--branch", + str(branch_id), + ) + self._created_branches.remove(branch_id) + + # Verify branch is gone + data = self._run_ok("branch", "list", "--project", self.alias) + branch_ids = [b["id"] for b in data["data"]["branches"]] + assert branch_id not in branch_ids + + def _test_permissions(self) -> None: + """Test permissions list, show, and check commands.""" + # permissions list -- returns array of operations + data = self._run_ok("permissions", "list") + operations = data["data"] + assert isinstance(operations, list) + assert len(operations) > 0 + # Each operation should have required fields + op = operations[0] + assert "name" in op + assert "category" in op + + # permissions show -- no policy set, should show inactive + data = self._run_ok("permissions", "show") + assert data["data"]["active"] is False + + # permissions check -- without policy, everything should be allowed + data = self._run_ok("permissions", "check", "branch.delete") + assert data["data"]["operation"] == "branch.delete" + assert data["data"]["allowed"] is True + + def _test_sharing_and_lineage(self) -> None: + """Test sharing list and lineage show (read-only, may be empty).""" + # sharing list + data = self._run_ok("sharing", "list", "--project", self.alias) + assert "shared_buckets" in data["data"] or "errors" in data["data"] + + # lineage show + data = self._run_ok("lineage", "show", "--project", self.alias) + # Lineage may be empty on a single-project setup + assert data["status"] == "ok" + + def _test_kai_commands(self) -> None: + """Test Kai AI Assistant commands (gracefully skip if not available).""" + # kai ping — check if Kai is available for this project + result = self._run("kai", "ping", "--project", self.alias) + if result.exit_code != 0: + output = result.output + if "KAI_NOT_ENABLED" in output or "KAI_ERROR" in output: + print( + f" {_YELLOW}SKIP: Kai not available for this project " + f"(exit {result.exit_code}){_RESET}" + ) + return + # Unexpected error — fail the test + assert result.exit_code == 0, f"kai ping failed unexpectedly: {result.output}" + + # Ping succeeded — verify structure + ping_data = json.loads(result.output) + assert ping_data["status"] == "ok" + assert "timestamp" in ping_data["data"] + assert "mcp_status" in ping_data["data"] + + # kai ask — one-shot question + result = self._run( + "kai", + "ask", + "--project", + self.alias, + "-m", + "Reply with just the word OK", + ) + if result.exit_code != 0: + # Auth issue (e.g. token type) — skip remaining kai tests + print( + f" {_YELLOW}SKIP: kai ask failed " + f"(exit {result.exit_code}), skipping chat/history{_RESET}" + ) + return + + ask_data = json.loads(result.output) + assert ask_data["status"] == "ok" + assert "response" in ask_data["data"] + assert "chat_id" in ask_data["data"] + assert len(ask_data["data"]["response"]) > 0 + + # kai history — list recent chats (at least the one we just created) + data = self._run_ok("kai", "history", "--project", self.alias, "--limit", "5") + assert "chats" in data["data"] + # We just chatted, so there should be at least 1 + assert len(data["data"]["chats"]) >= 1 + + def _test_job_commands(self) -> None: + """Verify job listing structure and detail (if jobs exist).""" + # job list + data = self._run_ok( + "job", + "list", + "--project", + self.alias, + "--limit", + "5", + ) + assert "jobs" in data["data"] + assert "errors" in data["data"] + assert data["data"]["errors"] == [] + + # job list with component filter + data = self._run_ok( + "job", + "list", + "--project", + self.alias, + "--component-id", + TEST_COMPONENT_ID, + "--limit", + "5", + ) + assert "jobs" in data["data"] + + # If any jobs exist, get detail for the first one + jobs = data["data"]["jobs"] + if jobs: + job_id = str(jobs[0]["id"]) + detail_data = self._run_ok( + "job", + "detail", + "--project", + self.alias, + "--job-id", + job_id, + ) + assert detail_data["data"]["id"] + + def _test_config_delete(self, config_id: str) -> None: + """Delete the test config via CLI.""" + data = self._run_ok( + "config", + "delete", + "--project", + self.alias, + "--component-id", + TEST_COMPONENT_ID, + "--config-id", + config_id, + ) + assert data["data"]["config_id"] == config_id + # Remove from cleanup since we deleted via CLI + self._created_config_ids.remove((TEST_COMPONENT_ID, config_id)) + + def _test_delete_column(self, table_id: str) -> None: + """Delete a column from a table: dry-run, actual delete, verify.""" + # Verify the table has 'value' column before we delete it + data = self._run_ok( + "storage", + "table-detail", + "--project", + self.alias, + "--table-id", + table_id, + ) + columns_before = data["data"]["columns"] + assert "value" in columns_before, f"Expected 'value' column, got {columns_before}" + + # delete-column dry-run + data = self._run_ok( + "storage", + "delete-column", + "--project", + self.alias, + "--table-id", + table_id, + "--column", + "value", + "--dry-run", + ) + assert data["data"]["dry_run"] is True + assert "value" in data["data"]["would_delete"] + assert data["data"]["table_id"] == table_id + + # delete-column (actual) + data = self._run_ok( + "storage", + "delete-column", + "--project", + self.alias, + "--table-id", + table_id, + "--column", + "value", + "--yes", + ) + assert "value" in data["data"]["deleted"] + assert data["data"]["failed"] == [] + assert data["data"]["table_id"] == table_id + + # Verify the column is gone + data = self._run_ok( + "storage", + "table-detail", + "--project", + self.alias, + "--table-id", + table_id, + ) + columns_after = data["data"]["columns"] + assert "value" not in columns_after, ( + f"'value' column should be deleted, got {columns_after}" + ) + assert "id" in columns_after + assert "name" in columns_after + + def _test_storage_cleanup(self, bucket_id: str, table_id: str) -> None: + """Delete table and bucket via CLI commands.""" + # delete-table (dry-run first) + data = self._run_ok( + "storage", + "delete-table", + "--project", + self.alias, + "--table-id", + table_id, + "--dry-run", + ) + assert table_id in data["data"]["would_delete"] + + # delete-table (actual) + data = self._run_ok( + "storage", + "delete-table", + "--project", + self.alias, + "--table-id", + table_id, + "--yes", + ) + assert table_id in data["data"]["deleted"] + + # delete-bucket (dry-run first) + data = self._run_ok( + "storage", + "delete-bucket", + "--project", + self.alias, + "--bucket-id", + bucket_id, + "--dry-run", + ) + assert bucket_id in data["data"]["would_delete"] + + # delete-bucket (actual) + data = self._run_ok( + "storage", + "delete-bucket", + "--project", + self.alias, + "--bucket-id", + bucket_id, + "--yes", + ) + assert bucket_id in data["data"]["deleted"] + self._created_buckets.remove(bucket_id) + + def _test_project_edit_and_remove(self) -> None: + """Edit project URL, then remove it.""" + # project edit -- change URL back to same (just verify command works) + data = self._run_ok( + "project", + "edit", + "--project", + self.alias, + "--url", + self.url, + ) + assert data["data"]["alias"] == self.alias + + # project remove + data = self._run_ok("project", "remove", "--project", self.alias) + assert data["data"]["message"] + + # Verify project is gone + data = self._run_ok("project", "list") + remaining = [p["alias"] for p in data["data"]] + assert self.alias not in remaining + + +# --------------------------------------------------------------------------- +# Error handling tests (separate from the main flow) +# --------------------------------------------------------------------------- + + +@skip_without_credentials +@pytest.mark.e2e +class TestE2EErrorHandling: + """Test error paths and edge cases.""" + + @pytest.fixture(autouse=True) + def setup(self, tmp_path: Path) -> None: + self.token = os.environ[ENV_TOKEN] + raw_url = os.environ.get(ENV_URL, "connection.keboola.com") + self.url = raw_url if raw_url.startswith("https://") else f"https://{raw_url}" + self.config_dir = tmp_path / "config" + self.config_dir.mkdir() + + def _run(self, *args: str) -> Any: + return _invoke(self.config_dir, ["--json", *args]) + + def test_add_with_invalid_token(self) -> None: + """Adding a project with an invalid token returns exit code 3.""" + result = self._run( + "project", + "add", + "--project", + "bad-project", + "--url", + self.url, + "--token", + "000-definitely-invalid-token", + ) + assert result.exit_code == 3 + data = json.loads(result.output) + assert data["status"] == "error" + + def test_status_of_nonexistent_project(self) -> None: + """Status of a project that doesn't exist returns exit code 5.""" + result = self._run("project", "status", "--project", "nonexistent") + assert result.exit_code == 5 + + def test_remove_nonexistent_project(self) -> None: + """Removing a nonexistent project returns exit code 5.""" + result = self._run("project", "remove", "--project", "nonexistent") + assert result.exit_code == 5 + + def test_config_detail_nonexistent(self) -> None: + """Config detail for nonexistent config returns error.""" + # First add a valid project + self._run( + "project", + "add", + "--project", + "err-test", + "--url", + self.url, + "--token", + self.token, + ) + result = self._run( + "config", + "detail", + "--project", + "err-test", + "--component-id", + "keboola.ex-db-snowflake", + "--config-id", + "999999999", + ) + assert result.exit_code != 0 + + def test_download_nonexistent_table(self) -> None: + """Downloading a nonexistent table returns error.""" + self._run( + "project", + "add", + "--project", + "err-test2", + "--url", + self.url, + "--token", + self.token, + ) + result = self._run( + "storage", + "download-table", + "--project", + "err-test2", + "--table-id", + "in.c-nonexistent.nonexistent", + ) + assert result.exit_code != 0 + + def test_delete_nonexistent_bucket(self) -> None: + """Deleting a nonexistent bucket returns error.""" + self._run( + "project", + "add", + "--project", + "err-test3", + "--url", + self.url, + "--token", + self.token, + ) + result = self._run( + "storage", + "delete-bucket", + "--project", + "err-test3", + "--bucket-id", + "in.c-nonexistent-bucket-xyz", + "--yes", + ) + assert result.exit_code != 0 + + +# --------------------------------------------------------------------------- +# JSON output consistency tests +# --------------------------------------------------------------------------- + + +@skip_without_credentials +@pytest.mark.e2e +class TestE2EJsonConsistency: + """Verify that all commands produce valid JSON with --json flag.""" + + @pytest.fixture(autouse=True) + def setup(self, tmp_path: Path) -> None: + self.token = os.environ[ENV_TOKEN] + raw_url = os.environ.get(ENV_URL, "connection.keboola.com") + self.url = raw_url if raw_url.startswith("https://") else f"https://{raw_url}" + self.alias = f"{RUN_ID}-json" + self.config_dir = tmp_path / "config" + self.config_dir.mkdir() + + # Add project + _invoke( + self.config_dir, + [ + "--json", + "project", + "add", + "--project", + self.alias, + "--url", + self.url, + "--token", + self.token, + ], + ) + + def _run(self, *args: str) -> Any: + return _invoke(self.config_dir, ["--json", *args]) + + def test_all_read_commands_return_valid_json(self) -> None: + """Every read command should return parseable JSON with status field.""" + commands = [ + ["project", "list"], + ["project", "status", "--project", self.alias], + ["config", "list", "--project", self.alias], + ["storage", "buckets", "--project", self.alias], + ["job", "list", "--project", self.alias, "--limit", "1"], + ["component", "list", "--project", self.alias], + ["branch", "list", "--project", self.alias], + ["sharing", "list", "--project", self.alias], + ["lineage", "show", "--project", self.alias], + ["doctor"], + ["permissions", "list"], + ["permissions", "show"], + ] + for cmd in commands: + result = self._run(*cmd) + assert result.exit_code == 0, ( + f"Command {' '.join(cmd)} failed (exit {result.exit_code}): {result.output}" + ) + try: + data = json.loads(result.output) + except json.JSONDecodeError: + pytest.fail( + f"Command {' '.join(cmd)} did not return valid JSON: {result.output[:200]}" + ) + assert "status" in data, f"Command {' '.join(cmd)} missing 'status' key: {data}" + + def test_token_never_appears_in_any_output(self) -> None: + """The full token should never appear in any command output.""" + commands = [ + ["project", "list"], + ["project", "status", "--project", self.alias], + ["doctor"], + ] + for cmd in commands: + result = self._run(*cmd) + assert self.token not in result.output, ( + f"Full token leaked in output of: {' '.join(cmd)}" + ) + + +# --------------------------------------------------------------------------- +# Sync workflow tests +# --------------------------------------------------------------------------- + + +@skip_without_credentials +@pytest.mark.e2e +class TestE2ESyncWorkflow: + """Test sync init/pull/diff/status/push in a temp git repo.""" + + @pytest.fixture(autouse=True) + def setup(self, tmp_path: Path) -> None: + """Set up config dir, project dir (as git repo), and register project.""" + self.token = os.environ[ENV_TOKEN] + raw_url = os.environ.get(ENV_URL, "connection.keboola.com") + self.url = raw_url if raw_url.startswith("https://") else f"https://{raw_url}" + self.alias = f"{RUN_ID}-sync" + + self.config_dir = tmp_path / "config" + self.config_dir.mkdir() + self.project_dir = tmp_path / "project" + self.project_dir.mkdir() + + # Register the project + result = _invoke( + self.config_dir, + [ + "--json", + "project", + "add", + "--project", + self.alias, + "--url", + self.url, + "--token", + self.token, + ], + ) + assert result.exit_code == 0, f"project add failed: {result.output}" + + # Initialize git repo + _git(self.project_dir, "init") + _git(self.project_dir, "config", "user.email", "e2e@test.local") + _git(self.project_dir, "config", "user.name", "E2E Test") + _git( + self.project_dir, + "commit", + "--allow-empty", + "-m", + "init", + ) + + def _run(self, *args: str) -> Any: + return _invoke(self.config_dir, ["--json", *args]) + + def _run_ok(self, *args: str) -> dict[str, Any]: + return _json_ok(self._run(*args)) + + def test_sync_workflow(self) -> None: + """Full sync lifecycle: init, pull, status, diff, push --dry-run.""" + + # 1. sync init + _step(1, "sync init") + data = self._run_ok( + "sync", + "init", + "--project", + self.alias, + "--directory", + str(self.project_dir), + ) + result = data["data"] + assert result["project_alias"] == self.alias + + # 2. sync pull + _step(2, "sync pull") + data = self._run_ok( + "sync", + "pull", + "--project", + self.alias, + "--directory", + str(self.project_dir), + ) + pull_result = data["data"] + # Should have configs_pulled key (may be 0 on empty project) + assert "configs_pulled" in pull_result + + # Commit pulled files so status/diff have a baseline + _git(self.project_dir, "add", "-A") + _git(self.project_dir, "commit", "-m", "pulled configs") + + # 3. sync status + _step(3, "sync status") + data = self._run_ok( + "sync", + "status", + "--directory", + str(self.project_dir), + ) + assert data["status"] == "ok" + + # 4. sync diff + _step(4, "sync diff") + data = self._run_ok( + "sync", + "diff", + "--project", + self.alias, + "--directory", + str(self.project_dir), + ) + assert data["status"] == "ok" + + # 5. sync push --dry-run + _step(5, "sync push --dry-run") + data = self._run_ok( + "sync", + "push", + "--project", + self.alias, + "--directory", + str(self.project_dir), + "--dry-run", + ) + assert data["status"] == "ok" + + +# --------------------------------------------------------------------------- +# Tool command tests (requires MCP server) +# --------------------------------------------------------------------------- + + +@skip_without_credentials +@skip_without_mcp +@pytest.mark.e2e +class TestE2EToolCommands: + """Test MCP tool list and call commands.""" + + @pytest.fixture(autouse=True) + def setup(self, tmp_path: Path) -> None: + """Register a project for tool tests.""" + self.token = os.environ[ENV_TOKEN] + raw_url = os.environ.get(ENV_URL, "connection.keboola.com") + self.url = raw_url if raw_url.startswith("https://") else f"https://{raw_url}" + self.alias = f"{RUN_ID}-tool" + self.config_dir = tmp_path / "config" + self.config_dir.mkdir() + + result = _invoke( + self.config_dir, + [ + "--json", + "project", + "add", + "--project", + self.alias, + "--url", + self.url, + "--token", + self.token, + ], + ) + assert result.exit_code == 0, f"project add failed: {result.output}" + + def _run(self, *args: str) -> Any: + return _invoke(self.config_dir, ["--json", *args]) + + def _run_ok(self, *args: str) -> dict[str, Any]: + return _json_ok(self._run(*args)) + + def test_tool_list(self) -> None: + """tool list should return a list of available MCP tools.""" + result = self._run("tool", "list", "--project", self.alias) + assert result.exit_code == 0 + + def test_tool_call_get_buckets(self) -> None: + """tool call get_buckets should return bucket data.""" + result = self._run( + "tool", + "call", + "get_buckets", + "--project", + self.alias, + ) + assert result.exit_code == 0 diff --git a/tests/test_kai_cli.py b/tests/test_kai_cli.py new file mode 100644 index 00000000..9273e657 --- /dev/null +++ b/tests/test_kai_cli.py @@ -0,0 +1,646 @@ +"""Tests for Kai CLI commands via CliRunner. + +Tests the `kbagent kai` subcommands: ping, ask, chat, history. +Each command is tested in both JSON and human output modes, plus error cases. +""" + +import json +from pathlib import Path +from unittest.mock import MagicMock + +import pytest +from typer.testing import CliRunner + +from helpers import setup_single_project +from keboola_agent_cli.cli import app +from keboola_agent_cli.errors import KeboolaApiError +from keboola_agent_cli.services.kai_service import KaiService + +runner = CliRunner() + + +class TestKaiPingCli: + """Tests for `kbagent kai ping` command.""" + + def test_kai_ping_json_output(self, tmp_config_dir: Path) -> None: + """kai ping --json returns structured JSON with server info.""" + setup_single_project(tmp_config_dir) + + mock_service = MagicMock(spec=KaiService) + mock_service.resolve_alias.return_value = "prod" + mock_service.ping.return_value = { + "project_alias": "prod", + "timestamp": "2025-01-15T10:30:00+00:00", + "app_name": "kai-api", + "app_version": "1.2.3", + "server_version": "2.0.0", + "mcp_status": "connected", + } + + with pytest.MonkeyPatch.context() as mp: + mp.setattr( + "keboola_agent_cli.commands.kai.get_service", + lambda ctx, name: mock_service, + ) + + result = runner.invoke( + app, + [ + "--json", + "--config-dir", + str(tmp_config_dir), + "kai", + "ping", + "--project", + "prod", + ], + ) + + assert result.exit_code == 0, f"Exit code {result.exit_code}: {result.output}" + output = json.loads(result.output) + assert output["status"] == "ok" + assert output["data"]["project_alias"] == "prod" + assert output["data"]["app_name"] == "kai-api" + assert output["data"]["mcp_status"] == "connected" + + def test_kai_ping_human_output(self, tmp_config_dir: Path) -> None: + """kai ping in human mode shows readable server info.""" + setup_single_project(tmp_config_dir) + + mock_service = MagicMock(spec=KaiService) + mock_service.resolve_alias.return_value = "prod" + mock_service.ping.return_value = { + "project_alias": "prod", + "timestamp": "2025-01-15T10:30:00+00:00", + "app_name": "kai-api", + "app_version": "1.2.3", + "server_version": "2.0.0", + "mcp_status": "connected", + } + + with pytest.MonkeyPatch.context() as mp: + mp.setattr( + "keboola_agent_cli.commands.kai.get_service", + lambda ctx, name: mock_service, + ) + + result = runner.invoke( + app, + [ + "--config-dir", + str(tmp_config_dir), + "kai", + "ping", + "--project", + "prod", + ], + ) + + assert result.exit_code == 0, f"Exit code {result.exit_code}: {result.output}" + assert "Kai is alive" in result.output + assert "kai-api" in result.output + assert "connected" in result.output + + def test_kai_ping_api_error(self, tmp_config_dir: Path) -> None: + """kai ping with API error returns structured error and exit code 1.""" + setup_single_project(tmp_config_dir) + + mock_service = MagicMock(spec=KaiService) + mock_service.resolve_alias.return_value = "prod" + mock_service.ping.side_effect = KeboolaApiError( + message="Kai ping failed: Connection refused", + status_code=0, + error_code="KAI_ERROR", + ) + + with pytest.MonkeyPatch.context() as mp: + mp.setattr( + "keboola_agent_cli.commands.kai.get_service", + lambda ctx, name: mock_service, + ) + + result = runner.invoke( + app, + [ + "--json", + "--config-dir", + str(tmp_config_dir), + "kai", + "ping", + "--project", + "prod", + ], + ) + + assert result.exit_code == 1 + output = json.loads(result.output) + assert output["status"] == "error" + assert "KAI_ERROR" in output["error"]["code"] + + def test_kai_ping_not_enabled(self, tmp_config_dir: Path) -> None: + """kai ping when Kai is not enabled returns KAI_NOT_ENABLED error.""" + setup_single_project(tmp_config_dir) + + mock_service = MagicMock(spec=KaiService) + mock_service.resolve_alias.return_value = "prod" + mock_service.ping.side_effect = KeboolaApiError( + message="Kai is not enabled for project 'prod'.", + status_code=0, + error_code="KAI_NOT_ENABLED", + ) + + with pytest.MonkeyPatch.context() as mp: + mp.setattr( + "keboola_agent_cli.commands.kai.get_service", + lambda ctx, name: mock_service, + ) + + result = runner.invoke( + app, + [ + "--json", + "--config-dir", + str(tmp_config_dir), + "kai", + "ping", + "--project", + "prod", + ], + ) + + assert result.exit_code == 1 + output = json.loads(result.output) + assert output["status"] == "error" + assert "KAI_NOT_ENABLED" in output["error"]["code"] + + def test_kai_ping_help(self) -> None: + """kai ping --help shows usage information.""" + result = runner.invoke(app, ["kai", "ping", "--help"]) + + assert result.exit_code == 0 + assert "Check Kai server health" in result.output + assert "--project" in result.output + + +class TestKaiAskCli: + """Tests for `kbagent kai ask` command.""" + + def test_kai_ask_json_output(self, tmp_config_dir: Path) -> None: + """kai ask --json returns structured JSON with response text.""" + setup_single_project(tmp_config_dir) + + mock_service = MagicMock(spec=KaiService) + mock_service.resolve_alias.return_value = "prod" + mock_service.ask.return_value = { + "project_alias": "prod", + "chat_id": "chat-xyz-789", + "response": "You have 5 transformations configured.", + } + + with pytest.MonkeyPatch.context() as mp: + mp.setattr( + "keboola_agent_cli.commands.kai.get_service", + lambda ctx, name: mock_service, + ) + + result = runner.invoke( + app, + [ + "--json", + "--config-dir", + str(tmp_config_dir), + "kai", + "ask", + "--message", + "How many transformations?", + "--project", + "prod", + ], + ) + + assert result.exit_code == 0, f"Exit code {result.exit_code}: {result.output}" + output = json.loads(result.output) + assert output["status"] == "ok" + assert output["data"]["chat_id"] == "chat-xyz-789" + assert output["data"]["response"] == "You have 5 transformations configured." + + def test_kai_ask_human_output(self, tmp_config_dir: Path) -> None: + """kai ask in human mode shows just the response text.""" + setup_single_project(tmp_config_dir) + + mock_service = MagicMock(spec=KaiService) + mock_service.resolve_alias.return_value = "prod" + mock_service.ask.return_value = { + "project_alias": "prod", + "chat_id": "chat-xyz-789", + "response": "You have 5 transformations configured.", + } + + with pytest.MonkeyPatch.context() as mp: + mp.setattr( + "keboola_agent_cli.commands.kai.get_service", + lambda ctx, name: mock_service, + ) + + result = runner.invoke( + app, + [ + "--config-dir", + str(tmp_config_dir), + "kai", + "ask", + "--message", + "How many transformations?", + "--project", + "prod", + ], + ) + + assert result.exit_code == 0, f"Exit code {result.exit_code}: {result.output}" + assert "You have 5 transformations configured." in result.output + + def test_kai_ask_api_error(self, tmp_config_dir: Path) -> None: + """kai ask with API error returns structured error.""" + setup_single_project(tmp_config_dir) + + mock_service = MagicMock(spec=KaiService) + mock_service.resolve_alias.return_value = "prod" + mock_service.ask.side_effect = KeboolaApiError( + message="Kai ask failed: timeout", + status_code=0, + error_code="KAI_ERROR", + ) + + with pytest.MonkeyPatch.context() as mp: + mp.setattr( + "keboola_agent_cli.commands.kai.get_service", + lambda ctx, name: mock_service, + ) + + result = runner.invoke( + app, + [ + "--json", + "--config-dir", + str(tmp_config_dir), + "kai", + "ask", + "--message", + "test", + "--project", + "prod", + ], + ) + + assert result.exit_code == 1 + output = json.loads(result.output) + assert output["status"] == "error" + + def test_kai_ask_help(self) -> None: + """kai ask --help shows usage information.""" + result = runner.invoke(app, ["kai", "ask", "--help"]) + + assert result.exit_code == 0 + assert "Ask Kai a one-shot question" in result.output + assert "--message" in result.output + assert "--project" in result.output + + +class TestKaiChatCli: + """Tests for `kbagent kai chat` command.""" + + def test_kai_chat_json_output(self, tmp_config_dir: Path) -> None: + """kai chat --json returns structured JSON with chat_id and response.""" + setup_single_project(tmp_config_dir) + + mock_service = MagicMock(spec=KaiService) + mock_service.resolve_alias.return_value = "prod" + mock_service.chat_message.return_value = { + "project_alias": "prod", + "chat_id": "chat-session-001", + "response": "I can help with that.", + } + + with pytest.MonkeyPatch.context() as mp: + mp.setattr( + "keboola_agent_cli.commands.kai.get_service", + lambda ctx, name: mock_service, + ) + + result = runner.invoke( + app, + [ + "--json", + "--config-dir", + str(tmp_config_dir), + "kai", + "chat", + "--message", + "Help me debug", + "--project", + "prod", + ], + ) + + assert result.exit_code == 0, f"Exit code {result.exit_code}: {result.output}" + output = json.loads(result.output) + assert output["status"] == "ok" + assert output["data"]["chat_id"] == "chat-session-001" + assert output["data"]["response"] == "I can help with that." + + def test_kai_chat_with_chat_id(self, tmp_config_dir: Path) -> None: + """kai chat --chat-id passes the ID to the service for continuation.""" + setup_single_project(tmp_config_dir) + + mock_service = MagicMock(spec=KaiService) + mock_service.resolve_alias.return_value = "prod" + mock_service.chat_message.return_value = { + "project_alias": "prod", + "chat_id": "existing-chat-42", + "response": "Continuing our conversation.", + } + + with pytest.MonkeyPatch.context() as mp: + mp.setattr( + "keboola_agent_cli.commands.kai.get_service", + lambda ctx, name: mock_service, + ) + + result = runner.invoke( + app, + [ + "--json", + "--config-dir", + str(tmp_config_dir), + "kai", + "chat", + "--message", + "What about now?", + "--chat-id", + "existing-chat-42", + "--project", + "prod", + ], + ) + + assert result.exit_code == 0, f"Exit code {result.exit_code}: {result.output}" + # Verify chat_id was passed through to the service + mock_service.chat_message.assert_called_once_with( + "prod", "What about now?", chat_id="existing-chat-42" + ) + + def test_kai_chat_human_output(self, tmp_config_dir: Path) -> None: + """kai chat in human mode shows response text and chat ID.""" + setup_single_project(tmp_config_dir) + + mock_service = MagicMock(spec=KaiService) + mock_service.resolve_alias.return_value = "prod" + mock_service.chat_message.return_value = { + "project_alias": "prod", + "chat_id": "chat-session-001", + "response": "Here is the answer.", + } + + with pytest.MonkeyPatch.context() as mp: + mp.setattr( + "keboola_agent_cli.commands.kai.get_service", + lambda ctx, name: mock_service, + ) + + result = runner.invoke( + app, + [ + "--config-dir", + str(tmp_config_dir), + "kai", + "chat", + "--message", + "question", + "--project", + "prod", + ], + ) + + assert result.exit_code == 0, f"Exit code {result.exit_code}: {result.output}" + assert "Here is the answer." in result.output + assert "chat-session-001" in result.output + + def test_kai_chat_help(self) -> None: + """kai chat --help shows usage information.""" + result = runner.invoke(app, ["kai", "chat", "--help"]) + + assert result.exit_code == 0 + assert "Send a message to Kai" in result.output + assert "--message" in result.output + assert "--chat-id" in result.output + + +class TestKaiHistoryCli: + """Tests for `kbagent kai history` command.""" + + def test_kai_history_json_output(self, tmp_config_dir: Path) -> None: + """kai history --json returns structured JSON with chat list.""" + setup_single_project(tmp_config_dir) + + mock_service = MagicMock(spec=KaiService) + mock_service.resolve_alias.return_value = "prod" + mock_service.get_history.return_value = { + "project_alias": "prod", + "chats": [ + { + "id": "chat-aaa-111", + "title": "Data pipeline question", + "created_at": "2025-01-10T08:00:00+00:00", + "visibility": "private", + }, + { + "id": "chat-bbb-222", + "title": "(untitled)", + "created_at": None, + "visibility": "public", + }, + ], + "has_more": False, + } + + with pytest.MonkeyPatch.context() as mp: + mp.setattr( + "keboola_agent_cli.commands.kai.get_service", + lambda ctx, name: mock_service, + ) + + result = runner.invoke( + app, + [ + "--json", + "--config-dir", + str(tmp_config_dir), + "kai", + "history", + "--project", + "prod", + ], + ) + + assert result.exit_code == 0, f"Exit code {result.exit_code}: {result.output}" + output = json.loads(result.output) + assert output["status"] == "ok" + assert len(output["data"]["chats"]) == 2 + assert output["data"]["chats"][0]["title"] == "Data pipeline question" + assert output["data"]["has_more"] is False + + def test_kai_history_with_limit(self, tmp_config_dir: Path) -> None: + """kai history --limit passes the limit to the service.""" + setup_single_project(tmp_config_dir) + + mock_service = MagicMock(spec=KaiService) + mock_service.resolve_alias.return_value = "prod" + mock_service.get_history.return_value = { + "project_alias": "prod", + "chats": [], + "has_more": False, + } + + with pytest.MonkeyPatch.context() as mp: + mp.setattr( + "keboola_agent_cli.commands.kai.get_service", + lambda ctx, name: mock_service, + ) + + result = runner.invoke( + app, + [ + "--json", + "--config-dir", + str(tmp_config_dir), + "kai", + "history", + "--limit", + "25", + "--project", + "prod", + ], + ) + + assert result.exit_code == 0 + mock_service.get_history.assert_called_once_with("prod", limit=25) + + def test_kai_history_human_output(self, tmp_config_dir: Path) -> None: + """kai history in human mode shows a table of chats.""" + setup_single_project(tmp_config_dir) + + mock_service = MagicMock(spec=KaiService) + mock_service.resolve_alias.return_value = "prod" + mock_service.get_history.return_value = { + "project_alias": "prod", + "chats": [ + { + "id": "chat-aaa-111-full-uuid", + "title": "Pipeline debugging", + "created_at": "2025-01-10T08:00:00+00:00", + "visibility": "private", + }, + ], + "has_more": True, + } + + with pytest.MonkeyPatch.context() as mp: + mp.setattr( + "keboola_agent_cli.commands.kai.get_service", + lambda ctx, name: mock_service, + ) + + result = runner.invoke( + app, + [ + "--config-dir", + str(tmp_config_dir), + "kai", + "history", + "--project", + "prod", + ], + ) + + assert result.exit_code == 0, f"Exit code {result.exit_code}: {result.output}" + assert "Pipeline debugging" in result.output + assert "More chats available" in result.output + + def test_kai_history_empty_human(self, tmp_config_dir: Path) -> None: + """kai history in human mode shows 'No chat history' when empty.""" + setup_single_project(tmp_config_dir) + + mock_service = MagicMock(spec=KaiService) + mock_service.resolve_alias.return_value = "prod" + mock_service.get_history.return_value = { + "project_alias": "prod", + "chats": [], + "has_more": False, + } + + with pytest.MonkeyPatch.context() as mp: + mp.setattr( + "keboola_agent_cli.commands.kai.get_service", + lambda ctx, name: mock_service, + ) + + result = runner.invoke( + app, + [ + "--config-dir", + str(tmp_config_dir), + "kai", + "history", + "--project", + "prod", + ], + ) + + assert result.exit_code == 0, f"Exit code {result.exit_code}: {result.output}" + assert "No chat history" in result.output + + def test_kai_history_help(self) -> None: + """kai history --help shows usage information.""" + result = runner.invoke(app, ["kai", "history", "--help"]) + + assert result.exit_code == 0 + assert "List recent Kai chat sessions" in result.output + assert "--project" in result.output + assert "--limit" in result.output + + +class TestKaiConfigError: + """Tests for ConfigError handling across kai commands.""" + + def test_kai_ping_config_error(self, tmp_config_dir: Path) -> None: + """kai ping with ConfigError returns exit code 5.""" + setup_single_project(tmp_config_dir) + + from keboola_agent_cli.errors import ConfigError + + mock_service = MagicMock(spec=KaiService) + mock_service.resolve_alias.side_effect = ConfigError("Project 'unknown' not found.") + + with pytest.MonkeyPatch.context() as mp: + mp.setattr( + "keboola_agent_cli.commands.kai.get_service", + lambda ctx, name: mock_service, + ) + + result = runner.invoke( + app, + [ + "--json", + "--config-dir", + str(tmp_config_dir), + "kai", + "ping", + "--project", + "unknown", + ], + ) + + assert result.exit_code == 5 + output = json.loads(result.output) + assert output["status"] == "error" + assert "CONFIG_ERROR" in output["error"]["code"] diff --git a/tests/test_kai_service.py b/tests/test_kai_service.py new file mode 100644 index 00000000..e11422b9 --- /dev/null +++ b/tests/test_kai_service.py @@ -0,0 +1,339 @@ +"""Tests for KaiService — Keboola AI Assistant business logic. + +Tests the sync wrapper methods (ping, ask, chat_message, get_history) +with mocked KaiClient and feature-flag detection. +""" + +from datetime import UTC, datetime +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from helpers import setup_single_project +from keboola_agent_cli.errors import ConfigError, KeboolaApiError +from keboola_agent_cli.models import TokenVerifyResponse +from keboola_agent_cli.services.kai_service import KaiService + + +def _make_kai_service(tmp_config_dir: Path, features: list[str] | None = None): + """Create a KaiService with a mock client that returns given features.""" + store = setup_single_project(tmp_config_dir) + mock_client = MagicMock() + mock_client.verify_token.return_value = TokenVerifyResponse( + token_id="t-123", + token_description="test token", + project_id=258, + project_name="Production", + owner_name="Production", + features=features or [], + ) + mock_client.close.return_value = None + + service = KaiService( + config_store=store, + client_factory=lambda url, token: mock_client, + ) + return service, mock_client + + +class TestKaiServicePing: + """Tests for KaiService.ping().""" + + def test_ping_success(self, tmp_config_dir: Path) -> None: + """ping returns server health info when Kai is enabled.""" + service, _ = _make_kai_service(tmp_config_dir, features=["agent-chat"]) + + mock_ping_resp = MagicMock() + mock_ping_resp.timestamp = datetime(2025, 1, 15, 10, 30, 0, tzinfo=UTC) + + mock_info_resp = MagicMock() + mock_info_resp.app_name = "kai-api" + mock_info_resp.app_version = "1.2.3" + mock_info_resp.server_version = "2.0.0" + mock_info_resp.connected_mcp = {"status": "connected"} + + mock_kai_client = AsyncMock() + mock_kai_client.ping.return_value = mock_ping_resp + mock_kai_client.info.return_value = mock_info_resp + mock_kai_client.__aenter__ = AsyncMock(return_value=mock_kai_client) + mock_kai_client.__aexit__ = AsyncMock(return_value=False) + + with patch.object(service, "_create_kai_client", return_value=mock_kai_client): + result = service.ping("prod") + + assert result["project_alias"] == "prod" + assert result["timestamp"] == "2025-01-15T10:30:00+00:00" + assert result["app_name"] == "kai-api" + assert result["app_version"] == "1.2.3" + assert result["server_version"] == "2.0.0" + assert result["mcp_status"] == "connected" + + def test_ping_kai_not_enabled(self, tmp_config_dir: Path) -> None: + """ping raises KeboolaApiError when agent-chat feature flag is missing.""" + service, _ = _make_kai_service(tmp_config_dir, features=[]) + + with pytest.raises(KeboolaApiError) as exc_info: + service.ping("prod") + + assert exc_info.value.error_code == "KAI_NOT_ENABLED" + assert "Kai is not enabled" in exc_info.value.message + + def test_ping_mcp_status_unknown(self, tmp_config_dir: Path) -> None: + """ping returns 'unknown' when connected_mcp is not a dict.""" + service, _ = _make_kai_service(tmp_config_dir, features=["agent-chat"]) + + mock_ping_resp = MagicMock() + mock_ping_resp.timestamp = datetime(2025, 1, 15, 10, 30, 0, tzinfo=UTC) + + mock_info_resp = MagicMock() + mock_info_resp.app_name = "kai-api" + mock_info_resp.app_version = "1.2.3" + mock_info_resp.server_version = "2.0.0" + mock_info_resp.connected_mcp = "not-a-dict" + + mock_kai_client = AsyncMock() + mock_kai_client.ping.return_value = mock_ping_resp + mock_kai_client.info.return_value = mock_info_resp + mock_kai_client.__aenter__ = AsyncMock(return_value=mock_kai_client) + mock_kai_client.__aexit__ = AsyncMock(return_value=False) + + with patch.object(service, "_create_kai_client", return_value=mock_kai_client): + result = service.ping("prod") + + assert result["mcp_status"] == "unknown" + + +class TestKaiServiceAsk: + """Tests for KaiService.ask().""" + + def test_ask_success(self, tmp_config_dir: Path) -> None: + """ask returns chat_id and response text.""" + service, _ = _make_kai_service(tmp_config_dir, features=["agent-chat"]) + + mock_kai_client = AsyncMock() + mock_kai_client.chat.return_value = ("chat-abc-123", "The answer is 42.") + mock_kai_client.__aenter__ = AsyncMock(return_value=mock_kai_client) + mock_kai_client.__aexit__ = AsyncMock(return_value=False) + + with patch.object(service, "_create_kai_client", return_value=mock_kai_client): + result = service.ask("prod", "What is the answer?") + + assert result["project_alias"] == "prod" + assert result["chat_id"] == "chat-abc-123" + assert result["response"] == "The answer is 42." + + def test_ask_api_error(self, tmp_config_dir: Path) -> None: + """ask wraps KaiError into KeboolaApiError.""" + from kai_client import KaiError + + service, _ = _make_kai_service(tmp_config_dir, features=["agent-chat"]) + + mock_kai_client = AsyncMock() + mock_kai_client.chat.side_effect = KaiError(message="Service unavailable") + mock_kai_client.__aenter__ = AsyncMock(return_value=mock_kai_client) + mock_kai_client.__aexit__ = AsyncMock(return_value=False) + + with ( + patch.object(service, "_create_kai_client", return_value=mock_kai_client), + pytest.raises(KeboolaApiError) as exc_info, + ): + service.ask("prod", "test question") + + assert exc_info.value.error_code == "KAI_ERROR" + assert "Kai ask failed" in exc_info.value.message + assert "Service unavailable" in exc_info.value.message + + def test_ask_kai_not_enabled(self, tmp_config_dir: Path) -> None: + """ask raises KAI_NOT_ENABLED when feature flag is missing.""" + service, _ = _make_kai_service(tmp_config_dir, features=[]) + + with pytest.raises(KeboolaApiError) as exc_info: + service.ask("prod", "some question") + + assert exc_info.value.error_code == "KAI_NOT_ENABLED" + + +class TestKaiServiceChat: + """Tests for KaiService.chat_message().""" + + def test_chat_new_session(self, tmp_config_dir: Path) -> None: + """chat_message without chat_id creates a new session.""" + service, _ = _make_kai_service(tmp_config_dir, features=["agent-chat"]) + + # Create a mock event with type="text" and text attribute + mock_event = MagicMock() + mock_event.type = "text" + mock_event.text = "Hello from Kai!" + + mock_kai_client = AsyncMock() + # new_chat_id() is called without await, so use MagicMock for it + mock_kai_client.new_chat_id = MagicMock(return_value="new-chat-id-456") + + # send_message returns an async iterator + async def mock_send_message(cid, msg): + yield mock_event + + mock_kai_client.send_message = mock_send_message + mock_kai_client.__aenter__ = AsyncMock(return_value=mock_kai_client) + mock_kai_client.__aexit__ = AsyncMock(return_value=False) + + with patch.object(service, "_create_kai_client", return_value=mock_kai_client): + result = service.chat_message("prod", "Hello!") + + assert result["project_alias"] == "prod" + assert result["chat_id"] == "new-chat-id-456" + assert result["response"] == "Hello from Kai!" + + def test_chat_continue(self, tmp_config_dir: Path) -> None: + """chat_message with chat_id continues an existing session.""" + service, _ = _make_kai_service(tmp_config_dir, features=["agent-chat"]) + + mock_event1 = MagicMock() + mock_event1.type = "text" + mock_event1.text = "Part one. " + + mock_event2 = MagicMock() + mock_event2.type = "text" + mock_event2.text = "Part two." + + # Non-text event should be skipped + mock_event_other = MagicMock() + mock_event_other.type = "tool_call" + + mock_kai_client = AsyncMock() + + async def mock_send_message(cid, msg): + yield mock_event1 + yield mock_event_other + yield mock_event2 + + mock_kai_client.send_message = mock_send_message + mock_kai_client.__aenter__ = AsyncMock(return_value=mock_kai_client) + mock_kai_client.__aexit__ = AsyncMock(return_value=False) + + with patch.object(service, "_create_kai_client", return_value=mock_kai_client): + result = service.chat_message("prod", "Continue please", chat_id="existing-chat-789") + + assert result["chat_id"] == "existing-chat-789" + assert result["response"] == "Part one. Part two." + + def test_chat_kai_error(self, tmp_config_dir: Path) -> None: + """chat_message wraps KaiError into KeboolaApiError.""" + from kai_client import KaiError + + service, _ = _make_kai_service(tmp_config_dir, features=["agent-chat"]) + + mock_kai_client = AsyncMock() + mock_kai_client.new_chat_id.return_value = "chat-err" + + async def mock_send_message(cid, msg): + raise KaiError(message="Chat session expired") + yield # needed to make this an async generator + + mock_kai_client.send_message = mock_send_message + mock_kai_client.__aenter__ = AsyncMock(return_value=mock_kai_client) + mock_kai_client.__aexit__ = AsyncMock(return_value=False) + + with ( + patch.object(service, "_create_kai_client", return_value=mock_kai_client), + pytest.raises(KeboolaApiError) as exc_info, + ): + service.chat_message("prod", "test") + + assert exc_info.value.error_code == "KAI_ERROR" + assert "Kai chat failed" in exc_info.value.message + + +class TestKaiServiceHistory: + """Tests for KaiService.get_history().""" + + def test_history_success(self, tmp_config_dir: Path) -> None: + """get_history returns a list of chat summaries.""" + service, _ = _make_kai_service(tmp_config_dir, features=["agent-chat"]) + + chat1 = MagicMock() + chat1.id = "chat-aaa" + chat1.title = "First chat" + chat1.created_at = datetime(2025, 1, 10, 8, 0, 0, tzinfo=UTC) + chat1.visibility = "private" + + chat2 = MagicMock() + chat2.id = "chat-bbb" + chat2.title = None # untitled + chat2.created_at = None + chat2.visibility = "public" + + mock_history = MagicMock() + mock_history.chats = [chat1, chat2] + mock_history.has_more = True + + mock_kai_client = AsyncMock() + mock_kai_client.get_history.return_value = mock_history + mock_kai_client.__aenter__ = AsyncMock(return_value=mock_kai_client) + mock_kai_client.__aexit__ = AsyncMock(return_value=False) + + with patch.object(service, "_create_kai_client", return_value=mock_kai_client): + result = service.get_history("prod", limit=5) + + assert result["project_alias"] == "prod" + assert result["has_more"] is True + assert len(result["chats"]) == 2 + + assert result["chats"][0]["id"] == "chat-aaa" + assert result["chats"][0]["title"] == "First chat" + assert result["chats"][0]["created_at"] == "2025-01-10T08:00:00+00:00" + assert result["chats"][0]["visibility"] == "private" + + assert result["chats"][1]["id"] == "chat-bbb" + assert result["chats"][1]["title"] == "(untitled)" + assert result["chats"][1]["created_at"] is None + assert result["chats"][1]["visibility"] == "public" + + def test_history_empty(self, tmp_config_dir: Path) -> None: + """get_history returns empty list when no chats exist.""" + service, _ = _make_kai_service(tmp_config_dir, features=["agent-chat"]) + + mock_history = MagicMock() + mock_history.chats = [] + mock_history.has_more = False + + mock_kai_client = AsyncMock() + mock_kai_client.get_history.return_value = mock_history + mock_kai_client.__aenter__ = AsyncMock(return_value=mock_kai_client) + mock_kai_client.__aexit__ = AsyncMock(return_value=False) + + with patch.object(service, "_create_kai_client", return_value=mock_kai_client): + result = service.get_history("prod") + + assert result["chats"] == [] + assert result["has_more"] is False + + def test_history_kai_not_enabled(self, tmp_config_dir: Path) -> None: + """get_history raises KAI_NOT_ENABLED when feature flag is missing.""" + service, _ = _make_kai_service(tmp_config_dir, features=[]) + + with pytest.raises(KeboolaApiError) as exc_info: + service.get_history("prod") + + assert exc_info.value.error_code == "KAI_NOT_ENABLED" + + +class TestKaiServiceResolveAlias: + """Tests for KaiService.resolve_alias().""" + + def test_resolve_explicit_alias(self, tmp_config_dir: Path) -> None: + """resolve_alias with explicit alias validates and returns it.""" + service, _ = _make_kai_service(tmp_config_dir) + assert service.resolve_alias("prod") == "prod" + + def test_resolve_default_alias(self, tmp_config_dir: Path) -> None: + """resolve_alias with None returns the first (default) project.""" + service, _ = _make_kai_service(tmp_config_dir) + assert service.resolve_alias(None) == "prod" + + def test_resolve_unknown_alias(self, tmp_config_dir: Path) -> None: + """resolve_alias raises ConfigError for unknown alias.""" + service, _ = _make_kai_service(tmp_config_dir) + with pytest.raises(ConfigError): + service.resolve_alias("nonexistent") diff --git a/tests/test_storage_delete.py b/tests/test_storage_delete.py index 6d37fcd3..abbc5fb6 100644 --- a/tests/test_storage_delete.py +++ b/tests/test_storage_delete.py @@ -396,6 +396,229 @@ def test_delete_bucket_linked_blocked(self, tmp_path: Path) -> None: assert result.exit_code == 1 +# --------------------------------------------------------------------------- +# delete-column: service layer +# --------------------------------------------------------------------------- + + +class TestDeleteColumnsService: + """Tests for StorageService.delete_columns().""" + + def test_single_column_success(self, tmp_path: Path) -> None: + store = _make_store(tmp_path) + mock_client = MagicMock() + mock_client.delete_column.return_value = None + service = _make_service(store, mock_client) + + result = service.delete_columns(alias="test", table_id="in.c-data.users", columns=["age"]) + + assert result["deleted"] == ["age"] + assert result["failed"] == [] + assert result["dry_run"] is False + assert result["table_id"] == "in.c-data.users" + mock_client.delete_column.assert_called_once_with("in.c-data.users", "age", branch_id=None) + + def test_batch_multiple_columns(self, tmp_path: Path) -> None: + store = _make_store(tmp_path) + mock_client = MagicMock() + mock_client.delete_column.return_value = None + service = _make_service(store, mock_client) + + result = service.delete_columns( + alias="test", table_id="in.c-data.users", columns=["age", "email"] + ) + + assert result["deleted"] == ["age", "email"] + assert result["failed"] == [] + assert mock_client.delete_column.call_count == 2 + + def test_batch_partial_failure(self, tmp_path: Path) -> None: + store = _make_store(tmp_path) + mock_client = MagicMock() + mock_client.delete_column.side_effect = [ + None, + KeboolaApiError("Column not found", status_code=404, error_code="NOT_FOUND"), + ] + service = _make_service(store, mock_client) + + result = service.delete_columns( + alias="test", table_id="in.c-data.users", columns=["age", "missing"] + ) + + assert result["deleted"] == ["age"] + assert len(result["failed"]) == 1 + assert result["failed"][0]["column"] == "missing" + + def test_dry_run(self, tmp_path: Path) -> None: + store = _make_store(tmp_path) + mock_client = MagicMock() + service = _make_service(store, mock_client) + + result = service.delete_columns( + alias="test", table_id="in.c-data.users", columns=["age"], dry_run=True + ) + + assert result["dry_run"] is True + assert result["would_delete"] == ["age"] + assert result["table_id"] == "in.c-data.users" + mock_client.delete_column.assert_not_called() + + def test_unknown_project(self, tmp_path: Path) -> None: + from keboola_agent_cli.errors import ConfigError + + store = _make_store(tmp_path) + mock_client = MagicMock() + service = _make_service(store, mock_client) + + with pytest.raises(ConfigError): + service.delete_columns(alias="nonexistent", table_id="in.c-data.t", columns=["col"]) + + +# --------------------------------------------------------------------------- +# delete-column: CLI layer +# --------------------------------------------------------------------------- + + +class TestDeleteColumnCLI: + """CLI tests for `kbagent storage delete-column`.""" + + def test_delete_column_json(self, tmp_path: Path) -> None: + store = _make_store(tmp_path) + with ( + patch("keboola_agent_cli.cli.ConfigStore") as MockStore, + patch("keboola_agent_cli.cli.StorageService") as MockSvc, + ): + MockStore.return_value = store + svc = MockSvc.return_value + svc.delete_columns.return_value = { + "deleted": ["age"], + "failed": [], + "dry_run": False, + "project_alias": "test", + "table_id": "in.c-data.users", + } + result = runner.invoke( + app, + [ + "--json", + "storage", + "delete-column", + "--project", + "test", + "--table-id", + "in.c-data.users", + "--column", + "age", + "--yes", + ], + ) + assert result.exit_code == 0 + data = json.loads(result.output)["data"] + assert data["deleted"] == ["age"] + assert data["table_id"] == "in.c-data.users" + + def test_delete_column_dry_run(self, tmp_path: Path) -> None: + store = _make_store(tmp_path) + with ( + patch("keboola_agent_cli.cli.ConfigStore") as MockStore, + patch("keboola_agent_cli.cli.StorageService") as MockSvc, + ): + MockStore.return_value = store + svc = MockSvc.return_value + svc.delete_columns.return_value = { + "deleted": [], + "failed": [], + "would_delete": ["age"], + "dry_run": True, + "project_alias": "test", + "table_id": "in.c-data.users", + } + result = runner.invoke( + app, + [ + "--json", + "storage", + "delete-column", + "--project", + "test", + "--table-id", + "in.c-data.users", + "--column", + "age", + "--dry-run", + ], + ) + assert result.exit_code == 0 + + def test_delete_column_multiple(self, tmp_path: Path) -> None: + store = _make_store(tmp_path) + with ( + patch("keboola_agent_cli.cli.ConfigStore") as MockStore, + patch("keboola_agent_cli.cli.StorageService") as MockSvc, + ): + MockStore.return_value = store + svc = MockSvc.return_value + svc.delete_columns.return_value = { + "deleted": ["age", "email"], + "failed": [], + "dry_run": False, + "project_alias": "test", + "table_id": "in.c-data.users", + } + result = runner.invoke( + app, + [ + "--json", + "storage", + "delete-column", + "--project", + "test", + "--table-id", + "in.c-data.users", + "--column", + "age", + "--column", + "email", + "--yes", + ], + ) + assert result.exit_code == 0 + data = json.loads(result.output)["data"] + assert data["deleted"] == ["age", "email"] + + def test_delete_column_exit_1_on_failure(self, tmp_path: Path) -> None: + store = _make_store(tmp_path) + with ( + patch("keboola_agent_cli.cli.ConfigStore") as MockStore, + patch("keboola_agent_cli.cli.StorageService") as MockSvc, + ): + MockStore.return_value = store + svc = MockSvc.return_value + svc.delete_columns.return_value = { + "deleted": [], + "failed": [{"column": "missing", "error": "not found"}], + "dry_run": False, + "project_alias": "test", + "table_id": "in.c-data.users", + } + result = runner.invoke( + app, + [ + "--json", + "storage", + "delete-column", + "--project", + "test", + "--table-id", + "in.c-data.users", + "--column", + "missing", + "--yes", + ], + ) + assert result.exit_code == 1 + + # --------------------------------------------------------------------------- # Branch support tests # --------------------------------------------------------------------------- @@ -619,3 +842,58 @@ def test_cli_branch_flag(self, tmp_path: Path) -> None: assert result.exit_code == 0 call_kwargs = svc.list_tables.call_args.kwargs assert call_kwargs["branch_id"] == 30 + + +class TestDeleteColumnBranch: + """Tests for --branch support in delete-column.""" + + def test_service_passes_branch_id(self, tmp_path: Path) -> None: + """delete_columns passes branch_id to client.delete_column.""" + store = _make_store(tmp_path) + mock_client = MagicMock() + mock_client.delete_column.return_value = None + service = _make_service(store, mock_client) + + result = service.delete_columns( + alias="test", table_id="in.c-data.users", columns=["age"], branch_id=42 + ) + + assert result["deleted"] == ["age"] + mock_client.delete_column.assert_called_once_with("in.c-data.users", "age", branch_id=42) + + def test_cli_branch_flag_json(self, tmp_path: Path) -> None: + """storage delete-column --branch 42 passes branch_id to service.""" + store = _make_store(tmp_path) + with ( + patch("keboola_agent_cli.cli.ConfigStore") as MockStore, + patch("keboola_agent_cli.cli.StorageService") as MockSvc, + ): + MockStore.return_value = store + svc = MockSvc.return_value + svc.delete_columns.return_value = { + "deleted": ["age"], + "failed": [], + "dry_run": False, + "project_alias": "test", + "table_id": "in.c-data.users", + } + result = runner.invoke( + app, + [ + "--json", + "storage", + "delete-column", + "--project", + "test", + "--table-id", + "in.c-data.users", + "--column", + "age", + "--branch", + "42", + "--yes", + ], + ) + assert result.exit_code == 0 + call_kwargs = svc.delete_columns.call_args.kwargs + assert call_kwargs["branch_id"] == 42 diff --git a/tests/test_storage_files.py b/tests/test_storage_files.py index 845490c5..c7a818de 100644 --- a/tests/test_storage_files.py +++ b/tests/test_storage_files.py @@ -449,8 +449,8 @@ def test_unload_with_tags(self, tmp_path: Path) -> None: ) assert mock_client.tag_file.call_count == 2 - mock_client.tag_file.assert_any_call(99, "export") - mock_client.tag_file.assert_any_call(99, "daily") + mock_client.tag_file.assert_any_call(99, "export", branch_id=None) + mock_client.tag_file.assert_any_call(99, "daily", branch_id=None) def test_unload_with_download(self, tmp_path: Path) -> None: store = _make_store(tmp_path) @@ -1071,3 +1071,121 @@ def test_gigabytes(self) -> None: from keboola_agent_cli.commands.storage import _format_file_size assert _format_file_size(2 * 1024 * 1024 * 1024) == "2.00 GB" + + +# ------------------------------------------------------------------ +# Branch-scoped file operations (issue #161) +# ------------------------------------------------------------------ + + +class TestClientGetFileInfoBranch: + """Verify get_file_info uses branch-scoped URL prefix.""" + + def test_get_file_info_without_branch(self, httpx_mock) -> None: + from keboola_agent_cli.client import KeboolaClient + + httpx_mock.add_response(json=SAMPLE_FILE) + client = KeboolaClient( + stack_url="https://connection.keboola.com", + token=TEST_TOKEN, + ) + client.get_file_info(12345) + request = httpx_mock.get_requests()[0] + assert "/v2/storage/files/12345" in str(request.url) + assert "/branch/" not in str(request.url) + client.close() + + def test_get_file_info_with_branch(self, httpx_mock) -> None: + from keboola_agent_cli.client import KeboolaClient + + httpx_mock.add_response(json=SAMPLE_FILE) + client = KeboolaClient( + stack_url="https://connection.keboola.com", + token=TEST_TOKEN, + ) + client.get_file_info(12345, branch_id=42) + request = httpx_mock.get_requests()[0] + assert "/v2/storage/branch/42/files/12345" in str(request.url) + client.close() + + +class TestClientDeleteFileBranch: + """Verify delete_file uses branch-scoped URL prefix.""" + + def test_delete_file_with_branch(self, httpx_mock) -> None: + from keboola_agent_cli.client import KeboolaClient + + httpx_mock.add_response(method="DELETE", status_code=204) + client = KeboolaClient( + stack_url="https://connection.keboola.com", + token=TEST_TOKEN, + ) + client.delete_file(12345, branch_id=42) + request = httpx_mock.get_requests()[0] + assert "/v2/storage/branch/42/files/12345" in str(request.url) + client.close() + + +class TestClientTagFileBranch: + """Verify tag_file and untag_file use branch-scoped URL prefix.""" + + def test_tag_file_with_branch(self, httpx_mock) -> None: + from keboola_agent_cli.client import KeboolaClient + + httpx_mock.add_response(method="POST", status_code=201) + client = KeboolaClient( + stack_url="https://connection.keboola.com", + token=TEST_TOKEN, + ) + client.tag_file(12345, "my-tag", branch_id=42) + request = httpx_mock.get_requests()[0] + assert "/v2/storage/branch/42/files/12345/tags" in str(request.url) + client.close() + + def test_untag_file_with_branch(self, httpx_mock) -> None: + from keboola_agent_cli.client import KeboolaClient + + httpx_mock.add_response(method="DELETE", status_code=204) + client = KeboolaClient( + stack_url="https://connection.keboola.com", + token=TEST_TOKEN, + ) + client.untag_file(12345, "old-tag", branch_id=42) + request = httpx_mock.get_requests()[0] + assert "/v2/storage/branch/42/files/12345/tags/old-tag" in str(request.url) + client.close() + + +class TestUnloadTableToFileBranchService: + """Verify unload_table_to_file passes branch_id to file operations.""" + + def test_unload_with_branch_passes_branch_to_file_ops(self, tmp_path: Path) -> None: + store = _make_store(tmp_path) + mock_client = MagicMock() + mock_client.export_table_async.return_value = { + "results": {"file": {"id": 99}}, + } + mock_client.get_file_info.return_value = { + "id": 99, + "name": "export.csv.gz", + "sizeBytes": 2048, + "isSliced": False, + "tags": ["export"], + } + service = _make_service(store, mock_client) + + service.unload_table_to_file( + alias="test", + table_id="in.c-data.users", + tags=["export"], + branch_id=33, + ) + + mock_client.export_table_async.assert_called_once_with( + table_id="in.c-data.users", + columns=None, + limit=None, + branch_id=33, + ) + mock_client.tag_file.assert_called_once_with(99, "export", branch_id=33) + mock_client.get_file_info.assert_called_once_with(99, branch_id=33) diff --git a/tests/test_storage_write.py b/tests/test_storage_write.py index 60b2b498..a1ca3d7e 100644 --- a/tests/test_storage_write.py +++ b/tests/test_storage_write.py @@ -1255,7 +1255,7 @@ def _fake_download(url, path): limit=None, branch_id=None, ) - mock_client.get_file_info.assert_called_once_with(42) + mock_client.get_file_info.assert_called_once_with(42, branch_id=None) mock_client.close.assert_called_once() def test_with_columns_and_limit(self, tmp_path: Path) -> None: @@ -1432,6 +1432,8 @@ def _fake_download(url, path): limit=None, branch_id=42, ) + # Issue #161: get_file_info must also receive branch_id + mock_client.get_file_info.assert_called_once_with(7, branch_id=42) # --------------------------------------------------------------------------- diff --git a/tests/test_sync_service.py b/tests/test_sync_service.py index 331f1631..ec8230d4 100644 --- a/tests/test_sync_service.py +++ b/tests/test_sync_service.py @@ -614,6 +614,87 @@ def test_pull_dry_run_preserves_directories(self, tmp_config_dir: Path, tmp_path # Directory must still exist after dry-run assert snowflake_dir.exists(), "Dry-run should NOT delete directories" + def test_pull_auto_renames_config_on_remote_name_change( + self, tmp_config_dir: Path, tmp_path: Path + ) -> None: + """pull auto-renames local directory when config name changed on remote.""" + project_root = tmp_path / "project" + project_root.mkdir() + + store = self._init_project(tmp_config_dir, project_root) + + # First pull: download config with original name "My HTTP Extractor" + pull_client = _make_sync_mock_client( + components_response=SAMPLE_COMPONENTS_NO_ROWS, + ) + svc = SyncService( + config_store=store, + client_factory=lambda url, token: pull_client, + ) + result1 = svc.pull(alias="prod", project_root=project_root) + assert result1["configs_pulled"] == 1 + + # Verify original directory exists at expected path + old_dir = project_root / "main" / "extractor" / "keboola.ex-http" / "my-http-extractor" + assert old_dir.exists(), "Original config directory should exist after first pull" + assert (old_dir / CONFIG_FILENAME).exists() + + # Verify manifest tracks the original path + manifest_before = load_manifest(project_root) + assert len(manifest_before.configurations) == 1 + assert ( + manifest_before.configurations[0].path == "extractor/keboola.ex-http/my-http-extractor" + ) + + # Second pull: same config ID but with renamed name + renamed_components = [ + { + "id": "keboola.ex-http", + "type": "extractor", + "configurations": [ + { + "id": "cfg-001", + "name": "Renamed HTTP Extractor", + "description": "Fetches data", + "configuration": { + "parameters": {"baseUrl": "https://api.example.com"}, + }, + "rows": [], + } + ], + }, + ] + pull_client2 = _make_sync_mock_client(components_response=renamed_components) + svc2 = SyncService( + config_store=store, + client_factory=lambda url, token: pull_client2, + ) + result2 = svc2.pull(alias="prod", project_root=project_root) + + # Verify the old directory no longer exists + assert not old_dir.exists(), "Old config directory should be gone after rename" + + # Verify the new directory exists + new_dir = project_root / "main" / "extractor" / "keboola.ex-http" / "renamed-http-extractor" + assert new_dir.exists(), "Renamed config directory should exist" + assert (new_dir / CONFIG_FILENAME).exists() + + # Verify manifest path was updated + manifest_after = load_manifest(project_root) + assert len(manifest_after.configurations) == 1 + assert ( + manifest_after.configurations[0].path + == "extractor/keboola.ex-http/renamed-http-extractor" + ) + + # Verify pull_details contains a "renamed" action + renamed_details = [d for d in result2["details"] if d["action"] == "renamed"] + assert len(renamed_details) == 1 + assert renamed_details[0]["component_id"] == "keboola.ex-http" + assert renamed_details[0]["config_name"] == "Renamed HTTP Extractor" + assert renamed_details[0]["old_path"] == "extractor/keboola.ex-http/my-http-extractor" + assert renamed_details[0]["path"] == "extractor/keboola.ex-http/renamed-http-extractor" + # =================================================================== # status tests diff --git a/uv.lock b/uv.lock index c503a0b0..a019bdda 100644 --- a/uv.lock +++ b/uv.lock @@ -406,13 +406,29 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/41/45/1a4ed80516f02155c51f51e8cedb3c1902296743db0bbc66608a0db2814f/jsonschema_specifications-2025.9.1-py3-none-any.whl", hash = "sha256:98802fee3a11ee76ecaca44429fda8a41bff98b00a0f2838151b113f210cc6fe", size = 18437, upload-time = "2025-09-08T01:34:57.871Z" }, ] +[[package]] +name = "kai-client" +version = "0.11.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click" }, + { name = "httpx" }, + { name = "pydantic" }, + { name = "python-dotenv" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/3d/ab/67459462a6fa561d8264fc2a615cf96a091755bc9328ff13c57c74778a59/kai_client-0.11.0.tar.gz", hash = "sha256:f3d42c96a8f92c56af784570d3b07c10701ffec29ce9cf0a4b8a4af374798d7a", size = 101197, upload-time = "2026-02-21T14:10:05.608Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/64/58/271197c664cee82c6775830945f46776ef74dc70f2e711b4087115ae8d88/kai_client-0.11.0-py3-none-any.whl", hash = "sha256:cbd173a805888cf2d46a841be1b8ae193a0f1237eafda0a8b0e589f8f006b81e", size = 26690, upload-time = "2026-02-21T14:10:04.144Z" }, +] + [[package]] name = "keboola-agent-cli" -version = "0.18.5" +version = "0.19.0" source = { editable = "." } dependencies = [ { name = "httpx" }, { name = "jsonschema" }, + { name = "kai-client" }, { name = "mcp" }, { name = "packaging" }, { name = "platformdirs" }, @@ -437,6 +453,7 @@ dev = [ requires-dist = [ { name = "httpx", specifier = ">=0.27" }, { name = "jsonschema", specifier = ">=4.20" }, + { name = "kai-client", specifier = ">=0.11.0" }, { name = "mcp", specifier = ">=1.0.0,<2.0.0" }, { name = "packaging", specifier = ">=23" }, { name = "platformdirs", specifier = ">=4" },