diff --git a/src/drs/cli.py b/src/drs/cli.py index 5ab21da..b5c4b5d 100644 --- a/src/drs/cli.py +++ b/src/drs/cli.py @@ -40,6 +40,7 @@ reflection, role, schema, + semantic, setup, space, tag, @@ -60,6 +61,7 @@ app.add_typer(query.app, name="query") app.add_typer(folder.app, name="folder") app.add_typer(schema.app, name="schema") +app.add_typer(semantic.app, name="semantic") app.add_typer(wiki.app, name="wiki") app.add_typer(tag.app, name="tag") app.add_typer(reflection.app, name="reflection") diff --git a/src/drs/client.py b/src/drs/client.py index 5762169..1c41a0e 100644 --- a/src/drs/client.py +++ b/src/drs/client.py @@ -70,6 +70,10 @@ def _v3(self, path: str) -> str: def _v1(self, path: str) -> str: return f"{self.config.uri}/v1{path}" + def _v1_project(self, path: str) -> str: + """Project-scoped v1 URL: /v1/projects/{pid}/...""" + return f"{self.config.uri}/v1/projects/{self.config.project_id}{path}" + # -- HTTP helpers with retry -- async def _request_with_retry(self, method: str, url: str, **kwargs: Any) -> httpx.Response: @@ -224,6 +228,10 @@ 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 format_catalog_table(self, entity_id: str, body: dict) -> dict: + """Format a file or folder as a physical dataset. POST /catalog/{id}.""" + return await self._post(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 @@ -252,6 +260,95 @@ async def set_tags(self, entity_id: str, tags: list[str], version: int | None = body["version"] = version return await self._post(self._v3(f"/catalog/{entity_id}/collaboration/tag"), json=body) + # -- Semantic layer (v4) -- + + async def delete_semantic_layer(self) -> dict: + return await self._delete(self._v1_project("/semantic-layer")) + + async def initialize_semantic_layer(self, body: dict) -> dict: + return await self._post(self._v1_project("/semantic-layer/initialize"), json=body) + + async def deploy_semantic_layer_entities(self, task_id: str) -> dict: + return await self._post(self._v1_project("/semantic-layer/deploy"), json={"taskId": task_id}) + + async def abandon_semantic_layer_entities(self, task_id: str) -> dict: + return await self._post(self._v1_project("/semantic-layer/abandon"), json={"taskId": task_id}) + + async def cancel_semantic_layer_task(self, task_id: str) -> dict: + return await self._post(self._v1_project("/semantic-layer/cancel"), json={"taskId": task_id}) + + async def list_semantic_layer_tasks(self, task_state: str | None = None) -> dict: + params = {"taskState": task_state} if task_state else None + return await self._get(self._v1_project("/semantic-layer/tasks"), params=params) + + async def add_semantic_layer_task(self, body: dict) -> dict: + return await self._post(self._v1_project("/semantic-layer/tasks"), json=body) + + async def get_semantic_layer_task(self, task_id: str) -> dict: + return await self._get(self._v1_project(f"/semantic-layer/tasks/{task_id}")) + + async def list_semantic_layer_entities( + self, + entity_type: str, + task_id: str | None = None, + page_token: str | None = None, + limit: int | None = None, + ) -> dict: + params: dict[str, Any] = {"EntityType": entity_type} + if task_id: + params["taskId"] = task_id + if page_token: + params["pageToken"] = page_token + if limit is not None: + params["limit"] = limit + return await self._get(self._v1_project("/semantic-layer/entities"), params=params) + + async def add_semantic_layer_entity(self, body: dict, task_id: str | None = None) -> dict: + url = self._v1_project("/semantic-layer/entities") + if task_id: + url = f"{url}?taskId={task_id}" + return await self._post(url, json=body) + + async def get_semantic_layer_entity( + self, + entity_id: str, + entity_type: str, + task_id: str | None = None, + include_relationships: bool = False, + ) -> dict: + params: dict[str, Any] = {"EntityType": entity_type} + if task_id: + params["taskId"] = task_id + if include_relationships: + params["includeRelationships"] = "true" + return await self._get(self._v1_project(f"/semantic-layer/entities/{entity_id}"), params=params) + + async def update_semantic_layer_entity(self, entity_id: str, body: dict, task_id: str | None = None) -> dict: + url = self._v1_project(f"/semantic-layer/entities/{entity_id}") + params = f"?taskId={task_id}" if task_id else "" + return await self._put(f"{url}{params}", json=body) + + async def delete_semantic_layer_entity(self, entity_id: str, entity_type: str, task_id: str | None = None) -> dict: + params: dict[str, Any] = {"EntityType": entity_type} + if task_id: + params["taskId"] = task_id + return await self._delete(self._v1_project(f"/semantic-layer/entities/{entity_id}"), params=params) + + async def get_semantic_layer_scope(self) -> dict: + return await self._get(self._v1_project("/semantic-layer/scope")) + + async def patch_semantic_layer_scope(self, body: dict) -> dict: + logger.debug("PATCH %s body=%s", self._v1_project("/semantic-layer/scope"), body) + resp = await self._request_with_retry("PATCH", self._v1_project("/semantic-layer/scope"), json=body) + logger.debug( + "PATCH %s → %d (%d bytes)", + self._v1_project("/semantic-layer/scope"), + resp.status_code, + len(resp.content), + ) + resp.raise_for_status() + return resp.json() + # -- Reflections (v3) -- async def get_reflection(self, reflection_id: str) -> dict: diff --git a/src/drs/commands/folder.py b/src/drs/commands/folder.py index 37f5250..3e9dae9 100644 --- a/src/drs/commands/folder.py +++ b/src/drs/commands/folder.py @@ -18,6 +18,7 @@ from __future__ import annotations import asyncio +from pathlib import Path import httpx import typer @@ -25,7 +26,14 @@ from drs.client import DremioClient from drs.commands.query import run_query from drs.output import OutputFormat, error, output, warn -from drs.utils import DremioAPIError, NestedPathUnsupported, handle_api_error, parse_path, quote_path_sql +from drs.utils import ( + DremioAPIError, + NestedPathUnsupported, + handle_api_error, + parse_path, + quote_path_sql, + sanitize_input, +) app = typer.Typer( help="Manage nested folders and list top-level catalog entities. Use `dremio space` for top-level spaces.", @@ -114,6 +122,91 @@ async def grants(client: DremioClient, path: str) -> dict: } +def _split_slash_path(path: str) -> list[str]: + """Split a slash-separated relative path into validated path parts.""" + sanitize_input(path, "path") + parts = [part.strip() for part in path.split("/") if part.strip()] + if not parts: + raise ValueError("Path is empty.") + if any(part in {".", ".."} for part in parts): + raise ValueError(f"Invalid path '{path}': '.' and '..' path segments are not allowed.") + return parts + + +async def promote_path_as_table(client: DremioClient, path_parts: list[str], format_type: str = "Delta") -> dict: + """Format a file or folder as a physical dataset using the Catalog API.""" + try: + entity = await client.get_catalog_by_path(path_parts) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + + body = { + "entityType": "dataset", + "type": "PHYSICAL_DATASET", + # Use the caller-resolved catalog path parts directly. The lookup response + # can normalize or collapse file-source segments containing dots, which + # breaks the subsequent format request for file/folder sources. + "path": path_parts, + "format": {"type": format_type}, + } + + try: + return await client.format_catalog_table(entity["id"], body) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + + +async def promote_folder(client: DremioClient, path: str, format_type: str = "Delta") -> dict: + """Format a dot-separated file or folder path as a physical dataset.""" + return await promote_path_as_table(client, parse_path(path), format_type=format_type) + + +async def promote_from_file( + client: DremioClient, + paths_file: Path, + source: str, + under: str | None = None, + format_type: str = "Delta", +) -> dict: + """Format multiple slash-separated relative paths from a file as datasets.""" + source_part = sanitize_input(source.strip(), "source") + if not source_part: + raise ValueError("Source is empty.") + + base_parts = [source_part] + if under: + base_parts.extend(_split_slash_path(under)) + + results: list[dict] = [] + for line_no, raw in enumerate(paths_file.read_text(encoding="utf-8").splitlines(), start=1): + raw = raw.strip() + if not raw or raw.startswith("#"): + continue + + rel_parts = _split_slash_path(raw) + full_parts = [*base_parts, *rel_parts] + result = await promote_path_as_table(client, full_parts, format_type=format_type) + results.append( + { + "line": line_no, + "input": raw, + "path": result.get("path", full_parts), + "id": result.get("id"), + "entityType": result.get("entityType"), + "type": result.get("type"), + "format": result.get("format"), + } + ) + + return { + "source": source_part, + "under": under, + "formatType": format_type, + "count": len(results), + "results": results, + } + + # -- CLI wrappers -- @@ -205,3 +298,55 @@ def cli_grants( """Show ACL grants on a catalog entity.""" client = _get_client() _run_command(grants(client, path), client, fmt) + + +@app.command("promote") +def cli_promote( + path: str = typer.Argument( + None, + help="Dot-separated file or folder path to format as a table (e.g., source.folder.table_dir)", + ), + paths_file: Path | None = typer.Option( + None, + "--paths-file", + exists=True, + file_okay=True, + dir_okay=False, + readable=True, + resolve_path=True, + help="File containing slash-separated relative paths, one per line", + ), + source: str | None = typer.Option( + None, + "--source", + help="Source name to prefix to every line from --paths-file", + ), + under: str | None = typer.Option( + None, + "--under", + help="Optional slash-separated base path under the source for every line from --paths-file", + ), + format_type: str = typer.Option("Delta", "--format-type", help="Dataset format type to promote as"), + 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: + """Format a file or folder as a table using the Catalog API. + + For batch promotion, pass --paths-file with slash-separated relative paths and + prefix them with --source and optional --under. + """ + if bool(path) == bool(paths_file): + error("Provide exactly one of PATH or --paths-file.") + raise typer.Exit(1) + + if paths_file and not source: + error("--source is required when using --paths-file.") + raise typer.Exit(1) + + client = _get_client() + if paths_file: + _run_command( + promote_from_file(client, paths_file, source, under=under, format_type=format_type), client, fmt, fields + ) + return + _run_command(promote_folder(client, path, format_type=format_type), client, fmt, fields=fields) diff --git a/src/drs/commands/query.py b/src/drs/commands/query.py index 9b115da..9fbaaa6 100644 --- a/src/drs/commands/query.py +++ b/src/drs/commands/query.py @@ -18,6 +18,7 @@ from __future__ import annotations import asyncio +import sys from pathlib import Path from typing import Annotated, Any @@ -134,17 +135,8 @@ async def _execute(): def cli_run( sql: str | None = typer.Argument(None, help="SQL query to execute (use '-' to read from stdin)"), file: Annotated[ - Path | None, - typer.Option( - "--file", - exists=True, - file_okay=True, - dir_okay=False, - writable=False, - readable=True, - allow_dash=True, - help="Path to a SQL file to execute (use '-' for stdin)", - ), + typer.FileText|None, + typer.Option(help="Path to a SQL file to execute (use '-' for stdin)"), ] = None, context: str = typer.Option(None, help="Dot-separated default schema context (e.g., myspace.folder)"), fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), @@ -171,9 +163,11 @@ def cli_run( if sql is not None: error("Cannot specify both a SQL argument and --file.") raise typer.Exit(1) - sql = file.read_text().strip() + sql = file.read().strip() elif sql is not None: sql = sql.strip() + if sql == "-": + sql = sys.stdin.read().strip() if not sql: error("SQL query is empty. Provide SQL as an argument, --file path, or pipe via stdin (use '-' for stdin).") diff --git a/src/drs/commands/semantic.py b/src/drs/commands/semantic.py new file mode 100644 index 0000000..89ec23e --- /dev/null +++ b/src/drs/commands/semantic.py @@ -0,0 +1,935 @@ +# +# 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 semantic — semantic layer operations and ingestion helpers.""" + +from __future__ import annotations + +import asyncio +import json +from pathlib import Path +from typing import Any + +import httpx +import typer +import yaml + +from drs.client import DremioClient +from drs.output import OutputFormat, error, output +from drs.utils import handle_api_error, parse_path + +app = typer.Typer( + help="Manage semantic layer lifecycle, entities, and ingestion helpers.", + context_settings={"help_option_names": ["-h", "--help"]}, +) +entity_app = typer.Typer( + help="Manage semantic layer entities, including wiki/glossary imports.", + context_settings={"help_option_names": ["-h", "--help"]}, +) +task_app = typer.Typer( + help="Manage semantic layer ingestion tasks.", + context_settings={"help_option_names": ["-h", "--help"]}, +) + +app.add_typer(task_app, name="task") +app.add_typer(entity_app, name="entity") + +VALID_ENTITY_TYPES = {"TABLE", "METRIC"} +VALID_AUTOPILOT_MODES = {"OFF", "OBSERVE", "ASSIST", "AUTOMATE", "CUSTOM"} +VALID_KEY_TYPES = {"PRIMARY", "FOREIGN", "NONE"} +MAX_INITIALIZE_ENTITIES = 10 + + +def _get_client() -> DremioClient: + from drs.cli import get_client + + return get_client() + + +def _run_command(coro, client, fmt: OutputFormat = OutputFormat.json) -> 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 | ValueError): + error(str(exc)) + raise typer.Exit(1) + if isinstance(exc, httpx.HTTPError): + message = str(exc) + if isinstance(exc, httpx.ConnectError): + message = f"Failed to connect to Dremio API at {client.config.uri}: {exc}" + error(message) + raise typer.Exit(1) + raise + output(result, fmt) + + +def _read_text_file(path: Path) -> str: + try: + return path.read_text(encoding="utf-8") + except OSError as exc: + raise ValueError(f"Failed to read file '{path}': {exc}") from exc + + +def _load_structured_file(path: Path) -> Any: + raw = _read_text_file(path) + try: + if path.suffix.lower() == ".json": + return json.loads(raw) + return yaml.safe_load(raw) + except Exception as exc: + raise ValueError(f"Failed to parse structured file '{path}': {exc}") from exc + + +def _load_dictionary_text(dictionary_file: Path | None, dictionary_text: str | None) -> str | None: + if dictionary_file and dictionary_text: + raise ValueError("Use either --dictionary-file or --dictionary-text, not both.") + if dictionary_file: + return _read_text_file(dictionary_file) + return dictionary_text + + +def _extract_path_string(value: Any) -> str: + if isinstance(value, str): + return value.strip() + if isinstance(value, list): + return ".".join(str(part) for part in value) + if isinstance(value, dict): + if "components" in value and isinstance(value["components"], list): + return ".".join(str(part) for part in value["components"]) + if "path" in value: + return _extract_path_string(value["path"]) + raise ValueError(f"Unsupported path entry: {value!r}") + + +def _load_paths(paths: list[str] | None, paths_file: Path | None) -> list[str]: + cli_paths = [path for path in (paths or []) if path.strip()] + if paths_file is None: + if not cli_paths: + raise ValueError("Provide at least one dataset path or use --paths-file.") + return cli_paths + if cli_paths: + raise ValueError("Use either positional paths or --paths-file, not both.") + + if paths_file.suffix.lower() in {".json", ".yaml", ".yml"}: + payload = _load_structured_file(paths_file) + if isinstance(payload, dict): + payload = payload.get("paths") or payload.get("entities") + if not isinstance(payload, list): + raise ValueError("Structured paths file must contain a list or an object with 'paths' or 'entities'.") + resolved = [_extract_path_string(entry) for entry in payload] + else: + resolved = [ + line.strip() + for line in _read_text_file(paths_file).splitlines() + if line.strip() and not line.lstrip().startswith("#") + ] + + cleaned = [path for path in resolved if path] + if not cleaned: + raise ValueError(f"No dataset paths found in '{paths_file}'.") + return cleaned + + +def _validate_entity_count(paths: list[str]) -> None: + if len(paths) > MAX_INITIALIZE_ENTITIES: + raise ValueError( + f"Semantic layer initialize/task add accepts at most {MAX_INITIALIZE_ENTITIES} entities per request; got {len(paths)}." + ) + + +def _chunk_paths(paths: list[str], chunk_size: int = MAX_INITIALIZE_ENTITIES) -> list[list[str]]: + return [paths[index : index + chunk_size] for index in range(0, len(paths), chunk_size)] + + +def _load_task_ids(task_ids: list[str] | None, task_ids_file: Path | None) -> list[str]: + cli_task_ids = [task_id.strip() for task_id in (task_ids or []) if task_id.strip()] + if task_ids_file is None: + if not cli_task_ids: + raise ValueError("Provide at least one task ID or use --task-ids-file.") + return cli_task_ids + if cli_task_ids: + raise ValueError("Use either positional task IDs or --task-ids-file, not both.") + + if task_ids_file.suffix.lower() in {".json", ".yaml", ".yml"}: + payload = _load_structured_file(task_ids_file) + if isinstance(payload, dict): + if isinstance(payload.get("results"), list): + resolved = [ + str(item.get("task", {}).get("taskId") or "").strip() + for item in payload["results"] + if isinstance(item, dict) + ] + else: + payload = payload.get("taskIds") or payload.get("tasks") + if not isinstance(payload, list): + raise ValueError("Structured task ID file must contain 'results', 'taskIds', or 'tasks'.") + resolved = [str(item).strip() for item in payload] + elif isinstance(payload, list): + resolved = [str(item).strip() for item in payload] + else: + raise ValueError("Structured task ID file must contain a list or supported object shape.") + else: + resolved = [ + line.strip() + for line in _read_text_file(task_ids_file).splitlines() + if line.strip() and not line.lstrip().startswith("#") + ] + + cleaned = [task_id for task_id in resolved if task_id] + if not cleaned: + raise ValueError(f"No task IDs found in '{task_ids_file}'.") + return cleaned + + +def _normalize_entity_type(entity_type: str) -> str: + normalized = entity_type.upper() + if normalized not in VALID_ENTITY_TYPES: + raise ValueError(f"Invalid entity type '{entity_type}'. Valid values: {', '.join(sorted(VALID_ENTITY_TYPES))}") + return normalized + + +def _normalize_autopilot_mode(mode: str) -> str: + normalized = mode.upper() + if normalized not in VALID_AUTOPILOT_MODES: + raise ValueError(f"Invalid autopilot mode '{mode}'. Valid values: {', '.join(sorted(VALID_AUTOPILOT_MODES))}") + return normalized + + +def _entity_reference_from_path(path: str) -> dict[str, Any]: + return {"type": "TABLE", "path": parse_path(path)} + + +def _coerce_attribute(entry: dict[str, Any]) -> dict[str, str]: + source_column = str(entry.get("sourceColumn") or entry.get("column") or entry.get("name") or "").strip() + if not source_column: + raise ValueError("Each glossary attribute must include sourceColumn, column, or name.") + + alias = str(entry.get("alias") or entry.get("displayName") or source_column).strip() + description = str(entry.get("description") or entry.get("wiki") or "").strip() + key_type = str(entry.get("keyType") or "NONE").upper() + if key_type not in VALID_KEY_TYPES: + raise ValueError(f"Invalid keyType '{key_type}' for column '{source_column}'.") + + return { + "sourceColumn": source_column, + "alias": alias, + "description": description, + "keyType": key_type, + } + + +def _normalize_glossary_entries(payload: Any) -> list[dict[str, str]]: + if payload is None: + return [] + if isinstance(payload, dict): + payload = payload.get("attributes") or payload.get("columns") or payload.get("glossary") + if not isinstance(payload, list): + raise ValueError("Glossary input must be a list or an object with attributes/columns/glossary.") + return [_coerce_attribute(item) for item in payload] + + +def _field_name(field: dict[str, Any]) -> str: + return str(field.get("name") or field.get("fieldName") or field.get("sourceColumn") or "").strip() + + +def _format_entity_path(path: list[str] | None) -> str: + return ".".join(path or []) + + +def _flatten_semantic_entity(entity: dict[str, Any]) -> dict[str, Any]: + attributes = entity.get("attributes") or [] + relationships = entity.get("relationships") or [] + return { + "id": entity.get("id", ""), + "type": entity.get("type", ""), + "name": entity.get("name", ""), + "path": _format_entity_path(entity.get("path")), + "description": entity.get("description", "") or "", + "confidenceScore": entity.get("confidenceScore", ""), + "attributeCount": len(attributes), + "relationshipCount": len(relationships), + "relatedMetricCount": entity.get("relatedMetricCount", ""), + } + + +def _build_default_attributes(catalog_entity: dict[str, Any]) -> list[dict[str, str]]: + fields = catalog_entity.get("fields") or [] + attributes: list[dict[str, str]] = [] + for field in fields: + name = _field_name(field) + if not name: + continue + attributes.append( + { + "sourceColumn": name, + "alias": name, + "description": "", + "keyType": "NONE", + } + ) + if not attributes: + raise ValueError("Dataset schema is unavailable. Provide a glossary file with column definitions.") + return attributes + + +def _merge_attributes( + catalog_entity: dict[str, Any], glossary_entries: list[dict[str, str]] | None = None +) -> list[dict[str, str]]: + merged = {entry["sourceColumn"]: entry for entry in _build_default_attributes(catalog_entity)} + for entry in glossary_entries or []: + merged[entry["sourceColumn"]] = entry + return list(merged.values()) + + +async def _get_catalog_entity_by_path(client: DremioClient, path: str) -> dict[str, Any]: + try: + return await client.get_catalog_by_path(parse_path(path)) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + + +async def _resolve_wiki_text(client: DremioClient, entity_id: str) -> str: + try: + wiki = await client.get_wiki(entity_id) + except httpx.HTTPStatusError as exc: + if exc.response.status_code == 404: + return "" + raise handle_api_error(exc) from exc + return str(wiki.get("text") or "").strip() + + +async def _build_table_payload( + client: DremioClient, + path: str, + *, + name: str | None, + description: str | None, + wiki_from_catalog: bool, + glossary_file: Path | None, + glossary_inline: list[dict[str, Any]] | None = None, +) -> dict[str, Any]: + catalog_entity = await _get_catalog_entity_by_path(client, path) + glossary_entries = glossary_inline + if glossary_file: + glossary_entries = _normalize_glossary_entries(_load_structured_file(glossary_file)) + resolved_description = (description or "").strip() + if wiki_from_catalog and not resolved_description: + resolved_description = await _resolve_wiki_text(client, catalog_entity["id"]) + + payload = { + "type": "TABLE", + "path": parse_path(path), + "name": name or parse_path(path)[-1], + "description": resolved_description, + "attributes": _merge_attributes(catalog_entity, glossary_entries), + } + return payload + + +async def initialize_semantic_layer( + client: DremioClient, + paths: list[str], + *, + dictionary_file: Path | None, + dictionary_text: str | None, + scoped_dictionaries_file: Path | None, + include_query_history: bool, + autopilot_mode: str | None, + max_jobs_to_process: int | None, +) -> dict[str, Any]: + _validate_entity_count(paths) + body: dict[str, Any] = { + "entities": [_entity_reference_from_path(path) for path in paths], + "includeQueryHistory": include_query_history, + } + if autopilot_mode: + body["autopilotMode"] = _normalize_autopilot_mode(autopilot_mode) + if max_jobs_to_process is not None: + body["maxJobsToProcess"] = max_jobs_to_process + + dictionary = _load_dictionary_text(dictionary_file, dictionary_text) + if dictionary and scoped_dictionaries_file: + raise ValueError("Use either a global dictionary input or --scoped-dictionaries-file, not both.") + if dictionary: + body["dataDictionaries"] = [{"dataDictionary": dictionary}] + elif scoped_dictionaries_file: + scoped = _load_structured_file(scoped_dictionaries_file) + if not isinstance(scoped, list): + raise ValueError("Scoped dictionaries file must contain a list of scoped dictionaries.") + body["dataDictionaries"] = scoped + + try: + return await client.initialize_semantic_layer(body) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + + +async def add_semantic_layer_task( + client: DremioClient, + paths: list[str], + *, + dictionary_file: Path | None, + dictionary_text: str | None, + include_query_history: bool, + max_jobs_to_process: int | None, +) -> dict[str, Any]: + _validate_entity_count(paths) + body: dict[str, Any] = {"entities": [_entity_reference_from_path(path) for path in paths]} + dictionary = _load_dictionary_text(dictionary_file, dictionary_text) + if dictionary: + body["dataDictionary"] = dictionary + if include_query_history: + body["includeQueryHistory"] = True + if max_jobs_to_process is not None: + body["maxJobsToProcess"] = max_jobs_to_process + + try: + return await client.add_semantic_layer_task(body) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + + +async def add_semantic_layer_task_batched( + client: DremioClient, + paths: list[str], + *, + dictionary_file: Path | None, + dictionary_text: str | None, + include_query_history: bool, + max_jobs_to_process: int | None, +) -> dict[str, Any]: + dictionary = _load_dictionary_text(dictionary_file, dictionary_text) + chunks = _chunk_paths(paths) + results = [] + for index, chunk in enumerate(chunks, start=1): + result = await add_semantic_layer_task( + client, + chunk, + dictionary_file=None, + dictionary_text=dictionary, + include_query_history=include_query_history, + max_jobs_to_process=max_jobs_to_process, + ) + results.append( + { + "batch": index, + "entitiesSubmitted": len(chunk), + "paths": chunk, + "task": result, + } + ) + return { + "batches": len(chunks), + "totalEntities": len(paths), + "results": results, + } + + +async def _run_task_action_batched( + client: DremioClient, + task_ids: list[str], + *, + action_name: str, + action_coro, +) -> dict[str, Any]: + results = [] + for task_id in task_ids: + result = await action_coro(task_id) + results.append({"taskId": task_id, "result": result}) + return {"action": action_name, "count": len(task_ids), "results": results} + + +async def deploy_semantic_layer_tasks_batched(client: DremioClient, task_ids: list[str]) -> dict[str, Any]: + return await _run_task_action_batched( + client, + task_ids, + action_name="deploy", + action_coro=client.deploy_semantic_layer_entities, + ) + + +async def abandon_semantic_layer_tasks_batched(client: DremioClient, task_ids: list[str]) -> dict[str, Any]: + return await _run_task_action_batched( + client, + task_ids, + action_name="abandon", + action_coro=client.abandon_semantic_layer_entities, + ) + + +async def list_semantic_entities( + client: DremioClient, + entity_type: str, + *, + task_id: str | None, + page_token: str | None, + limit: int, + fetch_all: bool, + fmt: OutputFormat, +) -> dict[str, Any] | list[dict[str, Any]]: + entities: list[dict[str, Any]] = [] + current_page_token = page_token + total_count: int | None = None + next_page_token: str | None = None + + while True: + response = await client.list_semantic_layer_entities( + entity_type, + task_id=task_id, + page_token=current_page_token, + limit=limit, + ) + page_entities = response.get("entities", []) + entities.extend(page_entities) + total_count = response.get("totalCount", total_count) + next_page_token = response.get("nextPageToken") + if not fetch_all or not next_page_token: + break + current_page_token = next_page_token + + if fmt in {OutputFormat.pretty, OutputFormat.csv}: + return [_flatten_semantic_entity(entity) for entity in entities] + + result: dict[str, Any] = { + "entities": entities, + "entityType": entity_type, + "count": len(entities), + } + if total_count is not None: + result["totalCount"] = total_count + if task_id: + result["taskId"] = task_id + if next_page_token and not fetch_all: + result["nextPageToken"] = next_page_token + if fetch_all: + result["allPagesFetched"] = True + return result + + +async def upsert_table_entity( + client: DremioClient, + path: str, + *, + name: str | None, + description: str | None, + wiki_from_catalog: bool, + glossary_file: Path | None, + task_id: str | None, +) -> dict[str, Any]: + payload = await _build_table_payload( + client, + path, + name=name, + description=description, + wiki_from_catalog=wiki_from_catalog, + glossary_file=glossary_file, + ) + try: + result = await client.update_semantic_layer_entity(path, payload, task_id=task_id) + action = "updated" + except httpx.HTTPStatusError as exc: + if exc.response.status_code != 404: + raise handle_api_error(exc) from exc + result = await client.add_semantic_layer_entity(payload, task_id=task_id) + action = "created" + return {"action": action, "path": path, "entity": result} + + +async def bulk_upsert_entities( + client: DremioClient, + spec_file: Path, + *, + wiki_from_catalog: bool, + task_id: str | None, +) -> dict[str, Any]: + payload = _load_structured_file(spec_file) + entities = payload.get("entities") if isinstance(payload, dict) else payload + if not isinstance(entities, list): + raise ValueError("Bulk entity spec must be a list or an object with an 'entities' list.") + + results = [] + for raw_entity in entities: + if not isinstance(raw_entity, dict): + raise ValueError("Each bulk entity entry must be an object.") + path = raw_entity.get("path") + dot_path = ".".join(path) if isinstance(path, list) else str(path or "").strip() + if not dot_path: + raise ValueError("Each bulk entity entry must include a path.") + glossary_file = raw_entity.get("glossaryFile") + glossary_path = (spec_file.parent / glossary_file) if glossary_file else None + glossary_inline = _normalize_glossary_entries(raw_entity.get("attributes") or raw_entity.get("glossary") or []) + entity_payload = await _build_table_payload( + client, + dot_path, + name=raw_entity.get("name"), + description=raw_entity.get("description"), + wiki_from_catalog=bool(raw_entity.get("wikiFromCatalog", wiki_from_catalog)), + glossary_file=glossary_path, + glossary_inline=glossary_inline, + ) + try: + entity_result = await client.update_semantic_layer_entity(dot_path, entity_payload, task_id=task_id) + action = "updated" + except httpx.HTTPStatusError as exc: + if exc.response.status_code != 404: + raise handle_api_error(exc) from exc + entity_result = await client.add_semantic_layer_entity(entity_payload, task_id=task_id) + action = "created" + results.append({"path": dot_path, "action": action, "entity": entity_result}) + return {"count": len(results), "results": results} + + +async def patch_semantic_layer_scope( + client: DremioClient, add_paths: list[str], remove_paths: list[str] +) -> dict[str, Any]: + body = { + "add": [{"components": parse_path(path)} for path in add_paths], + "remove": [{"components": parse_path(path)} for path in remove_paths], + } + try: + return await client.patch_semantic_layer_scope(body) + except httpx.HTTPStatusError as exc: + raise handle_api_error(exc) from exc + + +@app.command("delete") +def cli_delete( + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Delete the semantic layer.""" + client = _get_client() + _run_command(client.delete_semantic_layer(), client, fmt) + + +@app.command("initialize") +def cli_initialize( + paths: list[str] | None = typer.Argument( + None, help="One or more dot-separated dataset paths to initialize into the semantic layer." + ), + paths_file: Path | None = typer.Option( + None, "--paths-file", help="Text, JSON, or YAML file containing dataset paths." + ), + dictionary_file: Path | None = typer.Option( + None, "--dictionary-file", help="Path to a global data dictionary text file." + ), + dictionary_text: str | None = typer.Option(None, "--dictionary-text", help="Inline global data dictionary text."), + scoped_dictionaries_file: Path | None = typer.Option( + None, "--scoped-dictionaries-file", help="JSON/YAML file containing a list of scoped dictionaries." + ), + include_query_history: bool = typer.Option( + False, "--include-query-history", help="Analyze historical query jobs during initialization." + ), + autopilot_mode: str | None = typer.Option( + None, "--autopilot-mode", help="OFF, OBSERVE, ASSIST, AUTOMATE, or CUSTOM." + ), + max_jobs_to_process: int | None = typer.Option(None, "--max-jobs-to-process", min=1, help="Historical job limit."), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Initialize the semantic layer from one or more datasets.""" + client = _get_client() + resolved_paths = _load_paths(paths, paths_file) + _run_command( + initialize_semantic_layer( + client, + resolved_paths, + dictionary_file=dictionary_file, + dictionary_text=dictionary_text, + scoped_dictionaries_file=scoped_dictionaries_file, + include_query_history=include_query_history, + autopilot_mode=autopilot_mode, + max_jobs_to_process=max_jobs_to_process, + ), + client, + fmt, + ) + + +@task_app.command("add") +def cli_task_add( + paths: list[str] | None = typer.Argument( + None, help="One or more dot-separated dataset paths to hydrate into a draft semantic layer task." + ), + paths_file: Path | None = typer.Option( + None, "--paths-file", help="Text, JSON, or YAML file containing dataset paths." + ), + dictionary_file: Path | None = typer.Option(None, "--dictionary-file", help="Path to a data dictionary text file."), + dictionary_text: str | None = typer.Option(None, "--dictionary-text", help="Inline data dictionary text."), + include_query_history: bool = typer.Option(False, "--include-query-history", help="Analyze historical query jobs."), + max_jobs_to_process: int | None = typer.Option(None, "--max-jobs-to-process", min=1, help="Historical job limit."), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Create a semantic layer hydration task.""" + client = _get_client() + resolved_paths = _load_paths(paths, paths_file) + _run_command( + add_semantic_layer_task( + client, + resolved_paths, + dictionary_file=dictionary_file, + dictionary_text=dictionary_text, + include_query_history=include_query_history, + max_jobs_to_process=max_jobs_to_process, + ), + client, + fmt, + ) + + +@task_app.command("add-batch") +def cli_task_add_batch( + paths: list[str] | None = typer.Argument(None, help="One or more dot-separated dataset paths to hydrate."), + paths_file: Path | None = typer.Option( + None, "--paths-file", help="Text, JSON, or YAML file containing dataset paths." + ), + dictionary_file: Path | None = typer.Option(None, "--dictionary-file", help="Path to a data dictionary text file."), + dictionary_text: str | None = typer.Option(None, "--dictionary-text", help="Inline data dictionary text."), + include_query_history: bool = typer.Option(False, "--include-query-history", help="Analyze historical query jobs."), + max_jobs_to_process: int | None = typer.Option(None, "--max-jobs-to-process", min=1, help="Historical job limit."), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Split a large entity list into max-10 task-add requests and submit them sequentially.""" + client = _get_client() + resolved_paths = _load_paths(paths, paths_file) + _run_command( + add_semantic_layer_task_batched( + client, + resolved_paths, + dictionary_file=dictionary_file, + dictionary_text=dictionary_text, + include_query_history=include_query_history, + max_jobs_to_process=max_jobs_to_process, + ), + client, + fmt, + ) + + +@task_app.command("list") +def cli_task_list( + task_state: str | None = typer.Option(None, "--state", help="Optional task state filter."), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """List semantic layer tasks.""" + client = _get_client() + _run_command(client.list_semantic_layer_tasks(task_state), client, fmt) + + +@task_app.command("get") +def cli_task_get( + task_id: str = typer.Argument(help="Semantic layer task ID."), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Get semantic layer task status.""" + client = _get_client() + _run_command(client.get_semantic_layer_task(task_id), client, fmt) + + +@task_app.command("deploy") +def cli_task_deploy( + task_id: str = typer.Argument(help="Semantic layer task ID."), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Deploy a draft semantic layer task.""" + client = _get_client() + _run_command(client.deploy_semantic_layer_entities(task_id), client, fmt) + + +@task_app.command("deploy-batch") +def cli_task_deploy_batch( + task_ids: list[str] | None = typer.Argument(None, help="One or more semantic layer task IDs."), + task_ids_file: Path | None = typer.Option( + None, "--task-ids-file", help="Text, JSON, or YAML file containing task IDs." + ), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Deploy multiple semantic layer tasks sequentially.""" + client = _get_client() + resolved_task_ids = _load_task_ids(task_ids, task_ids_file) + _run_command(deploy_semantic_layer_tasks_batched(client, resolved_task_ids), client, fmt) + + +@task_app.command("abandon") +def cli_task_abandon( + task_id: str = typer.Argument(help="Semantic layer task ID."), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Abandon a draft semantic layer task.""" + client = _get_client() + _run_command(client.abandon_semantic_layer_entities(task_id), client, fmt) + + +@task_app.command("abandon-batch") +def cli_task_abandon_batch( + task_ids: list[str] | None = typer.Argument(None, help="One or more semantic layer task IDs."), + task_ids_file: Path | None = typer.Option( + None, "--task-ids-file", help="Text, JSON, or YAML file containing task IDs." + ), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Abandon multiple semantic layer tasks sequentially.""" + client = _get_client() + resolved_task_ids = _load_task_ids(task_ids, task_ids_file) + _run_command(abandon_semantic_layer_tasks_batched(client, resolved_task_ids), client, fmt) + + +@task_app.command("cancel") +def cli_task_cancel( + task_id: str = typer.Argument(help="Semantic layer task ID."), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Cancel an in-progress semantic layer task.""" + client = _get_client() + _run_command(client.cancel_semantic_layer_task(task_id), client, fmt) + + +@entity_app.command("list") +def cli_entity_list( + entity_type: str = typer.Argument(help="Entity type: TABLE or METRIC."), + task_id: str | None = typer.Option(None, "--task-id", help="Optional draft task ID."), + page_token: str | None = typer.Option(None, "--page-token", help="Pagination token."), + limit: int = typer.Option(50, "--limit", min=1, max=500, help="Page size."), + all_pages: bool = typer.Option(False, "--all", help="Fetch all pages by following nextPageToken."), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """List semantic layer entities.""" + client = _get_client() + _run_command( + list_semantic_entities( + client, + _normalize_entity_type(entity_type), + task_id=task_id, + page_token=page_token, + limit=limit, + fetch_all=all_pages, + fmt=fmt, + ), + client, + fmt, + ) + + +@entity_app.command("get") +def cli_entity_get( + entity_type: str = typer.Argument(help="Entity type: TABLE or METRIC."), + entity_id: str = typer.Argument(help="Entity ID or dot-separated table path."), + task_id: str | None = typer.Option(None, "--task-id", help="Optional draft task ID."), + include_relationships: bool = typer.Option(False, "--include-relationships", help="Include relationships."), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Get a semantic layer entity.""" + client = _get_client() + _run_command( + client.get_semantic_layer_entity( + entity_id, + _normalize_entity_type(entity_type), + task_id=task_id, + include_relationships=include_relationships, + ), + client, + fmt, + ) + + +@entity_app.command("delete") +def cli_entity_delete( + entity_type: str = typer.Argument(help="Entity type: TABLE or METRIC."), + entity_id: str = typer.Argument(help="Entity ID or dot-separated table path."), + task_id: str | None = typer.Option(None, "--task-id", help="Optional draft task ID."), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Delete a semantic layer entity.""" + client = _get_client() + _run_command( + client.delete_semantic_layer_entity(entity_id, _normalize_entity_type(entity_type), task_id), client, fmt + ) + + +@entity_app.command("upsert-table") +def cli_entity_upsert_table( + path: str = typer.Argument(help="Dot-separated dataset path."), + name: str | None = typer.Option(None, "--name", help="Display name to use in the semantic layer."), + description: str | None = typer.Option( + None, "--description", help="Table description to store in the semantic layer." + ), + wiki_from_catalog: bool = typer.Option( + False, "--wiki-from-catalog", help="Use the catalog wiki as the table description when --description is absent." + ), + glossary_file: Path | None = typer.Option( + None, "--glossary-file", help="JSON/YAML file with column glossary entries." + ), + task_id: str | None = typer.Option(None, "--task-id", help="Optional draft task ID."), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Create or update a TABLE entity from catalog schema, wiki, and glossary metadata.""" + client = _get_client() + _run_command( + upsert_table_entity( + client, + path, + name=name, + description=description, + wiki_from_catalog=wiki_from_catalog, + glossary_file=glossary_file, + task_id=task_id, + ), + client, + fmt, + ) + + +@entity_app.command("bulk-upsert") +def cli_entity_bulk_upsert( + spec_file: Path = typer.Argument(help="JSON/YAML file describing table entities to create or update."), + wiki_from_catalog: bool = typer.Option( + False, "--wiki-from-catalog", help="Use catalog wiki as the fallback description for entries that omit one." + ), + task_id: str | None = typer.Option(None, "--task-id", help="Optional draft task ID."), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Bulk create or update TABLE entities from a JSON/YAML spec.""" + client = _get_client() + _run_command( + bulk_upsert_entities(client, spec_file, wiki_from_catalog=wiki_from_catalog, task_id=task_id), client, fmt + ) + + +@app.command("scope") +def cli_scope_get( + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Get semantic layer scope.""" + client = _get_client() + _run_command(client.get_semantic_layer_scope(), client, fmt) + + +@app.command("scope-patch") +def cli_scope_patch( + add: list[str] = typer.Option([], "--add", help="Dataset path to add to semantic layer scope."), + remove: list[str] = typer.Option([], "--remove", help="Dataset path to remove from semantic layer scope."), + fmt: OutputFormat = typer.Option(OutputFormat.json, "--output", "-o", help="Output format"), +) -> None: + """Add or remove datasets from semantic layer scope.""" + if not add and not remove: + raise typer.BadParameter("Provide at least one --add or --remove dataset path.") + client = _get_client() + _run_command(patch_semantic_layer_scope(client, add, remove), client, fmt) diff --git a/src/drs/introspect.py b/src/drs/introspect.py index c88d174..2d1721d 100644 --- a/src/drs/introspect.py +++ b/src/drs/introspect.py @@ -154,6 +154,57 @@ {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, ], }, + "folder.promote": { + "group": "folder", + "command": "promote", + "description": "Format a file or folder as a physical dataset using the Catalog API. Supports either a single dot-separated path or a file of slash-separated relative paths prefixed by --source and optional --under.", + "mechanism": "REST", + "mutating": True, + "endpoints": [ + "GET /v0/projects/{pid}/catalog/by-path/{path}", + "POST /v0/projects/{pid}/catalog/{id}", + ], + "parameters": [ + { + "name": "path", + "type": "string", + "required": False, + "positional": True, + "description": "Dot-separated file or folder path to format as a table", + }, + { + "name": "paths_file", + "type": "path", + "required": False, + "flag": "--paths-file", + "description": "File with slash-separated relative paths, one per line", + }, + { + "name": "source", + "type": "string", + "required": False, + "flag": "--source", + "description": "Source name to prefix to every line from --paths-file", + }, + { + "name": "under", + "type": "string", + "required": False, + "flag": "--under", + "description": "Optional slash-separated base path under the source for every line from --paths-file", + }, + { + "name": "format_type", + "type": "string", + "required": False, + "default": "Delta", + "flag": "--format-type", + "description": "Dataset format type, for example Delta or Parquet", + }, + {"name": "output", "type": "enum", "required": False, "default": "json", "enum": ["json", "csv", "pretty"]}, + {"name": "fields", "type": "string", "required": False}, + ], + }, # -- Space -- "space.list": { "group": "space", diff --git a/src/drs/output.py b/src/drs/output.py index 0397a3f..4ea6d82 100644 --- a/src/drs/output.py +++ b/src/drs/output.py @@ -24,6 +24,10 @@ from enum import StrEnum from typing import Any +import pandas as pd +from rich.console import Console +from rich.table import Table + class OutputFormat(StrEnum): json = "json" @@ -100,22 +104,20 @@ def _list_table(rows: list) -> str: return "(no results)" if not isinstance(rows[0], dict): return "\n".join(str(r) for r in rows) + return _rich_table(rows) - cols = list(rows[0].keys()) - widths = {c: len(c) for c in cols} - str_rows = [] - for row in rows: - sr = {c: str(row.get(c, "")) for c in cols} - for c in cols: - widths[c] = max(widths[c], len(sr[c])) - str_rows.append(sr) - - header = " ".join(c.ljust(widths[c]) for c in cols) - sep = " ".join("-" * widths[c] for c in cols) - lines = [header, sep] - for sr in str_rows: - lines.append(" ".join(sr[c].ljust(widths[c]) for c in cols)) - return "\n".join(lines) + +def _rich_table(rows: list[dict[str, Any]]) -> str: + dataframe = pd.DataFrame(rows).fillna("") + table = Table(show_header=True, header_style="bold") + for column in dataframe.columns: + table.add_column(str(column), overflow="fold") + for row in dataframe.itertuples(index=False, name=None): + table.add_row(*[str(value) for value in row]) + + console = Console(record=True, force_terminal=False, width=120) + console.print(table) + return console.export_text().rstrip() def error(msg: str) -> None: diff --git a/tests/test_client.py b/tests/test_client.py index 4a38c6f..795a1da 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -74,6 +74,34 @@ def test_catalog_root_no_trailing_slash(self, client: DremioClient) -> None: 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" + @pytest.mark.asyncio + async def test_format_catalog_table_posts_expected_body(self, client: DremioClient) -> None: + captured: dict = {} + + async def _capture(request: httpx.Request) -> httpx.Response: + import json + + captured["url"] = str(request.url) + captured["body"] = json.loads(request.content) + return httpx.Response(200, json={"id": "ds-1", "path": ["dataproducts", "dre-unstructured", "orders"]}) + + client._client = httpx.AsyncClient(transport=httpx.MockTransport(_capture)) + + await client.format_catalog_table( + "folder-123", + { + "entityType": "dataset", + "type": "PHYSICAL_DATASET", + "path": ["dataproducts", "dre-unstructured", "orders"], + "format": {"type": "Delta"}, + }, + ) + + assert captured["url"] == "https://api.dremio.cloud/v0/projects/proj-123/catalog/folder-123" + assert captured["body"]["entityType"] == "dataset" + assert captured["body"]["type"] == "PHYSICAL_DATASET" + assert captured["body"]["format"]["type"] == "Delta" + class TestEngineURLs: def test_engines_list_url(self, client: DremioClient) -> None: diff --git a/tests/test_commands/test_folder.py b/tests/test_commands/test_folder.py index eb27858..c7fe139 100644 --- a/tests/test_commands/test_folder.py +++ b/tests/test_commands/test_folder.py @@ -17,11 +17,21 @@ from __future__ import annotations +from pathlib import Path from unittest.mock import AsyncMock import pytest -from drs.commands.folder import create_folder, delete_entity, delete_folder, get_entity, get_folder, grants +from drs.commands.folder import ( + create_folder, + delete_entity, + delete_folder, + get_entity, + get_folder, + grants, + promote_folder, + promote_from_file, +) from drs.utils import DremioAPIError @@ -156,3 +166,100 @@ async def test_grants(mock_client) -> None: result = await grants(mock_client, "myspace.table") assert result["path"] == "myspace.table" assert "accessControlList" in result + + +@pytest.mark.asyncio +async def test_promote_folder_formats_as_delta(mock_client) -> None: + mock_client.get_catalog_by_path = AsyncMock( + return_value={ + "id": "folder-1", + "path": ["dataproducts", "dre-unstructured", "sales_orders"], + "entityType": "folder", + } + ) + mock_client.format_catalog_table = AsyncMock( + return_value={ + "id": "ds-1", + "path": ["dataproducts", "dre-unstructured", "sales_orders"], + "entityType": "dataset", + "type": "PHYSICAL_DATASET", + "format": {"type": "Delta"}, + } + ) + + result = await promote_folder(mock_client, "dataproducts.dre-unstructured.sales_orders") + + mock_client.get_catalog_by_path.assert_called_once_with(["dataproducts", "dre-unstructured", "sales_orders"]) + mock_client.format_catalog_table.assert_called_once_with( + "folder-1", + { + "entityType": "dataset", + "type": "PHYSICAL_DATASET", + "path": ["dataproducts", "dre-unstructured", "sales_orders"], + "format": {"type": "Delta"}, + }, + ) + assert result["entityType"] == "dataset" + + +@pytest.mark.asyncio +async def test_promote_from_file_prefixes_source_and_under(mock_client, tmp_path: Path) -> None: + paths_file = tmp_path / "tables.txt" + paths_file.write_text("sales/orders\n# comment\n\ninventory/stock\n", encoding="utf-8") + + mock_client.get_catalog_by_path = AsyncMock( + side_effect=[ + { + "id": "folder-1", + "path": ["dataproducts", "dre-unstructured", "sales", "orders"], + }, + { + "id": "folder-2", + "path": ["dataproducts", "dre-unstructured", "inventory", "stock"], + }, + ] + ) + mock_client.format_catalog_table = AsyncMock( + side_effect=[ + { + "id": "ds-1", + "path": ["dataproducts", "dre-unstructured", "sales", "orders"], + "entityType": "dataset", + "type": "PHYSICAL_DATASET", + "format": {"type": "Delta"}, + }, + { + "id": "ds-2", + "path": ["dataproducts", "dre-unstructured", "inventory", "stock"], + "entityType": "dataset", + "type": "PHYSICAL_DATASET", + "format": {"type": "Delta"}, + }, + ] + ) + + result = await promote_from_file(mock_client, paths_file, "dataproducts", under="dre-unstructured") + + assert mock_client.get_catalog_by_path.call_args_list[0].args[0] == [ + "dataproducts", + "dre-unstructured", + "sales", + "orders", + ] + assert mock_client.get_catalog_by_path.call_args_list[1].args[0] == [ + "dataproducts", + "dre-unstructured", + "inventory", + "stock", + ] + assert result["count"] == 2 + assert result["results"][0]["path"] == ["dataproducts", "dre-unstructured", "sales", "orders"] + + +@pytest.mark.asyncio +async def test_promote_from_file_rejects_dotdot_segments(mock_client, tmp_path: Path) -> None: + paths_file = tmp_path / "tables.txt" + paths_file.write_text("../bad\n", encoding="utf-8") + + with pytest.raises(ValueError, match="not allowed"): + await promote_from_file(mock_client, paths_file, "dataproducts", under="dre-unstructured") diff --git a/tests/test_commands/test_semantic.py b/tests/test_commands/test_semantic.py new file mode 100644 index 0000000..a09173d --- /dev/null +++ b/tests/test_commands/test_semantic.py @@ -0,0 +1,448 @@ +# +# 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 semantic commands.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock + +import httpx +import pytest + +from drs.commands.semantic import ( + _chunk_paths, + _flatten_semantic_entity, + _load_paths, + _load_task_ids, + abandon_semantic_layer_tasks_batched, + add_semantic_layer_task, + add_semantic_layer_task_batched, + bulk_upsert_entities, + deploy_semantic_layer_tasks_batched, + initialize_semantic_layer, + list_semantic_entities, + patch_semantic_layer_scope, + upsert_table_entity, +) +from drs.output import OutputFormat + + +@pytest.mark.asyncio +async def test_add_semantic_layer_task_with_dictionary_file(mock_client, tmp_path) -> None: + dictionary = tmp_path / "dictionary.md" + dictionary.write_text("orders glossary", encoding="utf-8") + mock_client.add_semantic_layer_task = AsyncMock(return_value={"taskId": "t1"}) + + result = await add_semantic_layer_task( + mock_client, + ["Samples.sales.orders"], + dictionary_file=dictionary, + dictionary_text=None, + include_query_history=True, + max_jobs_to_process=25, + ) + + mock_client.add_semantic_layer_task.assert_called_once_with( + { + "entities": [{"type": "TABLE", "path": ["Samples", "sales", "orders"]}], + "dataDictionary": "orders glossary", + "includeQueryHistory": True, + "maxJobsToProcess": 25, + } + ) + assert result == {"taskId": "t1"} + + +@pytest.mark.asyncio +async def test_initialize_semantic_layer_with_scoped_dictionaries(mock_client, tmp_path) -> None: + scoped = tmp_path / "scoped.yaml" + scoped.write_text( + """ +- scope: + components: [Samples, sales] + dataDictionary: Sales glossary +- dataDictionary: Global glossary +""".strip(), + encoding="utf-8", + ) + mock_client.initialize_semantic_layer = AsyncMock(return_value={"taskId": "init-1"}) + + result = await initialize_semantic_layer( + mock_client, + ["Samples.sales.orders"], + dictionary_file=None, + dictionary_text=None, + scoped_dictionaries_file=scoped, + include_query_history=False, + autopilot_mode="assist", + max_jobs_to_process=None, + ) + + mock_client.initialize_semantic_layer.assert_called_once_with( + { + "entities": [{"type": "TABLE", "path": ["Samples", "sales", "orders"]}], + "includeQueryHistory": False, + "autopilotMode": "ASSIST", + "dataDictionaries": [ + {"scope": {"components": ["Samples", "sales"]}, "dataDictionary": "Sales glossary"}, + {"dataDictionary": "Global glossary"}, + ], + } + ) + assert result == {"taskId": "init-1"} + + +@pytest.mark.asyncio +async def test_upsert_table_entity_uses_catalog_wiki_and_glossary(mock_client, tmp_path) -> None: + glossary = tmp_path / "glossary.yaml" + glossary.write_text( + """ +attributes: + - sourceColumn: order_id + alias: Order ID + description: Unique order identifier + keyType: PRIMARY +""".strip(), + encoding="utf-8", + ) + mock_client.get_catalog_by_path = AsyncMock( + return_value={"id": "catalog-1", "fields": [{"name": "order_id"}, {"name": "customer_id"}]} + ) + mock_client.get_wiki = AsyncMock(return_value={"text": "Orders business definition"}) + request = httpx.Request("PUT", "https://example.com") + response = httpx.Response(404, request=request) + mock_client.update_semantic_layer_entity = AsyncMock( + side_effect=httpx.HTTPStatusError("Not Found", request=request, response=response) + ) + mock_client.add_semantic_layer_entity = AsyncMock(return_value={"id": "entity-1"}) + + result = await upsert_table_entity( + mock_client, + "Samples.sales.orders", + name=None, + description=None, + wiki_from_catalog=True, + glossary_file=glossary, + task_id="draft-1", + ) + + mock_client.add_semantic_layer_entity.assert_called_once_with( + { + "type": "TABLE", + "path": ["Samples", "sales", "orders"], + "name": "orders", + "description": "Orders business definition", + "attributes": [ + { + "sourceColumn": "order_id", + "alias": "Order ID", + "description": "Unique order identifier", + "keyType": "PRIMARY", + }, + { + "sourceColumn": "customer_id", + "alias": "customer_id", + "description": "", + "keyType": "NONE", + }, + ], + }, + task_id="draft-1", + ) + assert result["action"] == "created" + + +@pytest.mark.asyncio +async def test_bulk_upsert_entities_supports_relative_glossary_files(mock_client, tmp_path) -> None: + glossary = tmp_path / "customers-glossary.json" + glossary.write_text( + '[{"sourceColumn":"customer_id","alias":"Customer ID","description":"Business customer key","keyType":"PRIMARY"}]', + encoding="utf-8", + ) + spec = tmp_path / "entities.yaml" + spec.write_text( + """ +entities: + - path: Samples.sales.customers + description: Customer dimension + glossaryFile: customers-glossary.json +""".strip(), + encoding="utf-8", + ) + mock_client.get_catalog_by_path = AsyncMock(return_value={"id": "catalog-2", "fields": [{"name": "customer_id"}]}) + mock_client.update_semantic_layer_entity = AsyncMock(return_value={"id": "entity-2"}) + + result = await bulk_upsert_entities(mock_client, spec, wiki_from_catalog=False, task_id=None) + + assert result["count"] == 1 + mock_client.update_semantic_layer_entity.assert_called_once() + payload = mock_client.update_semantic_layer_entity.call_args.args[1] + assert payload["description"] == "Customer dimension" + assert payload["attributes"][0]["alias"] == "Customer ID" + + +@pytest.mark.asyncio +async def test_patch_semantic_layer_scope(mock_client) -> None: + mock_client.patch_semantic_layer_scope = AsyncMock(return_value={"datasets": []}) + + result = await patch_semantic_layer_scope( + mock_client, + ["Samples.sales.orders"], + ["Samples.sales.customers"], + ) + + mock_client.patch_semantic_layer_scope.assert_called_once_with( + { + "add": [{"components": ["Samples", "sales", "orders"]}], + "remove": [{"components": ["Samples", "sales", "customers"]}], + } + ) + assert result == {"datasets": []} + + +def test_chunk_paths() -> None: + paths = [f"Samples.sales.table_{i}" for i in range(23)] + + result = _chunk_paths(paths) + + assert [len(chunk) for chunk in result] == [10, 10, 3] + + +def test_load_paths_from_text_file(tmp_path) -> None: + paths_file = tmp_path / "paths.txt" + paths_file.write_text( + '\n# comment\n"Samples"."sales"."orders"\n"Samples"."sales"."customers"\n', + encoding="utf-8", + ) + + result = _load_paths(None, paths_file) + + assert result == ['"Samples"."sales"."orders"', '"Samples"."sales"."customers"'] + + +def test_load_paths_from_yaml_file(tmp_path) -> None: + paths_file = tmp_path / "paths.yaml" + paths_file.write_text( + """ +paths: + - path: [Samples, sales, orders] + - components: [Samples, sales, customers] +""".strip(), + encoding="utf-8", + ) + + result = _load_paths(None, paths_file) + + assert result == ["Samples.sales.orders", "Samples.sales.customers"] + + +def test_load_task_ids_from_add_batch_output(tmp_path) -> None: + task_file = tmp_path / "tasks.json" + task_file.write_text( + """ +{ + "batches": 2, + "results": [ + {"task": {"taskId": "11111111-1111-1111-1111-111111111111"}}, + {"task": {"taskId": "22222222-2222-2222-2222-222222222222"}} + ] +} +""".strip(), + encoding="utf-8", + ) + + result = _load_task_ids(None, task_file) + + assert result == [ + "11111111-1111-1111-1111-111111111111", + "22222222-2222-2222-2222-222222222222", + ] + + +@pytest.mark.asyncio +async def test_initialize_semantic_layer_rejects_more_than_ten_paths(mock_client) -> None: + paths = [f"Samples.sales.table_{i}" for i in range(11)] + + with pytest.raises(ValueError, match="at most 10 entities"): + await initialize_semantic_layer( + mock_client, + paths, + dictionary_file=None, + dictionary_text=None, + scoped_dictionaries_file=None, + include_query_history=False, + autopilot_mode=None, + max_jobs_to_process=None, + ) + + +@pytest.mark.asyncio +async def test_add_semantic_layer_task_batched_submits_multiple_requests(mock_client, tmp_path) -> None: + dictionary = tmp_path / "dictionary.md" + dictionary.write_text("global dictionary", encoding="utf-8") + mock_client.add_semantic_layer_task = AsyncMock( + side_effect=[ + {"taskId": "t1", "state": "PENDING", "entitiesSubmitted": 10}, + {"taskId": "t2", "state": "PENDING", "entitiesSubmitted": 2}, + ] + ) + + result = await add_semantic_layer_task_batched( + mock_client, + [f"Samples.sales.table_{i}" for i in range(12)], + dictionary_file=dictionary, + dictionary_text=None, + include_query_history=False, + max_jobs_to_process=None, + ) + + assert result["batches"] == 2 + assert result["totalEntities"] == 12 + assert [entry["task"]["taskId"] for entry in result["results"]] == ["t1", "t2"] + first_call = mock_client.add_semantic_layer_task.call_args_list[0].args[0] + second_call = mock_client.add_semantic_layer_task.call_args_list[1].args[0] + assert first_call["dataDictionary"] == "global dictionary" + assert second_call["dataDictionary"] == "global dictionary" + assert len(first_call["entities"]) == 10 + assert len(second_call["entities"]) == 2 + + +@pytest.mark.asyncio +async def test_deploy_semantic_layer_tasks_batched(mock_client) -> None: + mock_client.deploy_semantic_layer_entities = AsyncMock( + side_effect=[ + {"taskId": "t1", "success": True}, + {"taskId": "t2", "success": True}, + ] + ) + + result = await deploy_semantic_layer_tasks_batched(mock_client, ["t1", "t2"]) + + assert result["action"] == "deploy" + assert result["count"] == 2 + assert [entry["taskId"] for entry in result["results"]] == ["t1", "t2"] + + +@pytest.mark.asyncio +async def test_abandon_semantic_layer_tasks_batched(mock_client) -> None: + mock_client.abandon_semantic_layer_entities = AsyncMock( + side_effect=[ + {"taskId": "t1", "success": True}, + {"taskId": "t2", "success": True}, + ] + ) + + result = await abandon_semantic_layer_tasks_batched(mock_client, ["t1", "t2"]) + + assert result["action"] == "abandon" + assert result["count"] == 2 + assert [entry["taskId"] for entry in result["results"]] == ["t1", "t2"] + + +def test_flatten_semantic_entity() -> None: + result = _flatten_semantic_entity( + { + "id": "table-1", + "type": "TABLE", + "name": "orders", + "path": ["Samples", "sales", "orders"], + "description": "Orders fact table", + "confidenceScore": 92.5, + "attributes": [{"id": "a1"}, {"id": "a2"}], + "relationships": [{"id": "r1"}], + "relatedMetricCount": 4, + } + ) + + assert result["path"] == "Samples.sales.orders" + assert result["attributeCount"] == 2 + assert result["relationshipCount"] == 1 + + +@pytest.mark.asyncio +async def test_list_semantic_entities_fetch_all_pages_for_json(mock_client) -> None: + mock_client.list_semantic_layer_entities = AsyncMock( + side_effect=[ + { + "entities": [{"id": "e1", "type": "TABLE", "path": ["a"], "attributes": []}], + "nextPageToken": "p2", + "totalCount": 2, + }, + { + "entities": [{"id": "e2", "type": "TABLE", "path": ["b"], "attributes": []}], + "totalCount": 2, + }, + ] + ) + + result = await list_semantic_entities( + mock_client, + "TABLE", + task_id="draft-1", + page_token=None, + limit=50, + fetch_all=True, + fmt=OutputFormat.json, + ) + + assert result["count"] == 2 + assert result["allPagesFetched"] is True + assert [entity["id"] for entity in result["entities"]] == ["e1", "e2"] + assert mock_client.list_semantic_layer_entities.call_count == 2 + + +@pytest.mark.asyncio +async def test_list_semantic_entities_flattens_for_pretty(mock_client) -> None: + mock_client.list_semantic_layer_entities = AsyncMock( + return_value={ + "entities": [ + { + "id": "e1", + "type": "TABLE", + "name": "orders", + "path": ["Samples", "sales", "orders"], + "description": "Orders fact table", + "attributes": [{"id": "a1"}], + "relationships": [], + } + ], + "totalCount": 1, + } + ) + + result = await list_semantic_entities( + mock_client, + "TABLE", + task_id=None, + page_token=None, + limit=50, + fetch_all=False, + fmt=OutputFormat.pretty, + ) + + assert result == [ + { + "id": "e1", + "type": "TABLE", + "name": "orders", + "path": "Samples.sales.orders", + "description": "Orders fact table", + "confidenceScore": "", + "attributeCount": 1, + "relationshipCount": 0, + "relatedMetricCount": "", + } + ] diff --git a/tests/test_output.py b/tests/test_output.py index 20a6c1c..50567ca 100644 --- a/tests/test_output.py +++ b/tests/test_output.py @@ -50,7 +50,7 @@ def test_pretty_output_table() -> None: result = render(data, OutputFormat.pretty) assert "id" in result assert "alice" in result - assert "---" in result or "--" in result + assert "bob" in result def test_pretty_output_dict() -> None: diff --git a/uv.lock b/uv.lock index d89917f..25e2fce 100644 --- a/uv.lock +++ b/uv.lock @@ -1,6 +1,17 @@ version = 1 revision = 3 requires-python = ">=3.11" +resolution-markers = [ + "python_full_version >= '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.14' and sys_platform == 'emscripten'", + "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'win32'", + "python_full_version < '3.12' and sys_platform == 'win32'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'emscripten'", + "python_full_version < '3.12' and sys_platform == 'emscripten'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version < '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'", +] [[package]] name = "annotated-doc"