diff --git a/README.md b/README.md index 191b9a2..6e0ed15 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,4 @@ -# drs — Developer CLI for Dremio Cloud +# dremio — Developer CLI for Dremio Cloud A command-line tool for working with Dremio Cloud. Run SQL queries, browse the catalog, inspect table schemas, manage reflections, monitor jobs, and audit access — from your terminal or any automation pipeline. @@ -16,7 +16,7 @@ Dremio Cloud has a powerful REST API and rich system tables, but no official CLI - Scripting catalog operations means hand-rolling `curl` commands with auth headers - AI agents (Claude, GPT, etc.) need structured tool interfaces, not raw HTTP -`drs` wraps all of this into a single binary with consistent output formats, input validation, and structured error handling. +`dremio` wraps all of this into a single binary with consistent output formats, input validation, and structured error handling. ## Prerequisites @@ -29,18 +29,26 @@ Dremio Cloud has a powerful REST API and rich system tables, but no official CLI ### 1. Install ```bash -# Recommended — install as a standalone tool +# Clone the repo +git clone https://github.com/dremio/cli.git +cd cli + +# Install as a standalone tool (recommended) uv tool install . # Or with pip pip install . -# Or for development -git clone https://github.com/dremio/cli.git -cd cli +# Or for development (editable install) uv sync ``` +After install, verify the binary is available: + +```bash +dremio --help +``` + ### 2. Configure There are three ways to authenticate, in order of priority: @@ -48,10 +56,10 @@ There are three ways to authenticate, in order of priority: **Option A: CLI flags** (highest priority — override everything) ```bash -drs --token YOUR_PAT --project-id YOUR_PROJECT_ID query run "SELECT 1" +dremio --token YOUR_PAT --project-id YOUR_PROJECT_ID query run "SELECT 1" # EU region -drs --uri https://api.eu.dremio.cloud --token YOUR_PAT --project-id YOUR_PROJECT_ID query run "SELECT 1" +dremio --uri https://api.eu.dremio.cloud --token YOUR_PAT --project-id YOUR_PROJECT_ID query run "SELECT 1" ``` **Option B: Environment variables** @@ -81,7 +89,7 @@ chmod 600 ~/.config/dremioai/config.yaml ### 3. Verify ```bash -drs query run "SELECT 1 AS hello" +dremio query run "SELECT 1 AS hello" ``` If this returns `{"job_id": "...", "state": "COMPLETED", "rowCount": 1, "rows": [{"hello": "1"}]}`, you're set. @@ -92,33 +100,62 @@ If this returns `{"job_id": "...", "state": "COMPLETED", "rowCount": 1, "rows": | Group | Commands | What it does | |-------|----------|--------------| -| `drs query` | `run`, `status`, `cancel` | Execute SQL, check job status, cancel running jobs | -| `drs catalog` | `list`, `get`, `search` | Browse sources/spaces, get entity metadata, full-text search | -| `drs schema` | `describe`, `lineage`, `wiki`, `sample` | Column types, upstream/downstream deps, wiki docs, preview rows | -| `drs reflect` | `list`, `status`, `refresh`, `drop` | List reflections on a dataset, check freshness, trigger refresh | -| `drs jobs` | `list`, `get`, `profile` | Recent jobs with filters, job details, operator-level profiles | -| `drs access` | `grants`, `roles`, `whoami`, `audit` | ACLs on entities, org roles, user permission audit | +| `dremio query` | `run`, `status`, `cancel` | Execute SQL, check job status, cancel running jobs | +| `dremio folder` | `list`, `get`, `create`, `delete`, `grants` | Browse and manage spaces/folders, view ACLs | +| `dremio schema` | `describe`, `lineage`, `sample` | Column types, dependency graph, preview rows | +| `dremio wiki` | `get`, `update` | Read and update wiki documentation on entities | +| `dremio tag` | `get`, `update` | Read and update tags on entities | +| `dremio reflection` | `create`, `list`, `get`, `refresh`, `delete` | Full CRUD for reflections (materialized views) | +| `dremio job` | `list`, `get`, `profile` | Recent jobs with filters, job details, operator-level profiles | +| `dremio engine` | `list`, `get`, `create`, `update`, `delete`, `enable`, `disable` | Full CRUD for Dremio Cloud engines | +| `dremio user` | `list`, `get`, `create`, `delete`, `whoami`, `audit` | Manage org users, check identity, audit permissions | +| `dremio role` | `list`, `get`, `create`, `update`, `delete` | Full CRUD for organization roles | +| `dremio grant` | `get`, `update`, `delete` | Manage grants on projects, engines, org resources | +| `dremio search` | *(top-level)* | Full-text search across all catalog entities | +| `dremio describe` | *(top-level)* | Machine-readable schema for any command | ### Examples ```bash # Run a query and get results as a pretty table -drs query run "SELECT * FROM myspace.orders LIMIT 5" --output pretty +dremio query run "SELECT * FROM myspace.orders LIMIT 5" --output pretty # Search the catalog for anything matching "revenue" -drs catalog search "revenue" +dremio search "revenue" + +# Create a space, then a folder inside it +dremio folder create "Analytics" +dremio folder create Analytics.reports # Describe a table's columns -drs schema describe myspace.analytics.monthly_revenue +dremio schema describe myspace.analytics.monthly_revenue -# Check what reflections exist on a dataset -drs reflect list myspace.orders +# Read and update wiki docs on a table +dremio wiki get myspace.orders +dremio wiki update myspace.orders "Primary orders table. Refreshed daily from Salesforce." + +# Update tags on a table +dremio tag update myspace.orders "pii,finance,daily" + +# Create a raw reflection on a dataset +dremio reflection create myspace.orders --type raw +dremio reflection list myspace.orders + +# Manage engines +dremio engine list +dremio engine create "analytics-engine" --size LARGE +dremio engine disable eng-abc-123 + +# Manage users and roles +dremio user list --output pretty +dremio role create "data-analyst" +dremio grant update projects my-project-id role role-abc "MANAGE_GRANTS,CREATE_TABLE" # Find failed jobs from recent history -drs jobs list --status FAILED --output pretty +dremio job list --status FAILED --output pretty # Audit what roles and permissions a user has -drs access audit rahim.bhojani +dremio user audit rahim.bhojani ``` ### Output formats @@ -137,10 +174,10 @@ Reduce output to just the fields you need with `--fields` / `-f`. Supports dot n ```bash # Only show column names and types -drs schema describe myspace.orders --fields columns.name,columns.type +dremio schema describe myspace.orders --fields columns.name,columns.type # Only show job ID and state -drs jobs list --fields job_id,job_state +dremio job list --fields job_id,job_state ``` This is especially useful for AI agents to keep context windows small. @@ -150,19 +187,38 @@ This is especially useful for AI agents to keep context windows small. Discover parameters for any command programmatically: ```bash -drs describe query.run -drs describe reflect.list +dremio describe query.run +dremio describe reflection.list ``` -Returns a JSON schema with parameter names, types, required/optional, and descriptions. Useful for building automation on top of `drs`. +Returns a JSON schema with parameter names, types, required/optional, and descriptions. Useful for building automation on top of `dremio`. + +## CRUD design principle + +Every Dremio object has consistent CLI commands using standard CRUD verbs (`list`, `get`, `create`, `update`, `delete`): + +| Object | List | Get | Create | Update | Delete | +|--------|------|-----|--------|--------|--------| +| **Spaces/Folders** | `folder list` | `folder get` | `folder create` | — | `folder delete` | +| **Tables/Views** | `folder get` | `schema describe/sample` | `query run` (DDL) | `query run` (DDL) | `folder delete` | +| **Wiki** | — | `wiki get` | `wiki update` | `wiki update` | — | +| **Tags** | — | `tag get` | `tag update` | `tag update` | — | +| **Reflections** | `reflection list` | `reflection get` | `reflection create` | `reflection refresh` | `reflection delete` | +| **Engines** | `engine list` | `engine get` | `engine create` | `engine update` | `engine delete` | +| **Users** | `user list` | `user get` | `user create` | — | `user delete` | +| **Roles** | `role list` | `role get` | `role create` | `role update` | `role delete` | +| **Grants** | `grant get` | `grant get` | `grant update` | `grant update` | `grant delete` | +| **Jobs** | `job list` | `job get/profile` | `query run` | — | `query cancel` | + +`folder create` uses SQL under the hood (`CREATE SPACE` for top-level, `CREATE FOLDER` for nested paths). All other mutations use the REST API. ## How it works ``` -┌─────────────┐ ┌──────────────┐ ┌─────────────────┐ -│ drs CLI │────▶│ client.py │────▶│ Dremio Cloud API │ +┌──────────────┐ ┌──────────────┐ ┌─────────────────┐ +│ dremio CLI │────▶│ client.py │────▶│ Dremio Cloud API │ │ (typer) │ │ (httpx) │ │ (REST + SQL) │ -└─────────────┘ └──────────────┘ └─────────────────┘ +└──────────────┘ └──────────────┘ └─────────────────┘ ``` - **One HTTP layer** — `client.py` is the only file that makes network calls. Every command goes through it. @@ -177,26 +233,34 @@ All endpoints target `https://api.dremio.cloud`. See the [Dremio Cloud API refer | URL pattern | Used by | Docs | |-------------|---------|------| | `POST /v0/projects/{pid}/sql` | `query run` | [SQL](https://docs.dremio.com/dremio-cloud/api/sql) | -| `GET /v0/projects/{pid}/job/{id}` | `query status`, `jobs get` | [Job](https://docs.dremio.com/dremio-cloud/api/job/) | +| `GET /v0/projects/{pid}/job/{id}` | `query status`, `job get` | [Job](https://docs.dremio.com/dremio-cloud/api/job/) | | `GET /v0/projects/{pid}/job/{id}/results` | `query run` (result fetch) | [Job Results](https://docs.dremio.com/dremio-cloud/api/job/job-results/) | | `POST /v0/projects/{pid}/job/{id}/cancel` | `query cancel` | [Job](https://docs.dremio.com/dremio-cloud/api/job/) | -| `GET /v0/projects/{pid}/catalog` | `catalog list` | [Catalog](https://docs.dremio.com/dremio-cloud/api/catalog/) | -| `GET /v0/projects/{pid}/catalog/by-path/{path}` | `catalog get`, `schema describe`, `schema lineage`, `schema wiki`, `access grants` | [Catalog](https://docs.dremio.com/dremio-cloud/api/catalog/) | +| `GET /v0/projects/{pid}/catalog` | `folder list` | [Catalog](https://docs.dremio.com/dremio-cloud/api/catalog/) | +| `GET /v0/projects/{pid}/catalog/by-path/{path}` | `folder get`, `schema describe`, `wiki get`, `tag get`, `folder grants` | [Catalog](https://docs.dremio.com/dremio-cloud/api/catalog/) | +| `DELETE /v0/projects/{pid}/catalog/{id}` | `folder delete` | [Catalog](https://docs.dremio.com/dremio-cloud/api/catalog/) | | `GET /v0/projects/{pid}/catalog/{id}/graph` | `schema lineage` | [Lineage](https://docs.dremio.com/dremio-cloud/api/catalog/lineage) | -| `GET /v0/projects/{pid}/catalog/{id}/collaboration/wiki` | `schema wiki` | [Wiki](https://docs.dremio.com/dremio-cloud/api/catalog/wiki) | -| `GET /v0/projects/{pid}/catalog/{id}/collaboration/tag` | `schema wiki` | [Tag](https://docs.dremio.com/dremio-cloud/api/catalog/tag) | -| `POST /v0/projects/{pid}/search` | `catalog search` | [Search](https://docs.dremio.com/dremio-cloud/api/search) | -| `GET /v0/projects/{pid}/reflection/{id}` | `reflect status` | [Reflection](https://docs.dremio.com/dremio-cloud/api/reflection/) | -| `POST /v0/projects/{pid}/reflection/{id}/refresh` | `reflect refresh` | [Reflection](https://docs.dremio.com/dremio-cloud/api/reflection/) | -| `DELETE /v0/projects/{pid}/reflection/{id}` | `reflect drop` | [Reflection](https://docs.dremio.com/dremio-cloud/api/reflection/) | -| `GET /v1/users`, `GET /v1/users/name/{name}` | `access whoami`, `access audit` | — | -| `GET /v1/roles` | `access roles` | — | - -Commands that query system tables (`jobs list`, `jobs profile`, `reflect list`, `schema sample`) use `POST /v0/projects/{pid}/sql` to submit SQL against `sys.project.*` tables. +| `GET/PUT /v0/projects/{pid}/catalog/{id}/collaboration/wiki` | `wiki get`, `wiki update` | [Wiki](https://docs.dremio.com/dremio-cloud/api/catalog/wiki) | +| `GET/PUT /v0/projects/{pid}/catalog/{id}/collaboration/tag` | `tag get`, `tag update` | [Tag](https://docs.dremio.com/dremio-cloud/api/catalog/tag) | +| `POST /v0/projects/{pid}/search` | `search` | [Search](https://docs.dremio.com/dremio-cloud/api/search) | +| `POST /v0/projects/{pid}/reflection` | `reflection create` | [Reflection](https://docs.dremio.com/dremio-cloud/api/reflection/) | +| `GET /v0/projects/{pid}/reflection/{id}` | `reflection get` | [Reflection](https://docs.dremio.com/dremio-cloud/api/reflection/) | +| `POST /v0/projects/{pid}/reflection/{id}/refresh` | `reflection refresh` | [Reflection](https://docs.dremio.com/dremio-cloud/api/reflection/) | +| `DELETE /v0/projects/{pid}/reflection/{id}` | `reflection delete` | [Reflection](https://docs.dremio.com/dremio-cloud/api/reflection/) | +| `GET/POST/PUT/DELETE /v0/projects/{pid}/engines[/{id}]` | `engine list/get/create/update/delete` | [Engines](https://docs.dremio.com/dremio-cloud/api/) | +| `PUT /v0/projects/{pid}/engines/{id}/enable\|disable` | `engine enable`, `engine disable` | [Engines](https://docs.dremio.com/dremio-cloud/api/) | +| `GET /v1/users`, `GET /v1/users/name/{name}`, `GET /v1/users/{id}` | `user list/get`, `user whoami/audit` | [Users](https://docs.dremio.com/dremio-cloud/api/) | +| `POST /v1/users/invite` | `user create` | [Users](https://docs.dremio.com/dremio-cloud/api/) | +| `DELETE /v1/users/{id}` | `user delete` | [Users](https://docs.dremio.com/dremio-cloud/api/) | +| `GET /v1/roles[/{id}]`, `GET /v1/roles/name/{name}` | `role list/get` | [Roles](https://docs.dremio.com/dremio-cloud/api/) | +| `POST /v1/roles`, `PUT /v1/roles/{id}`, `DELETE /v1/roles/{id}` | `role create/update/delete` | [Roles](https://docs.dremio.com/dremio-cloud/api/) | +| `GET/PUT/DELETE /v1/{scope}/{id}/grants/{type}/{id}` | `grant get/update/delete` | [Grants](https://docs.dremio.com/dremio-cloud/api/) | + +Commands that query system tables (`job list`, `job profile`, `reflection list`, `schema sample`) use `POST /v0/projects/{pid}/sql` to submit SQL against `sys.project.*` tables. ## Configuration reference -`drs` resolves each setting using the first match (highest priority first): +`dremio` resolves each setting using the first match (highest priority first): | Priority | Token | Project ID | API URI | |----------|-------|------------|---------| @@ -210,20 +274,20 @@ The config file also accepts the legacy `dremio-mcp` format (`token`, `projectId ```bash # Custom config file -drs --config /path/to/my/config.yaml query run "SELECT 1" +dremio --config /path/to/my/config.yaml query run "SELECT 1" # EU region -drs --uri https://api.eu.dremio.cloud query run "SELECT 1" +dremio --uri https://api.eu.dremio.cloud query run "SELECT 1" ``` ## Claude Code Plugin -`drs` ships with a Claude Code plugin that adds Dremio-aware skills to your coding sessions: +`dremio` ships with a Claude Code plugin that adds Dremio-aware skills to your coding sessions: | Skill | What it does | |-------|-------------| | `dremio` | Core reference — SQL dialect, system tables, functions, REST patterns | -| `dremio-setup` | Interactive setup wizard for `drs` | +| `dremio-setup` | Interactive setup wizard for `dremio` | | `dremio-dbt` | dbt-dremio Cloud integration guide and patterns | | `investigate-slow-query` | Walks through job profile analysis and reflection recommendations | | `audit-dataset-access` | Traces grants, role inheritance, and effective permissions | @@ -233,15 +297,15 @@ drs --uri https://api.eu.dremio.cloud query run "SELECT 1" ## For AI agents -`drs` is designed to be agent-friendly: +`dremio` is designed to be agent-friendly: - **Structured JSON output** by default — no parsing needed -- **`drs describe `** lets agents self-discover parameter schemas at runtime +- **`dremio describe `** lets agents self-discover parameter schemas at runtime - **`--fields` filtering** reduces output size to fit context windows - **Input validation** catches hallucinated paths, malformed UUIDs, and injection attempts before they hit the API - **Consistent error format** — all API errors return `{"error": "...", "status_code": N}` rather than raw HTTP tracebacks -If you're building an agent that talks to Dremio, you can either shell out to `drs` commands or import the async functions directly: +If you're building an agent that talks to Dremio, you can either shell out to `dremio` commands or import the async functions directly: ```python from drs.auth import load_config @@ -277,22 +341,27 @@ src/drs/ client.py # The single HTTP layer (all API calls) output.py # JSON / CSV / pretty formatting utils.py # Path parsing, input validation, error handling - introspect.py # Command schema registry for drs describe + introspect.py # Command schema registry for dremio describe commands/ query.py # run, status, cancel - catalog.py # list, get, search - schema.py # describe, lineage, wiki, sample - reflect.py # list, status, refresh, drop - jobs.py # list, get, profile - access.py # grants, roles, whoami, audit + folder.py # list, get, create, delete, grants + schema.py # describe, lineage, sample + wiki.py # get, update + tag.py # get, update + reflection.py # create, list, get, refresh, delete + job.py # list, get, profile + engine.py # list, get, create, update, delete, enable, disable + user.py # list, get, create, delete, whoami, audit + role.py # list, get, create, update, delete + grant.py # get, update, delete ``` ## Related projects | Repo | Relationship | |------|-------------| -| `dremio/dremio-mcp` | Sibling — MCP server for AI agent integration. `drs` focuses on CLI; config format is shared. | -| `dremio/claude-plugins` | Predecessor — skills have been rewritten to use `drs` commands instead of raw curl. | +| `dremio/dremio-mcp` | Sibling — MCP server for AI agent integration. `dremio` focuses on CLI; config format is shared. | +| `dremio/claude-plugins` | Predecessor — skills have been rewritten to use `dremio` commands instead of raw curl. | ## License diff --git a/docs/ARCHITECTURE_REVIEW.md b/docs/ARCHITECTURE_REVIEW.md deleted file mode 100644 index f24c4d5..0000000 --- a/docs/ARCHITECTURE_REVIEW.md +++ /dev/null @@ -1,291 +0,0 @@ -# drs — Dremio Developer CLI: Architecture Review - -**Repo:** `github.com/dremio/cli` -**Binary:** `drs` -**Target:** Dremio Cloud only (PAT + project ID) -**Stack:** Python 3.11+, typer, httpx, FastMCP, pydantic - ---- - -## 1. What is this? - -A single repo that ships three distribution channels for the same Dremio Cloud operations: - -``` - ┌──────────────────────────────────┐ - │ dremio/cli repo │ - │ │ - │ ┌───────┐ ┌───────┐ ┌────────┐ │ - │ │ CLI │ │ MCP │ │ Plugin │ │ - │ │ (drs) │ │Server │ │(Skills)│ │ - │ └───┬───┘ └───┬───┘ └────────┘ │ - │ │ │ │ - │ ▼ ▼ │ - │ ┌─────────────────┐ │ - │ │ Command Layer │ │ - │ │ (async funcs) │ │ - │ └────────┬────────┘ │ - │ │ │ - │ ▼ │ - │ ┌─────────────────┐ │ - │ │ client.py │ │ - │ │ (single HTTP │ │ - │ │ layer — httpx) │ │ - │ └────────┬────────┘ │ - └───────────┼───────────────────────┘ - │ - ┌───────────▼───────────────────────┐ - │ Dremio Cloud REST API │ - │ v0 (SQL/Jobs) │ v3 (Catalog/ │ - │ │ Reflections) │ - │ v1 (Users/ │ │ - │ Roles/RBAC)│ │ - └───────────────────────────────────┘ -``` - -| Channel | What | Who | How | -|---------|------|-----|-----| -| **CLI** (`drs`) | Terminal commands | Developers, scripts, CI | `drs query run "SELECT 1"` | -| **MCP Server** (`drs mcp`) | Stdio tools for AI agents | Claude Desktop, Claude Code, any MCP client | `drs mcp` → 19 tools | -| **Plugin** (skills) | Claude Code workflow skills | Claude Code users | `/plugin install dremio@dremio-cli` → 8 skills | - -**Key constraint:** The CLI commands and MCP tools call the **exact same async functions**. Zero duplicated business logic. `client.py` is the only file that makes HTTP calls. - ---- - -## 2. Why not just dremio-mcp? - -| | `dremio-mcp` (existing) | `dremio/cli` (this) | -|---|---|---| -| CLI layer | None | Full CLI with 18 commands | -| Debugging | Requires MCP client | `drs query run "..."` directly | -| Output formats | JSON only | JSON, CSV, pretty table | -| Scripting | Not possible | Pipe-friendly (`drs jobs list --output csv \| ...`) | -| Skills | Separate repo (`claude-plugins`) | Same repo, references CLI commands | -| Organization | Flat tool list | Grouped by domain (query, catalog, schema, etc.) | - -`dremio-mcp` goes straight from API to MCP tools. `drs` puts a CLI layer in between, making every operation testable, scriptable, and debuggable without an AI agent. - ---- - -## 3. Command Inventory - -### CLI Commands (18) across 6 groups - -| Group | Subcommands | Mechanism | Description | -|-------|-------------|-----------|-------------| -| **query** | `run`, `status`, `cancel` | REST (v0) | Execute SQL, poll for results, cancel jobs | -| **catalog** | `list`, `get`, `search` | REST (v3) | Browse sources/spaces, get entity metadata, full-text search | -| **schema** | `describe`, `lineage`, `wiki`, `sample` | REST (v3) + SQL | Column types, dependency graph, wiki/tags, sample rows | -| **reflect** | `list`, `status`, `refresh`, `drop` | SQL + REST (v3) | List reflections (via sys.reflections), CRUD single reflection | -| **jobs** | `list`, `get`, `profile` | SQL + REST (v0) | Recent jobs (via sys.project.jobs_recent), job details, execution profile | -| **access** | `grants`, `roles`, `whoami`, `audit` | REST (v3, v1) | ACLs, role listing, user info, permission audit | - -### MCP Tools (19) - -Every CLI command plus `drs mcp` itself. Each tool has an LLM-optimized description — this is the #1 quality bar for agent usability. - -### Skills (8) - -| Tier | Skill | What it does | -|------|-------|-------------| -| **Core** | `dremio` | SQL reference, system tables, functions, REST patterns | -| **Setup** | `dremio-setup` | Install wizard: config file, PAT, verify, MCP setup | -| **Integration** | `dremio-dbt` | dbt-dremio profiles, materializations, troubleshooting | -| **Workflow** | `investigate-slow-query` | Job profile → reflection check → optimization recommendations | -| **Workflow** | `audit-dataset-access` | Grants → roles → effective permissions trace | -| **Workflow** | `document-dataset` | Schema + lineage + sample + wiki → markdown doc card | -| **Workflow** | `investigate-data-quality` | Null analysis, duplicates, outliers, freshness checks | -| **Workflow** | `onboard-new-source` | Discover → describe → reflect → verify access | - -Skill naming follows the Google Workspace CLI taxonomy: -- **Tier 1** (API): `dremio` (core reference) -- **Tier 2** (Action): `dremio-setup`, `dremio-dbt` -- **Tier 3** (Workflow): `investigate-*`, `audit-*`, `document-*`, `onboard-*` - ---- - -## 4. SQL vs REST: The Dual-Mechanism Design - -Some Dremio operations are only available via REST, others only via SQL system tables. The CLI abstracts this — users don't need to know which mechanism is used. - -``` -┌─────────────────────────────────────────────────┐ -│ Command Layer │ -│ │ -│ REST-based SQL-based │ -│ ───────── ───────── │ -│ catalog list/get/search jobs list │ -│ schema describe/lineage jobs profile │ -│ schema wiki reflect list (by path) │ -│ reflect status/refresh schema sample │ -│ reflect drop │ -│ access grants/roles │ -│ access whoami/audit │ -│ query run/status/cancel │ -│ │ │ │ -│ ▼ ▼ │ -│ ┌──────────────┐ ┌──────────────────┐ │ -│ │ client.py │ │ run_query() │ │ -│ │ (HTTP calls)│ │ (submits SQL via │ │ -│ │ │ │ client.py) │ │ -│ └──────────────┘ └──────────────────┘ │ -└─────────────────────────────────────────────────┘ -``` - -**Why SQL for some commands?** -- **No REST endpoint** to list reflections by dataset — must query `sys.reflections` -- **No REST endpoint** to list recent jobs with filtering — must query `sys.project.jobs_recent` -- **Sample data** is naturally a SQL operation (`SELECT * FROM ... LIMIT 10`) - -**SQL injection mitigation:** All SQL-interpolated parameters are validated before insertion. Job states are checked against a whitelist; job IDs must match UUID format. - ---- - -## 5. API Coverage - -### Source of truth: `dremio/js-sdk` - -The [dremio/js-sdk](https://github.com/dremio/js-sdk) TypeScript SDK is the canonical, auto-maintained API surface. `scripts/parse_jssdk.py` parses it to extract every endpoint: - -``` -js-sdk endpoints: 79 unique (97 total calls across 15 resources) -Covered by drs: 14 -Not in drs: 65 (engines, AI, projects, scripts, etc.) -drs-only (SQL-based): 5 (jobs list, profiles, reflection list, sample, lineage) -``` - -**Resources by coverage:** - -| Resource | js-sdk endpoints | drs coverage | Priority | -|----------|-----------------|--------------|----------| -| jobs | 6 | 4 covered | Core | -| catalog | 16 | 4 covered | Core | -| users | 8 | 3 covered | Core | -| roles | 5 | 2 covered | Core | -| grants | 1 | — (matched differently) | Core | -| ai | 18 | 0 | High — agent conversations, model providers | -| engines | 5 | 0 | Medium — list, start, stop | -| scripts | 3 | 0 | Medium — saved SQL | -| projects | 6 | 0 | Low — admin | -| organizations | 2 | 0 | Low — admin | -| arctic | 2 | 0 | Deferred — Open Catalog | - -Run coverage yourself: -```bash -git clone --depth 1 https://github.com/dremio/js-sdk.git /tmp/js-sdk -python scripts/parse_jssdk.py --sdk-path /tmp/js-sdk --compare -``` - ---- - -## 6. Discovery Service Roadmap - -Current state: hand-maintained client methods. Future: auto-generated from js-sdk. - -### Phase progression: - -``` -Phase 1 (now) Phase 2 (next) Phase 3 (future) -────────────── ─────────────── ──────────────── -Hand-written → js-sdk parsed → Runtime discovery -+ js-sdk parser + auto-generated + zero-release updates -+ coverage report client & CLI -``` - -**Phase 1 (done):** Regex parser proves the concept — `parse_jssdk.py` extracts 79 endpoints from js-sdk, compares against drs coverage. Not production-grade (can't extract types, breaks on refactors). - -**Phase 2 (next — requires js-sdk team):** Request `discovery.json` export from js-sdk build. The SDK already has typed Zod schemas and TypeScript interfaces for every endpoint — we need them serialized to JSON Schema. This gives us parameter types, enums, required fields, response schemas, deprecation markers. - -**Phase 3 (after discovery.json exists):** -1. `registry.py` loads `discovery.json` into typed `ApiOperation` objects -2. `executor.py` replaces hand-written client methods with generic HTTP execution -3. CLI and MCP tools auto-generated from registry — including `--help`, `--fields`, `--dry-run` -4. SQL-based commands coexist via explicit `SQL_COMMANDS` registry -5. CI fetches new `discovery.json` on js-sdk release → auto-detect new endpoints - -**Phase 4 (optional, requires server-side):** -1. Dremio Cloud serves `/api/openapi.json` -2. `drs` fetches at startup, caches 24h -3. New endpoints appear in CLI without a drs release - -Full plan: `DISCOVERY_SERVICE_PLAN.md` - ---- - -## 7. Relationship to Existing Repos - -``` -┌─────────────────────┐ -│ dremio/dremio-mcp │ ──── Predecessor. drs mcp supersedes it. -│ (existing MCP) │ Config format preserved for compatibility. -└─────────────────────┘ - -┌─────────────────────┐ -│ dremio/claude-plugins│ ──── Absorbed. 3 skills rewritten to use drs -│ (existing skills) │ commands. Repo can be archived. -└─────────────────────┘ - -┌──────────────────────────────┐ -│ developer-advocacy-dremio/ │ ──── Referenced. Wizard patterns -│ dremio-agent-skill │ informed skill design. -│ (knowledge pack) │ Not merged — different owner/scope. -└──────────────────────────────┘ -``` - ---- - -## 8. Project Structure - -``` -dremio/cli/ -├── .claude-plugin/marketplace.json # Claude Code marketplace registration -├── plugins/dremio/ -│ ├── .claude-plugin/plugin.json # Plugin manifest -│ └── skills/ # 8 SKILL.md files -├── src/drs/ -│ ├── cli.py # Typer entry point, registers all groups -│ ├── auth.py # PAT + config loading (env > file > defaults) -│ ├── client.py # ONLY file making HTTP calls (18 methods) -│ ├── output.py # JSON/CSV/pretty formatting -│ ├── utils.py # Path parsing, validation, error handling -│ ├── mcp_server.py # FastMCP adapter (19 tools, zero logic) -│ └── commands/ # 6 command modules -│ ├── query.py # run, status, cancel -│ ├── catalog.py # list, get, search -│ ├── schema.py # describe, lineage, wiki, sample -│ ├── reflect.py # list, status, refresh, drop -│ ├── jobs.py # list, get, profile -│ └── access.py # grants, roles, whoami, audit -├── scripts/ -│ ├── parse_jssdk.py # Parses dremio/js-sdk → api_registry.json -│ └── validate_api_coverage.py # Legacy OpenAPI spec validator -├── docs/ -│ ├── api_registry.json # Machine-readable endpoint catalog (79 endpoints) -│ ├── coverage_report.json # drs vs js-sdk coverage comparison -│ ├── ARCHITECTURE_REVIEW.md # This document -│ └── ONE_PAGER.md # Concise overview -├── tests/ # 56 unit tests (mocked HTTP) -├── README.md -├── TESTING.md -├── SPIKE_NOTES.md -└── DISCOVERY_SERVICE_PLAN.md -``` - ---- - -## 9. Open Questions for Review - -1. **Naming:** Repo is `dremio/cli`, binary is `drs`. Any concerns with the short name? - -2. **Scope:** Currently Dremio Cloud only. Should we design the auth layer now to accommodate Software later, or keep it strictly Cloud? - -3. **Write operations:** Current commands are read-only + reflection refresh/drop. Should we add catalog CRUD (create source, create VDS), or keep the first release read-heavy for safety? - -4. **Discovery service priority:** `parse_jssdk.py` already extracts 79 endpoints from js-sdk. Should we build the auto-generation pipeline (registry → executor → dynamic CLI) before the first release, or ship hand-maintained and iterate? - -5. **Skill taxonomy:** Following Google's pattern (Tier 1 API → Tier 2 Action → Tier 3 Workflow → Personas → Recipes). How deep should we go for v1? - -6. **Auth evolution:** PAT-only today. OAuth2 device flow for interactive login — priority for v1 or later? - -7. **Distribution:** PyPI (`pip install dremio-cli`), Homebrew tap, or both? diff --git a/docs/ONE_PAGER.md b/docs/ONE_PAGER.md deleted file mode 100644 index 221618a..0000000 --- a/docs/ONE_PAGER.md +++ /dev/null @@ -1,51 +0,0 @@ -# drs — Dremio Developer CLI - -**One CLI. Three channels. Zero duplicated logic.** - -## The Problem - -Dremio's developer tooling is fragmented across repos with no CLI layer: -- `dremio-mcp` — MCP server only, no CLI, can't script or debug without an AI agent -- `claude-plugins` — skills reference raw curl, not actionable commands -- `dremio-agent-skill` — knowledge pack, different owner, multi-tool targeting - -## The Solution - -`dremio/cli` is a single repo that ships: - -| Channel | Users | Example | -|---------|-------|---------| -| **CLI** (`drs`) | Developers, CI scripts | `drs query run "SELECT * FROM orders LIMIT 10"` | -| **MCP Server** | Claude Desktop, AI agents | `drs mcp` → 19 tools auto-available | -| **Plugin** | Claude Code | 8 workflow skills (slow query diagnosis, access audit, etc.) | - -All three call the same async command functions → `client.py` (single HTTP layer) → Dremio Cloud API. - -## What ships in v1 - -- **18 CLI commands** in 6 groups: query, catalog, schema, reflect, jobs, access -- **19 MCP tools** with LLM-optimized descriptions -- **8 Claude Code skills** (3 absorbed from claude-plugins + 5 new workflows) -- **3 output formats**: JSON (default), CSV, pretty table -- **Config compatibility** with existing dremio-mcp format -- **43 unit tests**, input validation, structured error handling - -## The Discovery Service Path - -Today: hand-maintained client methods + spec validator that catches drift against OpenAPI specs. - -Next: auto-generate the client layer from OpenAPI specs already in the Dremio monorepo (40+ YAML files). SQL-based commands (jobs, reflections by dataset) coexist alongside auto-generated REST commands. - -End state: Dremio Cloud serves `/api/openapi.json`, CLI fetches at startup — new endpoints appear without a release. - -## Key Design Decisions - -| Decision | Rationale | -|----------|-----------| -| CLI-first, MCP wraps | Every operation is testable/scriptable without an AI agent | -| Cloud-only for v1 | Software has different auth, URLs, API behavior — separate concern | -| SQL + REST hybrid | Some data only available via system tables (jobs, reflections by dataset) | -| Single `client.py` | One file makes all HTTP calls — easy to audit, replace, or auto-generate | -| Google CLI naming | Repo: `dremio/cli`, binary: `drs` (matches `googleworkspace/cli` → `gws`) | - -## Repo: `github.com/dremio/cli` diff --git a/pyproject.toml b/pyproject.toml index b8b51da..f1c6376 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "drs" version = "0.1.0" -description = "Developer CLI for Dremio — query, catalog, schema, reflections, jobs, and access management" +description = "Developer CLI for Dremio Cloud" readme = "README.md" requires-python = ">=3.11" license = "Apache-2.0" @@ -13,7 +13,7 @@ dependencies = [ ] [project.scripts] -drs = "drs.cli:app" +dremio = "drs.cli:app" [build-system] requires = ["hatchling"] diff --git a/src/drs/cli.py b/src/drs/cli.py index c16a5ff..ed9ce4e 100644 --- a/src/drs/cli.py +++ b/src/drs/cli.py @@ -13,34 +13,42 @@ # See the License for the specific language governing permissions and # limitations under the License. # -"""drs — Developer CLI for Dremio Cloud. Entry point and command registration.""" +"""dremio — Developer CLI for Dremio Cloud. Entry point and command registration.""" from __future__ import annotations +import asyncio import json import sys from pathlib import Path from typing import Optional +import httpx import typer from drs.auth import DrsConfig, load_config from drs.client import DremioClient -from drs.commands import query, catalog, schema, reflect, jobs, access +from drs.commands import query, schema, engine, user, role, grant +from drs.commands import folder, reflection, wiki, tag, job app = typer.Typer( - name="drs", - help="Developer CLI for Dremio Cloud — query, catalog, schema, reflections, jobs, and access.", + name="dremio", + help="Developer CLI for Dremio Cloud.", no_args_is_help=True, ) # Register command groups app.add_typer(query.app, name="query") -app.add_typer(catalog.app, name="catalog") +app.add_typer(folder.app, name="folder") app.add_typer(schema.app, name="schema") -app.add_typer(reflect.app, name="reflect") -app.add_typer(jobs.app, name="jobs") -app.add_typer(access.app, name="access") +app.add_typer(wiki.app, name="wiki") +app.add_typer(tag.app, name="tag") +app.add_typer(reflection.app, name="reflection") +app.add_typer(job.app, name="job") +app.add_typer(engine.app, name="engine") +app.add_typer(user.app, name="user") +app.add_typer(role.app, name="role") +app.add_typer(grant.app, name="grant") # Global state for config _config: DrsConfig | None = None @@ -54,7 +62,7 @@ def main( project_id: Optional[str] = typer.Option(None, "--project-id", help="Dremio Cloud project ID"), uri: Optional[str] = typer.Option(None, "--uri", help="Dremio API base URI (e.g., https://api.dremio.cloud, https://api.eu.dremio.cloud)"), ) -> None: - """Global options for drs CLI.""" + """Global options for dremio CLI.""" global _cli_opts _cli_opts = { "config_path": Path(config) if config else None, @@ -89,23 +97,47 @@ def get_client() -> DremioClient: return DremioClient(get_config()) +# -- Top-level commands -- + +@app.command("search") +def search_command( + term: str = typer.Argument(help="Search term (matches table names, view names, source names)"), + fmt: str = typer.Option("json", "--output", "-o", help="Output format: json, csv, pretty"), +) -> None: + """Full-text search across all catalog entities (tables, views, sources).""" + from drs.output import OutputFormat, output, error + from drs.utils import handle_api_error, DremioAPIError + + client = get_client() + + async def _execute(): + try: + try: + return await client.search(term) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + finally: + await client.close() + + try: + result = asyncio.run(_execute()) + except DremioAPIError as exc: + error(str(exc)) + raise typer.Exit(1) + output(result, OutputFormat(fmt)) + + @app.command("describe") def describe_command( - command: str = typer.Argument(help="Command to describe (e.g., 'query.run', 'catalog.get', 'reflect.drop')"), + command: str = typer.Argument(help="Command to describe (e.g., 'query.run', 'folder.get', 'reflection.delete')"), ) -> None: - """Show machine-readable schema for a command — parameters, types, and descriptions. - - Use this to discover what parameters a command accepts before calling it. - Outputs JSON with parameter names, types, required/optional, and descriptions. - """ - from drs.introspect import describe_command as _describe + """Show machine-readable schema for a command — parameters, types, and descriptions.""" + from drs.introspect import describe_command as _describe, list_commands result = _describe(command) if result is None: print(f"Unknown command: {command}", file=sys.stderr) - print("Available commands: query.run, query.status, query.cancel, catalog.list, catalog.get, " - "catalog.search, schema.describe, schema.lineage, schema.wiki, schema.sample, " - "reflect.list, reflect.status, reflect.refresh, reflect.drop, jobs.list, jobs.get, " - "jobs.profile, access.grants, access.roles, access.whoami, access.audit", file=sys.stderr) + commands = list_commands() + print(f"Available commands: {', '.join(commands)}", file=sys.stderr) raise typer.Exit(1) print(json.dumps(result, indent=2)) diff --git a/src/drs/client.py b/src/drs/client.py index 61a9090..e0f075e 100644 --- a/src/drs/client.py +++ b/src/drs/client.py @@ -17,18 +17,29 @@ from __future__ import annotations +import asyncio +import logging from typing import Any import httpx from drs.auth import DrsConfig +logger = logging.getLogger(__name__) + +_MAX_RETRIES = 3 +_RETRY_BACKOFF = (1.0, 2.0, 4.0) +_RETRYABLE_STATUS_CODES = {429, 502, 503, 504} + class DremioClient: """Async HTTP client for Dremio Cloud REST APIs. This is the ONLY file that makes HTTP calls. All commands and MCP tools call methods on this class. + + Transient failures (timeouts, 429, 502, 503, 504) are retried up to 3 + times with exponential backoff (1s, 2s, 4s). """ def __init__(self, config: DrsConfig) -> None: @@ -57,20 +68,55 @@ def _v3(self, path: str) -> str: def _v1(self, path: str) -> str: return f"{self.config.uri}/v1{path}" - # -- HTTP helpers -- + # -- HTTP helpers with retry -- + + async def _request_with_retry(self, method: str, url: str, **kwargs: Any) -> httpx.Response: + """Execute an HTTP request with retry on transient errors.""" + last_exc: Exception | None = None + for attempt in range(_MAX_RETRIES): + try: + resp = await self._client.request(method, url, **kwargs) + if resp.status_code in _RETRYABLE_STATUS_CODES and attempt < _MAX_RETRIES - 1: + delay = _RETRY_BACKOFF[attempt] + logger.warning( + "Retryable HTTP %d on %s %s — retrying in %.1fs (attempt %d/%d)", + resp.status_code, method, url, delay, attempt + 1, _MAX_RETRIES, + ) + await asyncio.sleep(delay) + continue + return resp + except httpx.TimeoutException as exc: + last_exc = exc + if attempt < _MAX_RETRIES - 1: + delay = _RETRY_BACKOFF[attempt] + logger.warning( + "Timeout on %s %s — retrying in %.1fs (attempt %d/%d)", + method, url, delay, attempt + 1, _MAX_RETRIES, + ) + await asyncio.sleep(delay) + continue + raise + raise last_exc # type: ignore[misc] async def _get(self, url: str, params: dict | None = None) -> Any: - resp = await self._client.get(url, params=params) + resp = await self._request_with_retry("GET", url, params=params) resp.raise_for_status() return resp.json() async def _post(self, url: str, json: dict | None = None) -> Any: - resp = await self._client.post(url, json=json) + resp = await self._request_with_retry("POST", url, json=json) resp.raise_for_status() return resp.json() - async def _delete(self, url: str) -> Any: - resp = await self._client.delete(url) + async def _put(self, url: str, json: dict | None = None) -> Any: + resp = await self._request_with_retry("PUT", url, json=json) + resp.raise_for_status() + if resp.content: + return resp.json() + return {"status": "ok"} + + async def _delete(self, url: str, params: dict | None = None) -> Any: + resp = await self._request_with_retry("DELETE", url, params=params) resp.raise_for_status() if resp.content: return resp.json() @@ -122,6 +168,19 @@ async def search(self, query: str, filter_: str | None = None) -> dict: body["filter"] = filter_ return await self._post(self._v0("/search"), json=body) + async def create_catalog_entity(self, body: dict) -> dict: + """Create a catalog entity (space, folder, etc.). POST /catalog.""" + return await self._post(self._v3("/catalog"), json=body) + + async def update_catalog_entity(self, entity_id: str, body: dict) -> dict: + """Update a catalog entity. PUT /catalog/{id}.""" + return await self._put(self._v3(f"/catalog/{entity_id}"), json=body) + + async def delete_catalog_entity(self, entity_id: str, tag: str | None = None) -> dict: + """Delete a catalog entity. DELETE /catalog/{id}.""" + params = {"tag": tag} if tag else None + return await self._delete(self._v3(f"/catalog/{entity_id}"), params=params) + async def get_lineage(self, entity_id: str) -> dict: return await self._get(self._v3(f"/catalog/{entity_id}/graph")) @@ -131,6 +190,20 @@ async def get_wiki(self, entity_id: str) -> dict: async def get_tags(self, entity_id: str) -> dict: return await self._get(self._v3(f"/catalog/{entity_id}/collaboration/tag")) + async def set_wiki(self, entity_id: str, text: str, version: int | None = None) -> dict: + """Set wiki text for a catalog entity. POST /catalog/{id}/collaboration/wiki.""" + body: dict[str, Any] = {"text": text} + if version is not None: + body["version"] = version + return await self._post(self._v3(f"/catalog/{entity_id}/collaboration/wiki"), json=body) + + async def set_tags(self, entity_id: str, tags: list[str], version: int | None = None) -> dict: + """Set tags for a catalog entity. POST /catalog/{id}/collaboration/tag.""" + body: dict[str, Any] = {"tags": tags} + if version is not None: + body["version"] = version + return await self._post(self._v3(f"/catalog/{entity_id}/collaboration/tag"), json=body) + # -- Reflections (v3) -- async def get_reflection(self, reflection_id: str) -> dict: @@ -139,9 +212,36 @@ async def get_reflection(self, reflection_id: str) -> dict: async def refresh_reflection(self, reflection_id: str) -> dict: return await self._post(self._v3(f"/reflection/{reflection_id}/refresh")) + async def create_reflection(self, body: dict) -> dict: + """Create a reflection. POST /reflection.""" + return await self._post(self._v3("/reflection"), json=body) + async def delete_reflection(self, reflection_id: str) -> dict: return await self._delete(self._v3(f"/reflection/{reflection_id}")) + # -- Engines (v0) -- + + async def list_engines(self) -> dict: + return await self._get(self._v0("/engines")) + + async def get_engine(self, engine_id: str) -> dict: + return await self._get(self._v0(f"/engines/{engine_id}")) + + async def create_engine(self, body: dict) -> dict: + return await self._post(self._v0("/engines"), json=body) + + async def update_engine(self, engine_id: str, body: dict) -> dict: + return await self._put(self._v0(f"/engines/{engine_id}"), json=body) + + async def delete_engine(self, engine_id: str) -> dict: + return await self._delete(self._v0(f"/engines/{engine_id}")) + + async def enable_engine(self, engine_id: str) -> dict: + return await self._put(self._v0(f"/engines/{engine_id}/enable")) + + async def disable_engine(self, engine_id: str) -> dict: + return await self._put(self._v0(f"/engines/{engine_id}/disable")) + # -- Users & Roles (v1) -- async def list_users(self, max_results: int = 100) -> dict: @@ -150,13 +250,59 @@ async def list_users(self, max_results: int = 100) -> dict: async def get_user_by_name(self, name: str) -> dict: return await self._get(self._v1(f"/users/name/{name}")) + async def get_user(self, user_id: str) -> dict: + return await self._get(self._v1(f"/users/{user_id}")) + + async def invite_user(self, body: dict) -> dict: + """Invite a user. POST /v1/users/invite.""" + return await self._post(self._v1("/users/invite"), json=body) + + async def update_user(self, user_id: str, body: dict) -> dict: + return await self._put(self._v1(f"/users/{user_id}"), json=body) + + async def delete_user(self, user_id: str) -> dict: + return await self._delete(self._v1(f"/users/{user_id}")) + async def list_roles(self, max_results: int = 100) -> dict: return await self._get(self._v1("/roles"), params={"maxResults": max_results}) + async def get_role(self, role_id: str) -> dict: + return await self._get(self._v1(f"/roles/{role_id}")) + + async def get_role_by_name(self, name: str) -> dict: + return await self._get(self._v1(f"/roles/name/{name}")) + + async def create_role(self, body: dict) -> dict: + return await self._post(self._v1("/roles"), json=body) + + async def update_role(self, role_id: str, body: dict) -> dict: + return await self._put(self._v1(f"/roles/{role_id}"), json=body) + + async def delete_role(self, role_id: str) -> dict: + return await self._delete(self._v1(f"/roles/{role_id}")) + + # -- Grants (v1) -- + async def get_grants( self, scope: str, scope_id: str, grantee_type: str, grantee_id: str ) -> dict: - """Get grants. scope is 'catalog' or 'org', grantee_type is 'user' or 'role'.""" + """Get grants. scope is 'projects', 'orgs', 'clouds', etc.""" return await self._get( self._v1(f"/{scope}/{scope_id}/grants/{grantee_type}/{grantee_id}") ) + + async def set_grants( + self, scope: str, scope_id: str, grantee_type: str, grantee_id: str, body: dict + ) -> dict: + """Set grants. PUT /v1/{scope}/{scopeId}/grants/{granteeType}/{granteeId}.""" + return await self._put( + self._v1(f"/{scope}/{scope_id}/grants/{grantee_type}/{grantee_id}"), json=body + ) + + async def delete_grants( + self, scope: str, scope_id: str, grantee_type: str, grantee_id: str + ) -> dict: + """Remove grants. DELETE /v1/{scope}/{scopeId}/grants/{granteeType}/{granteeId}.""" + return await self._delete( + self._v1(f"/{scope}/{scope_id}/grants/{grantee_type}/{grantee_id}") + ) diff --git a/src/drs/commands/access.py b/src/drs/commands/access.py deleted file mode 100644 index 64cae5f..0000000 --- a/src/drs/commands/access.py +++ /dev/null @@ -1,160 +0,0 @@ -# -# Copyright (C) 2017-2026 Dremio Corporation -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -"""drs access — inspect grants, roles, and user permissions.""" - -from __future__ import annotations - -import asyncio - -import httpx -import typer - -from drs.client import DremioClient -from drs.output import OutputFormat, output, error -from drs.utils import handle_api_error, parse_path - -app = typer.Typer(help="Inspect grants, roles, and user permissions.") - - -async def grants(client: DremioClient, path: str) -> dict: - """Get ACL grants on a catalog entity.""" - parts = parse_path(path) - try: - entity = await client.get_catalog_by_path(parts) - except httpx.HTTPStatusError as exc: - raise handle_api_error(exc) from exc - acl = entity.get("accessControlList", {}) - return { - "path": path, - "id": entity.get("id"), - "accessControlList": acl, - } - - -async def roles(client: DremioClient) -> dict: - """List all roles.""" - try: - return await client.list_roles() - except httpx.HTTPStatusError as exc: - raise handle_api_error(exc) from exc - - -async def whoami(client: DremioClient) -> dict: - """Get info about the current user. - - Note: Dremio Cloud lacks a dedicated 'whoami' endpoint. This returns - the first user from the user list, which may not be the PAT owner - in all configurations. Use 'drs access audit ' for reliable - user lookups. - """ - try: - return await client.list_users(max_results=1) - except httpx.HTTPStatusError as exc: - raise handle_api_error(exc) from exc - - -async def audit(client: DremioClient, username: str) -> dict: - """Audit a user's effective permissions: user -> roles.""" - try: - user = await client.get_user_by_name(username) - except httpx.HTTPStatusError as exc: - raise handle_api_error(exc) from exc - user_roles = user.get("roles", []) - - role_grants: list[dict] = [] - for role in user_roles: - role_id = role.get("id", role) if isinstance(role, dict) else role - role_name = role.get("name", role_id) if isinstance(role, dict) else role_id - role_grants.append({"role_id": role_id, "role_name": role_name}) - - return { - "username": username, - "user_id": user.get("id"), - "roles": role_grants, - } - - -# -- CLI wrappers -- - -def _get_client() -> DremioClient: - from drs.cli import get_client - return get_client() - - -def _run_command(coro, client, fmt: OutputFormat = OutputFormat.json, fields: str | None = None) -> None: - async def _execute(): - try: - return await coro - finally: - await client.close() - - try: - result = asyncio.run(_execute()) - except Exception as exc: - from drs.utils import DremioAPIError - if isinstance(exc, DremioAPIError): - error(str(exc)) - raise typer.Exit(1) - if isinstance(exc, ValueError): - error(str(exc)) - raise typer.Exit(1) - raise - output(result, fmt, fields=fields) - - -@app.command("grants") -def cli_grants( - path: str = typer.Argument(help="Dot-separated entity path (e.g., myspace.mytable)"), - fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), -) -> None: - """Show ACL grants on a catalog entity. - - Displays which users and roles have permissions on the specified - table, view, source, or space. - """ - client = _get_client() - _run_command(grants(client, path), client, fmt) - - -@app.command("roles") -def cli_roles( - fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), -) -> None: - """List all roles in the organization.""" - client = _get_client() - _run_command(roles(client), client, fmt) - - -@app.command("whoami") -def cli_whoami( - fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), -) -> None: - """Show current authenticated user info (best-effort, see 'audit' for reliable lookups).""" - client = _get_client() - _run_command(whoami(client), client, fmt) - - -@app.command("audit") -def cli_audit( - username: str = typer.Argument(help="Username to look up and audit"), - fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), -) -> None: - """Audit a user's roles and effective permissions. - - Looks up the user by name, then lists all roles they belong to. - """ - client = _get_client() - _run_command(audit(client, username), client, fmt) diff --git a/src/drs/commands/engine.py b/src/drs/commands/engine.py new file mode 100644 index 0000000..607d8f9 --- /dev/null +++ b/src/drs/commands/engine.py @@ -0,0 +1,201 @@ +# +# Copyright (C) 2017-2026 Dremio Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""drs engine — manage Dremio Cloud engines.""" + +from __future__ import annotations + +import asyncio + +import httpx +import typer + +from drs.client import DremioClient +from drs.output import OutputFormat, output, error +from drs.utils import handle_api_error + +app = typer.Typer(help="Manage Dremio Cloud engines.") + + +async def list_engines(client: DremioClient) -> dict: + """List all engines in the project.""" + try: + return await client.list_engines() + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + + +async def get_engine(client: DremioClient, engine_id: str) -> dict: + """Get engine details by ID.""" + try: + return await client.get_engine(engine_id) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + + +async def create_engine(client: DremioClient, name: str, size: str = "SMALL") -> dict: + """Create a new engine.""" + body = {"name": name, "size": size.upper()} + try: + return await client.create_engine(body) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + + +async def update_engine(client: DremioClient, engine_id: str, name: str | None = None, size: str | None = None) -> dict: + """Update engine configuration.""" + try: + existing = await client.get_engine(engine_id) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + body = dict(existing) + if name: + body["name"] = name + if size: + body["size"] = size.upper() + try: + return await client.update_engine(engine_id, body) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + + +async def delete_engine(client: DremioClient, engine_id: str) -> dict: + """Delete an engine.""" + try: + return await client.delete_engine(engine_id) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + + +async def enable_engine(client: DremioClient, engine_id: str) -> dict: + """Enable a disabled engine.""" + try: + return await client.enable_engine(engine_id) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + + +async def disable_engine(client: DremioClient, engine_id: str) -> dict: + """Disable a running engine.""" + try: + return await client.disable_engine(engine_id) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + + +# -- CLI wrappers -- + +def _get_client() -> DremioClient: + from drs.cli import get_client + return get_client() + + +def _run_command(coro, client, fmt: OutputFormat = OutputFormat.json, fields: str | None = None) -> None: + async def _execute(): + try: + return await coro + finally: + await client.close() + + try: + result = asyncio.run(_execute()) + except Exception as exc: + from drs.utils import DremioAPIError + if isinstance(exc, DremioAPIError): + error(str(exc)) + raise typer.Exit(1) + if isinstance(exc, ValueError): + error(str(exc)) + raise typer.Exit(1) + raise + output(result, fmt, fields=fields) + + +@app.command("list") +def cli_list( + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), + fields: str = typer.Option(None, "--fields", "-f", help="Comma-separated fields to include"), +) -> None: + """List all engines in the project.""" + client = _get_client() + _run_command(list_engines(client), client, fmt, fields=fields) + + +@app.command("get") +def cli_get( + engine_id: str = typer.Argument(help="Engine ID"), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), + fields: str = typer.Option(None, "--fields", "-f", help="Comma-separated fields to include"), +) -> None: + """Get details for a specific engine.""" + client = _get_client() + _run_command(get_engine(client, engine_id), client, fmt, fields=fields) + + +@app.command("create") +def cli_create( + name: str = typer.Argument(help="Name for the new engine"), + size: str = typer.Option("SMALL", "--size", "-s", help="Engine size (e.g., SMALL, MEDIUM, LARGE, XLARGE, XXLARGE)"), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Create a new engine.""" + client = _get_client() + _run_command(create_engine(client, name, size=size), client, fmt) + + +@app.command("update") +def cli_update( + engine_id: str = typer.Argument(help="Engine ID to update"), + name: str = typer.Option(None, "--name", help="New engine name"), + size: str = typer.Option(None, "--size", "-s", help="New engine size"), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Update engine configuration (name, size).""" + client = _get_client() + _run_command(update_engine(client, engine_id, name=name, size=size), client, fmt) + + +@app.command("delete") +def cli_delete( + engine_id: str = typer.Argument(help="Engine ID to delete"), + dry_run: bool = typer.Option(False, "--dry-run", help="Show engine details without deleting"), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Delete an engine. Cannot be undone.""" + client = _get_client() + if dry_run: + _run_command(get_engine(client, engine_id), client, fmt) + return + _run_command(delete_engine(client, engine_id), client, fmt) + + +@app.command("enable") +def cli_enable( + engine_id: str = typer.Argument(help="Engine ID to enable"), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Enable a disabled engine.""" + client = _get_client() + _run_command(enable_engine(client, engine_id), client, fmt) + + +@app.command("disable") +def cli_disable( + engine_id: str = typer.Argument(help="Engine ID to disable"), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Disable a running engine.""" + client = _get_client() + _run_command(disable_engine(client, engine_id), client, fmt) diff --git a/src/drs/commands/folder.py b/src/drs/commands/folder.py new file mode 100644 index 0000000..308392f --- /dev/null +++ b/src/drs/commands/folder.py @@ -0,0 +1,177 @@ +# +# Copyright (C) 2017-2026 Dremio Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""dremio folder — manage spaces and folders in the Dremio catalog.""" + +from __future__ import annotations + +import asyncio + +import httpx +import typer + +from drs.client import DremioClient +from drs.commands.query import run_query +from drs.output import OutputFormat, output, error +from drs.utils import handle_api_error, parse_path, quote_path_sql + +app = typer.Typer(help="Manage spaces and folders in the Dremio catalog.") + + +async def list_catalog(client: DremioClient) -> dict: + """List top-level catalog entities (sources, spaces, home).""" + try: + root = await client.get_catalog_entity("") + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + children = root.get("data", root.get("children", [])) + return {"entities": children} + + +async def get_entity(client: DremioClient, path: str) -> dict: + """Get a catalog entity by dot-separated path.""" + parts = parse_path(path) + try: + return await client.get_catalog_by_path(parts) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + + +async def create_folder(client: DremioClient, path: str) -> dict: + """Create a space (single component) or folder (nested path) using SQL.""" + parts = parse_path(path) + if len(parts) == 1: + sql = f'CREATE SPACE "{parts[0]}"' + else: + quoted = quote_path_sql(path) + sql = f"CREATE FOLDER {quoted}" + return await run_query(client, sql) + + +async def delete_entity(client: DremioClient, path: str) -> dict: + """Delete a catalog entity by path.""" + parts = parse_path(path) + try: + entity = await client.get_catalog_by_path(parts) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + entity_id = entity["id"] + tag = entity.get("tag") + try: + return await client.delete_catalog_entity(entity_id, tag=tag) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + + +async def grants(client: DremioClient, path: str) -> dict: + """Get ACL grants on a catalog entity.""" + parts = parse_path(path) + try: + entity = await client.get_catalog_by_path(parts) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + acl = entity.get("accessControlList", {}) + return { + "path": path, + "id": entity.get("id"), + "accessControlList": acl, + } + + +# -- CLI wrappers -- + +def _get_client() -> DremioClient: + from drs.cli import get_client + return get_client() + + +def _run_command(coro, client, fmt: OutputFormat = OutputFormat.json, fields: str | None = None) -> None: + async def _execute(): + try: + return await coro + finally: + await client.close() + + try: + result = asyncio.run(_execute()) + except Exception as exc: + from drs.utils import DremioAPIError + if isinstance(exc, DremioAPIError): + error(str(exc)) + raise typer.Exit(1) + if isinstance(exc, ValueError): + error(str(exc)) + raise typer.Exit(1) + raise + output(result, fmt, fields=fields) + + +@app.command("list") +def cli_list( + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), + fields: str = typer.Option(None, "--fields", "-f", help="Comma-separated fields to include"), +) -> None: + """List top-level catalog entities: sources, spaces, and home folder.""" + client = _get_client() + _run_command(list_catalog(client), client, fmt, fields=fields) + + +@app.command("get") +def cli_get( + path: str = typer.Argument(help='Dot-separated entity path (e.g., myspace.folder.table)'), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), + fields: str = typer.Option(None, "--fields", "-f", help="Comma-separated fields to include"), +) -> None: + """Get full metadata for a catalog entity by path.""" + client = _get_client() + _run_command(get_entity(client, path), client, fmt, fields=fields) + + +@app.command("create") +def cli_create( + path: str = typer.Argument(help="Space name (single component) or dot-separated folder path (e.g., myspace.newfolder)"), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Create a space or folder. + + Single path component (e.g., 'Analytics') creates a space. + Nested path (e.g., 'Analytics.reports') creates a folder. + """ + client = _get_client() + _run_command(create_folder(client, path), client, fmt) + + +@app.command("delete") +def cli_delete( + path: str = typer.Argument(help="Dot-separated entity path to delete"), + dry_run: bool = typer.Option(False, "--dry-run", help="Show what would be deleted without deleting"), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Delete a catalog entity (space, folder, view, etc.). Cannot be undone.""" + client = _get_client() + if dry_run: + _run_command(get_entity(client, path), client, fmt) + return + _run_command(delete_entity(client, path), client, fmt) + + +@app.command("grants") +def cli_grants( + path: str = typer.Argument(help="Dot-separated entity path"), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Show ACL grants on a catalog entity.""" + client = _get_client() + _run_command(grants(client, path), client, fmt) diff --git a/src/drs/commands/grant.py b/src/drs/commands/grant.py new file mode 100644 index 0000000..e46a75c --- /dev/null +++ b/src/drs/commands/grant.py @@ -0,0 +1,137 @@ +# +# Copyright (C) 2017-2026 Dremio Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""drs grant — manage grants on Dremio Cloud resources.""" + +from __future__ import annotations + +import asyncio + +import httpx +import typer + +from drs.client import DremioClient +from drs.output import OutputFormat, output, error +from drs.utils import handle_api_error + +app = typer.Typer(help="Manage grants on projects, engines, and org resources.") + + +async def get_grants( + client: DremioClient, scope: str, scope_id: str, grantee_type: str, grantee_id: str +) -> dict: + """Get grants for a grantee on a resource.""" + try: + return await client.get_grants(scope, scope_id, grantee_type, grantee_id) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + + +async def set_grants( + client: DremioClient, scope: str, scope_id: str, grantee_type: str, grantee_id: str, + privileges: list[str], +) -> dict: + """Set grants for a grantee on a resource.""" + body = {"privileges": privileges} + try: + return await client.set_grants(scope, scope_id, grantee_type, grantee_id, body) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + + +async def remove_grants( + client: DremioClient, scope: str, scope_id: str, grantee_type: str, grantee_id: str +) -> dict: + """Remove all grants for a grantee on a resource.""" + try: + return await client.delete_grants(scope, scope_id, grantee_type, grantee_id) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + + +# -- CLI wrappers -- + +def _get_client() -> DremioClient: + from drs.cli import get_client + return get_client() + + +def _run_command(coro, client, fmt: OutputFormat = OutputFormat.json, fields: str | None = None) -> None: + async def _execute(): + try: + return await coro + finally: + await client.close() + + try: + result = asyncio.run(_execute()) + except Exception as exc: + from drs.utils import DremioAPIError + if isinstance(exc, DremioAPIError): + error(str(exc)) + raise typer.Exit(1) + if isinstance(exc, ValueError): + error(str(exc)) + raise typer.Exit(1) + raise + output(result, fmt, fields=fields) + + +@app.command("get") +def cli_get( + scope: str = typer.Argument(help="Resource scope (projects, orgs, clouds)"), + scope_id: str = typer.Argument(help="Resource ID (project ID, org ID, etc.)"), + grantee_type: str = typer.Argument(help="Grantee type (user or role)"), + grantee_id: str = typer.Argument(help="Grantee ID (user ID or role ID)"), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Get grants for a user or role on a resource.""" + client = _get_client() + _run_command(get_grants(client, scope, scope_id, grantee_type, grantee_id), client, fmt) + + +@app.command("update") +def cli_update( + scope: str = typer.Argument(help="Resource scope (projects, orgs, clouds)"), + scope_id: str = typer.Argument(help="Resource ID"), + grantee_type: str = typer.Argument(help="Grantee type (user or role)"), + grantee_id: str = typer.Argument(help="Grantee ID"), + privileges: str = typer.Argument(help='Comma-separated privileges (e.g., "MANAGE_GRANTS,CREATE_TABLE")'), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Set grants (privileges) for a user or role on a resource. + + Replaces any existing grants for this grantee on the resource. + """ + client = _get_client() + priv_list = [p.strip() for p in privileges.split(",") if p.strip()] + _run_command(set_grants(client, scope, scope_id, grantee_type, grantee_id, priv_list), client, fmt) + + +@app.command("delete") +def cli_delete( + scope: str = typer.Argument(help="Resource scope (projects, orgs, clouds)"), + scope_id: str = typer.Argument(help="Resource ID"), + grantee_type: str = typer.Argument(help="Grantee type (user or role)"), + grantee_id: str = typer.Argument(help="Grantee ID"), + dry_run: bool = typer.Option(False, "--dry-run", help="Show current grants without removing"), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Remove all grants for a user or role on a resource. Cannot be undone.""" + client = _get_client() + if dry_run: + _run_command(get_grants(client, scope, scope_id, grantee_type, grantee_id), client, fmt) + return + _run_command(remove_grants(client, scope, scope_id, grantee_type, grantee_id), client, fmt) diff --git a/src/drs/commands/jobs.py b/src/drs/commands/job.py similarity index 87% rename from src/drs/commands/jobs.py rename to src/drs/commands/job.py index d000b2a..4c70183 100644 --- a/src/drs/commands/jobs.py +++ b/src/drs/commands/job.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # -"""drs jobs — list and inspect query jobs.""" +"""dremio job — list and inspect query jobs.""" from __future__ import annotations @@ -97,13 +97,9 @@ def cli_list( status_filter: Optional[str] = typer.Option(None, "--status", "-s", help="Filter by job state: COMPLETED, FAILED, RUNNING, CANCELED, PLANNING, ENQUEUED"), limit: int = typer.Option(25, "--limit", "-n", help="Max jobs to return (default 25)"), fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), - fields: str = typer.Option(None, "--fields", "-f", help="Comma-separated fields to include in output"), + fields: str = typer.Option(None, "--fields", "-f", help="Comma-separated fields to include"), ) -> None: - """List recent query jobs from sys.project.jobs. - - Shows job ID, user, query type, state, and timing. Results are ordered - by start time (most recent first). - """ + """List recent query jobs from sys.project.jobs.""" client = _get_client() _run_command(list_jobs(client, status_filter=status_filter, limit=limit), client, fmt, fields=fields) @@ -112,7 +108,7 @@ def cli_list( def cli_get( job_id: str = typer.Argument(help="Job ID (UUID)"), fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), - fields: str = typer.Option(None, "--fields", "-f", help="Comma-separated fields to include in output"), + fields: str = typer.Option(None, "--fields", "-f", help="Comma-separated fields to include"), ) -> None: """Get detailed status and metadata for a specific job.""" client = _get_client() @@ -121,13 +117,9 @@ def cli_get( @app.command("profile") def cli_profile( - job_id: str = typer.Argument(help="Job ID (UUID) to get execution profile for"), + job_id: str = typer.Argument(help="Job ID (UUID)"), fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), ) -> None: - """Show execution profile for a completed job. - - Queries sys.project.jobs for cost estimates, scan stats, - timing breakdown, and acceleration info. Useful for diagnosing slow queries. - """ + """Show execution profile for a completed job.""" client = _get_client() _run_command(profile(client, job_id), client, fmt) diff --git a/src/drs/commands/reflect.py b/src/drs/commands/reflection.py similarity index 59% rename from src/drs/commands/reflect.py rename to src/drs/commands/reflection.py index 4d8f282..9af7066 100644 --- a/src/drs/commands/reflect.py +++ b/src/drs/commands/reflection.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # -"""drs reflect — manage Dremio reflections (materialized views).""" +"""dremio reflection — manage Dremio reflections (materialized views).""" from __future__ import annotations @@ -30,6 +30,34 @@ app = typer.Typer(help="Manage reflections (materialized views).") +async def create(client: DremioClient, path: str, rtype: str, display_fields: list[str] | None = None) -> dict: + """Create a reflection on a dataset.""" + parts = parse_path(path) + try: + entity = await client.get_catalog_by_path(parts) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + + dataset_id = entity["id"] + body: dict = { + "type": rtype.upper(), + "datasetId": dataset_id, + } + + fields = entity.get("fields", []) + if rtype.lower() == "raw": + display = display_fields or [f["name"] for f in fields] + body["displayFields"] = [{"name": n} for n in display] + elif rtype.lower() == "aggregation": + if display_fields: + body["dimensionFields"] = [{"name": n, "granularity": "DATE"} for n in display_fields] + + try: + return await client.create_reflection(body) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + + async def list_reflections(client: DremioClient, path: str) -> dict: """List reflections on a dataset via sys.project.reflections.""" parts = parse_path(path) @@ -42,7 +70,8 @@ async def list_reflections(client: DremioClient, path: str) -> dict: return await run_query(client, sql) -async def status(client: DremioClient, reflection_id: str) -> dict: +async def get_reflection(client: DremioClient, reflection_id: str) -> dict: + """Get detailed status of a reflection.""" try: return await client.get_reflection(reflection_id) except httpx.HTTPStatusError as exc: @@ -50,13 +79,15 @@ async def status(client: DremioClient, reflection_id: str) -> dict: async def refresh(client: DremioClient, reflection_id: str) -> dict: + """Trigger an immediate refresh of a reflection.""" try: return await client.refresh_reflection(reflection_id) except httpx.HTTPStatusError as exc: raise handle_api_error(exc) from exc -async def drop(client: DremioClient, reflection_id: str) -> dict: +async def delete(client: DremioClient, reflection_id: str) -> dict: + """Delete a reflection.""" try: return await client.delete_reflection(reflection_id) except httpx.HTTPStatusError as exc: @@ -91,64 +122,64 @@ async def _execute(): output(result, fmt, fields=fields) +@app.command("create") +def cli_create( + path: str = typer.Argument(help="Dot-separated dataset path to create a reflection on"), + rtype: str = typer.Option("raw", "--type", "-t", help="Reflection type: raw or aggregation"), + fields_list: str = typer.Option(None, "--fields", "-f", help="Comma-separated field names to include"), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Create a new reflection on a dataset.""" + client = _get_client() + display = [f.strip() for f in fields_list.split(",") if f.strip()] if fields_list else None + _run_command(create(client, path, rtype, display_fields=display), client, fmt) + + @app.command("list") def cli_list( - path: str = typer.Argument(help="Dot-separated dataset path to list reflections for"), + path: str = typer.Argument(help="Dot-separated dataset path"), fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), ) -> None: - """List all reflections defined on a dataset. - - Queries sys.project.reflections filtered by the dataset ID. Shows reflection type, - status, and configuration. - """ + """List all reflections defined on a dataset.""" client = _get_client() _run_command(list_reflections(client, path), client, fmt) -@app.command("status") -def cli_status( - reflection_id: str = typer.Argument(help="Reflection ID (get IDs from 'drs reflect list')"), +@app.command("get") +def cli_get( + reflection_id: str = typer.Argument(help="Reflection ID"), fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), ) -> None: - """Show detailed status of a reflection. - - Includes freshness, staleness, size, last refresh time, and configuration. - """ + """Get detailed status of a reflection.""" client = _get_client() - _run_command(status(client, reflection_id), client, fmt) + _run_command(get_reflection(client, reflection_id), client, fmt) @app.command("refresh") def cli_refresh( reflection_id: str = typer.Argument(help="Reflection ID to refresh"), - dry_run: bool = typer.Option(False, "--dry-run", help="Validate the request without executing it"), + dry_run: bool = typer.Option(False, "--dry-run", help="Validate without executing"), fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), ) -> None: - """Trigger an immediate refresh of a reflection. - - Use --dry-run to validate the reflection ID without triggering the refresh. - """ + """Trigger an immediate refresh of a reflection.""" if dry_run: client = _get_client() - _run_command(status(client, reflection_id), client, fmt) + _run_command(get_reflection(client, reflection_id), client, fmt) return client = _get_client() _run_command(refresh(client, reflection_id), client, fmt) -@app.command("drop") -def cli_drop( +@app.command("delete") +def cli_delete( reflection_id: str = typer.Argument(help="Reflection ID to delete"), - dry_run: bool = typer.Option(False, "--dry-run", help="Validate the request without executing it"), + dry_run: bool = typer.Option(False, "--dry-run", help="Validate without executing"), fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), ) -> None: - """Permanently delete a reflection. Cannot be undone. - - Use --dry-run to verify the reflection exists before deleting. - """ + """Permanently delete a reflection. Cannot be undone.""" if dry_run: client = _get_client() - _run_command(status(client, reflection_id), client, fmt) + _run_command(get_reflection(client, reflection_id), client, fmt) return client = _get_client() - _run_command(drop(client, reflection_id), client, fmt) + _run_command(delete(client, reflection_id), client, fmt) diff --git a/src/drs/commands/role.py b/src/drs/commands/role.py new file mode 100644 index 0000000..cf139e9 --- /dev/null +++ b/src/drs/commands/role.py @@ -0,0 +1,164 @@ +# +# Copyright (C) 2017-2026 Dremio Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""drs role — manage Dremio Cloud roles.""" + +from __future__ import annotations + +import asyncio + +import httpx +import typer + +from drs.client import DremioClient +from drs.output import OutputFormat, output, error +from drs.utils import handle_api_error + +app = typer.Typer(help="Manage Dremio Cloud roles.") + + +async def list_roles(client: DremioClient) -> dict: + """List all roles.""" + try: + return await client.list_roles() + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + + +async def get_role(client: DremioClient, identifier: str) -> dict: + """Get role by name or ID. Tries name first, falls back to ID.""" + try: + return await client.get_role_by_name(identifier) + except httpx.HTTPStatusError as exc: + if exc.response.status_code == 404: + try: + return await client.get_role(identifier) + except httpx.HTTPStatusError as exc2: + raise handle_api_error(exc2) from exc2 + raise handle_api_error(exc) from exc + + +async def create_role(client: DremioClient, name: str) -> dict: + """Create a new role.""" + try: + return await client.create_role({"name": name}) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + + +async def update_role(client: DremioClient, role_id: str, name: str) -> dict: + """Update a role's name.""" + try: + existing = await client.get_role(role_id) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + body = dict(existing) + body["name"] = name + try: + return await client.update_role(role_id, body) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + + +async def delete_role(client: DremioClient, role_id: str) -> dict: + """Delete a role.""" + try: + return await client.delete_role(role_id) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + + +# -- CLI wrappers -- + +def _get_client() -> DremioClient: + from drs.cli import get_client + return get_client() + + +def _run_command(coro, client, fmt: OutputFormat = OutputFormat.json, fields: str | None = None) -> None: + async def _execute(): + try: + return await coro + finally: + await client.close() + + try: + result = asyncio.run(_execute()) + except Exception as exc: + from drs.utils import DremioAPIError + if isinstance(exc, DremioAPIError): + error(str(exc)) + raise typer.Exit(1) + if isinstance(exc, ValueError): + error(str(exc)) + raise typer.Exit(1) + raise + output(result, fmt, fields=fields) + + +@app.command("list") +def cli_list( + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), + fields: str = typer.Option(None, "--fields", "-f", help="Comma-separated fields to include"), +) -> None: + """List all roles in the organization.""" + client = _get_client() + _run_command(list_roles(client), client, fmt, fields=fields) + + +@app.command("get") +def cli_get( + identifier: str = typer.Argument(help="Role name or role ID"), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), + fields: str = typer.Option(None, "--fields", "-f", help="Comma-separated fields to include"), +) -> None: + """Get role details by name or ID.""" + client = _get_client() + _run_command(get_role(client, identifier), client, fmt, fields=fields) + + +@app.command("create") +def cli_create( + name: str = typer.Argument(help="Name for the new role"), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Create a new role.""" + client = _get_client() + _run_command(create_role(client, name), client, fmt) + + +@app.command("update") +def cli_update( + role_id: str = typer.Argument(help="Role ID to update"), + name: str = typer.Option(..., "--name", help="New role name"), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Update a role's name.""" + client = _get_client() + _run_command(update_role(client, role_id, name), client, fmt) + + +@app.command("delete") +def cli_delete( + role_id: str = typer.Argument(help="Role ID to delete"), + dry_run: bool = typer.Option(False, "--dry-run", help="Show role details without deleting"), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Delete a role. Cannot be undone.""" + client = _get_client() + if dry_run: + _run_command(get_role(client, role_id), client, fmt) + return + _run_command(delete_role(client, role_id), client, fmt) diff --git a/src/drs/commands/schema.py b/src/drs/commands/schema.py index 501774d..141def1 100644 --- a/src/drs/commands/schema.py +++ b/src/drs/commands/schema.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # -"""drs schema — describe tables, trace lineage, sample data.""" +"""dremio schema — describe tables, trace lineage, sample data.""" from __future__ import annotations @@ -27,7 +27,7 @@ from drs.output import OutputFormat, output, error from drs.utils import handle_api_error, parse_path, quote_path_sql -app = typer.Typer(help="Describe schemas, trace lineage, and sample data.") +app = typer.Typer(help="Describe table schemas, trace lineage, and sample data.") async def describe(client: DremioClient, path: str) -> dict: @@ -61,42 +61,6 @@ async def lineage(client: DremioClient, path: str) -> dict: return {"path": path, "id": entity_id, "graph": graph} -async def wiki(client: DremioClient, path: str) -> dict: - """Get wiki description and tags for an entity.""" - parts = parse_path(path) - try: - entity = await client.get_catalog_by_path(parts) - except httpx.HTTPStatusError as exc: - raise handle_api_error(exc) from exc - entity_id = entity["id"] - - wiki_text = "" - tags_list: list[str] = [] - try: - wiki_data = await client.get_wiki(entity_id) - wiki_text = wiki_data.get("text", "") - except httpx.HTTPStatusError as exc: - if exc.response.status_code == 404: - pass # no wiki exists for this entity - else: - raise handle_api_error(exc) from exc - try: - tags_data = await client.get_tags(entity_id) - tags_list = tags_data.get("tags", []) - except httpx.HTTPStatusError as exc: - if exc.response.status_code == 404: - pass # no tags exist for this entity - else: - raise handle_api_error(exc) from exc - - return { - "path": path, - "id": entity_id, - "wiki": wiki_text, - "tags": tags_list, - } - - async def sample(client: DremioClient, path: str, limit: int = 10) -> dict: """Return sample rows from a table/view.""" quoted = quote_path_sql(path) @@ -153,16 +117,6 @@ def cli_lineage( _run_command(lineage(client, path), client, fmt) -@app.command("wiki") -def cli_wiki( - path: str = typer.Argument(help='Dot-separated entity path'), - fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), -) -> None: - """Show wiki description and tags for a catalog entity.""" - client = _get_client() - _run_command(wiki(client, path), client, fmt) - - @app.command("sample") def cli_sample( path: str = typer.Argument(help='Dot-separated table/view path'), diff --git a/src/drs/commands/tag.py b/src/drs/commands/tag.py new file mode 100644 index 0000000..a00ad24 --- /dev/null +++ b/src/drs/commands/tag.py @@ -0,0 +1,132 @@ +# +# Copyright (C) 2017-2026 Dremio Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""dremio tag — get and update tags on catalog entities.""" + +from __future__ import annotations + +import asyncio + +import httpx +import typer + +from drs.client import DremioClient +from drs.output import OutputFormat, output, error +from drs.utils import handle_api_error, parse_path + +app = typer.Typer(help="Get and update tags on catalog entities.") + + +async def get_tags(client: DremioClient, path: str) -> dict: + """Get tags for an entity.""" + parts = parse_path(path) + try: + entity = await client.get_catalog_by_path(parts) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + entity_id = entity["id"] + + tags_list: list[str] = [] + try: + tags_data = await client.get_tags(entity_id) + tags_list = tags_data.get("tags", []) + except httpx.HTTPStatusError as exc: + if exc.response.status_code == 404: + pass # no tags exist for this entity + else: + raise handle_api_error(exc) from exc + + return { + "path": path, + "id": entity_id, + "tags": tags_list, + } + + +async def update_tags(client: DremioClient, path: str, tags: list[str]) -> dict: + """Set tags for an entity.""" + parts = parse_path(path) + try: + entity = await client.get_catalog_by_path(parts) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + entity_id = entity["id"] + + # Try to get existing tags for version number + version = None + try: + existing = await client.get_tags(entity_id) + version = existing.get("version") + except httpx.HTTPStatusError as exc: + if exc.response.status_code == 404: + pass # no tags exist yet + else: + raise handle_api_error(exc) from exc + + try: + result = await client.set_tags(entity_id, tags, version=version) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + return {"path": path, "id": entity_id, "tags": tags, "result": result} + + +# -- CLI wrappers -- + +def _get_client() -> DremioClient: + from drs.cli import get_client + return get_client() + + +def _run_command(coro, client, fmt: OutputFormat = OutputFormat.json, fields: str | None = None) -> None: + async def _execute(): + try: + return await coro + finally: + await client.close() + + try: + result = asyncio.run(_execute()) + except Exception as exc: + from drs.utils import DremioAPIError + if isinstance(exc, DremioAPIError): + error(str(exc)) + raise typer.Exit(1) + if isinstance(exc, ValueError): + error(str(exc)) + raise typer.Exit(1) + raise + output(result, fmt, fields=fields) + + +@app.command("get") +def cli_get( + path: str = typer.Argument(help='Dot-separated entity path'), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Get tags for a catalog entity.""" + client = _get_client() + _run_command(get_tags(client, path), client, fmt) + + +@app.command("update") +def cli_update( + path: str = typer.Argument(help='Dot-separated entity path'), + tags: str = typer.Argument(help='Comma-separated list of tags (e.g., "pii,finance,daily")'), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Set tags on a catalog entity. Replaces all existing tags.""" + client = _get_client() + tag_list = [t.strip() for t in tags.split(",") if t.strip()] + _run_command(update_tags(client, path, tag_list), client, fmt) diff --git a/src/drs/commands/user.py b/src/drs/commands/user.py new file mode 100644 index 0000000..20308b3 --- /dev/null +++ b/src/drs/commands/user.py @@ -0,0 +1,218 @@ +# +# Copyright (C) 2017-2026 Dremio Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""dremio user — manage Dremio Cloud users.""" + +from __future__ import annotations + +import asyncio + +import httpx +import typer + +from drs.client import DremioClient +from drs.output import OutputFormat, output, error +from drs.utils import handle_api_error + +app = typer.Typer(help="Manage Dremio Cloud users.") + + +async def list_users(client: DremioClient, max_results: int = 100) -> dict: + """List all users.""" + try: + return await client.list_users(max_results=max_results) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + + +async def get_user(client: DremioClient, identifier: str) -> dict: + """Get user by name or ID. Tries name first, falls back to ID.""" + try: + return await client.get_user_by_name(identifier) + except httpx.HTTPStatusError as exc: + if exc.response.status_code == 404: + try: + return await client.get_user(identifier) + except httpx.HTTPStatusError as exc2: + raise handle_api_error(exc2) from exc2 + raise handle_api_error(exc) from exc + + +async def create_user(client: DremioClient, email: str, role_id: str | None = None) -> dict: + """Create (invite) a user by email.""" + body: dict = {"email": email} + if role_id: + body["roleId"] = role_id + try: + return await client.invite_user(body) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + + +async def whoami(client: DremioClient) -> dict: + """Get info about the current user.""" + try: + return await client.list_users(max_results=1) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + + +async def audit(client: DremioClient, username: str) -> dict: + """Audit a user's effective permissions: user -> roles.""" + try: + user = await client.get_user_by_name(username) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + user_roles = user.get("roles", []) + + role_grants: list[dict] = [] + for role in user_roles: + role_id = role.get("id", role) if isinstance(role, dict) else role + role_name = role.get("name", role_id) if isinstance(role, dict) else role_id + role_grants.append({"role_id": role_id, "role_name": role_name}) + + return { + "username": username, + "user_id": user.get("id"), + "roles": role_grants, + } + + +async def update_user(client: DremioClient, user_id: str, name: str | None = None) -> dict: + """Update a user.""" + try: + existing = await client.get_user(user_id) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + body = dict(existing) + if name: + body["name"] = name + try: + return await client.update_user(user_id, body) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + + +async def delete_user(client: DremioClient, user_id: str) -> dict: + """Delete a user.""" + try: + return await client.delete_user(user_id) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + + +# -- CLI wrappers -- + +def _get_client() -> DremioClient: + from drs.cli import get_client + return get_client() + + +def _run_command(coro, client, fmt: OutputFormat = OutputFormat.json, fields: str | None = None) -> None: + async def _execute(): + try: + return await coro + finally: + await client.close() + + try: + result = asyncio.run(_execute()) + except Exception as exc: + from drs.utils import DremioAPIError + if isinstance(exc, DremioAPIError): + error(str(exc)) + raise typer.Exit(1) + if isinstance(exc, ValueError): + error(str(exc)) + raise typer.Exit(1) + raise + output(result, fmt, fields=fields) + + +@app.command("list") +def cli_list( + limit: int = typer.Option(100, "--limit", "-n", help="Max users to return"), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), + fields: str = typer.Option(None, "--fields", "-f", help="Comma-separated fields to include"), +) -> None: + """List all users in the organization.""" + client = _get_client() + _run_command(list_users(client, max_results=limit), client, fmt, fields=fields) + + +@app.command("get") +def cli_get( + identifier: str = typer.Argument(help="Username or user ID to look up"), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), + fields: str = typer.Option(None, "--fields", "-f", help="Comma-separated fields to include"), +) -> None: + """Get user details by name or ID.""" + client = _get_client() + _run_command(get_user(client, identifier), client, fmt, fields=fields) + + +@app.command("create") +def cli_create( + email: str = typer.Argument(help="Email address to invite"), + role_id: str = typer.Option(None, "--role-id", help="Role ID to assign to the new user"), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Create (invite) a new user by email address.""" + client = _get_client() + _run_command(create_user(client, email, role_id=role_id), client, fmt) + + +@app.command("update") +def cli_update( + user_id: str = typer.Argument(help="User ID to update"), + name: str = typer.Option(None, "--name", help="New display name"), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Update a user's properties.""" + client = _get_client() + _run_command(update_user(client, user_id, name=name), client, fmt) + + +@app.command("delete") +def cli_delete( + user_id: str = typer.Argument(help="User ID to delete"), + dry_run: bool = typer.Option(False, "--dry-run", help="Show user details without deleting"), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Delete a user from the organization. Cannot be undone.""" + client = _get_client() + if dry_run: + _run_command(get_user(client, user_id), client, fmt) + return + _run_command(delete_user(client, user_id), client, fmt) + + +@app.command("whoami") +def cli_whoami( + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Show current authenticated user info (best-effort).""" + client = _get_client() + _run_command(whoami(client), client, fmt) + + +@app.command("audit") +def cli_audit( + username: str = typer.Argument(help="Username to look up and audit"), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Audit a user's roles and effective permissions.""" + client = _get_client() + _run_command(audit(client, username), client, fmt) diff --git a/src/drs/commands/catalog.py b/src/drs/commands/wiki.py similarity index 51% rename from src/drs/commands/catalog.py rename to src/drs/commands/wiki.py index 762a5f1..947d846 100644 --- a/src/drs/commands/catalog.py +++ b/src/drs/commands/wiki.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # -"""drs catalog — browse and search the Dremio catalog.""" +"""dremio wiki — get and update wiki descriptions on catalog entities.""" from __future__ import annotations @@ -26,34 +26,60 @@ from drs.output import OutputFormat, output, error from drs.utils import handle_api_error, parse_path -app = typer.Typer(help="Browse and search the Dremio catalog.") +app = typer.Typer(help="Get and update wiki descriptions on catalog entities.") -async def list_catalog(client: DremioClient) -> dict: - """List top-level catalog entities (sources, spaces, home).""" +async def get_wiki(client: DremioClient, path: str) -> dict: + """Get wiki description for an entity.""" + parts = parse_path(path) try: - root = await client.get_catalog_entity("") + entity = await client.get_catalog_by_path(parts) except httpx.HTTPStatusError as exc: raise handle_api_error(exc) from exc - children = root.get("data", root.get("children", [])) - return {"entities": children} + entity_id = entity["id"] + + wiki_text = "" + try: + wiki_data = await client.get_wiki(entity_id) + wiki_text = wiki_data.get("text", "") + except httpx.HTTPStatusError as exc: + if exc.response.status_code == 404: + pass # no wiki exists for this entity + else: + raise handle_api_error(exc) from exc + return { + "path": path, + "id": entity_id, + "wiki": wiki_text, + } -async def get_entity(client: DremioClient, path: str) -> dict: - """Get a catalog entity by dot-separated path.""" + +async def update_wiki(client: DremioClient, path: str, text: str) -> dict: + """Set or update wiki description text for an entity.""" parts = parse_path(path) try: - return await client.get_catalog_by_path(parts) + entity = await client.get_catalog_by_path(parts) except httpx.HTTPStatusError as exc: raise handle_api_error(exc) from exc + entity_id = entity["id"] + # Try to get existing wiki for version number (optimistic concurrency) + version = None + try: + existing = await client.get_wiki(entity_id) + version = existing.get("version") + except httpx.HTTPStatusError as exc: + if exc.response.status_code == 404: + pass # no wiki exists yet + else: + raise handle_api_error(exc) from exc -async def search_catalog(client: DremioClient, term: str) -> dict: - """Full-text search for catalog entities.""" try: - return await client.search(term) + result = await client.set_wiki(entity_id, text, version=version) except httpx.HTTPStatusError as exc: raise handle_api_error(exc) from exc + return {"path": path, "id": entity_id, "wiki": text, "result": result} # -- CLI wrappers -- @@ -84,36 +110,22 @@ async def _execute(): output(result, fmt, fields=fields) -@app.command("list") -def cli_list( - fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), - fields: str = typer.Option(None, "--fields", "-f", help="Comma-separated fields to include in output"), -) -> None: - """List top-level catalog entities: sources, spaces, and home folder.""" - client = _get_client() - _run_command(list_catalog(client), client, fmt, fields=fields) - - @app.command("get") def cli_get( - path: str = typer.Argument(help='Dot-separated entity path (e.g., myspace.folder.table). Quote components with dots: \'"My Source".table\''), + path: str = typer.Argument(help='Dot-separated entity path'), fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), - fields: str = typer.Option(None, "--fields", "-f", help="Comma-separated fields to include in output"), ) -> None: - """Get full metadata for a catalog entity by path. - - Returns entity type, ID, children (for containers), fields (for datasets), - and access control information. - """ + """Get wiki description for a catalog entity.""" client = _get_client() - _run_command(get_entity(client, path), client, fmt, fields=fields) + _run_command(get_wiki(client, path), client, fmt) -@app.command("search") -def cli_search( - term: str = typer.Argument(help="Search term (matches table names, view names, source names)"), +@app.command("update") +def cli_update( + path: str = typer.Argument(help='Dot-separated entity path'), + text: str = typer.Argument(help='Wiki text to set (Markdown supported)'), fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), ) -> None: - """Full-text search across all catalog entities (tables, views, sources).""" + """Set or update the wiki description for a catalog entity.""" client = _get_client() - _run_command(search_catalog(client, term), client, fmt) + _run_command(update_wiki(client, path, text), client, fmt) diff --git a/src/drs/introspect.py b/src/drs/introspect.py index 7edffc8..dad1336 100644 --- a/src/drs/introspect.py +++ b/src/drs/introspect.py @@ -13,84 +13,108 @@ # See the License for the specific language governing permissions and # limitations under the License. # -"""Runtime schema introspection for drs commands. +"""Runtime schema introspection for dremio commands. Provides machine-readable descriptions of command parameters, types, and -constraints. Agents call `drs describe ` to self-serve schema info -instead of relying on pre-stuffed documentation (cheaper on token budget, -always up to date). +constraints. Agents call `dremio describe ` to self-serve schema info. """ from __future__ import annotations from drs.utils import VALID_JOB_STATES -# Command registry — one entry per CLI command with full parameter schema. -# This is the canonical source of truth for what each command accepts. - COMMAND_SCHEMAS: dict[str, dict] = { + # -- Query -- "query.run": { - "group": "query", - "command": "run", - "description": "Execute a SQL query against Dremio Cloud, wait for completion, return results as JSON.", + "group": "query", "command": "run", + "description": "Execute a SQL query against Dremio Cloud, wait for completion, return results.", "mechanism": "REST", "endpoints": ["POST /v0/projects/{pid}/sql", "GET /v0/projects/{pid}/job/{id}", "GET /v0/projects/{pid}/job/{id}/results"], "parameters": [ {"name": "sql", "type": "string", "required": True, "positional": True, "description": "SQL query to execute"}, - {"name": "context", "type": "string", "required": False, "description": "Dot-separated default schema context (e.g., myspace.folder)"}, - {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"], "description": "Output format"}, - {"name": "fields", "type": "string", "required": False, "flag": "--fields/-f", "description": "Comma-separated fields to include in output (reduces context window usage)"}, + {"name": "context", "type": "string", "required": False, "description": "Dot-separated default schema context"}, + {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + {"name": "fields", "type": "string", "required": False, "flag": "--fields/-f", "description": "Comma-separated fields to include"}, ], }, "query.status": { - "group": "query", - "command": "status", + "group": "query", "command": "status", "description": "Check the status of a Dremio job by UUID.", "mechanism": "REST", "endpoints": ["GET /v0/projects/{pid}/job/{id}"], "parameters": [ - {"name": "job_id", "type": "string", "required": True, "positional": True, "format": "uuid", "description": "Job ID (UUID)"}, + {"name": "job_id", "type": "string", "required": True, "positional": True, "format": "uuid"}, {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, ], }, "query.cancel": { - "group": "query", - "command": "cancel", - "description": "Cancel a running Dremio job. No effect if already completed.", - "mechanism": "REST", + "group": "query", "command": "cancel", + "description": "Cancel a running Dremio job.", + "mechanism": "REST", "mutating": True, "endpoints": ["POST /v0/projects/{pid}/job/{id}/cancel"], - "mutating": True, "parameters": [ - {"name": "job_id", "type": "string", "required": True, "positional": True, "format": "uuid", "description": "Job ID (UUID) to cancel"}, - {"name": "dry_run", "type": "boolean", "required": False, "default": False, "flag": "--dry-run", "description": "Check job status without cancelling"}, + {"name": "job_id", "type": "string", "required": True, "positional": True, "format": "uuid"}, + {"name": "dry_run", "type": "boolean", "required": False, "default": False, "flag": "--dry-run"}, ], }, - "catalog.list": { - "group": "catalog", - "command": "list", + + # -- Folder -- + "folder.list": { + "group": "folder", "command": "list", "description": "List top-level catalog entities: sources, spaces, and home folder.", "mechanism": "REST", "endpoints": ["GET /v0/projects/{pid}/catalog"], "parameters": [ {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, - {"name": "fields", "type": "string", "required": False, "description": "Comma-separated fields to include"}, + {"name": "fields", "type": "string", "required": False}, ], }, - "catalog.get": { - "group": "catalog", - "command": "get", + "folder.get": { + "group": "folder", "command": "get", "description": "Get full metadata for a catalog entity by dot-separated path.", "mechanism": "REST", "endpoints": ["GET /v0/projects/{pid}/catalog/by-path/{path}"], "parameters": [ - {"name": "path", "type": "string", "required": True, "positional": True, "description": "Dot-separated entity path (e.g., myspace.folder.table)"}, + {"name": "path", "type": "string", "required": True, "positional": True, "description": "Dot-separated entity path"}, + {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + {"name": "fields", "type": "string", "required": False}, + ], + }, + "folder.create": { + "group": "folder", "command": "create", + "description": "Create a space (single name) or folder (nested path) using SQL.", + "mechanism": "SQL", "mutating": True, + "sql_template": "CREATE SPACE \"{name}\" / CREATE FOLDER {path}", + "parameters": [ + {"name": "path", "type": "string", "required": True, "positional": True, "description": "Space name or dot-separated folder path"}, + {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + ], + }, + "folder.delete": { + "group": "folder", "command": "delete", + "description": "Delete a catalog entity (space, folder, view, etc.). Cannot be undone.", + "mechanism": "REST", "mutating": True, + "endpoints": ["GET /v0/projects/{pid}/catalog/by-path/{path}", "DELETE /v0/projects/{pid}/catalog/{id}"], + "parameters": [ + {"name": "path", "type": "string", "required": True, "positional": True}, + {"name": "dry_run", "type": "boolean", "required": False, "default": False}, + {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + ], + }, + "folder.grants": { + "group": "folder", "command": "grants", + "description": "Show ACL grants on a catalog entity.", + "mechanism": "REST", + "endpoints": ["GET /v0/projects/{pid}/catalog/by-path/{path}"], + "parameters": [ + {"name": "path", "type": "string", "required": True, "positional": True}, {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, - {"name": "fields", "type": "string", "required": False, "description": "Comma-separated fields to include"}, ], }, - "catalog.search": { - "group": "catalog", - "command": "search", + + # -- Search (top-level) -- + "search": { + "group": None, "command": "search", "description": "Full-text search for tables, views, and sources by keyword.", "mechanism": "REST", "endpoints": ["POST /v0/projects/{pid}/search"], @@ -99,174 +123,424 @@ {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, ], }, + + # -- Schema -- "schema.describe": { - "group": "schema", - "command": "describe", + "group": "schema", "command": "describe", "description": "Get column names, data types, and nullability for a table or view.", "mechanism": "REST", "endpoints": ["GET /v0/projects/{pid}/catalog/by-path/{path}"], "parameters": [ - {"name": "path", "type": "string", "required": True, "positional": True, "description": "Dot-separated table/view path"}, + {"name": "path", "type": "string", "required": True, "positional": True}, {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, - {"name": "fields", "type": "string", "required": False, "description": "Comma-separated fields to include (e.g., 'columns.name,columns.type')"}, + {"name": "fields", "type": "string", "required": False}, ], }, "schema.lineage": { - "group": "schema", - "command": "lineage", + "group": "schema", "command": "lineage", "description": "Get upstream and downstream dependency graph for a table or view.", "mechanism": "REST", "endpoints": ["GET /v0/projects/{pid}/catalog/by-path/{path}", "GET /v0/projects/{pid}/catalog/{id}/graph"], "parameters": [ - {"name": "path", "type": "string", "required": True, "positional": True, "description": "Dot-separated table/view path"}, + {"name": "path", "type": "string", "required": True, "positional": True}, + {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + ], + }, + "schema.sample": { + "group": "schema", "command": "sample", + "description": "Return sample rows from a table or view.", + "mechanism": "SQL", + "sql_template": "SELECT * FROM {path} LIMIT {limit}", + "parameters": [ + {"name": "path", "type": "string", "required": True, "positional": True}, + {"name": "limit", "type": "integer", "required": False, "default": 10}, {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + {"name": "fields", "type": "string", "required": False}, ], }, - "schema.wiki": { - "group": "schema", - "command": "wiki", - "description": "Get wiki documentation text and tags for a catalog entity.", + + # -- Wiki -- + "wiki.get": { + "group": "wiki", "command": "get", + "description": "Get wiki description for a catalog entity.", "mechanism": "REST", - "endpoints": ["GET /v0/projects/{pid}/catalog/by-path/{path}", "GET /v0/projects/{pid}/catalog/{id}/collaboration/wiki", "GET /v0/projects/{pid}/catalog/{id}/collaboration/tag"], + "endpoints": ["GET /v0/projects/{pid}/catalog/by-path/{path}", "GET /v0/projects/{pid}/catalog/{id}/collaboration/wiki"], "parameters": [ - {"name": "path", "type": "string", "required": True, "positional": True, "description": "Dot-separated entity path"}, + {"name": "path", "type": "string", "required": True, "positional": True}, {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, ], }, - "schema.sample": { - "group": "schema", - "command": "sample", - "description": "Return sample rows from a table or view.", - "mechanism": "SQL", - "sql_template": "SELECT * FROM {path} LIMIT {limit}", + "wiki.update": { + "group": "wiki", "command": "update", + "description": "Set or update the wiki description for a catalog entity.", + "mechanism": "REST", "mutating": True, + "endpoints": ["GET /v0/projects/{pid}/catalog/by-path/{path}", "POST /v0/projects/{pid}/catalog/{id}/collaboration/wiki"], "parameters": [ - {"name": "path", "type": "string", "required": True, "positional": True, "description": "Dot-separated table/view path"}, - {"name": "limit", "type": "integer", "required": False, "default": 10, "description": "Number of sample rows"}, + {"name": "path", "type": "string", "required": True, "positional": True}, + {"name": "text", "type": "string", "required": True, "positional": True, "description": "Wiki text (Markdown supported)"}, {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, - {"name": "fields", "type": "string", "required": False, "description": "Comma-separated fields to include"}, ], }, - "reflect.list": { - "group": "reflect", - "command": "list", + + # -- Tag -- + "tag.get": { + "group": "tag", "command": "get", + "description": "Get tags for a catalog entity.", + "mechanism": "REST", + "endpoints": ["GET /v0/projects/{pid}/catalog/by-path/{path}", "GET /v0/projects/{pid}/catalog/{id}/collaboration/tag"], + "parameters": [ + {"name": "path", "type": "string", "required": True, "positional": True}, + {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + ], + }, + "tag.update": { + "group": "tag", "command": "update", + "description": "Set tags on a catalog entity. Replaces all existing tags.", + "mechanism": "REST", "mutating": True, + "endpoints": ["GET /v0/projects/{pid}/catalog/by-path/{path}", "POST /v0/projects/{pid}/catalog/{id}/collaboration/tag"], + "parameters": [ + {"name": "path", "type": "string", "required": True, "positional": True}, + {"name": "tags", "type": "string", "required": True, "positional": True, "description": "Comma-separated tags"}, + {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + ], + }, + + # -- Reflection -- + "reflection.create": { + "group": "reflection", "command": "create", + "description": "Create a new reflection on a dataset.", + "mechanism": "REST", "mutating": True, + "endpoints": ["GET /v0/projects/{pid}/catalog/by-path/{path}", "POST /v0/projects/{pid}/reflection"], + "parameters": [ + {"name": "path", "type": "string", "required": True, "positional": True}, + {"name": "type", "type": "enum", "required": False, "default": "raw", "enum": ["raw", "aggregation"], "flag": "--type/-t"}, + {"name": "fields", "type": "string", "required": False, "flag": "--fields/-f"}, + {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + ], + }, + "reflection.list": { + "group": "reflection", "command": "list", "description": "List all reflections defined on a dataset.", "mechanism": "SQL", "sql_template": "SELECT * FROM sys.project.reflections WHERE dataset_id = '{dataset_id}'", "parameters": [ - {"name": "path", "type": "string", "required": True, "positional": True, "description": "Dot-separated dataset path"}, + {"name": "path", "type": "string", "required": True, "positional": True}, {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, ], }, - "reflect.status": { - "group": "reflect", - "command": "status", + "reflection.get": { + "group": "reflection", "command": "get", "description": "Get detailed status of a reflection by ID.", "mechanism": "REST", "endpoints": ["GET /v0/projects/{pid}/reflection/{id}"], "parameters": [ - {"name": "reflection_id", "type": "string", "required": True, "positional": True, "description": "Reflection ID (get from 'drs reflect list')"}, + {"name": "reflection_id", "type": "string", "required": True, "positional": True}, {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, ], }, - "reflect.refresh": { - "group": "reflect", - "command": "refresh", + "reflection.refresh": { + "group": "reflection", "command": "refresh", "description": "Trigger an immediate refresh of a reflection.", - "mechanism": "REST", + "mechanism": "REST", "mutating": True, "endpoints": ["POST /v0/projects/{pid}/reflection/{id}/refresh"], - "mutating": True, "parameters": [ - {"name": "reflection_id", "type": "string", "required": True, "positional": True, "description": "Reflection ID to refresh"}, - {"name": "dry_run", "type": "boolean", "required": False, "default": False, "description": "Validate without executing"}, + {"name": "reflection_id", "type": "string", "required": True, "positional": True}, + {"name": "dry_run", "type": "boolean", "required": False, "default": False}, ], }, - "reflect.drop": { - "group": "reflect", - "command": "drop", + "reflection.delete": { + "group": "reflection", "command": "delete", "description": "Permanently delete a reflection. Cannot be undone.", - "mechanism": "REST", + "mechanism": "REST", "mutating": True, "endpoints": ["DELETE /v0/projects/{pid}/reflection/{id}"], - "mutating": True, "parameters": [ - {"name": "reflection_id", "type": "string", "required": True, "positional": True, "description": "Reflection ID to delete"}, - {"name": "dry_run", "type": "boolean", "required": False, "default": False, "description": "Validate without executing"}, + {"name": "reflection_id", "type": "string", "required": True, "positional": True}, + {"name": "dry_run", "type": "boolean", "required": False, "default": False}, ], }, - "jobs.list": { - "group": "jobs", - "command": "list", + + # -- Job -- + "job.list": { + "group": "job", "command": "list", "description": "List recent query jobs, optionally filtered by status.", "mechanism": "SQL", - "sql_template": "SELECT job_id, user_name, query_type, status, submitted_ts, final_state_ts FROM sys.project.jobs WHERE ... ORDER BY submitted_ts DESC LIMIT {limit}", + "sql_template": "SELECT ... FROM sys.project.jobs WHERE ... ORDER BY submitted_ts DESC LIMIT {limit}", "parameters": [ - {"name": "status", "type": "enum", "required": False, "enum": sorted(VALID_JOB_STATES), "description": "Filter by job state"}, - {"name": "limit", "type": "integer", "required": False, "default": 25, "description": "Max jobs to return"}, + {"name": "status", "type": "enum", "required": False, "enum": sorted(VALID_JOB_STATES)}, + {"name": "limit", "type": "integer", "required": False, "default": 25}, {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, - {"name": "fields", "type": "string", "required": False, "description": "Comma-separated fields to include"}, + {"name": "fields", "type": "string", "required": False}, ], }, - "jobs.get": { - "group": "jobs", - "command": "get", + "job.get": { + "group": "job", "command": "get", "description": "Get detailed status and metadata for a specific job.", "mechanism": "REST", "endpoints": ["GET /v0/projects/{pid}/job/{id}"], "parameters": [ - {"name": "job_id", "type": "string", "required": True, "positional": True, "format": "uuid", "description": "Job ID (UUID)"}, + {"name": "job_id", "type": "string", "required": True, "positional": True, "format": "uuid"}, {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, - {"name": "fields", "type": "string", "required": False, "flag": "--fields/-f", "description": "Comma-separated fields to include in output"}, + {"name": "fields", "type": "string", "required": False, "flag": "--fields/-f"}, ], }, - "jobs.profile": { - "group": "jobs", - "command": "profile", - "description": "Get operator-level execution profile for a completed job.", + "job.profile": { + "group": "job", "command": "profile", + "description": "Get execution profile for a completed job.", "mechanism": "SQL", - "sql_template": "SELECT job_id, status, query_type, query, planner_estimated_cost, rows_scanned, bytes_scanned, rows_returned, bytes_returned, accelerated, engine, submitted_ts, final_state_ts, error_msg FROM sys.project.jobs WHERE job_id = '{job_id}'", + "sql_template": "SELECT ... FROM sys.project.jobs WHERE job_id = '{job_id}'", "parameters": [ - {"name": "job_id", "type": "string", "required": True, "positional": True, "format": "uuid", "description": "Job ID (UUID)"}, + {"name": "job_id", "type": "string", "required": True, "positional": True, "format": "uuid"}, {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, ], }, - "access.grants": { - "group": "access", - "command": "grants", - "description": "Get ACL grants on a catalog entity.", + + # -- Engine -- + "engine.list": { + "group": "engine", "command": "list", + "description": "List all engines in the project.", "mechanism": "REST", - "endpoints": ["GET /v0/projects/{pid}/catalog/by-path/{path}"], + "endpoints": ["GET /v0/projects/{pid}/engines"], "parameters": [ - {"name": "path", "type": "string", "required": True, "positional": True, "description": "Dot-separated entity path"}, {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + {"name": "fields", "type": "string", "required": False}, ], }, - "access.roles": { - "group": "access", - "command": "roles", - "description": "List all roles in the organization.", + "engine.get": { + "group": "engine", "command": "get", + "description": "Get details for a specific engine.", "mechanism": "REST", - "endpoints": ["GET /v1/roles"], + "endpoints": ["GET /v0/projects/{pid}/engines/{id}"], "parameters": [ + {"name": "engine_id", "type": "string", "required": True, "positional": True}, {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + {"name": "fields", "type": "string", "required": False}, ], }, - "access.whoami": { - "group": "access", - "command": "whoami", - "description": "Get current authenticated user info (best-effort).", + "engine.create": { + "group": "engine", "command": "create", + "description": "Create a new engine.", + "mechanism": "REST", "mutating": True, + "endpoints": ["POST /v0/projects/{pid}/engines"], + "parameters": [ + {"name": "name", "type": "string", "required": True, "positional": True}, + {"name": "size", "type": "enum", "required": False, "default": "SMALL", "enum": ["SMALL", "MEDIUM", "LARGE", "XLARGE", "XXLARGE"]}, + {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + ], + }, + "engine.update": { + "group": "engine", "command": "update", + "description": "Update engine configuration (name, size).", + "mechanism": "REST", "mutating": True, + "endpoints": ["GET /v0/projects/{pid}/engines/{id}", "PUT /v0/projects/{pid}/engines/{id}"], + "parameters": [ + {"name": "engine_id", "type": "string", "required": True, "positional": True}, + {"name": "name", "type": "string", "required": False, "flag": "--name"}, + {"name": "size", "type": "string", "required": False, "flag": "--size/-s"}, + {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + ], + }, + "engine.delete": { + "group": "engine", "command": "delete", + "description": "Delete an engine. Cannot be undone.", + "mechanism": "REST", "mutating": True, + "endpoints": ["DELETE /v0/projects/{pid}/engines/{id}"], + "parameters": [ + {"name": "engine_id", "type": "string", "required": True, "positional": True}, + {"name": "dry_run", "type": "boolean", "required": False, "default": False}, + {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + ], + }, + "engine.enable": { + "group": "engine", "command": "enable", + "description": "Enable a disabled engine.", + "mechanism": "REST", "mutating": True, + "endpoints": ["PUT /v0/projects/{pid}/engines/{id}/enable"], + "parameters": [ + {"name": "engine_id", "type": "string", "required": True, "positional": True}, + {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + ], + }, + "engine.disable": { + "group": "engine", "command": "disable", + "description": "Disable a running engine.", + "mechanism": "REST", "mutating": True, + "endpoints": ["PUT /v0/projects/{pid}/engines/{id}/disable"], + "parameters": [ + {"name": "engine_id", "type": "string", "required": True, "positional": True}, + {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + ], + }, + + # -- User -- + "user.list": { + "group": "user", "command": "list", + "description": "List all users in the organization.", "mechanism": "REST", "endpoints": ["GET /v1/users"], "parameters": [ + {"name": "limit", "type": "integer", "required": False, "default": 100, "flag": "--limit/-n"}, + {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + {"name": "fields", "type": "string", "required": False}, + ], + }, + "user.get": { + "group": "user", "command": "get", + "description": "Get user details by name or ID.", + "mechanism": "REST", + "endpoints": ["GET /v1/users/name/{userName}", "GET /v1/users/{userId}"], + "parameters": [ + {"name": "identifier", "type": "string", "required": True, "positional": True}, {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + {"name": "fields", "type": "string", "required": False}, ], }, - "access.audit": { - "group": "access", - "command": "audit", + "user.create": { + "group": "user", "command": "create", + "description": "Create (invite) a new user by email address.", + "mechanism": "REST", "mutating": True, + "endpoints": ["POST /v1/users/invite"], + "parameters": [ + {"name": "email", "type": "string", "required": True, "positional": True}, + {"name": "role_id", "type": "string", "required": False, "flag": "--role-id"}, + {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + ], + }, + "user.update": { + "group": "user", "command": "update", + "description": "Update a user's properties.", + "mechanism": "REST", "mutating": True, + "endpoints": ["GET /v1/users/{userId}", "PUT /v1/users/{userId}"], + "parameters": [ + {"name": "user_id", "type": "string", "required": True, "positional": True}, + {"name": "name", "type": "string", "required": False, "flag": "--name"}, + {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + ], + }, + "user.delete": { + "group": "user", "command": "delete", + "description": "Delete a user from the organization. Cannot be undone.", + "mechanism": "REST", "mutating": True, + "endpoints": ["DELETE /v1/users/{userId}"], + "parameters": [ + {"name": "user_id", "type": "string", "required": True, "positional": True}, + {"name": "dry_run", "type": "boolean", "required": False, "default": False}, + {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + ], + }, + "user.whoami": { + "group": "user", "command": "whoami", + "description": "Show current authenticated user info (best-effort).", + "mechanism": "REST", + "endpoints": ["GET /v1/users"], + "parameters": [ + {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + ], + }, + "user.audit": { + "group": "user", "command": "audit", "description": "Audit a user's roles and effective permissions by username.", "mechanism": "REST", "endpoints": ["GET /v1/users/name/{userName}"], "parameters": [ - {"name": "username", "type": "string", "required": True, "positional": True, "description": "Username to audit"}, + {"name": "username", "type": "string", "required": True, "positional": True}, + {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + ], + }, + + # -- Role -- + "role.list": { + "group": "role", "command": "list", + "description": "List all roles in the organization.", + "mechanism": "REST", + "endpoints": ["GET /v1/roles"], + "parameters": [ + {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + {"name": "fields", "type": "string", "required": False}, + ], + }, + "role.get": { + "group": "role", "command": "get", + "description": "Get role details by name or ID.", + "mechanism": "REST", + "endpoints": ["GET /v1/roles/name/{roleName}", "GET /v1/roles/{roleId}"], + "parameters": [ + {"name": "identifier", "type": "string", "required": True, "positional": True}, + {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + {"name": "fields", "type": "string", "required": False}, + ], + }, + "role.create": { + "group": "role", "command": "create", + "description": "Create a new role.", + "mechanism": "REST", "mutating": True, + "endpoints": ["POST /v1/roles"], + "parameters": [ + {"name": "name", "type": "string", "required": True, "positional": True}, + {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + ], + }, + "role.update": { + "group": "role", "command": "update", + "description": "Update a role's name.", + "mechanism": "REST", "mutating": True, + "endpoints": ["GET /v1/roles/{roleId}", "PUT /v1/roles/{roleId}"], + "parameters": [ + {"name": "role_id", "type": "string", "required": True, "positional": True}, + {"name": "name", "type": "string", "required": True, "flag": "--name"}, + {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + ], + }, + "role.delete": { + "group": "role", "command": "delete", + "description": "Delete a role. Cannot be undone.", + "mechanism": "REST", "mutating": True, + "endpoints": ["DELETE /v1/roles/{roleId}"], + "parameters": [ + {"name": "role_id", "type": "string", "required": True, "positional": True}, + {"name": "dry_run", "type": "boolean", "required": False, "default": False}, + {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + ], + }, + + # -- Grant -- + "grant.get": { + "group": "grant", "command": "get", + "description": "Get grants for a user or role on a resource.", + "mechanism": "REST", + "endpoints": ["GET /v1/{scope}/{scopeId}/grants/{granteeType}/{granteeId}"], + "parameters": [ + {"name": "scope", "type": "string", "required": True, "positional": True, "description": "Resource scope (projects, orgs, clouds)"}, + {"name": "scope_id", "type": "string", "required": True, "positional": True}, + {"name": "grantee_type", "type": "string", "required": True, "positional": True, "description": "user or role"}, + {"name": "grantee_id", "type": "string", "required": True, "positional": True}, + {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + ], + }, + "grant.update": { + "group": "grant", "command": "update", + "description": "Set grants (privileges) for a user or role on a resource.", + "mechanism": "REST", "mutating": True, + "endpoints": ["PUT /v1/{scope}/{scopeId}/grants/{granteeType}/{granteeId}"], + "parameters": [ + {"name": "scope", "type": "string", "required": True, "positional": True}, + {"name": "scope_id", "type": "string", "required": True, "positional": True}, + {"name": "grantee_type", "type": "string", "required": True, "positional": True}, + {"name": "grantee_id", "type": "string", "required": True, "positional": True}, + {"name": "privileges", "type": "string", "required": True, "positional": True, "description": "Comma-separated privileges"}, + {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + ], + }, + "grant.delete": { + "group": "grant", "command": "delete", + "description": "Remove all grants for a user or role on a resource.", + "mechanism": "REST", "mutating": True, + "endpoints": ["DELETE /v1/{scope}/{scopeId}/grants/{granteeType}/{granteeId}"], + "parameters": [ + {"name": "scope", "type": "string", "required": True, "positional": True}, + {"name": "scope_id", "type": "string", "required": True, "positional": True}, + {"name": "grantee_type", "type": "string", "required": True, "positional": True}, + {"name": "grantee_id", "type": "string", "required": True, "positional": True}, + {"name": "dry_run", "type": "boolean", "required": False, "default": False}, {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, ], }, diff --git a/tests/conftest.py b/tests/conftest.py index c28bc0b..92e4538 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -40,5 +40,6 @@ def mock_client(config: DrsConfig) -> DremioClient: client = DremioClient(config) client._get = AsyncMock() client._post = AsyncMock() + client._put = AsyncMock() client._delete = AsyncMock() return client diff --git a/tests/test_client.py b/tests/test_client.py index 20a8040..9c5799c 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -76,6 +76,24 @@ def test_catalog_entity_with_id(self, client: DremioClient) -> None: assert client._v3("/catalog/abc-123") == "https://api.dremio.cloud/v0/projects/proj-123/catalog/abc-123" +class TestEngineURLs: + def test_engines_list_url(self, client: DremioClient) -> None: + assert client._v0("/engines") == "https://api.dremio.cloud/v0/projects/proj-123/engines" + + def test_engine_enable_url(self, client: DremioClient) -> None: + assert client._v0("/engines/eng-1/enable") == "https://api.dremio.cloud/v0/projects/proj-123/engines/eng-1/enable" + + +class TestGrantURLs: + def test_project_grant_url(self, client: DremioClient) -> None: + url = client._v1("/projects/proj-1/grants/role/role-1") + assert url == "https://api.dremio.cloud/v1/projects/proj-1/grants/role/role-1" + + def test_org_grant_url(self, client: DremioClient) -> None: + url = client._v1("/orgs/org-1/grants/user/user-1") + assert url == "https://api.dremio.cloud/v1/orgs/org-1/grants/user/user-1" + + class TestClientHeaders: def test_auth_header(self, client: DremioClient) -> None: assert client._client.headers["authorization"] == "Bearer test-token" diff --git a/tests/test_client_retry.py b/tests/test_client_retry.py new file mode 100644 index 0000000..b2449b1 --- /dev/null +++ b/tests/test_client_retry.py @@ -0,0 +1,154 @@ +# +# Copyright (C) 2017-2026 Dremio Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Tests for retry logic in DremioClient.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, patch + +import httpx +import pytest + +from drs.client import DremioClient + + +@pytest.mark.asyncio +async def test_retry_on_timeout(config) -> None: + """Should retry on timeout and succeed on second attempt.""" + client = DremioClient(config) + ok_response = httpx.Response(200, json={"ok": True}, request=httpx.Request("GET", "https://example.com")) + + client._client.request = AsyncMock(side_effect=[ + httpx.TimeoutException("timed out"), + ok_response, + ]) + + with patch("drs.client.asyncio.sleep", new_callable=AsyncMock): + result = await client._get("https://example.com/test") + + assert result == {"ok": True} + assert client._client.request.call_count == 2 + + +@pytest.mark.asyncio +async def test_retry_on_429(config) -> None: + """Should retry on 429 Too Many Requests.""" + client = DremioClient(config) + rate_limited = httpx.Response(429, request=httpx.Request("GET", "https://example.com")) + ok_response = httpx.Response(200, json={"ok": True}, request=httpx.Request("GET", "https://example.com")) + + client._client.request = AsyncMock(side_effect=[rate_limited, ok_response]) + + with patch("drs.client.asyncio.sleep", new_callable=AsyncMock): + result = await client._get("https://example.com/test") + + assert result == {"ok": True} + assert client._client.request.call_count == 2 + + +@pytest.mark.asyncio +async def test_retry_on_503(config) -> None: + """Should retry on 503 Service Unavailable.""" + client = DremioClient(config) + unavailable = httpx.Response(503, request=httpx.Request("POST", "https://example.com")) + ok_response = httpx.Response(200, json={"data": 1}, request=httpx.Request("POST", "https://example.com")) + + client._client.request = AsyncMock(side_effect=[unavailable, ok_response]) + + with patch("drs.client.asyncio.sleep", new_callable=AsyncMock): + result = await client._post("https://example.com/test", json={"sql": "SELECT 1"}) + + assert result == {"data": 1} + assert client._client.request.call_count == 2 + + +@pytest.mark.asyncio +async def test_no_retry_on_400(config) -> None: + """Should NOT retry on 400 Bad Request.""" + client = DremioClient(config) + bad_request = httpx.Response(400, json={"error": "bad"}, request=httpx.Request("GET", "https://example.com")) + + client._client.request = AsyncMock(return_value=bad_request) + + with pytest.raises(httpx.HTTPStatusError): + await client._get("https://example.com/test") + + assert client._client.request.call_count == 1 + + +@pytest.mark.asyncio +async def test_no_retry_on_404(config) -> None: + """Should NOT retry on 404 Not Found.""" + client = DremioClient(config) + not_found = httpx.Response(404, json={"error": "not found"}, request=httpx.Request("GET", "https://example.com")) + + client._client.request = AsyncMock(return_value=not_found) + + with pytest.raises(httpx.HTTPStatusError): + await client._get("https://example.com/test") + + assert client._client.request.call_count == 1 + + +@pytest.mark.asyncio +async def test_exhausted_retries_raises_timeout(config) -> None: + """Should raise TimeoutException after all retries exhausted.""" + client = DremioClient(config) + + client._client.request = AsyncMock(side_effect=httpx.TimeoutException("timed out")) + + with patch("drs.client.asyncio.sleep", new_callable=AsyncMock) as mock_sleep: + with pytest.raises(httpx.TimeoutException): + await client._get("https://example.com/test") + + assert client._client.request.call_count == 3 + assert mock_sleep.call_count == 2 + + +@pytest.mark.asyncio +async def test_exhausted_retries_returns_last_status(config) -> None: + """Should raise HTTPStatusError if all retries return retryable status.""" + client = DremioClient(config) + unavailable = httpx.Response(503, request=httpx.Request("GET", "https://example.com")) + + client._client.request = AsyncMock(return_value=unavailable) + + with patch("drs.client.asyncio.sleep", new_callable=AsyncMock): + with pytest.raises(httpx.HTTPStatusError): + await client._get("https://example.com/test") + + assert client._client.request.call_count == 3 + + +@pytest.mark.asyncio +async def test_retry_backoff_delays(config) -> None: + """Should use exponential backoff delays (1s, 2s).""" + client = DremioClient(config) + ok_response = httpx.Response(200, json={"ok": True}, request=httpx.Request("GET", "https://example.com")) + + client._client.request = AsyncMock(side_effect=[ + httpx.TimeoutException("timed out"), + httpx.TimeoutException("timed out"), + ok_response, + ]) + + with patch("drs.client.asyncio.sleep", new_callable=AsyncMock) as mock_sleep: + result = await client._get("https://example.com/test") + + assert result == {"ok": True} + assert mock_sleep.call_count == 2 + mock_sleep.assert_any_call(1.0) + mock_sleep.assert_any_call(2.0) diff --git a/tests/test_commands/test_access.py b/tests/test_commands/test_access.py deleted file mode 100644 index 697a192..0000000 --- a/tests/test_commands/test_access.py +++ /dev/null @@ -1,92 +0,0 @@ -# -# Copyright (C) 2017-2026 Dremio Corporation -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -"""Tests for drs access commands — user lookup endpoint and audit.""" - -from __future__ import annotations - -from unittest.mock import AsyncMock - -import httpx -import pytest - -from drs.commands.access import audit -from drs.utils import DremioAPIError - - -@pytest.mark.asyncio -async def test_audit_calls_correct_user_endpoint(mock_client) -> None: - """Verify get_user_by_name hits /v1/users/name/{userName} (not /v1/user/by-name/).""" - mock_client.get_user_by_name = AsyncMock(return_value={ - "id": "user-abc", - "roles": [{"id": "role-1", "name": "admin"}], - }) - - result = await audit(mock_client, "rahim") - - mock_client.get_user_by_name.assert_called_once_with("rahim") - assert result["username"] == "rahim" - assert result["user_id"] == "user-abc" - assert result["roles"] == [{"role_id": "role-1", "role_name": "admin"}] - - -@pytest.mark.asyncio -async def test_audit_user_not_found(mock_client) -> None: - """Verify 404 from user lookup raises DremioAPIError, not raw httpx traceback.""" - request = httpx.Request("GET", "https://api.dremio.cloud/v1/users/name/nobody") - response = httpx.Response(404, request=request) - mock_client.get_user_by_name = AsyncMock( - side_effect=httpx.HTTPStatusError("Not Found", request=request, response=response) - ) - - with pytest.raises(DremioAPIError) as exc_info: - await audit(mock_client, "nobody") - - assert exc_info.value.status_code == 404 - assert "Not found" in exc_info.value.message - - -@pytest.mark.asyncio -async def test_audit_user_with_multiple_roles(mock_client) -> None: - """Verify audit correctly maps multiple roles.""" - mock_client.get_user_by_name = AsyncMock(return_value={ - "id": "user-xyz", - "roles": [ - {"id": "role-1", "name": "admin"}, - {"id": "role-2", "name": "analyst"}, - {"id": "role-3", "name": "viewer"}, - ], - }) - - result = await audit(mock_client, "testuser") - - assert len(result["roles"]) == 3 - assert result["roles"][1] == {"role_id": "role-2", "role_name": "analyst"} - - -@pytest.mark.asyncio -async def test_audit_expired_pat(mock_client) -> None: - """Verify 401 from user lookup raises structured auth error.""" - request = httpx.Request("GET", "https://api.dremio.cloud/v1/users/name/rahim") - response = httpx.Response(401, request=request) - mock_client.get_user_by_name = AsyncMock( - side_effect=httpx.HTTPStatusError("Unauthorized", request=request, response=response) - ) - - with pytest.raises(DremioAPIError) as exc_info: - await audit(mock_client, "rahim") - - assert exc_info.value.status_code == 401 - assert "Authentication failed" in exc_info.value.message diff --git a/tests/test_commands/test_catalog.py b/tests/test_commands/test_catalog.py deleted file mode 100644 index e9ae186..0000000 --- a/tests/test_commands/test_catalog.py +++ /dev/null @@ -1,64 +0,0 @@ -# -# Copyright (C) 2017-2026 Dremio Corporation -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -"""Tests for drs catalog commands.""" - -from __future__ import annotations - -from unittest.mock import AsyncMock - -import pytest - -from drs.commands.catalog import get_entity, search_catalog - - -@pytest.mark.asyncio -async def test_get_entity_splits_path(mock_client) -> None: - mock_client.get_catalog_by_path = AsyncMock(return_value={ - "id": "abc", "entityType": "dataset", "path": ["space", "table"], - }) - - result = await get_entity(mock_client, "space.table") - - mock_client.get_catalog_by_path.assert_called_once_with(["space", "table"]) - assert result["id"] == "abc" - - -@pytest.mark.asyncio -async def test_get_entity_strips_quotes(mock_client) -> None: - mock_client.get_catalog_by_path = AsyncMock(return_value={"id": "abc"}) - - await get_entity(mock_client, '"space"."table"') - - mock_client.get_catalog_by_path.assert_called_once_with(["space", "table"]) - - -@pytest.mark.asyncio -async def test_get_entity_handles_dots_in_quotes(mock_client) -> None: - mock_client.get_catalog_by_path = AsyncMock(return_value={"id": "abc"}) - - await get_entity(mock_client, '"My Source"."my.table"') - - mock_client.get_catalog_by_path.assert_called_once_with(["My Source", "my.table"]) - - -@pytest.mark.asyncio -async def test_search_catalog(mock_client) -> None: - mock_client.search = AsyncMock(return_value={"data": [{"name": "orders"}]}) - - result = await search_catalog(mock_client, "orders") - - mock_client.search.assert_called_once_with("orders") - assert result["data"][0]["name"] == "orders" diff --git a/tests/test_commands/test_engine.py b/tests/test_commands/test_engine.py new file mode 100644 index 0000000..d5fc759 --- /dev/null +++ b/tests/test_commands/test_engine.py @@ -0,0 +1,79 @@ +# +# Copyright (C) 2017-2026 Dremio Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Tests for drs engine commands.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock + +import pytest + +from drs.commands.engine import list_engines, get_engine, create_engine, update_engine, delete_engine, enable_engine, disable_engine + + +@pytest.mark.asyncio +async def test_list_engines(mock_client) -> None: + mock_client.list_engines = AsyncMock(return_value={"data": [{"id": "eng-1", "name": "default"}]}) + result = await list_engines(mock_client) + mock_client.list_engines.assert_called_once() + assert result["data"][0]["name"] == "default" + + +@pytest.mark.asyncio +async def test_get_engine(mock_client) -> None: + mock_client.get_engine = AsyncMock(return_value={"id": "eng-1", "name": "default", "size": "SMALL"}) + result = await get_engine(mock_client, "eng-1") + mock_client.get_engine.assert_called_once_with("eng-1") + assert result["size"] == "SMALL" + + +@pytest.mark.asyncio +async def test_create_engine(mock_client) -> None: + mock_client.create_engine = AsyncMock(return_value={"id": "eng-2", "name": "analytics", "size": "LARGE"}) + result = await create_engine(mock_client, "analytics", size="LARGE") + mock_client.create_engine.assert_called_once_with({"name": "analytics", "size": "LARGE"}) + assert result["name"] == "analytics" + + +@pytest.mark.asyncio +async def test_update_engine(mock_client) -> None: + mock_client.get_engine = AsyncMock(return_value={"id": "eng-1", "name": "old", "size": "SMALL"}) + mock_client.update_engine = AsyncMock(return_value={"id": "eng-1", "name": "new", "size": "MEDIUM"}) + result = await update_engine(mock_client, "eng-1", name="new", size="MEDIUM") + call_body = mock_client.update_engine.call_args[0][1] + assert call_body["name"] == "new" + assert call_body["size"] == "MEDIUM" + + +@pytest.mark.asyncio +async def test_delete_engine(mock_client) -> None: + mock_client.delete_engine = AsyncMock(return_value={"status": "ok"}) + result = await delete_engine(mock_client, "eng-1") + mock_client.delete_engine.assert_called_once_with("eng-1") + + +@pytest.mark.asyncio +async def test_enable_engine(mock_client) -> None: + mock_client.enable_engine = AsyncMock(return_value={"id": "eng-1", "state": "ACTIVE"}) + result = await enable_engine(mock_client, "eng-1") + mock_client.enable_engine.assert_called_once_with("eng-1") + + +@pytest.mark.asyncio +async def test_disable_engine(mock_client) -> None: + mock_client.disable_engine = AsyncMock(return_value={"id": "eng-1", "state": "DISABLED"}) + result = await disable_engine(mock_client, "eng-1") + mock_client.disable_engine.assert_called_once_with("eng-1") diff --git a/tests/test_commands/test_folder.py b/tests/test_commands/test_folder.py new file mode 100644 index 0000000..9aa7c0d --- /dev/null +++ b/tests/test_commands/test_folder.py @@ -0,0 +1,84 @@ +# +# Copyright (C) 2017-2026 Dremio Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Tests for dremio folder commands.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock + +import pytest + +from drs.commands.folder import get_entity, create_folder, delete_entity, grants + + +@pytest.mark.asyncio +async def test_get_entity_splits_path(mock_client) -> None: + mock_client.get_catalog_by_path = AsyncMock(return_value={ + "id": "abc", "entityType": "dataset", "path": ["space", "table"], + }) + result = await get_entity(mock_client, "space.table") + mock_client.get_catalog_by_path.assert_called_once_with(["space", "table"]) + assert result["id"] == "abc" + + +@pytest.mark.asyncio +async def test_get_entity_handles_dots_in_quotes(mock_client) -> None: + mock_client.get_catalog_by_path = AsyncMock(return_value={"id": "abc"}) + await get_entity(mock_client, '"My Source"."my.table"') + mock_client.get_catalog_by_path.assert_called_once_with(["My Source", "my.table"]) + + +@pytest.mark.asyncio +async def test_create_folder_single_creates_space(mock_client) -> None: + """Single-component path should CREATE SPACE.""" + mock_client.submit_sql = AsyncMock(return_value={"id": "job-1"}) + mock_client.get_job_status = AsyncMock(return_value={"jobState": "COMPLETED", "rowCount": 0}) + mock_client.get_job_results = AsyncMock(return_value={"rows": []}) + result = await create_folder(mock_client, "Analytics") + sql = mock_client.submit_sql.call_args[0][0] + assert 'CREATE SPACE "Analytics"' in sql + + +@pytest.mark.asyncio +async def test_create_folder_nested_creates_folder(mock_client) -> None: + """Nested path should CREATE FOLDER.""" + mock_client.submit_sql = AsyncMock(return_value={"id": "job-1"}) + mock_client.get_job_status = AsyncMock(return_value={"jobState": "COMPLETED", "rowCount": 0}) + mock_client.get_job_results = AsyncMock(return_value={"rows": []}) + result = await create_folder(mock_client, "Analytics.reports") + sql = mock_client.submit_sql.call_args[0][0] + assert "CREATE FOLDER" in sql + assert '"Analytics"."reports"' in sql + + +@pytest.mark.asyncio +async def test_delete_entity(mock_client) -> None: + mock_client.get_catalog_by_path = AsyncMock(return_value={ + "id": "entity-1", "tag": "v1", "entityType": "space" + }) + mock_client.delete_catalog_entity = AsyncMock(return_value={"status": "ok"}) + result = await delete_entity(mock_client, "myspace") + mock_client.delete_catalog_entity.assert_called_once_with("entity-1", tag="v1") + + +@pytest.mark.asyncio +async def test_grants(mock_client) -> None: + mock_client.get_catalog_by_path = AsyncMock(return_value={ + "id": "entity-1", "accessControlList": {"users": []} + }) + result = await grants(mock_client, "myspace.table") + assert result["path"] == "myspace.table" + assert "accessControlList" in result diff --git a/tests/test_commands/test_grant.py b/tests/test_commands/test_grant.py new file mode 100644 index 0000000..0bad5e4 --- /dev/null +++ b/tests/test_commands/test_grant.py @@ -0,0 +1,47 @@ +# +# Copyright (C) 2017-2026 Dremio Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Tests for drs grant commands.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock + +import pytest + +from drs.commands.grant import get_grants, set_grants, remove_grants + + +@pytest.mark.asyncio +async def test_get_grants(mock_client) -> None: + mock_client.get_grants = AsyncMock(return_value={"privileges": ["MANAGE_GRANTS"]}) + result = await get_grants(mock_client, "projects", "proj-1", "role", "role-1") + mock_client.get_grants.assert_called_once_with("projects", "proj-1", "role", "role-1") + + +@pytest.mark.asyncio +async def test_set_grants(mock_client) -> None: + mock_client.set_grants = AsyncMock(return_value={"privileges": ["MANAGE_GRANTS", "CREATE_TABLE"]}) + result = await set_grants(mock_client, "projects", "proj-1", "role", "role-1", ["MANAGE_GRANTS", "CREATE_TABLE"]) + mock_client.set_grants.assert_called_once_with( + "projects", "proj-1", "role", "role-1", {"privileges": ["MANAGE_GRANTS", "CREATE_TABLE"]} + ) + + +@pytest.mark.asyncio +async def test_remove_grants(mock_client) -> None: + mock_client.delete_grants = AsyncMock(return_value={"status": "ok"}) + await remove_grants(mock_client, "projects", "proj-1", "user", "user-1") + mock_client.delete_grants.assert_called_once_with("projects", "proj-1", "user", "user-1") diff --git a/tests/test_commands/test_reflection.py b/tests/test_commands/test_reflection.py new file mode 100644 index 0000000..36b1f0a --- /dev/null +++ b/tests/test_commands/test_reflection.py @@ -0,0 +1,45 @@ +# +# Copyright (C) 2017-2026 Dremio Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Tests for dremio reflection commands.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock + +import pytest + +from drs.commands.reflection import get_reflection, refresh, delete + + +@pytest.mark.asyncio +async def test_get_reflection(mock_client) -> None: + mock_client.get_reflection = AsyncMock(return_value={"id": "r1", "status": "CAN_ACCELERATE"}) + result = await get_reflection(mock_client, "r1") + assert result["status"] == "CAN_ACCELERATE" + + +@pytest.mark.asyncio +async def test_refresh(mock_client) -> None: + mock_client.refresh_reflection = AsyncMock(return_value={"status": "ok"}) + await refresh(mock_client, "r1") + mock_client.refresh_reflection.assert_called_once_with("r1") + + +@pytest.mark.asyncio +async def test_delete(mock_client) -> None: + mock_client.delete_reflection = AsyncMock(return_value={"status": "ok"}) + await delete(mock_client, "r1") + mock_client.delete_reflection.assert_called_once_with("r1") diff --git a/tests/test_commands/test_role.py b/tests/test_commands/test_role.py new file mode 100644 index 0000000..dee3234 --- /dev/null +++ b/tests/test_commands/test_role.py @@ -0,0 +1,76 @@ +# +# Copyright (C) 2017-2026 Dremio Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Tests for drs role commands.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock + +import httpx +import pytest + +from drs.commands.role import list_roles, get_role, create_role, update_role, delete_role + + +@pytest.mark.asyncio +async def test_list_roles(mock_client) -> None: + mock_client.list_roles = AsyncMock(return_value={"data": [{"id": "r1", "name": "admin"}]}) + result = await list_roles(mock_client) + mock_client.list_roles.assert_called_once() + + +@pytest.mark.asyncio +async def test_get_role_by_name(mock_client) -> None: + mock_client.get_role_by_name = AsyncMock(return_value={"id": "r1", "name": "admin"}) + result = await get_role(mock_client, "admin") + mock_client.get_role_by_name.assert_called_once_with("admin") + assert result["name"] == "admin" + + +@pytest.mark.asyncio +async def test_get_role_falls_back_to_id(mock_client) -> None: + request = httpx.Request("GET", "https://api.dremio.cloud/v1/roles/name/r1") + response = httpx.Response(404, request=request) + mock_client.get_role_by_name = AsyncMock( + side_effect=httpx.HTTPStatusError("Not Found", request=request, response=response) + ) + mock_client.get_role = AsyncMock(return_value={"id": "r1", "name": "admin"}) + result = await get_role(mock_client, "r1") + mock_client.get_role.assert_called_once_with("r1") + + +@pytest.mark.asyncio +async def test_create_role(mock_client) -> None: + mock_client.create_role = AsyncMock(return_value={"id": "r2", "name": "analyst"}) + result = await create_role(mock_client, "analyst") + mock_client.create_role.assert_called_once_with({"name": "analyst"}) + assert result["name"] == "analyst" + + +@pytest.mark.asyncio +async def test_update_role(mock_client) -> None: + mock_client.get_role = AsyncMock(return_value={"id": "r1", "name": "old"}) + mock_client.update_role = AsyncMock(return_value={"id": "r1", "name": "new"}) + result = await update_role(mock_client, "r1", "new") + call_body = mock_client.update_role.call_args[0][1] + assert call_body["name"] == "new" + + +@pytest.mark.asyncio +async def test_delete_role(mock_client) -> None: + mock_client.delete_role = AsyncMock(return_value={"status": "ok"}) + await delete_role(mock_client, "r1") + mock_client.delete_role.assert_called_once_with("r1") diff --git a/tests/test_commands/test_tag.py b/tests/test_commands/test_tag.py new file mode 100644 index 0000000..8bf2c60 --- /dev/null +++ b/tests/test_commands/test_tag.py @@ -0,0 +1,54 @@ +# +# Copyright (C) 2017-2026 Dremio Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Tests for dremio tag commands.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock + +import httpx +import pytest + +from drs.commands.tag import get_tags, update_tags + + +@pytest.mark.asyncio +async def test_get_tags(mock_client) -> None: + mock_client.get_catalog_by_path = AsyncMock(return_value={"id": "e1"}) + mock_client.get_tags = AsyncMock(return_value={"tags": ["pii", "finance"]}) + result = await get_tags(mock_client, "myspace.table") + assert result["tags"] == ["pii", "finance"] + + +@pytest.mark.asyncio +async def test_get_tags_404(mock_client) -> None: + mock_client.get_catalog_by_path = AsyncMock(return_value={"id": "e1"}) + request = httpx.Request("GET", "https://example.com") + response = httpx.Response(404, request=request) + mock_client.get_tags = AsyncMock( + side_effect=httpx.HTTPStatusError("Not Found", request=request, response=response) + ) + result = await get_tags(mock_client, "myspace.table") + assert result["tags"] == [] + + +@pytest.mark.asyncio +async def test_update_tags(mock_client) -> None: + mock_client.get_catalog_by_path = AsyncMock(return_value={"id": "e1"}) + mock_client.get_tags = AsyncMock(return_value={"tags": ["old"], "version": 1}) + mock_client.set_tags = AsyncMock(return_value={"tags": ["pii", "finance"], "version": 2}) + result = await update_tags(mock_client, "myspace.table", ["pii", "finance"]) + mock_client.set_tags.assert_called_once_with("e1", ["pii", "finance"], version=1) diff --git a/tests/test_commands/test_user.py b/tests/test_commands/test_user.py new file mode 100644 index 0000000..27def6d --- /dev/null +++ b/tests/test_commands/test_user.py @@ -0,0 +1,83 @@ +# +# Copyright (C) 2017-2026 Dremio Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Tests for drs user commands.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock + +import httpx +import pytest + +from drs.commands.user import list_users, get_user, create_user, delete_user, whoami, audit + + +@pytest.mark.asyncio +async def test_list_users(mock_client) -> None: + mock_client.list_users = AsyncMock(return_value={"data": [{"id": "u1", "name": "alice"}]}) + result = await list_users(mock_client) + mock_client.list_users.assert_called_once_with(max_results=100) + + +@pytest.mark.asyncio +async def test_get_user_by_name(mock_client) -> None: + mock_client.get_user_by_name = AsyncMock(return_value={"id": "u1", "name": "alice"}) + result = await get_user(mock_client, "alice") + mock_client.get_user_by_name.assert_called_once_with("alice") + assert result["name"] == "alice" + + +@pytest.mark.asyncio +async def test_get_user_falls_back_to_id(mock_client) -> None: + request = httpx.Request("GET", "https://api.dremio.cloud/v1/users/name/u1") + response = httpx.Response(404, request=request) + mock_client.get_user_by_name = AsyncMock( + side_effect=httpx.HTTPStatusError("Not Found", request=request, response=response) + ) + mock_client.get_user = AsyncMock(return_value={"id": "u1", "name": "alice"}) + result = await get_user(mock_client, "u1") + mock_client.get_user.assert_called_once_with("u1") + + +@pytest.mark.asyncio +async def test_create_user(mock_client) -> None: + mock_client.invite_user = AsyncMock(return_value={"id": "u2", "email": "bob@example.com"}) + result = await create_user(mock_client, "bob@example.com", role_id="role-1") + mock_client.invite_user.assert_called_once_with({"email": "bob@example.com", "roleId": "role-1"}) + + +@pytest.mark.asyncio +async def test_delete_user(mock_client) -> None: + mock_client.delete_user = AsyncMock(return_value={"status": "ok"}) + await delete_user(mock_client, "u1") + mock_client.delete_user.assert_called_once_with("u1") + + +@pytest.mark.asyncio +async def test_whoami(mock_client) -> None: + mock_client.list_users = AsyncMock(return_value={"data": [{"id": "u1", "name": "me"}]}) + result = await whoami(mock_client) + mock_client.list_users.assert_called_once_with(max_results=1) + + +@pytest.mark.asyncio +async def test_audit(mock_client) -> None: + mock_client.get_user_by_name = AsyncMock(return_value={ + "id": "u1", "roles": [{"id": "r1", "name": "admin"}] + }) + result = await audit(mock_client, "alice") + assert result["username"] == "alice" + assert result["roles"] == [{"role_id": "r1", "role_name": "admin"}] diff --git a/tests/test_commands/test_wiki.py b/tests/test_commands/test_wiki.py new file mode 100644 index 0000000..e8f3017 --- /dev/null +++ b/tests/test_commands/test_wiki.py @@ -0,0 +1,55 @@ +# +# Copyright (C) 2017-2026 Dremio Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Tests for dremio wiki commands.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock + +import httpx +import pytest + +from drs.commands.wiki import get_wiki, update_wiki + + +@pytest.mark.asyncio +async def test_get_wiki(mock_client) -> None: + mock_client.get_catalog_by_path = AsyncMock(return_value={"id": "e1"}) + mock_client.get_wiki = AsyncMock(return_value={"text": "Hello", "version": 1}) + result = await get_wiki(mock_client, "myspace.table") + assert result["wiki"] == "Hello" + + +@pytest.mark.asyncio +async def test_get_wiki_404(mock_client) -> None: + mock_client.get_catalog_by_path = AsyncMock(return_value={"id": "e1"}) + request = httpx.Request("GET", "https://example.com") + response = httpx.Response(404, request=request) + mock_client.get_wiki = AsyncMock( + side_effect=httpx.HTTPStatusError("Not Found", request=request, response=response) + ) + result = await get_wiki(mock_client, "myspace.table") + assert result["wiki"] == "" + + +@pytest.mark.asyncio +async def test_update_wiki(mock_client) -> None: + mock_client.get_catalog_by_path = AsyncMock(return_value={"id": "e1"}) + mock_client.get_wiki = AsyncMock(return_value={"text": "old", "version": 2}) + mock_client.set_wiki = AsyncMock(return_value={"text": "new", "version": 3}) + result = await update_wiki(mock_client, "myspace.table", "new") + mock_client.set_wiki.assert_called_once_with("e1", "new", version=2) + assert result["wiki"] == "new"