From 2a30bfb6516fb78f5c1ace8ed4c83d6d7da8b197 Mon Sep 17 00:00:00 2001 From: yustme Date: Thu, 16 Jul 2026 14:47:13 +0200 Subject: [PATCH 1/3] feat(semantic-layer): resolve alias/linked-bucket column types via workspace Alias and linked-bucket tables carry no per-column datatype metadata in Storage -- the types live on the source table, so `columnMetadata` on the alias is empty (verified against both BigQuery and Snowflake; it is an alias trait, not a backend one). The build heuristic therefore saw empty basetypes and classified every field as `dimension`. Add an opt-in `kbagent semantic-layer build --types-workspace `: for tables whose Storage metadata carries no column types, query the warehouse `INFORMATION_SCHEMA.COLUMNS` through the given workspace (Query Service) and backfill the real types before generation. Backend-aware SQL for BigQuery and Snowflake; identifiers are validated (not escaped) against a strict charset before interpolation. Unresolved tables are reported non-fatally in the build result as `type_resolution_errors`. Without the flag, behaviour is unchanged. Complements the field-role classifier fix (#486): this makes the types available, that maps a numeric measure column to the `measure` role. --- .../commands/semantic_layer.py | 18 ++ .../services/_semantic_layer_internals.py | 124 ++++++++++ .../services/semantic_layer_service.py | 23 ++ tests/test_semantic_layer_service.py | 228 ++++++++++++++++++ 4 files changed, 393 insertions(+) diff --git a/src/keboola_agent_cli/commands/semantic_layer.py b/src/keboola_agent_cli/commands/semantic_layer.py index e61a242c..655f476b 100644 --- a/src/keboola_agent_cli/commands/semantic_layer.py +++ b/src/keboola_agent_cli/commands/semantic_layer.py @@ -335,6 +335,13 @@ def _print_build_result(console: Console, data: dict) -> None: console.print(f" [red]✗[/red] {e['type']} {e['item']} — {e['detail']}") if warns: console.print(f"\n[bold yellow]Validation: {len(warns)} warning(s)[/bold yellow]") + type_errs = data.get("type_resolution_errors") or [] + if type_errs: + console.print( + f"\n[bold yellow]Type resolution: {len(type_errs)} table(s) unresolved[/bold yellow]" + ) + for e in type_errs: + console.print(f" [yellow]![/yellow] {e['table_id']} — {e['error']}") if data.get("created"): console.print(f"\nCreated: {data['created']}") elif data.get("dry_run"): @@ -434,6 +441,16 @@ def semantic_layer_build( output: Path | None = typer.Option( None, "--output", help="Also write the generated JSON to this file." ), + types_workspace: int | None = typer.Option( + None, + "--types-workspace", + help=( + "Workspace ID used to read real column types from the warehouse " + "INFORMATION_SCHEMA for tables whose Storage metadata carries none " + "(alias / linked-bucket tables). Without it such tables reach the " + "heuristic untyped and every field becomes a dimension." + ), + ), ) -> None: """Build a semantic-layer model from a list of storage tables (non-interactive). @@ -471,6 +488,7 @@ def semantic_layer_build( dry_run=dry_run, keep_on_failure=keep_on_failure, output_path=output, + types_workspace_id=types_workspace, ) formatter.output(result, _print_build_result) diff --git a/src/keboola_agent_cli/services/_semantic_layer_internals.py b/src/keboola_agent_cli/services/_semantic_layer_internals.py index cd461046..2b8209bb 100644 --- a/src/keboola_agent_cli/services/_semantic_layer_internals.py +++ b/src/keboola_agent_cli/services/_semantic_layer_internals.py @@ -25,6 +25,8 @@ from __future__ import annotations +import csv +import io import json import logging import os @@ -807,6 +809,128 @@ def _worker(tid: str) -> WorkerResult: return schemas_by_tid, fetch_errors +# ── build: warehouse type resolution (opt-in `--types-workspace`) ──── + +# Keboola bucket/table identifiers are restricted to this character set. We +# interpolate them into INFORMATION_SCHEMA SQL, so an identifier containing +# anything else is rejected rather than escaped -- defense in depth even +# though the Storage API would never mint such an id. +_SAFE_IDENTIFIER_RE = re.compile(r"^[A-Za-z0-9_.\-]+$") + + +def build_information_schema_sql(backend: str, bucket_id: str, table_name: str) -> str | None: + """Build an ``INFORMATION_SCHEMA.COLUMNS`` query for a table's real types. + + Returns a query projecting ``(column_name, data_type)`` for the given + backend, or ``None`` when the backend is unsupported or an identifier + contains an unexpected character (see ``_SAFE_IDENTIFIER_RE``). + """ + if not ( + _SAFE_IDENTIFIER_RE.match(bucket_id or "") and _SAFE_IDENTIFIER_RE.match(table_name or "") + ): + return None + normalized_backend = (backend or "").lower() + if normalized_backend == "bigquery": + # Keboola maps a bucket to a BigQuery dataset by replacing `.` and `-` + # with `_` (e.g. `in.c-Foo-Bar` -> `in_c_Foo_Bar`). + dataset = bucket_id.replace(".", "_").replace("-", "_") + return ( + "SELECT column_name, data_type " + f"FROM `{dataset}`.INFORMATION_SCHEMA.COLUMNS " + f"WHERE table_name = '{table_name}'" + ) + if normalized_backend == "snowflake": + # Keboola's Snowflake mapping keeps the bucketId verbatim as the schema + # name inside the KEBOOLA database (mirrors ``_derive_fqn``). + return ( + "SELECT COLUMN_NAME AS column_name, DATA_TYPE AS data_type " + 'FROM "KEBOOLA".INFORMATION_SCHEMA.COLUMNS ' + f"WHERE TABLE_SCHEMA = '{bucket_id}' AND TABLE_NAME = '{table_name}'" + ) + return None + + +def _parse_information_schema_result(result: dict[str, Any]) -> dict[str, str]: + """Extract a ``{column_name: data_type}`` map from an execute_query result. + + Reads the first statement's synthesized ``csv_data`` (always present on the + fast inline path), which carries a ``column_name,data_type`` header. + """ + statements = result.get("statements") or [] + if not statements: + return {} + csv_data = statements[0].get("csv_data") + if not csv_data: + return {} + type_map: dict[str, str] = {} + for row in csv.DictReader(io.StringIO(csv_data)): + name = (row.get("column_name") or row.get("COLUMN_NAME") or "").strip() + dtype = (row.get("data_type") or row.get("DATA_TYPE") or "").strip() + if name and dtype: + type_map[name] = dtype + return type_map + + +def enrich_schemas_with_workspace_types( + *, + schemas_by_tid: dict[str, dict[str, Any]], + alias: str, + workspace_id: int, + execute_query: Callable[..., dict[str, Any]], +) -> list[dict[str, str]]: + """Backfill missing column ``type``s from the warehouse INFORMATION_SCHEMA. + + Alias / linked-bucket tables carry no per-column datatype metadata in + Storage -- the types live on the source table -- so the build heuristic + sees empty basetypes and classifies every field as ``dimension``. When the + caller supplies a workspace, query the warehouse for the real types and + fill them in on each ``column_details`` entry in place. + + Returns a list of ``{"table_id", "error"}`` for tables that could not be + (fully) resolved. These are non-fatal and surfaced in the build result. + """ + unresolved: list[dict[str, str]] = [] + for tid, detail in schemas_by_tid.items(): + cols = detail.get("column_details") or [] + if not any(not col.get("type") for col in cols): + continue # already fully typed (or no columns) -- nothing to do + sql = build_information_schema_sql( + detail.get("backend", ""), + detail.get("bucket_id", ""), + detail.get("name", "") or tid.split(".")[-1], + ) + if sql is None: + unresolved.append( + { + "table_id": tid, + "error": f"type resolution unsupported for backend {detail.get('backend', '')!r}", + } + ) + continue + try: + result = execute_query(alias, workspace_id, sql) + except (KeboolaApiError, ConfigError) as exc: + unresolved.append({"table_id": tid, "error": str(exc)}) + continue + type_map = _parse_information_schema_result(result) + if not type_map: + unresolved.append( + {"table_id": tid, "error": "INFORMATION_SCHEMA returned no column types"} + ) + continue + for col in cols: + if not col.get("type"): + dtype = type_map.get(col.get("name", "")) + if dtype: + col["type"] = dtype + missing = [col.get("name", "") for col in cols if not col.get("type")] + if missing: + unresolved.append( + {"table_id": tid, "error": f"no type found for column(s): {', '.join(missing)}"} + ) + return unresolved + + def resolve_model_uuid( client: Any, model_name_or_uuid: str | None, diff --git a/src/keboola_agent_cli/services/semantic_layer_service.py b/src/keboola_agent_cli/services/semantic_layer_service.py index 3d9a4878..f527318f 100644 --- a/src/keboola_agent_cli/services/semantic_layer_service.py +++ b/src/keboola_agent_cli/services/semantic_layer_service.py @@ -39,6 +39,9 @@ from ._semantic_layer_internals import collect_side_from_file from ._semantic_layer_internals import default_export_path as _default_export_path from ._semantic_layer_internals import diff_one_type as _diff_one_type_helper +from ._semantic_layer_internals import ( + enrich_schemas_with_workspace_types as _enrich_schemas_with_workspace_types, +) from ._semantic_layer_internals import fetch_table_schemas as _fetch_table_schemas from ._semantic_layer_internals import heuristic_generate_model as _heuristic_generate_helper from ._semantic_layer_internals import push_built_model as _push_built_model @@ -58,6 +61,7 @@ from .base import BaseService, ClientFactory from .encrypt_service import EncryptService from .storage_service import StorageService +from .workspace_service import WorkspaceService # Constraint name regex enforced by the metastore server. CONSTRAINT_NAME_RE = re.compile(r"^[a-z][a-z0-9_]*$") @@ -1360,6 +1364,7 @@ def build_model( dry_run: bool = False, keep_on_failure: bool = False, output_path: Path | None = None, + types_workspace_id: int | None = None, ) -> dict[str, Any]: """Build (or update) a semantic-layer model from a list of tableIds. @@ -1397,6 +1402,23 @@ def build_model( ) schemas_by_tid, fetch_errors = _fetch_table_schemas(storage, alias, table_ids) + # Opt-in: backfill column types that Storage does not carry locally + # (alias / linked-bucket tables) by querying the warehouse + # INFORMATION_SCHEMA via the supplied workspace. Without this, such + # tables reach the heuristic with empty basetypes and every field is + # classified as `dimension`. + type_resolution_errors: list[dict[str, str]] = [] + if types_workspace_id is not None: + workspace = WorkspaceService( + config_store=self._config_store, client_factory=self._client_factory + ) + type_resolution_errors = _enrich_schemas_with_workspace_types( + schemas_by_tid=schemas_by_tid, + alias=alias, + workspace_id=types_workspace_id, + execute_query=workspace.execute_query, + ) + # Generate model JSON (heuristic). The internals helper takes the # FQN-derivation and role-classification callables as kwargs so its # module stays free of import-time coupling to this module. @@ -1430,6 +1452,7 @@ def build_model( "keep_on_failure": keep_on_failure, "fallback_used": "heuristic", # no AI endpoint shipped yet "fetch_errors": fetch_errors, + "type_resolution_errors": type_resolution_errors, "generated": generated, "validation": {"errors": errors, "warnings": warnings}, "validated": len(errors) == 0, diff --git a/tests/test_semantic_layer_service.py b/tests/test_semantic_layer_service.py index e463ff0f..67ca3604 100644 --- a/tests/test_semantic_layer_service.py +++ b/tests/test_semantic_layer_service.py @@ -2870,3 +2870,231 @@ def test_registry_entries(self) -> None: assert OPERATION_REGISTRY["semantic-layer.reference-data.get"] == "read" assert OPERATION_REGISTRY["semantic-layer.reference-data.set"] == "write" assert OPERATION_REGISTRY["semantic-layer.reference-data.delete"] == "destructive" + + +# --------------------------------------------------------------------------- +# build --types-workspace: warehouse type resolution for alias/linked tables +# --------------------------------------------------------------------------- + + +class TestBuildInformationSchemaSql: + """SQL generation for the INFORMATION_SCHEMA type-resolution query.""" + + def test_bigquery_dataset_derivation(self) -> None: + from keboola_agent_cli.services._semantic_layer_internals import ( + build_information_schema_sql, + ) + + sql = build_information_schema_sql( + "bigquery", "in.c-OUT_Shop_Category_Metrics", "out_category_metrics" + ) + assert sql is not None + # bucketId `.`/`-` map to `_` for the BigQuery dataset name. + assert "`in_c_OUT_Shop_Category_Metrics`.INFORMATION_SCHEMA.COLUMNS" in sql + assert "table_name = 'out_category_metrics'" in sql + assert "column_name, data_type" in sql + + def test_snowflake_keeps_bucket_id_as_schema(self) -> None: + from keboola_agent_cli.services._semantic_layer_internals import ( + build_information_schema_sql, + ) + + sql = build_information_schema_sql("snowflake", "out.c-fake-ecommerce", "orders") + assert sql is not None + assert '"KEBOOLA".INFORMATION_SCHEMA.COLUMNS' in sql + assert "TABLE_SCHEMA = 'out.c-fake-ecommerce'" in sql + assert "TABLE_NAME = 'orders'" in sql + + def test_unsupported_backend_returns_none(self) -> None: + from keboola_agent_cli.services._semantic_layer_internals import ( + build_information_schema_sql, + ) + + assert build_information_schema_sql("redshift", "in.c-x", "t") is None + assert build_information_schema_sql("", "in.c-x", "t") is None + + def test_unsafe_identifier_rejected(self) -> None: + from keboola_agent_cli.services._semantic_layer_internals import ( + build_information_schema_sql, + ) + + # A quote in the table name must not be escaped-and-passed; it is + # rejected outright (defense in depth against SQL injection). + assert build_information_schema_sql("bigquery", "in.c-x", "t'; DROP TABLE") is None + assert build_information_schema_sql("bigquery", "in.c-x`", "t") is None + + +class TestEnrichSchemasWithWorkspaceTypes: + """Backfilling missing column types from the warehouse INFORMATION_SCHEMA.""" + + @staticmethod + def _query_result(csv_data: str) -> dict[str, Any]: + return {"statements": [{"csv_data": csv_data}]} + + def test_fills_missing_types_in_place(self) -> None: + from keboola_agent_cli.services._semantic_layer_internals import ( + enrich_schemas_with_workspace_types, + ) + + schemas = { + "in.c-linked.metrics": { + "backend": "bigquery", + "bucket_id": "in.c-linked", + "name": "metrics", + "column_details": [ + {"name": "total_revenue"}, # no type (alias table) + {"name": "order_date"}, + ], + } + } + execute_query = MagicMock( + return_value=self._query_result( + "column_name,data_type\r\ntotal_revenue,NUMERIC\r\norder_date,DATE\r\n" + ) + ) + errors = enrich_schemas_with_workspace_types( + schemas_by_tid=schemas, alias="prod", workspace_id=42, execute_query=execute_query + ) + assert errors == [] + cols = {c["name"]: c["type"] for c in schemas["in.c-linked.metrics"]["column_details"]} + assert cols == {"total_revenue": "NUMERIC", "order_date": "DATE"} + execute_query.assert_called_once() + args = execute_query.call_args.args + assert args[0] == "prod" and args[1] == 42 + + def test_fully_typed_table_is_skipped(self) -> None: + from keboola_agent_cli.services._semantic_layer_internals import ( + enrich_schemas_with_workspace_types, + ) + + schemas = { + "out.c-native.t": { + "backend": "snowflake", + "bucket_id": "out.c-native", + "name": "t", + "column_details": [{"name": "x", "type": "NUMERIC"}], + } + } + execute_query = MagicMock() + errors = enrich_schemas_with_workspace_types( + schemas_by_tid=schemas, alias="prod", workspace_id=1, execute_query=execute_query + ) + assert errors == [] + execute_query.assert_not_called() + + def test_query_error_is_non_fatal(self) -> None: + from keboola_agent_cli.services._semantic_layer_internals import ( + enrich_schemas_with_workspace_types, + ) + + schemas = { + "in.c-linked.t": { + "backend": "bigquery", + "bucket_id": "in.c-linked", + "name": "t", + "column_details": [{"name": "a"}], + } + } + execute_query = MagicMock( + side_effect=KeboolaApiError(message="boom", error_code=ErrorCode.API_ERROR) + ) + errors = enrich_schemas_with_workspace_types( + schemas_by_tid=schemas, alias="prod", workspace_id=1, execute_query=execute_query + ) + assert len(errors) == 1 + assert errors[0]["table_id"] == "in.c-linked.t" + # column stays untyped -- build still proceeds, field defaults to dimension + assert "type" not in schemas["in.c-linked.t"]["column_details"][0] + + def test_columns_absent_from_result_reported(self) -> None: + from keboola_agent_cli.services._semantic_layer_internals import ( + enrich_schemas_with_workspace_types, + ) + + schemas = { + "in.c-linked.t": { + "backend": "bigquery", + "bucket_id": "in.c-linked", + "name": "t", + "column_details": [{"name": "a"}, {"name": "b"}], + } + } + execute_query = MagicMock( + return_value=self._query_result("column_name,data_type\r\na,INT64\r\n") + ) + errors = enrich_schemas_with_workspace_types( + schemas_by_tid=schemas, alias="prod", workspace_id=1, execute_query=execute_query + ) + assert schemas["in.c-linked.t"]["column_details"][0]["type"] == "INT64" + assert len(errors) == 1 + assert "b" in errors[0]["error"] + + +class TestBuildModelTypesWorkspace: + """`build_model(types_workspace_id=...)` wiring end to end (mocked).""" + + def test_types_workspace_resolves_alias_types(self, tmp_path: Path) -> None: + store = _make_store(tmp_path) + service, mock = _make_service(store) + mock.list_items.return_value = [] + mock.post_item.side_effect = [ + {"id": "new-model"}, + {"id": "d1"}, + {"id": "m1"}, + {"id": "g1"}, + ] + with ( + patch( + "keboola_agent_cli.services.semantic_layer_service.StorageService" + ) as MockStorageCls, + patch( + "keboola_agent_cli.services.semantic_layer_service.WorkspaceService" + ) as MockWorkspaceCls, + ): + MockStorageCls.return_value.get_table_detail.return_value = { + "display_name": "metrics", + "backend": "bigquery", + "bucket_id": "in.c-linked", + "name": "metrics", + # Alias table: columns carry no type. + "column_details": [{"name": "total_revenue"}], + } + MockWorkspaceCls.return_value.execute_query.return_value = { + "statements": [{"csv_data": "column_name,data_type\r\ntotal_revenue,NUMERIC\r\n"}] + } + result = service.build_model( + "prod", table_ids=["in.c-linked.metrics"], types_workspace_id=7, dry_run=True + ) + assert result["type_resolution_errors"] == [] + fields = result["generated"]["datasets"][0]["fields"] + # The alias column, untyped in Storage, is now resolved: the warehouse + # NUMERIC normalizes to the metastore `decimal`. (Whether a numeric + # measure column then becomes a `measure` role is the classifier's job, + # fixed in the companion field-role PR.) + assert fields[0]["type"] == "decimal" + + def test_without_workspace_alias_types_stay_empty(self, tmp_path: Path) -> None: + store = _make_store(tmp_path) + service, mock = _make_service(store) + mock.list_items.return_value = [] + with ( + patch( + "keboola_agent_cli.services.semantic_layer_service.StorageService" + ) as MockStorageCls, + patch( + "keboola_agent_cli.services.semantic_layer_service.WorkspaceService" + ) as MockWorkspaceCls, + ): + MockStorageCls.return_value.get_table_detail.return_value = { + "display_name": "metrics", + "backend": "bigquery", + "bucket_id": "in.c-linked", + "name": "metrics", + "column_details": [{"name": "total_revenue"}], + } + result = service.build_model("prod", table_ids=["in.c-linked.metrics"], dry_run=True) + MockWorkspaceCls.return_value.execute_query.assert_not_called() + fields = result["generated"]["datasets"][0]["fields"] + # Untyped -> string, and the measure-named column defaults to dimension. + assert fields[0]["type"] == "string" + assert fields[0]["role"] == "dimension" From 92f69e878b3859f8cc97c904b027ce8dc2cd0b2d Mon Sep 17 00:00:00 2001 From: yustme Date: Thu, 16 Jul 2026 14:53:24 +0200 Subject: [PATCH 2/3] fix(semantic-layer): surface bucket backend from get_table_detail for type resolution The workspace type-resolution path keys the INFORMATION_SCHEMA dialect off the table's storage backend, but get_table_detail did not expose it, so --types-workspace failed with "unsupported backend ''" and every field still defaulted to dimension. Add `backend` (from the owning bucket) to the get_table_detail response and cover it with a test. --- .../services/storage_service.py | 4 ++++ tests/test_storage_describe_service.py | 23 +++++++++++++++++++ 2 files changed, 27 insertions(+) diff --git a/src/keboola_agent_cli/services/storage_service.py b/src/keboola_agent_cli/services/storage_service.py index d4ee80d9..47c4264f 100644 --- a/src/keboola_agent_cli/services/storage_service.py +++ b/src/keboola_agent_cli/services/storage_service.py @@ -595,6 +595,10 @@ def get_table_detail( "name": table.get("name", ""), "display_name": table.get("displayName", ""), "bucket_id": table.get("bucket", {}).get("id", ""), + # Storage backend of the owning bucket (e.g. "snowflake", + # "bigquery"); needed to pick the right INFORMATION_SCHEMA dialect + # when resolving column types for alias / linked tables. + "backend": table.get("bucket", {}).get("backend", ""), "description": description, "columns": columns, "column_details": column_details, diff --git a/tests/test_storage_describe_service.py b/tests/test_storage_describe_service.py index 8098546c..e24f5743 100644 --- a/tests/test_storage_describe_service.py +++ b/tests/test_storage_describe_service.py @@ -689,6 +689,29 @@ def test_extracts_table_and_column_descriptions(self, tmp_path: Path) -> None: assert isinstance(result["metadata"], list) assert len(result["metadata"]) == 3 + def test_surfaces_bucket_backend(self, tmp_path: Path) -> None: + """The owning bucket's storage backend is exposed on the response. + + Consumed by semantic-layer build to pick the INFORMATION_SCHEMA + dialect when resolving column types for alias / linked tables. + """ + store = _make_store(tmp_path) + mock_client = MagicMock() + mock_client.get_table_detail.return_value = { + "id": "in.c-sales.orders", + "name": "orders", + "displayName": "orders", + "bucket": {"id": "in.c-sales", "backend": "bigquery"}, + "columns": ["order_id"], + "columnMetadata": {}, + "metadata": [], + } + service = _make_service(store, mock_client) + + result = service.get_table_detail(alias="prod", table_id="in.c-sales.orders") + + assert result["backend"] == "bigquery" + def test_columns_without_description_have_no_description_key(self, tmp_path: Path) -> None: """Columns without a matching KBC.column.{name}.description entry omit 'description'.""" store = _make_store(tmp_path) From 1872bf99994799538c66271e75aca86888f2b755 Mon Sep 17 00:00:00 2001 From: yustme Date: Thu, 16 Jul 2026 15:40:10 +0200 Subject: [PATCH 3/3] feat(semantic-layer): auto-resolve a workspace for build type resolution Wire type resolution through the server /build endpoint and make it work from the UI without any extra input: BuildRequest gains `types_workspace_id` (explicit) and `auto_resolve_types` (default True), and the endpoint passes both to build_model. In auto mode the service picks a Query-Service-capable, read-only workspace whose backend matches each table (via list_workspaces), cached per backend. The enrichment helper now takes a `resolve_workspace_id(backend)` callable instead of a fixed id, so the explicit-id and auto paths share one code path. No workspace found -> reported non-fatally, build still proceeds. CLI keeps explicit `--types-workspace`; adds `--auto-types-workspace` for parity. Because the server default is auto, an unchanged UI build now resolves alias/linked-bucket column types instead of defaulting every field to dimension. --- .../commands/semantic_layer.py | 9 ++ .../server/routers/semantic_layer.py | 7 ++ .../services/_semantic_layer_internals.py | 27 ++++-- .../services/semantic_layer_service.py | 39 +++++++- tests/test_semantic_layer_service.py | 95 ++++++++++++++++++- 5 files changed, 165 insertions(+), 12 deletions(-) diff --git a/src/keboola_agent_cli/commands/semantic_layer.py b/src/keboola_agent_cli/commands/semantic_layer.py index 655f476b..91e73a81 100644 --- a/src/keboola_agent_cli/commands/semantic_layer.py +++ b/src/keboola_agent_cli/commands/semantic_layer.py @@ -451,6 +451,14 @@ def semantic_layer_build( "heuristic untyped and every field becomes a dimension." ), ), + auto_types_workspace: bool = typer.Option( + False, + "--auto-types-workspace", + help=( + "Auto-pick a read-only workspace per backend for column-type " + "resolution (instead of passing --types-workspace explicitly)." + ), + ), ) -> None: """Build a semantic-layer model from a list of storage tables (non-interactive). @@ -489,6 +497,7 @@ def semantic_layer_build( keep_on_failure=keep_on_failure, output_path=output, types_workspace_id=types_workspace, + auto_resolve_types=auto_types_workspace, ) formatter.output(result, _print_build_result) diff --git a/src/keboola_agent_cli/server/routers/semantic_layer.py b/src/keboola_agent_cli/server/routers/semantic_layer.py index 6c20a5b7..b34b8428 100644 --- a/src/keboola_agent_cli/server/routers/semantic_layer.py +++ b/src/keboola_agent_cli/server/routers/semantic_layer.py @@ -188,6 +188,11 @@ class BuildRequest(BaseModel): name: str | None = None dry_run: bool = False keep_on_failure: bool = False + # Column-type resolution for alias / linked tables (empty Storage + # metadata). Provide a specific workspace, or leave auto-resolve on (the + # UI default) to have the server pick a read-only workspace per backend. + types_workspace_id: int | None = None + auto_resolve_types: bool = True class TokenEncryptRequest(BaseModel): @@ -587,6 +592,8 @@ def build(body: BuildRequest, registry: ServiceRegistry = Depends(get_registry)) model_name_or_uuid=body.model, dry_run=body.dry_run, keep_on_failure=body.keep_on_failure, + types_workspace_id=body.types_workspace_id, + auto_resolve_types=body.auto_resolve_types, ) diff --git a/src/keboola_agent_cli/services/_semantic_layer_internals.py b/src/keboola_agent_cli/services/_semantic_layer_internals.py index 2b8209bb..4d194bea 100644 --- a/src/keboola_agent_cli/services/_semantic_layer_internals.py +++ b/src/keboola_agent_cli/services/_semantic_layer_internals.py @@ -875,16 +875,21 @@ def enrich_schemas_with_workspace_types( *, schemas_by_tid: dict[str, dict[str, Any]], alias: str, - workspace_id: int, + resolve_workspace_id: Callable[[str], int | None], execute_query: Callable[..., dict[str, Any]], ) -> list[dict[str, str]]: """Backfill missing column ``type``s from the warehouse INFORMATION_SCHEMA. Alias / linked-bucket tables carry no per-column datatype metadata in Storage -- the types live on the source table -- so the build heuristic - sees empty basetypes and classifies every field as ``dimension``. When the - caller supplies a workspace, query the warehouse for the real types and - fill them in on each ``column_details`` entry in place. + sees empty basetypes and classifies every field as ``dimension``. When a + workspace is available, query the warehouse for the real types and fill + them in on each ``column_details`` entry in place. + + ``resolve_workspace_id`` maps a storage backend (e.g. ``"bigquery"``) to a + workspace ID that can run SQL against it -- either a fixed id supplied by + the caller, or one auto-picked per backend. Returning ``None`` means no + suitable workspace, and the table is reported as unresolved. Returns a list of ``{"table_id", "error"}`` for tables that could not be (fully) resolved. These are non-fatal and surfaced in the build result. @@ -894,8 +899,9 @@ def enrich_schemas_with_workspace_types( cols = detail.get("column_details") or [] if not any(not col.get("type") for col in cols): continue # already fully typed (or no columns) -- nothing to do + backend = detail.get("backend", "") sql = build_information_schema_sql( - detail.get("backend", ""), + backend, detail.get("bucket_id", ""), detail.get("name", "") or tid.split(".")[-1], ) @@ -903,7 +909,16 @@ def enrich_schemas_with_workspace_types( unresolved.append( { "table_id": tid, - "error": f"type resolution unsupported for backend {detail.get('backend', '')!r}", + "error": f"type resolution unsupported for backend {backend!r}", + } + ) + continue + workspace_id = resolve_workspace_id(backend) + if workspace_id is None: + unresolved.append( + { + "table_id": tid, + "error": f"no workspace available to read {backend!r} types", } ) continue diff --git a/src/keboola_agent_cli/services/semantic_layer_service.py b/src/keboola_agent_cli/services/semantic_layer_service.py index f527318f..20736a9b 100644 --- a/src/keboola_agent_cli/services/semantic_layer_service.py +++ b/src/keboola_agent_cli/services/semantic_layer_service.py @@ -1354,6 +1354,26 @@ def promote_model( # Phase 7 — build (AI-assisted / heuristic greenfield) # ------------------------------------------------------------------ + @staticmethod + def _pick_types_workspace(workspace: WorkspaceService, alias: str, backend: str) -> int | None: + """Auto-pick a workspace able to read ``backend`` types via the Query + Service. + + Prefers a Query-Service-compatible, read-only workspace whose backend + matches the table's (so ``INFORMATION_SCHEMA`` runs in the right + dialect). Returns its ID, or ``None`` if the project has none. + """ + try: + listing = workspace.list_workspaces(aliases=[alias], qs_compatible_only=True) + except (KeboolaApiError, ConfigError): + return None + candidates = listing.get("workspaces", []) + for ws in candidates: + if not backend or (ws.get("backend", "") or "").lower() == backend.lower(): + ws_id = ws.get("id") + return int(ws_id) if ws_id is not None else None + return None + def build_model( self, alias: str, @@ -1365,6 +1385,7 @@ def build_model( keep_on_failure: bool = False, output_path: Path | None = None, types_workspace_id: int | None = None, + auto_resolve_types: bool = False, ) -> dict[str, Any]: """Build (or update) a semantic-layer model from a list of tableIds. @@ -1408,14 +1429,28 @@ def build_model( # tables reach the heuristic with empty basetypes and every field is # classified as `dimension`. type_resolution_errors: list[dict[str, str]] = [] - if types_workspace_id is not None: + if types_workspace_id is not None or auto_resolve_types: workspace = WorkspaceService( config_store=self._config_store, client_factory=self._client_factory ) + if types_workspace_id is not None: + # Explicit workspace: use it for every backend. + def resolve_workspace_id(_backend: str) -> int | None: + return types_workspace_id + else: + # Auto mode (UI / server default): pick a Query-Service-capable + # read-only workspace per backend, cached across tables. + ws_cache: dict[str, int | None] = {} + + def resolve_workspace_id(backend: str) -> int | None: + if backend not in ws_cache: + ws_cache[backend] = self._pick_types_workspace(workspace, alias, backend) + return ws_cache[backend] + type_resolution_errors = _enrich_schemas_with_workspace_types( schemas_by_tid=schemas_by_tid, alias=alias, - workspace_id=types_workspace_id, + resolve_workspace_id=resolve_workspace_id, execute_query=workspace.execute_query, ) diff --git a/tests/test_semantic_layer_service.py b/tests/test_semantic_layer_service.py index 67ca3604..d8507d6f 100644 --- a/tests/test_semantic_layer_service.py +++ b/tests/test_semantic_layer_service.py @@ -2953,7 +2953,10 @@ def test_fills_missing_types_in_place(self) -> None: ) ) errors = enrich_schemas_with_workspace_types( - schemas_by_tid=schemas, alias="prod", workspace_id=42, execute_query=execute_query + schemas_by_tid=schemas, + alias="prod", + resolve_workspace_id=lambda _backend: 42, + execute_query=execute_query, ) assert errors == [] cols = {c["name"]: c["type"] for c in schemas["in.c-linked.metrics"]["column_details"]} @@ -2977,7 +2980,10 @@ def test_fully_typed_table_is_skipped(self) -> None: } execute_query = MagicMock() errors = enrich_schemas_with_workspace_types( - schemas_by_tid=schemas, alias="prod", workspace_id=1, execute_query=execute_query + schemas_by_tid=schemas, + alias="prod", + resolve_workspace_id=lambda _backend: 1, + execute_query=execute_query, ) assert errors == [] execute_query.assert_not_called() @@ -2999,7 +3005,10 @@ def test_query_error_is_non_fatal(self) -> None: side_effect=KeboolaApiError(message="boom", error_code=ErrorCode.API_ERROR) ) errors = enrich_schemas_with_workspace_types( - schemas_by_tid=schemas, alias="prod", workspace_id=1, execute_query=execute_query + schemas_by_tid=schemas, + alias="prod", + resolve_workspace_id=lambda _backend: 1, + execute_query=execute_query, ) assert len(errors) == 1 assert errors[0]["table_id"] == "in.c-linked.t" @@ -3023,7 +3032,10 @@ def test_columns_absent_from_result_reported(self) -> None: return_value=self._query_result("column_name,data_type\r\na,INT64\r\n") ) errors = enrich_schemas_with_workspace_types( - schemas_by_tid=schemas, alias="prod", workspace_id=1, execute_query=execute_query + schemas_by_tid=schemas, + alias="prod", + resolve_workspace_id=lambda _backend: 1, + execute_query=execute_query, ) assert schemas["in.c-linked.t"]["column_details"][0]["type"] == "INT64" assert len(errors) == 1 @@ -3098,3 +3110,78 @@ def test_without_workspace_alias_types_stay_empty(self, tmp_path: Path) -> None: # Untyped -> string, and the measure-named column defaults to dimension. assert fields[0]["type"] == "string" assert fields[0]["role"] == "dimension" + + def test_auto_resolve_picks_workspace_per_backend(self, tmp_path: Path) -> None: + """`auto_resolve_types=True` (the server/UI default) picks a read-only + workspace itself -- no explicit id needed.""" + store = _make_store(tmp_path) + service, mock = _make_service(store) + mock.list_items.return_value = [] + mock.post_item.side_effect = [ + {"id": "new-model"}, + {"id": "d1"}, + {"id": "m1"}, + {"id": "g1"}, + ] + with ( + patch( + "keboola_agent_cli.services.semantic_layer_service.StorageService" + ) as MockStorageCls, + patch( + "keboola_agent_cli.services.semantic_layer_service.WorkspaceService" + ) as MockWorkspaceCls, + ): + MockStorageCls.return_value.get_table_detail.return_value = { + "display_name": "metrics", + "backend": "bigquery", + "bucket_id": "in.c-linked", + "name": "metrics", + "column_details": [{"name": "total_revenue"}], + } + ws = MockWorkspaceCls.return_value + # Two workspaces; only the BigQuery one matches the table backend. + ws.list_workspaces.return_value = { + "workspaces": [ + {"id": 111, "backend": "snowflake"}, + {"id": 222, "backend": "bigquery"}, + ] + } + ws.execute_query.return_value = { + "statements": [{"csv_data": "column_name,data_type\r\ntotal_revenue,NUMERIC\r\n"}] + } + result = service.build_model( + "prod", table_ids=["in.c-linked.metrics"], auto_resolve_types=True, dry_run=True + ) + # The backend-matching workspace (222) was used, not 111. + assert ws.execute_query.call_args.args[1] == 222 + assert result["type_resolution_errors"] == [] + assert result["generated"]["datasets"][0]["fields"][0]["type"] == "decimal" + + def test_auto_resolve_no_workspace_reports_error(self, tmp_path: Path) -> None: + """Auto mode with no matching workspace is non-fatal: build proceeds, + the table is reported unresolved.""" + store = _make_store(tmp_path) + service, mock = _make_service(store) + mock.list_items.return_value = [] + with ( + patch( + "keboola_agent_cli.services.semantic_layer_service.StorageService" + ) as MockStorageCls, + patch( + "keboola_agent_cli.services.semantic_layer_service.WorkspaceService" + ) as MockWorkspaceCls, + ): + MockStorageCls.return_value.get_table_detail.return_value = { + "display_name": "metrics", + "backend": "bigquery", + "bucket_id": "in.c-linked", + "name": "metrics", + "column_details": [{"name": "total_revenue"}], + } + MockWorkspaceCls.return_value.list_workspaces.return_value = {"workspaces": []} + result = service.build_model( + "prod", table_ids=["in.c-linked.metrics"], auto_resolve_types=True, dry_run=True + ) + MockWorkspaceCls.return_value.execute_query.assert_not_called() + assert len(result["type_resolution_errors"]) == 1 + assert "no workspace" in result["type_resolution_errors"][0]["error"]