Compare commits

...

11 commits

Author SHA1 Message Date
Andrey Avtomonov
49a4ae6f5b
fix(ci): repair slow sl test and sync uv.lock to 0.16.0 (#345)
* test(sl): assert predefined_measures_only in slow sl query test

The slow-only sl query test asserted the outgoing compute.query() request
with an exact nested query object, so the predefined_measures_only flag
added by the semantic-layer-only query policy (#334) broke the match and
turned the Slow TypeScript tests CI job red on main.

* chore: sync uv.lock ktx-sl/ktx-daemon to 0.16.0

The 0.16.0 release bumped the pyproject versions but left uv.lock pinning
the editable ktx-sl/ktx-daemon packages at 0.15.0.
2026-07-07 12:38:33 +00:00
semantic-release-bot
a6dd8cf730 chore(release): 0.16.0 [skip ci]
## [0.16.0](https://github.com/Kaelio/ktx/compare/v0.15.0...v0.16.0) (2026-07-03)

### Features

* Add duckdb connector ([#308](https://github.com/Kaelio/ktx/issues/308)) ([3c4fcc2](3c4fcc27c7))
* **athena:** first-class AWS Athena warehouse identity (SL + BI mapping) ([#332](https://github.com/Kaelio/ktx/issues/332)) ([f310391](f310391da5))
* **connector:** add Amazon Athena connector via Glue Data Catalog ([#309](https://github.com/Kaelio/ktx/issues/309)) ([fe7e6bd](fe7e6bd1fa))
* query_policy semantic-layer-only restricts agents to predefined semantic-layer measures ([#334](https://github.com/Kaelio/ktx/issues/334)) ([a651b82](a651b82e2f))

### Bug Fixes

* **deps:** patch 22 Dependabot security alerts ([#328](https://github.com/Kaelio/ktx/issues/328)) ([6d01030](6d01030745))
* **sl:** classify semantic-query request rejections as expected, not faults ([#339](https://github.com/Kaelio/ktx/issues/339)) ([a0d19ba](a0d19ba26f))
* **sl:** correct reserved-word/week-grain SQL and classify sl_query errors ([#340](https://github.com/Kaelio/ktx/issues/340)) ([4ebce75](4ebce75449))
* **telemetry:** classify daemon query rejections as expected, not faults ([#335](https://github.com/Kaelio/ktx/issues/335)) ([5d17469](5d17469601))

### Documentation

* reflect recently added connectors in ingestion diagram and README ([#333](https://github.com/Kaelio/ktx/issues/333)) ([66768fe](66768fe009))
2026-07-03 21:22:57 +00:00
Andrey Avtomonov
4ebce75449
fix(sl): correct reserved-word/week-grain SQL and classify sl_query errors (#340)
Reserved-word columns (like, default, ...) referenced as source.col were
quoted with postgres double quotes even on BigQuery/MySQL, where a
double-quoted token is a string literal, not an identifier -- the
"Unexpected string literal" semantic-layer errors. quote_reserved_identifiers
now uses the identifier quote char of the dialect it will be parsed in
(backtick for BigQuery/MySQL), threaded through the planner and generator
parse sites; week_<weekday> granularity now emits WEEK(<weekday>) on BigQuery
instead of the invalid WEEK_MONDAY.

On the telemetry side, warehouse rejections from the sl_query execution path
are classified as expected (KtxQueryError) via a new shared markExpected()
helper, so routine agent/warehouse query failures stop reaching PostHog Error
Tracking as ktx faults; the sql_execution catch is refactored onto the same
helper. The daemon-compile boundary is deliberately left unclassified here so
genuine daemon crashes stay visible.
2026-07-03 23:19:33 +02:00
Andrey Avtomonov
a0d19ba26f
fix(sl): classify semantic-query request rejections as expected, not faults (#339)
The daemon rejects an invalid semantic-query request (unknown source,
ambiguous measure, no join path) with a plain ValueError; the Node compute
port now maps the daemon's exit code 3 / HTTP 400 to KtxExpectedError so these
routine, caller-driven rejections stay out of Error Tracking.

A dedicated SemanticLayerRequestError(ValueError) is raised only for engine
rejections and routed through both daemon transports and the HTTP handler.
Because pydantic v2 ValidationError subclasses ValueError, malformed sources or
responses (contract faults) are kept as faults: they are reported and mapped to
exit 1 / HTTP 500 / plain Error on every path. Non-object stdin is likewise a
fault (exit 1), not exit 3.
2026-07-03 23:17:33 +02:00
Andrey Avtomonov
5d17469601
fix(telemetry): classify daemon query rejections as expected, not faults (#335)
* fix(telemetry): classify daemon query rejections as expected, not faults

Semantic-layer query rejections and warehouse-execution rejections from the
sl_query MCP tool were wrapped as generic Errors, so reportException filed them
as PostHog $exception faults indistinguishable from real ktx bugs.

The daemon already separates a caller rejection (planner ValueError -> exit 3 /
HTTP 400) from a crash. The Node runner now carries that distinction as a typed
KtxDaemonComputeError, and a shared throwClassifiedQueryError promotes daemon
input-rejections and warehouse rejections to KtxQueryError while daemon crashes
and native JS faults still reach Error Tracking. query_semantic_layer stops
report_exception-ing expected ValueErrors, and a missing 'file:' secret now
raises KtxExpectedError so absent .ktx/secrets/<conn>-password stops filing
faults.

* chore: sync uv.lock to ktx-daemon/ktx-sl 0.15.0
2026-07-03 22:39:34 +02:00
Luca Martial
a651b82e2f
feat: query_policy semantic-layer-only restricts agents to predefined semantic-layer measures (#334)
* feat(sl): add predefined_measures_only guard to semantic query planning

SemanticQuery gains a predefined_measures_only flag; the planner rejects
any measure resolved with Provenance.COMPOSED (runtime aggregate
expressions and query-time derivations) while predefined measures,
predefined derived chains, dimensions, filters, and segments pass.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* feat(config): add per-connection query_policy to warehouse connections

query_policy: semantic-layer-only | read-only-sql (default) on the
warehouse connection schema, plus a policy module with the raw-SQL
guard, federated member restriction lookup, and the project-level
predicate used to gate sql_execution registration.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* feat(cli): enforce query_policy on raw SQL through one shared executor

ktx sql and the MCP sql_execution tool now share executeProjectRawSql
(resolve, policy check, read-only validation, execute), collapsing
their duplicated validate-then-execute paths. Restricted connections
are rejected before validation; federated raw SQL is rejected when any
member is restricted. sql_execution is not registered when every SQL
connection is restricted, and connection_list marks restricted
connections so agents route to sl_query. executeProjectReadOnlySql
stays generic for ktx-internal SQL (scan, ingest, SL-generated).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* feat(sl): compile queries with predefined_measures_only from query_policy

compileLocalSlQuery injects the flag from the connection's query_policy,
never from caller input, covering both ktx sl query and the MCP
sl_query tool through the daemon compile path.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* docs: document query_policy semantic-layer-only

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* fix(sl): close semantic-layer-only bypasses via filters and federated hint

The predefined_measures_only guard only inspected query.measures, so a
composed aggregate written into `filters` slipped through _classify_filters
into a HAVING clause untouched — letting a restricted agent evaluate
arbitrary aggregates (e.g. threshold-probing `sum(x) BETWEEN a AND b`).
Reject filter clauses that compose an aggregate function; a HAVING that
compares a predefined measure by name (`orders.revenue > 100`) still works.

Also make the federated sl_query error policy-aware: when a member is
restricted, raw federated SQL is disabled too, so stop directing the agent
to `ktx sql -c _ktx_federated` / sql_execution (a guaranteed failure) and
point to per-connection semantic-layer queries instead.

---------

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Andrey Avtomonov <andreybavt@gmail.com>
2026-07-03 08:54:17 +00:00
Luca Martial
66768fe009
docs: reflect recently added connectors in ingestion diagram and README (#333)
* docs(diagram): show recently added connectors in ingestion flow

Add Amazon Athena and MongoDB (plus a "& more" chip), Sigma, and Google
Drive to the ingestion diagram source cards so it reflects the current
connector set. Trim the Databases card body to one line so the added
chips fit the fixed-height card without clipping.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* docs(readme): list newly added connectors in the "Works with" summary

Include DuckDB, Amazon Athena, MongoDB, Sigma, and Google Drive so the
README matches the integrations reference pages.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-03 01:07:22 -04:00
Andrey Avtomonov
f310391da5
feat(athena): first-class AWS Athena warehouse identity (SL + BI mapping) (#332)
* feat: support Athena warehouse identity

* fix: support Athena and DuckDB BI table parsing
2026-07-02 15:10:29 +02:00
Patel Dhrit
fe7e6bd1fa
feat(connector): add Amazon Athena connector via Glue Data Catalog (#309)
* feat(connector): add Amazon Athena connector via Glue Data Catalog

* fix(athena): address reviewer feedback

* fix(athena): wire scope discovery, fix normalizeDriver, tighten types and tests

* fix(athena): honor databases scope, wire sql-analysis dialect, harden config resolution

- introspect() limits to the configured `databases` scope instead of scanning
  every Glue database in the account (docs promised this; connector ignored it)
- add athena -> athena to sql-analysis SQLGLOT_DIALECTS so `ktx sql` and MCP
  read-only validation parse Athena SQL under the Trino grammar, not postgres
- stringConfigValue coerces a resolved-empty `env:` reference to undefined so
  optional fields fall back to their defaults (workgroup 'primary', catalog
  'AwsDataCatalog') instead of ''
- drop trailing whitespace in dialect.test.ts

* fix(athena): integrate with main's SQL/non-SQL dialect split and add dialect notes

Rebase onto main, which introduced the KtxDialect (core) vs KtxSqlDialect
(SQL-only) split for MongoDB:
- KtxAthenaDialect implements KtxSqlDialect; the connector resolves it via
  getSqlDialectForDriver so SQL-generation methods stay in scope
- add authored athena.md SQL notes for the sql_dialect_notes MCP tool, required
  now that athena resolves to the athena sqlglot dialect (dialect-notes coverage
  is derived from the warehouse-driver registry)

---------

Co-authored-by: Andrey Avtomonov <andreybavt@gmail.com>
2026-07-02 15:00:26 +02:00
Andrey Avtomonov
6d01030745
fix(deps): patch 22 Dependabot security alerts (#328)
Bump transitive dependencies to their patched versions to clear all open
Dependabot advisories. npm fixes go through the pnpm-workspace.yaml
overrides block; the Python fix goes through uv.lock.

npm: undici 6.27.0/7.28.0, hono 4.12.25, form-data 4.0.6, ws 8.21.0,
vite 8.0.16, esbuild 0.28.1, js-yaml 4.2.0.
pip: starlette 1.3.1.
2026-07-02 09:24:18 +00:00
Kevin Messiaen
3c4fcc27c7
feat: Add duckdb connector (#308)
* refactor(duckdb): extract shared json-safe bigint helper

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* feat(duckdb): add and register the duckdb primary connector

Add KtxDuckDbDialect, KtxDuckDbScanConnector (local file-backed, read-only,
never-create, main-schema introspection via information_schema and
duckdb_constraints() for foreign keys), and register the duckdb driver across
the dialect factory, driver registry, connection-type enum, warehouse descriptor,
config schema, scan normalization, connection test drivers, and status display.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* feat(duckdb): route live-database ingest through the DuckDB connector

Add the DuckDB live-database introspection bridge and dispatch duckdb
connections to it in local-adapters, matching the SQLite path. Repoint the
config-rejection test off duckdb (now a valid driver) onto the no-driver case.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* feat(duckdb): add duckdb to the setup database flow

Offer DuckDB in the interactive checklist and via ktx setup --database duckdb,
with a file-path prompt and duckdb-local default connection id, parallel to SQLite.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* feat(duckdb): attach native duckdb files in federation

Native .duckdb members ATTACH with (READ_ONLY) and no TYPE/INSTALL/LOAD, since
the duckdb format needs no extension. attachTypeForDriver returns null for the
native case; buildAttachStatements builds load statements from non-null types
only and emits a conditional ATTACH clause.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* docs(duckdb): document the duckdb primary-source connector

Add a DuckDB section to the primary-sources integration page (config, read-only
never-create behavior, main-schema scope, federation) and update the
supported-driver assertion in dialects.test.ts to include duckdb.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* fix(duckdb): use single-namespace display shape for main-only refs

DuckDB v1 introspects the main schema and sets db=null on every table, so its
display refs are single-namespace like SQLite. The ansi shape emitted a 1-part
table display it then refused to parse, breaking column-level display resolution.
Switch the dialect to the sqlite display shape and add a round-trip test plus a
composite-foreign-key test that were missing.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* refactor(duckdb): resolve connector dialect via getDialectForDriver

Route the connector's dialect through the shared factory like every other
connector, now that duckdb is registered. Single construction path.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* fix(duckdb): skip schema picker for single-file duckdb setup

DuckDB is a single-file, single-namespace ('main') database like SQLite,
but the setup scope step only skipped the schema picker for sqlite. DuckDB
fell into the multi-schema path with an empty schema list, rendering a
broken picker ("No matches found" for main). Extend the file-based-driver
early-return to cover duckdb so it ingests every table directly.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* refactor(duckdb): reuse shared config helper and derive scope skip

Route duckdb path resolution through the shared resolveStringReference
helper instead of a local third copy of env:/file: handling. Derive the
setup scope-picker skip from SCOPE_DISCOVERY_SPECS membership rather than
a hardcoded sqlite/duckdb driver list.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* test(duckdb): use a genuinely unknown driver in the rejection test

The merged "rejects unknown drivers" test used `driver: duckdb` as its
unknown-driver stand-in, which stopped being unknown once this branch
added the duckdb connector. Switch to `nonsense` so it again exercises
the unsupported-driver config error.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* test(duckdb): cover dialect, connector, and live-introspection branches

Codecov flagged uncovered branches as dead code; all are real connector,
dialect, and live-ingest behavior. Add unit tests instead of removing them.

- dialect: precedence ladder, sample/clause builders, profiling expressions
- connector: url/env config forms, error throws, never-create guard,
  cardinality cap branches, table-scope empty/non-empty paths
- live-introspection: full-schema and table-scope extraction

Functions 100%, lines ~99% across the duckdb connector dir.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* docs: add DuckDB to supported-driver references

The DuckDB connector PR documented the connector itself but left the
scattered supported-driver enumerations stale. Add duckdb to the
federation concept page (participation table, activation, table naming,
limitations), the ktx setup CLI reference, the ktx.yaml warehouse-driver
table, the primary-sources field reference, and the quickstart driver
list (which also restores the missing ClickHouse entry).

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Andrey Avtomonov <andreybavt@gmail.com>
2026-07-01 12:06:02 +00:00
102 changed files with 5365 additions and 479 deletions

View file

@ -103,8 +103,9 @@ upkeep and don't absorb the rest of your company's knowledge.
- You don't have a SQL warehouse - **ktx** sits on top of one - You don't have a SQL warehouse - **ktx** sits on top of one
- You only need one ad-hoc query - `psql` or a notebook will do - You only need one ad-hoc query - `psql` or a notebook will do
Works with PostgreSQL, Snowflake, BigQuery, ClickHouse, MySQL, SQL Server, and Works with PostgreSQL, Snowflake, BigQuery, ClickHouse, MySQL, SQL Server,
SQLite. Integrates with dbt, MetricFlow, LookML, Looker, Metabase, and Notion. SQLite, DuckDB, Amazon Athena, and MongoDB. Integrates with dbt, MetricFlow,
LookML, Looker, Metabase, Sigma, Notion, and Google Drive.
## Quick Start ## Quick Start

View file

@ -68,14 +68,14 @@ const EDGE_STROKE = "#94a3b8";
const sourceData: SourceNodeData[] = [ const sourceData: SourceNodeData[] = [
{ {
title: "Databases", title: "Databases",
body: "Schemas, columns, keys, row counts, and query history.", body: "Schemas, keys, row counts, query history.",
items: ["PostgreSQL", "Snowflake", "BigQuery", "SQLite"], items: ["PostgreSQL", "Snowflake", "BigQuery", "Athena", "MongoDB", "& more"],
accent: "#3b82f6", accent: "#3b82f6",
}, },
{ {
title: "BI tools", title: "BI tools",
body: "Dashboards, questions, explores, usage, and trusted examples.", body: "Dashboards, questions, explores, usage, and trusted examples.",
items: ["Metabase", "Looker"], items: ["Metabase", "Looker", "Sigma"],
accent: "#f97316", accent: "#f97316",
}, },
{ {
@ -87,7 +87,7 @@ const sourceData: SourceNodeData[] = [
{ {
title: "Docs and notes", title: "Docs and notes",
body: "Policies, caveats, team definitions, and analyst context.", body: "Policies, caveats, team definitions, and analyst context.",
items: ["Notion", "Any text"], items: ["Notion", "Google Drive", "Any text"],
accent: "#10b981", accent: "#10b981",
}, },
]; ];

View file

@ -120,9 +120,9 @@ runtime features are missing.
| Flag | Description | | Flag | Description |
|------|-------------| |------|-------------|
| `--database <driver>` | Database driver to configure; repeatable. Choices: `sqlite`, `postgres`, `mysql`, `clickhouse`, `sqlserver`, `bigquery`, `snowflake` | | `--database <driver>` | Database driver to configure; repeatable. Choices: `sqlite`, `duckdb`, `postgres`, `mysql`, `clickhouse`, `sqlserver`, `bigquery`, `snowflake` |
| `--database-connection-id <id>` | Existing selected connection id; repeatable. With `--database` or `--database-url`, connection id for the new connection. | | `--database-connection-id <id>` | Existing selected connection id; repeatable. With `--database` or `--database-url`, connection id for the new connection. |
| `--database-url <url>` | URL, `env:NAME`, or `file:/path` for one new URL-style database connection; also used as the SQLite path | | `--database-url <url>` | URL, `env:NAME`, or `file:/path` for one new URL-style database connection; also used as the SQLite or DuckDB path |
| `--database-schema <schema>` | Database schema or dataset to include; repeatable | | `--database-schema <schema>` | Database schema or dataset to include; repeatable |
| `--skip-databases` | Leave database setup incomplete | | `--skip-databases` | Leave database setup incomplete |

View file

@ -1,6 +1,6 @@
--- ---
title: Cross-database federation title: Cross-database federation
description: How ktx federates postgres, mysql, and sqlite connections so a single read-only SQL query can join across them without copying data. description: How ktx federates postgres, mysql, sqlite, and duckdb connections so a single read-only SQL query can join across them without copying data.
--- ---
Cross-database federation lets a single read-only SQL query join tables that Cross-database federation lets a single read-only SQL query join tables that
@ -20,13 +20,14 @@ block to add. With zero or one compatible connection the behavior is unchanged.
## Which connections participate ## Which connections participate
The v1 federation engine supports three drivers: The v1 federation engine supports four drivers:
| Driver | Participates in federation | | Driver | Participates in federation |
|--------|---------------------------| |--------|---------------------------|
| `postgres` | Yes | | `postgres` | Yes |
| `mysql` | Yes | | `mysql` | Yes |
| `sqlite` | Yes | | `sqlite` | Yes |
| `duckdb` | Yes |
| `snowflake` | No — standalone connection | | `snowflake` | No — standalone connection |
| `bigquery` | No — standalone connection | | `bigquery` | No — standalone connection |
| `clickhouse` | No — standalone connection | | `clickhouse` | No — standalone connection |
@ -38,7 +39,7 @@ queried independently; they do not appear as federation members.
## How it activates ## How it activates
**ktx** inspects the connections in `ktx.yaml` at startup. When it finds two or **ktx** inspects the connections in `ktx.yaml` at startup. When it finds two or
more connections whose driver is `postgres`, `mysql`, or `sqlite`, it more connections whose driver is `postgres`, `mysql`, `sqlite`, or `duckdb`, it
instantiates the DuckDB federation engine and attaches each one read-only. instantiates the DuckDB federation engine and attaches each one read-only.
There is no `federation:` key, no opt-in flag, and no connection-level setting There is no `federation:` key, no opt-in flag, and no connection-level setting
to enable. The engine is derived entirely from what is already declared. to enable. The engine is derived entirely from what is already declared.
@ -60,9 +61,10 @@ Two attach-compatible connections are present, so federation is active.
## Table naming in federated queries ## Table naming in federated queries
Inside a federated query, postgres and mysql tables use a three-part name: Inside a federated query, postgres and mysql tables use a three-part name:
`connectionId.schema.table`. SQLite tables, which have no schema layer in `connectionId.schema.table`. SQLite and DuckDB tables use the two-part form
DuckDB, use the two-part form `connectionId.table`. In both cases the `connectionId.table`, since ktx addresses both as single-namespace members. In
connection's `id` field in `ktx.yaml` becomes the catalog name inside DuckDB. both cases the connection's `id` field in `ktx.yaml` becomes the catalog name
inside DuckDB.
If a connection `id` is not a bare SQL identifier — for example it contains a If a connection `id` is not a bare SQL identifier — for example it contains a
hyphen, like `books-db` — double-quote it in the query the same way DuckDB hyphen, like `books-db` — double-quote it in the query the same way DuckDB
@ -131,12 +133,17 @@ ktx sql -c _ktx_federated \
Table names follow the rules from Table names follow the rules from
[Table naming in federated queries](#table-naming-in-federated-queries): [Table naming in federated queries](#table-naming-in-federated-queries):
three-part `connectionId.schema.table` for postgres and mysql, two-part three-part `connectionId.schema.table` for postgres and mysql, two-part
`connectionId.table` for sqlite. The `_ktx_federated` id is virtual — it is `connectionId.table` for sqlite and duckdb. The `_ktx_federated` id is virtual —
never written to `ktx.yaml` and only exists when two or more attach-compatible it is never written to `ktx.yaml` and only exists when two or more attach-compatible
connections are declared. It surfaces in `ktx connection` and in the agent's connections are declared. It surfaces in `ktx connection` and in the agent's
connection list so the id is discoverable. Querying a single member database connection list so the id is discoverable. Querying a single member database
directly with its own connection id (`ktx sql -c pg_books ...`) is unchanged. directly with its own connection id (`ktx sql -c pg_books ...`) is unchanged.
If any member connection sets
[`query_policy: semantic-layer-only`](/docs/configuration/ktx-yaml#query-policy),
raw SQL against `_ktx_federated` is rejected as a whole: a federated query can
touch any member's tables, so one restricted member restricts the federation.
## Federated queries are read-only ## Federated queries are read-only
DuckDB attaches every member database with read-only access. Federated queries DuckDB attaches every member database with read-only access. Federated queries
@ -149,6 +156,6 @@ database through the federation engine.
them in a source's `joins:` block and automatic discovery of cross-database them in a source's `joins:` block and automatic discovery of cross-database
relationships are not available yet. Intra-database relationship discovery for relationships are not available yet. Intra-database relationship discovery for
each member connection is unchanged. each member connection is unchanged.
- **postgres, mysql, and sqlite only.** Other drivers (snowflake, bigquery, - **postgres, mysql, sqlite, and duckdb only.** Other drivers (snowflake,
clickhouse, sqlserver) do not participate in federation in this version. They bigquery, clickhouse, sqlserver) do not participate in federation in this
remain usable as standalone connections. version. They remain usable as standalone connections.

View file

@ -109,6 +109,7 @@ context-source drivers share the map.
| `postgres` | Warehouse | `driver` | `url`, `enabled_tables`, `historicSql`, `context.queryHistory` | | `postgres` | Warehouse | `driver` | `url`, `enabled_tables`, `historicSql`, `context.queryHistory` |
| `mysql` | Warehouse | `driver` | `url`, `enabled_tables` | | `mysql` | Warehouse | `driver` | `url`, `enabled_tables` |
| `sqlite` | Warehouse | `driver` | `url` or `path`, `enabled_tables` | | `sqlite` | Warehouse | `driver` | `url` or `path`, `enabled_tables` |
| `duckdb` | Warehouse | `driver` | `url` or `path`, `enabled_tables` |
| `sqlserver` | Warehouse | `driver` | `url`, `enabled_tables` | | `sqlserver` | Warehouse | `driver` | `url`, `enabled_tables` |
| `bigquery` | Warehouse | `driver` | `credentials_json`, `dataset_ids`, `enabled_tables`, `historicSql` | | `bigquery` | Warehouse | `driver` | `credentials_json`, `dataset_ids`, `enabled_tables`, `historicSql` |
| `snowflake` | Warehouse | `driver` | `schema_names`, `enabled_tables`, `historicSql` | | `snowflake` | Warehouse | `driver` | `schema_names`, `enabled_tables`, `historicSql` |
@ -216,6 +217,41 @@ connections:
observed in-scope query history. The block uses `mode: exclude` and remains observed in-scope query history. The block uses `mode: exclude` and remains
hand-editable. hand-editable.
### Query policy
Set `query_policy: semantic-layer-only` on a warehouse connection to stop
agents from authoring SQL against it. The default, `read-only-sql`, allows
parser-validated read-only SQL through `ktx sql` and the `sql_execution` MCP
tool alongside semantic-layer queries.
```yaml
connections:
warehouse:
driver: snowflake
query_policy: semantic-layer-only
```
With `semantic-layer-only`:
- `ktx sql` and the `sql_execution` MCP tool reject the connection with a
clear error. When every SQL connection in the project is restricted, the
`sql_execution` tool is not registered at all.
- Raw SQL against the federated connection (`_ktx_federated`) is rejected
when any member connection is restricted.
- Semantic-layer queries (`ktx sl query`, the `sl_query` tool) accept only
measures predefined in the semantic-layer sources. Composed aggregate
expressions such as `sum(orders.amount)` are rejected wherever they appear,
including inside `filters` (a `HAVING`-style clause may only compare a
predefined measure by name, e.g. `orders.revenue > 100`). Grouping by
declared dimensions, filtering on columns, and segments remain available.
- `connection_list` marks the connection as restricted so agents route to
`sl_query` instead of burning a failed call.
The policy governs agent-facing query authorship, not data access: **ktx**'s
own scan, ingest, and semantic-layer-generated SQL still run, and context
tools such as `entity_details` and `dictionary_search` still expose schema
metadata and sampled values.
### Metabase ### Metabase
```yaml ```yaml

View file

@ -218,7 +218,8 @@ The wizard walks you through everything **ktx** needs in one pass:
3. **Embeddings** - picks an embeddings backend. Choose OpenAI for hosted 3. **Embeddings** - picks an embeddings backend. Choose OpenAI for hosted
embeddings or `sentence-transformers` to run locally without an API key. embeddings or `sentence-transformers` to run locally without an API key.
4. **Database** - adds at least one primary connection. Supported drivers: 4. **Database** - adds at least one primary connection. Supported drivers:
SQLite, PostgreSQL, MySQL, SQL Server, BigQuery, and Snowflake. PostgreSQL, Snowflake, BigQuery, MySQL, ClickHouse, SQL Server, SQLite, and
DuckDB.
5. **Context sources** - optionally adds dbt, MetricFlow, LookML, Looker, 5. **Context sources** - optionally adds dbt, MetricFlow, LookML, Looker,
Metabase, or Notion. You can skip and add them later. Metabase, or Notion. You can skip and add them later.
6. **Build** - offers to run the first ingest so semantic sources and wiki 6. **Build** - offers to run the first ingest so semantic sources and wiki

View file

@ -1,6 +1,6 @@
--- ---
title: Primary Sources title: Primary Sources
description: Connect ktx to PostgreSQL, Snowflake, BigQuery, MySQL, ClickHouse, SQL Server, SQLite, or MongoDB. description: Connect ktx to PostgreSQL, Snowflake, BigQuery, MySQL, ClickHouse, SQL Server, SQLite, DuckDB, MongoDB, or Amazon Athena.
--- ---
**ktx** connects to your data warehouse or database to build schema context, **ktx** connects to your data warehouse or database to build schema context,
@ -26,17 +26,23 @@ Agents should prefer environment or file references over literal secrets.
| Field | Required | Applies to | Description | | Field | Required | Applies to | Description |
|-------|----------|------------|-------------| |-------|----------|------------|-------------|
| `driver` | Yes | all connections | Connector driver such as `postgres`, `snowflake`, `bigquery`, `mysql`, `clickhouse`, `sqlserver`, `sqlite`, or `mongodb` | | `driver` | Yes | all connections | Connector driver such as `postgres`, `snowflake`, `bigquery`, `mysql`, `clickhouse`, `sqlserver`, `sqlite`, `duckdb`, `mongodb`, or `athena` |
| `url` | One of the connection methods | URL-style connectors | Database URL, `env:NAME`, or `file:/path/to/secret` | | `url` | One of the connection methods | URL-style connectors | Database URL, `env:NAME`, or `file:/path/to/secret` |
| `host`, `port`, `database`, `username`, `password` | One of the connection methods | PostgreSQL, MySQL, SQL Server | Field-by-field connection values | | `host`, `port`, `database`, `username`, `password` | One of the connection methods | PostgreSQL, MySQL, SQL Server | Field-by-field connection values |
| `schema` or `schemas` | No | schema-aware warehouses | Single schema or list of schemas to scan | | `schema` or `schemas` | No | schema-aware warehouses | Single schema or list of schemas to scan |
| `databases` | No | ClickHouse, MongoDB | List of databases to scan | | `databases` | No | ClickHouse, MongoDB, Athena | List of databases to scan |
| `sample_size`, `order_by` | No | MongoDB | Schema-inference sampling controls (recent documents, sort field) | | `sample_size`, `order_by` | No | MongoDB | Schema-inference sampling controls (recent documents, sort field) |
| `context.queryHistory` | No | PostgreSQL, Snowflake, BigQuery | Enables query-history ingestion when the warehouse supports it | | `context.queryHistory` | No | PostgreSQL, Snowflake, BigQuery | Enables query-history ingestion when the warehouse supports it |
| `path` | Yes for path-style SQLite | SQLite | Local SQLite database path or `env:NAME` reference | | `path` | Yes for path-style SQLite/DuckDB | SQLite, DuckDB | Local SQLite or DuckDB database path or `env:NAME` reference |
| `max_bytes_billed` | No | BigQuery | Maximum bytes billed per query job | | `max_bytes_billed` | No | BigQuery | Maximum bytes billed per query job |
| `query_timeout_ms` | No | all warehouses | Maximum execution time for a single read-only query, in milliseconds (default 30000). A query exceeding it is cancelled server-side (or, for SQLite, by terminating the off-process executor) and returns a `query exceeded Ns` error so the agent can revise. | | `query_timeout_ms` | No | all warehouses | Maximum execution time for a single read-only query, in milliseconds (default 30000). A query exceeding it is cancelled server-side (or, for SQLite, by terminating the off-process executor) and returns a `query exceeded Ns` error so the agent can revise. |
| `project_id` | No | BigQuery | Optional local descriptor and mapping metadata; not used for BigQuery authentication | | `project_id` | No | BigQuery | Optional local descriptor and mapping metadata; not used for BigQuery authentication |
| `region` | Yes | Athena | AWS region where the Athena workgroup and Glue catalog reside (e.g. `us-east-1`) |
| `s3_staging_dir` | Yes | Athena | S3 URI for Athena query result storage (e.g. `s3://my-bucket/athena-results/`) |
| `workgroup` | No | Athena | Athena workgroup name (default `primary`) |
| `catalog` | No | Athena | Glue Data Catalog name (default `AwsDataCatalog`) |
| `database` | No | Athena | Default Glue database name passed as the query execution context |
| `databases` | No | Athena | Glue databases to include in schema scans; written by `ktx setup` and read by `ktx ingest` |
## PostgreSQL ## PostgreSQL
@ -304,6 +310,76 @@ staged artifact shape as Postgres and Snowflake.
--- ---
## Amazon Athena
Connects to Amazon Athena using the AWS Glue Data Catalog for schema introspection and the Athena query API for read-only SQL execution. Authentication uses the standard AWS credential chain — no credentials are embedded in `ktx.yaml`.
### Connection config
```yaml title="ktx.yaml"
connections:
my-athena:
driver: athena
region: us-east-1
s3_staging_dir: s3://my-bucket/athena-results/
```
With optional fields:
```yaml title="ktx.yaml"
connections:
my-athena:
driver: athena
region: us-east-1
s3_staging_dir: env:ATHENA_S3_STAGING_DIR
workgroup: analytics
catalog: AwsDataCatalog
database: my_default_database
databases:
- analytics
- raw
```
`ktx setup` writes the `databases` array when you select Glue databases during setup. `ktx scan` reads it to limit introspection to those databases.
### Authentication
**ktx** uses the AWS SDK default credential chain — no credentials appear in `ktx.yaml`. The chain resolves credentials in this order:
| Method | How to configure |
|--------|-----------------|
| Environment variables | Set `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, and optionally `AWS_SESSION_TOKEN` |
| Shared credentials file | Configure `~/.aws/credentials` with a `[default]` or named profile; set `AWS_PROFILE` to select a non-default profile |
| IAM instance profile | Attach an IAM role to the EC2 instance or ECS task — no local configuration needed |
| IAM roles for service accounts (EKS) | Annotate the pod's service account with the IAM role ARN |
The IAM principal must have `athena:StartQueryExecution`, `athena:GetQueryExecution`, `athena:GetQueryResults`, `glue:GetDatabases`, and `glue:GetTables` permissions, plus read access to the S3 results bucket.
### Features
| Feature | Supported | Notes |
|---------|-----------|-------|
| Tables & views | Yes | Via AWS Glue Data Catalog |
| Primary keys | No | Glue does not expose constraint metadata |
| Foreign keys | No | Not available in Glue/Athena |
| Row count estimates | No | Glue table statistics are often stale |
| Column statistics | No | - |
| Query history | No | - |
| Table sampling | Yes | `SELECT ... LIMIT n` |
### Dialect notes
- SQL dialect is Presto/Trino; identifiers are quoted with double-quotes
- Table names use three-part format: `"catalog"."database"."table"` (e.g. `"AwsDataCatalog"."analytics"."orders"`)
- Partition columns (`PartitionKeys` in Glue) are included after regular columns in the schema and are fully queryable
- Athena does not support `TABLESAMPLE`; random sampling uses `ORDER BY rand()`
- Query execution is asynchronous: **ktx** starts the query, polls until completion, then fetches results from S3
- Results are stored in `s3_staging_dir`; the IAM principal needs write access to that bucket
- Use `workgroup` to apply per-workgroup cost controls and result configuration
- The connector always uses your account's default Glue Data Catalog; cross-account catalog access (`CatalogId` pointing to another account) is not supported
---
## MySQL ## MySQL
Standard MySQL/MariaDB connector with full foreign key support and schema introspection. Standard MySQL/MariaDB connector with full foreign key support and schema introspection.
@ -545,6 +621,52 @@ No authentication required - SQLite is file-based. The file must be readable by
--- ---
## DuckDB
File-based connector using the DuckDB Node.js API. Ideal for local analytics, embedded warehouses, and cross-database federation.
### Connection config
```yaml title="ktx.yaml"
connections:
warehouse:
driver: duckdb
path: data/warehouse.duckdb
```
`path` is resolved relative to the project directory. The `.duckdb` file must already exist — **ktx** never creates a missing database file.
### Authentication
No authentication required — DuckDB is file-based. The `.duckdb` file must be readable by the process running **ktx**.
### Features
| Feature | Supported | Notes |
|---------|-----------|-------|
| Tables & views | Yes | Via `information_schema` on the `main` schema |
| Primary keys | Yes | Via `information_schema.table_constraints` |
| Foreign keys | Yes | Via DuckDB's `duckdb_constraints()` catalog function |
| Row count estimates | Yes | Exact count via `SELECT COUNT(*)` |
| Column statistics | No | - |
| Query history | No | - |
| Table sampling | Yes | - |
| Nested analysis | No | - |
### Dialect notes
- Introspection scans the `main` schema only
- Execution is read-only; **ktx** opens the file without write access
- Parameter binding uses positional `?` placeholders
- Uses `LIMIT X OFFSET Y` for pagination
- Database file must exist before `ktx connection test` or ingest runs
### Cross-database federation
When a project declares two or more attach-compatible connections — any combination of `postgres`, `mysql`, `sqlite`, and `duckdb` — **ktx** derives a cross-database federation connection. That connection can ATTACH a native `.duckdb` file, allowing semantic queries to join across sources without manually copying data.
---
## MongoDB ## MongoDB
Connects to MongoDB as a primary context source. **ktx** treats each collection Connects to MongoDB as a primary context source. **ktx** treats each collection
@ -629,7 +751,9 @@ nullability from how often the field is present:
| Error or symptom | Likely cause | Recovery | | Error or symptom | Likely cause | Recovery |
|------------------|--------------|----------| |------------------|--------------|----------|
| Connection URL appears in git diff | A literal credential URL was written to `ktx.yaml` | Replace it with `env:NAME` or `file:/path/to/secret` and rotate exposed credentials | | Connection URL appears in git diff | A literal credential URL was written to `ktx.yaml` | Replace it with `env:NAME` or `file:/path/to/secret` and rotate exposed credentials |
| Database ingest returns no tables | Schema, database, or project filter is wrong, or the user lacks metadata permissions | Verify the schema list and grant metadata read permissions | | Database ingest returns no tables | Schema, database, or project filter is wrong, or the user lacks metadata permissions | Verify the schema list and grant metadata read permissions. For Athena, confirm the IAM principal has `glue:GetDatabases` and `glue:GetTables` permissions |
| Query history is empty | Query history extension or warehouse history view is unavailable | Enable the warehouse-specific history feature, then rerun `ktx ingest <connectionId> --query-history` or `ktx setup` | | Query history is empty | Query history extension or warehouse history view is unavailable | Enable the warehouse-specific history feature, then rerun `ktx ingest <connectionId> --query-history` or `ktx setup` |
| Column statistics are missing | Connector cannot access stats tables or the warehouse does not expose them | Grant stats permissions where supported; otherwise rely on schema-level context without column statistics | | Column statistics are missing | Connector cannot access stats tables or the warehouse does not expose them | Grant stats permissions where supported; otherwise rely on schema-level context without column statistics |
| Semantic query execution fails | Connection is missing, unreachable, or query execution is disabled | Run `ktx connection test <id>` and check the `ktx sl query` flags | | Semantic query execution fails | Connection is missing, unreachable, or query execution is disabled | Run `ktx connection test <id>` and check the `ktx sl query` flags |
| Athena query fails with `ACCESS_DENIED` | IAM principal lacks `athena:StartQueryExecution` or S3 write access to `s3_staging_dir` | Attach a policy granting Athena query permissions and `s3:PutObject` on the staging bucket |
| Athena ingest finds databases but no tables | IAM principal has `glue:GetDatabases` but not `glue:GetTables` | Grant `glue:GetTables` on the relevant Glue catalog resources |

View file

@ -1,6 +1,6 @@
{ {
"name": "ktx-workspace", "name": "ktx-workspace",
"version": "0.15.0", "version": "0.16.0",
"description": "Workspace root for ktx packages", "description": "Workspace root for ktx packages",
"private": true, "private": true,
"type": "module", "type": "module",

View file

@ -1,6 +1,6 @@
{ {
"name": "@kaelio/ktx", "name": "@kaelio/ktx",
"version": "0.15.0", "version": "0.16.0",
"description": "Standalone ktx context layer for data agents", "description": "Standalone ktx context layer for data agents",
"author": { "author": {
"name": "Kaelio", "name": "Kaelio",
@ -51,6 +51,8 @@
"@ai-sdk/devtools": "0.0.18", "@ai-sdk/devtools": "0.0.18",
"@ai-sdk/google-vertex": "^4.0.134", "@ai-sdk/google-vertex": "^4.0.134",
"@anthropic-ai/claude-agent-sdk": "0.3.146", "@anthropic-ai/claude-agent-sdk": "0.3.146",
"@aws-sdk/client-athena": "^3.1068.0",
"@aws-sdk/client-glue": "^3.1068.0",
"@clack/core": "1.3.1", "@clack/core": "1.3.1",
"@clack/prompts": "1.4.0", "@clack/prompts": "1.4.0",
"@clickhouse/client": "^1.18.5", "@clickhouse/client": "^1.18.5",

View file

@ -38,6 +38,7 @@ function llmBackend(value: string): KtxSetupLlmBackend {
function databaseDriver(value: string): KtxSetupDatabaseDriver { function databaseDriver(value: string): KtxSetupDatabaseDriver {
if ( if (
value === 'sqlite' || value === 'sqlite' ||
value === 'duckdb' ||
value === 'postgres' || value === 'postgres' ||
value === 'mysql' || value === 'mysql' ||
value === 'clickhouse' || value === 'clickhouse' ||

View file

@ -3,12 +3,14 @@ import type { KtxProjectConnectionConfig } from './context/project/config.js';
/** @internal Canonical SQL-warehouse driver ids; the dialect-notes coverage test derives its required coverage from this set. */ /** @internal Canonical SQL-warehouse driver ids; the dialect-notes coverage test derives its required coverage from this set. */
export const KTX_DATABASE_DRIVER_IDS = [ export const KTX_DATABASE_DRIVER_IDS = [
'sqlite', 'sqlite',
'duckdb',
'postgres', 'postgres',
'mysql', 'mysql',
'clickhouse', 'clickhouse',
'sqlserver', 'sqlserver',
'bigquery', 'bigquery',
'snowflake', 'snowflake',
'athena',
] as const; ] as const;
// mongodb is a database driver but has no SQL dialect, so it sits outside the // mongodb is a database driver but has no SQL dialect, so it sits outside the

View file

@ -51,6 +51,7 @@ export interface KtxConnectionDeps {
const SUPPORTED_TEST_DRIVERS = [ const SUPPORTED_TEST_DRIVERS = [
'sqlite', 'sqlite',
'duckdb',
'postgres', 'postgres',
'mysql', 'mysql',
'clickhouse', 'clickhouse',

View file

@ -0,0 +1,555 @@
import { AthenaClient, StartQueryExecutionCommand, GetQueryExecutionCommand, GetQueryResultsCommand } from '@aws-sdk/client-athena';
import { GlueClient, GetDatabasesCommand, GetTablesCommand } from '@aws-sdk/client-glue';
import { getSqlDialectForDriver } from '../../context/connections/dialects.js';
import { assertReadOnlySql, limitSqlForExecution } from '../../context/connections/read-only-sql.js';
import {
connectorTestFailure,
createKtxConnectorCapabilities,
type KtxConnectorTestResult,
type KtxColumnSampleInput,
type KtxColumnSampleResult,
type KtxColumnStatsInput,
type KtxColumnStatsResult,
type KtxQueryResult,
type KtxReadOnlyQueryInput,
type KtxScanConnector,
type KtxScanContext,
type KtxScanInput,
type KtxSchemaColumn,
type KtxSchemaSnapshot,
type KtxSchemaTable,
type KtxTableListEntry,
type KtxTableRef,
type KtxTableSampleInput,
type KtxTableSampleResult,
} from '../../context/scan/types.js';
import { scopedTableNames } from '../../context/scan/table-ref.js';
import { resolveStringReference } from '../shared/string-reference.js';
export interface KtxAthenaConnectionConfig {
driver?: string;
region?: string;
s3_staging_dir?: string;
workgroup?: string;
catalog?: string;
database?: string;
databases?: string[];
[key: string]: unknown;
}
export interface KtxAthenaResolvedConnectionConfig {
region: string;
s3StagingDir: string;
workgroup: string;
catalog: string;
database: string | undefined;
databases: string[];
}
interface KtxAthenaQueryExecutionStatus {
State?: string;
StateChangeReason?: string;
}
interface KtxAthenaQueryExecution {
Status?: KtxAthenaQueryExecutionStatus;
}
interface KtxAthenaColumnInfo {
Name?: string;
Type?: string;
}
interface KtxAthenaDatum {
VarCharValue?: string;
}
interface KtxAthenaRow {
Data?: KtxAthenaDatum[];
}
interface KtxAthenaResultSet {
Rows?: KtxAthenaRow[];
ResultSetMetadata?: { ColumnInfo?: KtxAthenaColumnInfo[] };
}
/** @internal */
export interface KtxAthenaClient {
startQueryExecution(input: {
QueryString: string;
ResultConfiguration: { OutputLocation: string };
WorkGroup: string;
QueryExecutionContext?: { Database?: string; Catalog?: string };
}): Promise<{ QueryExecutionId?: string }>;
getQueryExecution(input: { QueryExecutionId: string }): Promise<{ QueryExecution?: KtxAthenaQueryExecution }>;
getQueryResults(input: { QueryExecutionId: string; NextToken?: string }): Promise<{
ResultSet?: KtxAthenaResultSet;
NextToken?: string;
}>;
}
interface KtxGlueColumnDef {
Name?: string;
Type?: string;
Comment?: string;
}
interface KtxGlueStorageDescriptor {
Columns?: KtxGlueColumnDef[];
}
/** @internal */
export interface KtxGlueTable {
Name?: string;
TableType?: string;
StorageDescriptor?: KtxGlueStorageDescriptor;
PartitionKeys?: KtxGlueColumnDef[];
Description?: string;
Parameters?: Record<string, string>;
}
/** @internal */
export interface KtxGlueClient {
getDatabases(input: { CatalogId?: string; NextToken?: string }): Promise<{
DatabaseList?: Array<{ Name?: string }>;
NextToken?: string;
}>;
getTables(input: { DatabaseName: string; CatalogId?: string; NextToken?: string }): Promise<{
TableList?: KtxGlueTable[];
NextToken?: string;
}>;
}
export interface KtxAthenaClientFactory {
createAthenaClient(region: string): KtxAthenaClient;
createGlueClient(region: string): KtxGlueClient;
}
class DefaultAthenaClientFactory implements KtxAthenaClientFactory {
createAthenaClient(region: string): KtxAthenaClient {
const client = new AthenaClient({ region });
return {
startQueryExecution: async (input) => {
const result = await client.send(
new StartQueryExecutionCommand({
QueryString: input.QueryString,
ResultConfiguration: { OutputLocation: input.ResultConfiguration.OutputLocation },
WorkGroup: input.WorkGroup,
QueryExecutionContext: input.QueryExecutionContext,
}),
);
return { QueryExecutionId: result.QueryExecutionId };
},
getQueryExecution: async (input) => {
const result = await client.send(new GetQueryExecutionCommand({ QueryExecutionId: input.QueryExecutionId }));
return {
QueryExecution: result.QueryExecution
? {
Status: {
State: result.QueryExecution.Status?.State,
StateChangeReason: result.QueryExecution.Status?.StateChangeReason,
},
}
: undefined,
};
},
getQueryResults: async (input) => {
const result = await client.send(
new GetQueryResultsCommand({ QueryExecutionId: input.QueryExecutionId, NextToken: input.NextToken }),
);
return {
ResultSet: result.ResultSet as KtxAthenaResultSet | undefined,
NextToken: result.NextToken,
};
},
};
}
createGlueClient(region: string): KtxGlueClient {
const client = new GlueClient({ region });
return {
getDatabases: async (input) => {
const result = await client.send(new GetDatabasesCommand({ CatalogId: input.CatalogId, NextToken: input.NextToken }));
return {
DatabaseList: result.DatabaseList?.map((db) => ({ Name: db.Name })),
NextToken: result.NextToken,
};
},
getTables: async (input) => {
const result = await client.send(
new GetTablesCommand({ DatabaseName: input.DatabaseName, CatalogId: input.CatalogId, NextToken: input.NextToken }),
);
return {
TableList: result.TableList as KtxGlueTable[] | undefined,
NextToken: result.NextToken,
};
},
};
}
}
function stringConfigValue(
connection: KtxAthenaConnectionConfig | undefined,
key: keyof KtxAthenaConnectionConfig,
env: NodeJS.ProcessEnv,
): string | undefined {
const value = connection?.[key];
if (typeof value !== 'string' || value.trim().length === 0) return undefined;
// Resolve before checking emptiness: an unset `env:` reference resolves to '',
// which must become undefined so `?? default` applies instead of keeping ''.
const resolved = resolveStringReference(value.trim(), env).trim();
return resolved.length > 0 ? resolved : undefined;
}
function configuredAthenaDatabases(connection: KtxAthenaConnectionConfig): string[] {
if (!Array.isArray(connection.databases)) return [];
const selected = connection.databases
.filter((database): database is string => typeof database === 'string' && database.trim().length > 0)
.map((database) => database.trim());
return [...new Set(selected)];
}
export function isKtxAthenaConnectionConfig(
connection: unknown,
): connection is KtxAthenaConnectionConfig {
return (
typeof connection === 'object' &&
connection !== null &&
String((connection as { driver?: unknown }).driver ?? '').toLowerCase() === 'athena'
);
}
/** @internal */
export function athenaConnectionConfigFromConfig(input: {
connectionId: string;
connection: KtxAthenaConnectionConfig | undefined;
env?: NodeJS.ProcessEnv;
}): KtxAthenaResolvedConnectionConfig {
const inputDriver = input.connection?.driver ?? 'unknown';
if (!isKtxAthenaConnectionConfig(input.connection)) {
throw new Error(`Native Athena connector cannot run driver "${String(inputDriver)}"`);
}
const env = input.env ?? process.env;
const region = stringConfigValue(input.connection, 'region', env);
if (!region) {
throw new Error(`Native Athena connector requires connections.${input.connectionId}.region`);
}
const s3StagingDir = stringConfigValue(input.connection, 's3_staging_dir', env);
if (!s3StagingDir) {
throw new Error(`Native Athena connector requires connections.${input.connectionId}.s3_staging_dir`);
}
return {
region,
s3StagingDir,
workgroup: stringConfigValue(input.connection, 'workgroup', env) ?? 'primary',
catalog: stringConfigValue(input.connection, 'catalog', env) ?? 'AwsDataCatalog',
database: stringConfigValue(input.connection, 'database', env),
databases: configuredAthenaDatabases(input.connection),
};
}
function glueTableKind(tableType: string | undefined): 'table' | 'view' {
const t = String(tableType ?? '').toUpperCase();
if (t === 'VIRTUAL_VIEW') return 'view';
return 'table';
}
const POLL_INTERVAL_MS = 250;
const QUERY_TIMEOUT_MS = 5 * 60 * 1000;
export interface KtxAthenaScanConnectorOptions {
connectionId: string;
connection: KtxAthenaConnectionConfig | undefined;
clientFactory?: KtxAthenaClientFactory;
env?: NodeJS.ProcessEnv;
now?: () => Date;
}
export class KtxAthenaScanConnector implements KtxScanConnector {
readonly id: string;
readonly driver = 'athena' as const;
readonly capabilities = createKtxConnectorCapabilities({
tableSampling: true,
columnSampling: true,
columnStats: false,
readOnlySql: true,
nestedAnalysis: false,
formalForeignKeys: false,
estimatedRowCounts: false,
});
private readonly connectionId: string;
private readonly resolved: KtxAthenaResolvedConnectionConfig;
private readonly clientFactory: KtxAthenaClientFactory;
private readonly now: () => Date;
private readonly dialect = getSqlDialectForDriver('athena');
private athenaClient: KtxAthenaClient | null = null;
private glueClient: KtxGlueClient | null = null;
constructor(options: KtxAthenaScanConnectorOptions) {
this.connectionId = options.connectionId;
this.resolved = athenaConnectionConfigFromConfig({
connectionId: options.connectionId,
connection: options.connection,
env: options.env,
});
this.clientFactory = options.clientFactory ?? new DefaultAthenaClientFactory();
this.now = options.now ?? (() => new Date());
this.id = `athena:${options.connectionId}`;
}
async testConnection(): Promise<KtxConnectorTestResult> {
try {
await this.listDatabasesPaginated({ maxResults: 1 });
return { success: true };
} catch (error) {
return connectorTestFailure(error);
}
}
async introspect(input: KtxScanInput, _ctx: KtxScanContext): Promise<KtxSchemaSnapshot> {
this.assertConnection(input.connectionId);
// Honor the configured `databases` scope (written by `ktx setup`); fall back
// to every Glue database only when the scope is unset.
const databases =
this.resolved.databases.length > 0 ? this.resolved.databases : await this.listDatabasesPaginated({});
const tables: KtxSchemaTable[] = [];
for (const database of databases) {
const scopedNames = input.tableScope
? scopedTableNames(input.tableScope, { catalog: this.resolved.catalog, db: database })
: null;
tables.push(...(await this.introspectDatabase(database, scopedNames)));
}
return {
connectionId: this.connectionId,
driver: 'athena',
extractedAt: this.now().toISOString(),
scope: { catalogs: [this.resolved.catalog], datasets: databases },
metadata: {
catalog: this.resolved.catalog,
databases,
table_count: tables.length,
total_columns: tables.reduce((sum, t) => sum + t.columns.length, 0),
},
tables,
warnings: [],
};
}
async sampleTable(input: KtxTableSampleInput, _ctx: KtxScanContext): Promise<KtxTableSampleResult & { headerTypes?: string[] }> {
this.assertConnection(input.connectionId);
const result = await this.query(this.dialect.generateSampleQuery(this.qTableName(input.table), input.limit, input.columns));
return { headers: result.headers, headerTypes: result.headerTypes, rows: result.rows, totalRows: result.totalRows };
}
async sampleColumn(input: KtxColumnSampleInput, _ctx: KtxScanContext): Promise<KtxColumnSampleResult> {
this.assertConnection(input.connectionId);
const result = await this.query(
this.dialect.generateColumnSampleQuery(this.qTableName(input.table), input.column, input.limit),
);
return {
values: result.rows.filter((row) => row.length > 0 && row[0] !== null).map((row) => row[0]),
nullCount: null,
distinctCount: null,
};
}
async columnStats(_input: KtxColumnStatsInput, _ctx: KtxScanContext): Promise<KtxColumnStatsResult | null> {
return null;
}
async executeReadOnly(input: KtxReadOnlyQueryInput, _ctx: KtxScanContext): Promise<KtxQueryResult> {
this.assertConnection(input.connectionId);
const limitedSql = limitSqlForExecution(assertReadOnlySql(input.sql), input.maxRows);
const result = await this.query(limitedSql);
return { ...result, rowCount: result.rows.length };
}
async listSchemas(): Promise<string[]> {
return this.listDatabasesPaginated({});
}
async listTables(databases?: string[]): Promise<KtxTableListEntry[]> {
const targetDatabases = databases && databases.length > 0 ? databases : await this.listDatabasesPaginated({});
const entries: KtxTableListEntry[] = [];
for (const database of targetDatabases) {
const glueTables = await this.listGlueTablesPaginated(database);
for (const t of glueTables) {
if (!t.Name) continue;
entries.push({
catalog: this.resolved.catalog,
schema: database,
name: t.Name,
kind: glueTableKind(t.TableType),
});
}
}
return entries;
}
async cleanup(): Promise<void> {
this.athenaClient = null;
this.glueClient = null;
}
qTableName(table: Pick<KtxTableRef, 'name'> & Partial<Pick<KtxTableRef, 'catalog' | 'db'>>): string {
return this.dialect.formatTableName(table);
}
private getAthenaClient(): KtxAthenaClient {
if (!this.athenaClient) {
this.athenaClient = this.clientFactory.createAthenaClient(this.resolved.region);
}
return this.athenaClient;
}
private getGlueClient(): KtxGlueClient {
if (!this.glueClient) {
this.glueClient = this.clientFactory.createGlueClient(this.resolved.region);
}
return this.glueClient;
}
private async listDatabasesPaginated(opts: { maxResults?: number }): Promise<string[]> {
const names: string[] = [];
let nextToken: string | undefined;
do {
const result = await this.getGlueClient().getDatabases({ NextToken: nextToken });
for (const db of result.DatabaseList ?? []) {
if (db.Name) names.push(db.Name);
if (opts.maxResults && names.length >= opts.maxResults) return names;
}
nextToken = result.NextToken;
} while (nextToken);
return names;
}
private async listGlueTablesPaginated(database: string): Promise<KtxGlueTable[]> {
const tables: KtxGlueTable[] = [];
let nextToken: string | undefined;
do {
const result = await this.getGlueClient().getTables({ DatabaseName: database, NextToken: nextToken });
tables.push(...(result.TableList ?? []));
nextToken = result.NextToken;
} while (nextToken);
return tables;
}
private async introspectDatabase(database: string, scopedNames: readonly string[] | null): Promise<KtxSchemaTable[]> {
if (scopedNames && scopedNames.length === 0) return [];
const glueTables = await this.listGlueTablesPaginated(database);
const scopeSet = scopedNames ? new Set(scopedNames) : null;
return glueTables
.filter((t): t is KtxGlueTable & { Name: string } => Boolean(t.Name) && (!scopeSet || scopeSet.has(t.Name!)))
.map((t) => ({
catalog: this.resolved.catalog,
db: database,
name: t.Name,
kind: glueTableKind(t.TableType),
comment: t.Description ?? null,
estimatedRows: null,
columns: this.toSchemaColumns(t),
foreignKeys: [],
}));
}
private toSchemaColumns(table: KtxGlueTable): KtxSchemaColumn[] {
const columns = [...(table.StorageDescriptor?.Columns ?? []), ...(table.PartitionKeys ?? [])];
return columns
.filter((col): col is KtxGlueColumnDef & { Name: string } => Boolean(col.Name))
.map((col) => {
const nativeType = String(col.Type ?? 'string').toLowerCase();
return {
name: col.Name,
nativeType,
normalizedType: this.dialect.mapDataType(nativeType),
dimensionType: this.dialect.mapToDimensionType(nativeType),
nullable: true,
primaryKey: false,
comment: col.Comment ?? null,
};
});
}
private async query(sql: string): Promise<KtxQueryResult> {
const athena = this.getAthenaClient();
const { QueryExecutionId } = await athena.startQueryExecution({
QueryString: sql,
ResultConfiguration: { OutputLocation: this.resolved.s3StagingDir },
WorkGroup: this.resolved.workgroup,
...(this.resolved.database || this.resolved.catalog
? {
QueryExecutionContext: {
...(this.resolved.database ? { Database: this.resolved.database } : {}),
...(this.resolved.catalog ? { Catalog: this.resolved.catalog } : {}),
},
}
: {}),
});
if (!QueryExecutionId) {
throw new Error('Athena did not return a QueryExecutionId');
}
await this.waitForQueryCompletion(athena, QueryExecutionId);
const rows: unknown[][] = [];
let headers: string[] = [];
let headerTypes: string[] = [];
let nextToken: string | undefined;
let firstPage = true;
do {
const result = await athena.getQueryResults({ QueryExecutionId, NextToken: nextToken });
const resultSet = result.ResultSet;
if (firstPage) {
const columnInfo = resultSet?.ResultSetMetadata?.ColumnInfo ?? [];
headers = columnInfo.map((col) => col.Name ?? '');
headerTypes = columnInfo.map((col) => String(col.Type ?? 'varchar').toUpperCase());
firstPage = false;
}
const pageRows = resultSet?.Rows ?? [];
// Athena includes the header row as the first row of the first page — skip it.
const dataRows = nextToken === undefined ? pageRows.slice(1) : pageRows;
for (const row of dataRows) {
rows.push((row.Data ?? []).map((d) => d.VarCharValue ?? null));
}
nextToken = result.NextToken;
} while (nextToken);
return {
headers,
headerTypes: headerTypes.length > 0 ? headerTypes : undefined,
rows,
totalRows: rows.length,
rowCount: rows.length,
};
}
private async waitForQueryCompletion(athena: KtxAthenaClient, queryExecutionId: string): Promise<void> {
const terminalStates = new Set(['SUCCEEDED', 'FAILED', 'CANCELLED']);
const deadline = this.now().getTime() + QUERY_TIMEOUT_MS;
for (;;) {
const { QueryExecution } = await athena.getQueryExecution({ QueryExecutionId: queryExecutionId });
const state = QueryExecution?.Status?.State ?? '';
if (state === 'SUCCEEDED') return;
if (terminalStates.has(state)) {
const reason = QueryExecution?.Status?.StateChangeReason ?? state;
throw new Error(`Athena query ${state}: ${reason}`);
}
if (this.now().getTime() >= deadline) {
throw new Error(`Athena query ${queryExecutionId} timed out after ${QUERY_TIMEOUT_MS / 1000}s`);
}
await new Promise((resolve) => setTimeout(resolve, POLL_INTERVAL_MS));
}
}
private assertConnection(connectionId: string): void {
if (connectionId !== this.connectionId) {
throw new Error(`Athena connector ${this.connectionId} cannot scan connection ${connectionId}`);
}
}
}

View file

@ -0,0 +1,175 @@
import type { KtxSqlDialect } from '../../context/connections/dialects.js';
import {
columnDisplayPartCount,
formatDialectDisplayRef,
formatDialectTableName,
parseDialectDisplayRef,
} from '../../context/connections/dialect-helpers.js';
import type { KtxSchemaDimensionType, KtxTableRef } from '../../context/scan/types.js';
type AthenaTableNameRef = Pick<KtxTableRef, 'name'> & Partial<Pick<KtxTableRef, 'catalog' | 'db'>>;
/** @internal */
export class KtxAthenaDialect implements KtxSqlDialect {
readonly type = 'athena' as const;
private readonly dimensionTypeMappings: Record<string, KtxSchemaDimensionType> = {
timestamp: 'time',
date: 'time',
bigint: 'number',
int: 'number',
integer: 'number',
tinyint: 'number',
smallint: 'number',
double: 'number',
float: 'number',
real: 'number',
boolean: 'boolean',
};
quoteIdentifier(identifier: string): string {
return `"${identifier.replace(/"/g, '""')}"`;
}
formatTableName(table: AthenaTableNameRef): string {
return formatDialectTableName(table, this.quoteIdentifier.bind(this), 'ansi');
}
formatDisplayRef(table: AthenaTableNameRef): string {
return formatDialectDisplayRef(table, 'ansi');
}
parseDisplayRef(display: string): KtxTableRef | null {
return parseDialectDisplayRef(display, 'ansi');
}
columnDisplayTablePartCount(): 1 | 2 | 3 {
return columnDisplayPartCount('ansi');
}
mapDataType(nativeType: string): string {
const base = nativeType.toLowerCase().trim().split('<')[0]!.split('(')[0]!.trim();
const typeMap: Record<string, string> = {
string: 'VARCHAR',
varchar: 'VARCHAR',
char: 'CHAR',
binary: 'VARBINARY',
bigint: 'BIGINT',
int: 'INTEGER',
integer: 'INTEGER',
tinyint: 'TINYINT',
smallint: 'SMALLINT',
double: 'DOUBLE',
float: 'FLOAT',
real: 'REAL',
decimal: 'DECIMAL',
boolean: 'BOOLEAN',
timestamp: 'TIMESTAMP',
date: 'DATE',
array: 'ARRAY',
map: 'MAP',
struct: 'STRUCT',
uniontype: 'UNION',
};
return typeMap[base] ?? nativeType.toUpperCase();
}
mapToDimensionType(nativeType: string): KtxSchemaDimensionType {
const base = nativeType.toLowerCase().trim().split('<')[0]!.split('(')[0]!.trim();
const mapped = this.dimensionTypeMappings[base];
if (mapped) return mapped;
if (base.includes('timestamp') || base.includes('date')) return 'time';
if (base.includes('int') || base.includes('float') || base.includes('double') || base.includes('decimal') || base.includes('real')) return 'number';
if (base.includes('bool')) return 'boolean';
return 'string';
}
generateSampleQuery(tableName: string, limit: number, columns?: string[]): string {
const columnList =
columns && columns.length > 0 ? columns.map((c) => this.quoteIdentifier(c)).join(', ') : '*';
return `SELECT ${columnList} FROM ${tableName} LIMIT ${limit}`;
}
generateColumnSampleQuery(tableName: string, columnName: string, limit: number): string {
const quoted = this.quoteIdentifier(columnName);
return `SELECT ${quoted} FROM ${tableName} WHERE ${quoted} IS NOT NULL LIMIT ${limit}`;
}
generateCardinalitySampleQuery(tableName: string, columnName: string, sampleSize: number): string {
return `
SELECT approx_distinct(${columnName}) AS cardinality
FROM (
SELECT ${columnName}
FROM ${tableName}
WHERE ${columnName} IS NOT NULL
LIMIT ${sampleSize}
)
`;
}
generateRandomizedCardinalitySampleQuery(tableName: string, columnName: string, sampleSize: number): string {
return `
SELECT approx_distinct(${columnName}) AS cardinality
FROM (
SELECT ${columnName}
FROM ${tableName}
WHERE ${columnName} IS NOT NULL
ORDER BY rand()
LIMIT ${sampleSize}
)
`;
}
generateDistinctValuesQuery(tableName: string, columnName: string, limit: number): string {
return `
SELECT DISTINCT CAST(${columnName} AS VARCHAR) AS val
FROM ${tableName}
WHERE ${columnName} IS NOT NULL
ORDER BY val
LIMIT ${limit}
`;
}
generateColumnStatisticsQuery(_schemaName: string, _tableName: string): string | null {
return null;
}
getNullCountExpression(column: string): string {
return `COUNT_IF(${column} IS NULL)`;
}
getDistinctCountExpression(column: string): string {
return `approx_distinct(${column})`;
}
textLengthExpression(columnSql: string): string {
return `LENGTH(CAST(${columnSql} AS VARCHAR))`;
}
castToText(columnSql: string): string {
return `CAST(${columnSql} AS VARCHAR)`;
}
getSampleValueAggregation(innerSql: string): string {
return `(SELECT array_join(array_agg(CAST(value AS VARCHAR)), '\u001f') FROM (${innerSql}) AS relationship_profile_values)`;
}
getLimitOffsetClause(limit: number, offset?: number): string {
const safeLimit = Math.max(1, Math.floor(limit));
const safeOffset = offset !== undefined ? Math.floor(offset) : 0;
return safeOffset > 0 ? `OFFSET ${safeOffset} LIMIT ${safeLimit}` : `LIMIT ${safeLimit}`;
}
getTopClause(_limit: number): string {
return '';
}
getTableSampleClause(_samplePct: number): string {
return '';
}
getRandomSampleFilter(samplePct: number): string {
if (samplePct <= 0 || samplePct >= 1) return '';
return `rand() < ${samplePct}`;
}
}

View file

@ -0,0 +1,44 @@
import type {
LiveDatabaseIntrospectionOptions,
LiveDatabaseIntrospectionPort,
} from '../../context/ingest/adapters/live-database/types.js';
import type { KtxProjectConnectionConfig } from '../../context/project/config.js';
import {
KtxAthenaScanConnector,
type KtxAthenaClientFactory,
type KtxAthenaConnectionConfig,
} from './connector.js';
interface CreateAthenaLiveDatabaseIntrospectionOptions {
connections: Record<string, KtxProjectConnectionConfig>;
clientFactory?: KtxAthenaClientFactory;
now?: () => Date;
}
export function createAthenaLiveDatabaseIntrospection(
options: CreateAthenaLiveDatabaseIntrospectionOptions,
): LiveDatabaseIntrospectionPort {
return {
async extractSchema(connectionId: string, introspectionOptions?: LiveDatabaseIntrospectionOptions) {
const connection = options.connections[connectionId] as KtxAthenaConnectionConfig | undefined;
const connector = new KtxAthenaScanConnector({
connectionId,
connection,
clientFactory: options.clientFactory,
now: options.now,
});
try {
return await connector.introspect(
{
connectionId,
driver: 'athena',
...(introspectionOptions?.tableScope ? { tableScope: introspectionOptions.tableScope } : {}),
},
{ runId: `athena-${connectionId}` },
);
} finally {
await connector.cleanup();
}
},
};
}

View file

@ -0,0 +1,395 @@
import { DuckDBInstance, type DuckDBConnection } from '@duckdb/node-api';
import { existsSync, statSync } from 'node:fs';
import { isAbsolute, resolve } from 'node:path';
import { fileURLToPath } from 'node:url';
import { resolveStringReference } from '../shared/string-reference.js';
import { getSqlDialectForDriver } from '../../context/connections/dialects.js';
import { assertReadOnlySql, limitSqlForExecution } from '../../context/connections/read-only-sql.js';
import { normalizeQueryRows } from '../../context/connections/query-executor.js';
import { toJsonSafeRows } from '../shared/duckdb-json-safe.js';
import {
connectorTestFailure,
createKtxConnectorCapabilities,
type KtxColumnSampleInput,
type KtxColumnSampleResult,
type KtxColumnStatsInput,
type KtxColumnStatsResult,
type KtxConnectorTestResult,
type KtxQueryResult,
type KtxReadOnlyQueryInput,
type KtxScanConnector,
type KtxScanContext,
type KtxScanInput,
type KtxSchemaForeignKey,
type KtxSchemaSnapshot,
type KtxSchemaTable,
type KtxTableListEntry,
type KtxTableRef,
type KtxTableSampleInput,
type KtxTableSampleResult,
} from '../../context/scan/types.js';
import { scopedTableNames } from '../../context/scan/table-ref.js';
const MAIN_SCHEMA = 'main';
export interface KtxDuckDbConnectionConfig {
driver?: string;
path?: string;
url?: string;
[key: string]: unknown;
}
/** @internal */
export interface DuckDbDatabasePathInput {
connectionId: string;
projectDir?: string;
connection: KtxDuckDbConnectionConfig | undefined;
}
export interface KtxDuckDbScanConnectorOptions extends DuckDbDatabasePathInput {
now?: () => Date;
}
export interface KtxDuckDbColumnDistinctValuesOptions {
maxCardinality: number;
limit: number;
sampleSize?: number;
}
export interface KtxDuckDbColumnDistinctValuesResult {
values: string[] | null;
cardinality: number;
}
interface InfoSchemaTableRow {
table_name: string;
table_type: string;
}
interface InfoSchemaColumnRow {
column_name: string;
data_type: string;
is_nullable: string;
}
// `path` may be an env:/file: reference; `url` resolves env: only, since file:
// on a url is a native URI form (handled by duckDbPathFromUrl), not a file read.
function stringConfigValue(
connection: KtxDuckDbConnectionConfig | undefined,
key: 'path' | 'url',
): string | undefined {
const value = connection?.[key];
if (typeof value !== 'string' || value.trim().length === 0) {
return undefined;
}
const trimmed = value.trim();
if (key === 'url') {
return trimmed.startsWith('env:') ? (process.env[trimmed.slice('env:'.length)] ?? '') : trimmed;
}
return resolveStringReference(trimmed, process.env);
}
function duckDbPathFromUrl(url: string): string {
if (url.startsWith('file:')) {
return fileURLToPath(url);
}
if (url.startsWith('duckdb:')) {
const parsed = new URL(url);
return decodeURIComponent(parsed.pathname);
}
return url;
}
export function isKtxDuckDbConnectionConfig(
connection: KtxDuckDbConnectionConfig | undefined,
): connection is KtxDuckDbConnectionConfig {
return String(connection?.driver ?? '').toLowerCase() === 'duckdb';
}
/** @internal */
export function duckDbDatabasePathFromConfig(input: DuckDbDatabasePathInput): string {
const inputDriver = input.connection?.driver ?? 'unknown';
if (!isKtxDuckDbConnectionConfig(input.connection)) {
throw new Error(`Native DuckDB connector cannot run driver "${inputDriver}"`);
}
const configuredPath =
stringConfigValue(input.connection, 'path') ?? duckDbPathFromUrl(stringConfigValue(input.connection, 'url') ?? '');
if (!configuredPath) {
throw new Error(`Native DuckDB connector requires connections.${input.connectionId}.path or url`);
}
return isAbsolute(configuredPath) ? configuredPath : resolve(input.projectDir ?? process.cwd(), configuredPath);
}
export class KtxDuckDbScanConnector implements KtxScanConnector {
readonly id: string;
readonly driver = 'duckdb' as const;
readonly capabilities = createKtxConnectorCapabilities({
tableSampling: true,
columnSampling: true,
columnStats: false,
readOnlySql: true,
nestedAnalysis: false,
formalForeignKeys: true,
estimatedRowCounts: true,
});
private readonly connectionId: string;
private readonly dbPath: string;
private readonly now: () => Date;
private readonly dialect = getSqlDialectForDriver('duckdb');
private instance: DuckDBInstance | null = null;
private connection: DuckDBConnection | null = null;
constructor(options: KtxDuckDbScanConnectorOptions) {
this.connectionId = options.connectionId;
this.dbPath = duckDbDatabasePathFromConfig(options);
this.now = options.now ?? (() => new Date());
this.id = `duckdb:${options.connectionId}`;
}
async testConnection(): Promise<KtxConnectorTestResult> {
try {
if (!existsSync(this.dbPath) || !statSync(this.dbPath).isFile()) {
return { success: false, error: `File not found: ${this.dbPath}` };
}
await this.query('SELECT 1');
return { success: true };
} catch (error) {
return connectorTestFailure(error);
}
}
async introspect(input: KtxScanInput, _ctx: KtxScanContext): Promise<KtxSchemaSnapshot> {
this.assertConnection(input.connectionId);
const scopedNames = input.tableScope ? scopedTableNames(input.tableScope, { catalog: null, db: null }) : null;
const tableRows = await this.readTableRows(scopedNames);
const tables: KtxSchemaTable[] = [];
for (const row of tableRows) {
tables.push(await this.readTable(row));
}
return {
connectionId: this.connectionId,
driver: 'duckdb' as const,
extractedAt: this.now().toISOString(),
scope: {},
metadata: {
file_path: this.dbPath,
table_count: tables.length,
total_columns: tables.reduce((sum, table) => sum + table.columns.length, 0),
},
tables,
};
}
async listSchemas(): Promise<string[]> {
return [MAIN_SCHEMA];
}
async listTables(_schemas?: string[]): Promise<KtxTableListEntry[]> {
const rows = await this.readTableRows(null);
return rows.map((row) => ({
catalog: null,
schema: MAIN_SCHEMA,
name: row.table_name,
kind: row.table_type === 'VIEW' ? ('view' as const) : ('table' as const),
}));
}
async sampleTable(input: KtxTableSampleInput, _ctx: KtxScanContext): Promise<KtxTableSampleResult> {
this.assertConnection(input.connectionId);
const result = await this.query(
this.dialect.generateSampleQuery(this.qTableName(input.table), input.limit, input.columns),
);
return { headers: result.headers, rows: result.rows, totalRows: result.totalRows };
}
async sampleColumn(input: KtxColumnSampleInput, _ctx: KtxScanContext): Promise<KtxColumnSampleResult> {
this.assertConnection(input.connectionId);
const result = await this.query(
this.dialect.generateColumnSampleQuery(this.qTableName(input.table), input.column, input.limit),
);
const values = result.rows.filter((row) => row.length > 0 && row[0] !== null).map((row) => row[0]);
return { values, nullCount: null, distinctCount: null };
}
async columnStats(_input: KtxColumnStatsInput, _ctx: KtxScanContext): Promise<KtxColumnStatsResult | null> {
return null;
}
async executeReadOnly(input: KtxReadOnlyQueryInput, _ctx: KtxScanContext): Promise<KtxQueryResult> {
this.assertConnection(input.connectionId);
const result = await this.query(limitSqlForExecution(input.sql, input.maxRows));
return { ...result, rowCount: result.rows.length };
}
async getColumnDistinctValues(
table: KtxTableRef,
columnName: string,
options: KtxDuckDbColumnDistinctValuesOptions,
): Promise<KtxDuckDbColumnDistinctValuesResult | null> {
const sampleSize = options.sampleSize ?? 10000;
const tableName = this.qTableName(table);
const quotedColumn = this.dialect.quoteIdentifier(columnName);
const cardinalityResult = await this.query(
this.dialect.generateCardinalitySampleQuery(tableName, quotedColumn, sampleSize),
);
if (cardinalityResult.rows.length === 0) {
return null;
}
const cardinality = Number(cardinalityResult.rows[0][0]);
if (Number.isNaN(cardinality)) {
return null;
}
if (cardinality === 0) {
return { values: [], cardinality: 0 };
}
if (cardinality > options.maxCardinality) {
return { values: null, cardinality };
}
const valuesResult = await this.query(
this.dialect.generateDistinctValuesQuery(tableName, quotedColumn, options.limit),
);
return {
values: valuesResult.rows.filter((row) => row.length > 0 && row[0] !== null).map((row) => String(row[0])),
cardinality,
};
}
async getTableRowCount(tableName: string): Promise<number> {
const result = await this.query(`SELECT COUNT(*) AS count FROM ${this.dialect.quoteIdentifier(tableName)}`);
return Number(result.rows[0]?.[0] ?? 0);
}
qTableName(table: Pick<KtxTableRef, 'name'>): string {
return this.dialect.formatTableName(table);
}
quoteIdentifier(identifier: string): string {
return this.dialect.quoteIdentifier(identifier);
}
async cleanup(): Promise<void> {
this.connection?.closeSync();
this.instance?.closeSync();
this.connection = null;
this.instance = null;
}
private async db(): Promise<DuckDBConnection> {
if (!this.connection) {
// DuckDBInstance.create() creates the file if missing, so this pre-check
// enforces the never-create rule. Do not remove it.
if (!existsSync(this.dbPath) || !statSync(this.dbPath).isFile()) {
throw new Error(`File not found: ${this.dbPath}`);
}
this.instance = await DuckDBInstance.create(this.dbPath, { access_mode: 'read_only' });
this.connection = await this.instance.connect();
}
return this.connection;
}
private async query(sql: string): Promise<Omit<KtxQueryResult, 'rowCount'>> {
const connection = await this.db();
const reader = await connection.runAndReadAll(assertReadOnlySql(sql));
const rows = toJsonSafeRows(normalizeQueryRows(reader.getRows()));
return {
headers: reader.columnNames(),
rows,
totalRows: rows.length,
};
}
private async readTableRows(scopedNames: string[] | null): Promise<InfoSchemaTableRow[]> {
if (scopedNames && scopedNames.length === 0) {
return [];
}
const scopeClause = scopedNames
? `AND table_name IN (${scopedNames.map((name) => `'${name.replaceAll("'", "''")}'`).join(', ')})`
: '';
const result = await this.query(
`SELECT table_name, table_type
FROM information_schema.tables
WHERE table_schema = '${MAIN_SCHEMA}' ${scopeClause}
ORDER BY table_name`,
);
return result.rows.map((row) => ({ table_name: String(row[0]), table_type: String(row[1]) }));
}
private async readTable(table: InfoSchemaTableRow): Promise<KtxSchemaTable> {
const columnsResult = await this.query(
`SELECT column_name, data_type, is_nullable
FROM information_schema.columns
WHERE table_schema = '${MAIN_SCHEMA}' AND table_name = '${table.table_name.replaceAll("'", "''")}'
ORDER BY ordinal_position`,
);
const columns = columnsResult.rows.map<InfoSchemaColumnRow>((row) => ({
column_name: String(row[0]),
data_type: String(row[1]),
is_nullable: String(row[2]),
}));
const primaryKeys = await this.readPrimaryKeyColumns(table.table_name);
const isView = table.table_type === 'VIEW';
const estimatedRows = isView ? null : await this.getTableRowCount(table.table_name);
return {
catalog: null,
db: null,
name: table.table_name,
kind: isView ? 'view' : 'table',
comment: null,
estimatedRows,
columns: columns.map((column) => ({
name: column.column_name,
nativeType: column.data_type,
normalizedType: this.dialect.mapDataType(column.data_type),
dimensionType: this.dialect.mapToDimensionType(column.data_type),
nullable: column.is_nullable === 'YES' && !primaryKeys.has(column.column_name),
primaryKey: primaryKeys.has(column.column_name),
comment: null,
})),
foreignKeys: await this.readForeignKeys(table.table_name),
};
}
private async readPrimaryKeyColumns(tableName: string): Promise<Set<string>> {
const result = await this.query(
`SELECT kcu.column_name
FROM information_schema.table_constraints tc
JOIN information_schema.key_column_usage kcu
ON tc.constraint_name = kcu.constraint_name AND tc.table_schema = kcu.table_schema
WHERE tc.table_schema = '${MAIN_SCHEMA}'
AND tc.table_name = '${tableName.replaceAll("'", "''")}'
AND tc.constraint_type = 'PRIMARY KEY'`,
);
return new Set(result.rows.map((row) => String(row[0])));
}
private async readForeignKeys(tableName: string): Promise<KtxSchemaForeignKey[]> {
// information_schema.constraint_column_usage in DuckDB returns the constrained
// columns (source), not the referenced columns. Use duckdb_constraints() which
// exposes constraint_column_names and referenced_column_names directly.
const result = await this.query(
`SELECT unnest(constraint_column_names) AS from_column,
referenced_table,
unnest(referenced_column_names) AS to_column,
constraint_name
FROM duckdb_constraints()
WHERE schema_name = '${MAIN_SCHEMA}'
AND table_name = '${tableName.replaceAll("'", "''")}'
AND constraint_type = 'FOREIGN KEY'`,
);
return result.rows.map((row) => ({
fromColumn: String(row[0]),
toCatalog: null,
toDb: null,
toTable: String(row[1]),
toColumn: String(row[2]),
constraintName: row[3] === null ? null : String(row[3]),
}));
}
private assertConnection(connectionId: string): void {
if (connectionId !== this.connectionId) {
throw new Error(`ktx DuckDB connector ${this.id} cannot serve connection ${connectionId}`);
}
}
}

View file

@ -0,0 +1,192 @@
import type { KtxSqlDialect } from '../../context/connections/dialects.js';
import {
columnDisplayPartCount,
formatDialectDisplayRef,
formatDialectTableName,
limitOffsetClause,
parseDialectDisplayRef,
} from '../../context/connections/dialect-helpers.js';
import type { KtxSchemaDimensionType, KtxTableRef } from '../../context/scan/types.js';
type DuckDbTableNameRef = Pick<KtxTableRef, 'name'> & Partial<Pick<KtxTableRef, 'catalog' | 'db'>>;
/** @internal */
export class KtxDuckDbDialect implements KtxSqlDialect {
readonly type = 'duckdb' as const;
private readonly typeMappings: Record<string, KtxSchemaDimensionType> = {
TIMESTAMP: 'time',
'TIMESTAMP WITH TIME ZONE': 'time',
TIMESTAMPTZ: 'time',
DATE: 'time',
TIME: 'time',
BIGINT: 'number',
INTEGER: 'number',
SMALLINT: 'number',
TINYINT: 'number',
HUGEINT: 'number',
UBIGINT: 'number',
UINTEGER: 'number',
DECIMAL: 'number',
NUMERIC: 'number',
REAL: 'number',
FLOAT: 'number',
DOUBLE: 'number',
VARCHAR: 'string',
CHAR: 'string',
TEXT: 'string',
BLOB: 'string',
UUID: 'string',
BOOLEAN: 'boolean',
BOOL: 'boolean',
};
quoteIdentifier(identifier: string): string {
return `"${identifier.replace(/"/g, '""')}"`;
}
// v1 introspects the `main` schema only and sets db=null on every table, so
// refs are single-namespace like SQLite — use the matching display shape.
formatTableName(table: DuckDbTableNameRef): string {
return formatDialectTableName(table, this.quoteIdentifier.bind(this), 'sqlite');
}
formatDisplayRef(table: DuckDbTableNameRef): string {
return formatDialectDisplayRef(table, 'sqlite');
}
parseDisplayRef(display: string): KtxTableRef | null {
return parseDialectDisplayRef(display, 'sqlite');
}
columnDisplayTablePartCount(): 1 | 2 | 3 {
return columnDisplayPartCount('sqlite');
}
mapDataType(nativeType: string): string {
return nativeType;
}
mapToDimensionType(nativeType: string): KtxSchemaDimensionType {
if (!nativeType) {
return 'string';
}
let normalized = nativeType.toUpperCase().trim();
if (normalized.includes('(')) {
normalized = normalized.split('(')[0].trim();
}
if (this.typeMappings[normalized]) {
return this.typeMappings[normalized];
}
if (normalized.includes('TIME') || normalized.includes('DATE')) {
return 'time';
}
if (
normalized.includes('INT') ||
normalized.includes('DEC') ||
normalized.includes('NUM') ||
normalized.includes('REAL') ||
normalized.includes('FLOAT') ||
normalized.includes('DOUBLE')
) {
return 'number';
}
if (normalized.includes('BOOL')) {
return 'boolean';
}
return 'string';
}
generateSampleQuery(tableName: string, limit: number, columns?: string[]): string {
const columnList =
columns && columns.length > 0 ? columns.map((column) => this.quoteIdentifier(column)).join(', ') : '*';
return `SELECT ${columnList} FROM ${tableName} LIMIT ${limit}`;
}
generateColumnSampleQuery(tableName: string, columnName: string, limit: number): string {
const quoted = this.quoteIdentifier(columnName);
return `SELECT ${quoted} FROM ${tableName} WHERE ${quoted} IS NOT NULL AND TRIM(CAST(${quoted} AS VARCHAR)) != '' LIMIT ${limit}`;
}
getRandomSampleFilter(samplePct: number): string {
if (samplePct <= 0 || samplePct >= 1) {
return '';
}
return `RANDOM() < ${samplePct}`;
}
getTableSampleClause(samplePct: number): string {
if (samplePct <= 0 || samplePct >= 1) {
return '';
}
return `USING SAMPLE ${Math.round(samplePct * 100)} PERCENT (bernoulli)`;
}
getLimitOffsetClause(limit: number, offset?: number): string {
return limitOffsetClause(limit, offset);
}
getTopClause(_limit: number): string {
return '';
}
getNullCountExpression(column: string): string {
return `SUM(CASE WHEN ${column} IS NULL THEN 1 ELSE 0 END)`;
}
getDistinctCountExpression(column: string): string {
return `COUNT(DISTINCT ${column})`;
}
textLengthExpression(columnSql: string): string {
return `LENGTH(CAST(${columnSql} AS VARCHAR))`;
}
castToText(columnSql: string): string {
return `CAST(${columnSql} AS VARCHAR)`;
}
getSampleValueAggregation(innerSql: string): string {
return `(SELECT STRING_AGG(CAST(value AS VARCHAR), chr(31)) FROM (${innerSql}) AS relationship_profile_values)`;
}
generateCardinalitySampleQuery(tableName: string, columnName: string, sampleSize: number): string {
return `
WITH sampled AS (
SELECT ${columnName} AS val
FROM ${tableName}
WHERE ${columnName} IS NOT NULL
LIMIT ${sampleSize}
)
SELECT COUNT(DISTINCT val) AS cardinality
FROM sampled
`;
}
generateDistinctValuesQuery(tableName: string, columnName: string, limit: number): string {
return `
SELECT DISTINCT CAST(${columnName} AS VARCHAR) AS val
FROM ${tableName}
WHERE ${columnName} IS NOT NULL
ORDER BY val
LIMIT ${limit}
`;
}
generateColumnStatisticsQuery(_schemaName: string, _tableName: string): string | null {
return null;
}
generateRandomizedCardinalitySampleQuery(tableName: string, columnName: string, sampleSize: number): string {
return `
WITH sampled AS (
SELECT ${columnName} AS val
FROM ${tableName}
WHERE ${columnName} IS NOT NULL
USING SAMPLE ${sampleSize} ROWS
)
SELECT COUNT(DISTINCT val) AS cardinality
FROM sampled
`;
}
}

View file

@ -4,6 +4,7 @@ import {
mysqlConnectionPoolConfigFromConfig, mysqlConnectionPoolConfigFromConfig,
type KtxMysqlConnectionConfig, type KtxMysqlConnectionConfig,
} from '../mysql/connector.js'; } from '../mysql/connector.js';
import { duckDbDatabasePathFromConfig, type KtxDuckDbConnectionConfig } from '../../connectors/duckdb/connector.js';
import type { FederatedMember } from '../../context/connections/federation.js'; import type { FederatedMember } from '../../context/connections/federation.js';
function kvKeyword(value: string): string { function kvKeyword(value: string): string {
@ -74,6 +75,12 @@ function mysqlAttachString(member: FederatedMember, env: NodeJS.ProcessEnv): str
*/ */
export function federatedAttachTarget(member: FederatedMember, env: NodeJS.ProcessEnv): string { export function federatedAttachTarget(member: FederatedMember, env: NodeJS.ProcessEnv): string {
switch (member.driver.toLowerCase()) { switch (member.driver.toLowerCase()) {
case 'duckdb':
return duckDbDatabasePathFromConfig({
connectionId: member.connectionId,
projectDir: member.projectDir,
connection: member.connection as KtxDuckDbConnectionConfig,
});
case 'sqlite': case 'sqlite':
return sqliteDatabasePathFromConfig({ return sqliteDatabasePathFromConfig({
connectionId: member.connectionId, connectionId: member.connectionId,

View file

@ -7,25 +7,12 @@ import type {
import { normalizeQueryRows } from '../../context/connections/query-executor.js'; import { normalizeQueryRows } from '../../context/connections/query-executor.js';
import { assertReadOnlySql, limitSqlForExecution } from '../../context/connections/read-only-sql.js'; import { assertReadOnlySql, limitSqlForExecution } from '../../context/connections/read-only-sql.js';
import { attachTypeForDriver, type FederatedMember } from '../../context/connections/federation.js'; import { attachTypeForDriver, type FederatedMember } from '../../context/connections/federation.js';
import { toJsonSafeRows } from '../shared/duckdb-json-safe.js';
function quoteDuckdbIdentifier(id: string): string { function quoteDuckdbIdentifier(id: string): string {
return `"${id.replaceAll('"', '""')}"`; return `"${id.replaceAll('"', '""')}"`;
} }
const MIN_SAFE_BIGINT = BigInt(Number.MIN_SAFE_INTEGER);
const MAX_SAFE_BIGINT = BigInt(Number.MAX_SAFE_INTEGER);
// DuckDB returns integer columns as JS bigint (unserializable by JSON). Values
// in Number's safe range become Number; larger magnitudes become strings so a
// BIGINT beyond 2^53 keeps its exact value instead of silently rounding.
function jsonSafeBigint(value: bigint): number | string {
return value >= MIN_SAFE_BIGINT && value <= MAX_SAFE_BIGINT ? Number(value) : value.toString();
}
function toJsonSafeRows(rows: unknown[][]): unknown[][] {
return rows.map((row) => row.map((cell) => (typeof cell === 'bigint' ? jsonSafeBigint(cell) : cell)));
}
/** @internal */ /** @internal */
export function buildAttachStatements(members: FederatedMember[], env: NodeJS.ProcessEnv): string[] { export function buildAttachStatements(members: FederatedMember[], env: NodeJS.ProcessEnv): string[] {
const attachments = members.map((member) => ({ const attachments = members.map((member) => ({
@ -34,13 +21,12 @@ export function buildAttachStatements(members: FederatedMember[], env: NodeJS.Pr
alias: member.connectionId, alias: member.connectionId,
})); }));
const loadStatements = [...new Set(attachments.map((a) => a.type))].map( const loadTypes = [...new Set(attachments.map((a) => a.type).filter((type): type is string => type !== null))];
(type) => `INSTALL ${type}; LOAD ${type};`, const loadStatements = loadTypes.map((type) => `INSTALL ${type}; LOAD ${type};`);
); const attachStatements = attachments.map(({ type, url, alias }) => {
const attachStatements = attachments.map( const options = type === null ? 'READ_ONLY' : `TYPE ${type}, READ_ONLY`;
({ type, url, alias }) => return `ATTACH '${url.replaceAll("'", "''")}' AS ${quoteDuckdbIdentifier(alias)} (${options});`;
`ATTACH '${url.replaceAll("'", "''")}' AS ${quoteDuckdbIdentifier(alias)} (TYPE ${type}, READ_ONLY);`, });
);
return [...loadStatements, ...attachStatements]; return [...loadStatements, ...attachStatements];
} }

View file

@ -0,0 +1,40 @@
import type {
LiveDatabaseIntrospectionOptions,
LiveDatabaseIntrospectionPort,
} from '../../context/ingest/adapters/live-database/types.js';
import type { KtxProjectConnectionConfig } from '../../context/project/config.js';
import { KtxDuckDbScanConnector, type KtxDuckDbConnectionConfig } from './connector.js';
export interface CreateDuckDbLiveDatabaseIntrospectionOptions {
projectDir?: string;
connections: Record<string, KtxProjectConnectionConfig>;
now?: () => Date;
}
export function createDuckDbLiveDatabaseIntrospection(
options: CreateDuckDbLiveDatabaseIntrospectionOptions,
): LiveDatabaseIntrospectionPort {
return {
async extractSchema(connectionId: string, introspectionOptions?: LiveDatabaseIntrospectionOptions) {
const connection = options.connections[connectionId] as KtxDuckDbConnectionConfig | undefined;
const connector = new KtxDuckDbScanConnector({
connectionId,
connection,
projectDir: options.projectDir,
now: options.now,
});
try {
return await connector.introspect(
{
connectionId,
driver: 'duckdb',
...(introspectionOptions?.tableScope ? { tableScope: introspectionOptions.tableScope } : {}),
},
{ runId: `duckdb-${connectionId}` },
);
} finally {
await connector.cleanup();
}
},
};
}

View file

@ -0,0 +1,14 @@
const MIN_SAFE_BIGINT = BigInt(Number.MIN_SAFE_INTEGER);
const MAX_SAFE_BIGINT = BigInt(Number.MAX_SAFE_INTEGER);
// DuckDB returns integer columns as JS bigint (unserializable by JSON). Values
// in Number's safe range become Number; larger magnitudes become strings so a
// BIGINT beyond 2^53 keeps its exact value instead of silently rounding.
/** @internal */
export function jsonSafeBigint(value: bigint): number | string {
return value >= MIN_SAFE_BIGINT && value <= MAX_SAFE_BIGINT ? Number(value) : value.toString();
}
export function toJsonSafeRows(rows: unknown[][]): unknown[][] {
return rows.map((row) => row.map((cell) => (typeof cell === 'bigint' ? jsonSafeBigint(cell) : cell)));
}

View file

@ -0,0 +1,24 @@
import type { ConnectionType } from './connection-type.js';
const CONNECTION_TYPE_TO_SQLGLOT = {
POSTGRESQL: 'postgres',
SQLITE: 'sqlite',
DUCKDB: 'duckdb',
SQLSERVER: 'tsql',
BIGQUERY: 'bigquery',
SNOWFLAKE: 'snowflake',
MYSQL: 'mysql',
CLICKHOUSE: 'clickhouse',
ATHENA: 'athena',
METABASE: null,
LOOKER: null,
NOTION: null,
} satisfies Record<ConnectionType, string | null>;
export function dialectForConnectionType(connectionType: string): string {
return CONNECTION_TYPE_TO_SQLGLOT[connectionType.toUpperCase() as ConnectionType] ?? 'postgres';
}
export function warehouseTargetDialect(connectionType: string): string | null {
return CONNECTION_TYPE_TO_SQLGLOT[connectionType.toUpperCase() as ConnectionType] ?? null;
}

View file

@ -3,25 +3,16 @@ import { z } from 'zod';
export const connectionTypeSchema = z.enum([ export const connectionTypeSchema = z.enum([
'POSTGRESQL', 'POSTGRESQL',
'SQLITE', 'SQLITE',
'DUCKDB',
'SQLSERVER', 'SQLSERVER',
'BIGQUERY', 'BIGQUERY',
'SNOWFLAKE', 'SNOWFLAKE',
'CENTRALREACH',
'EPIC',
'CERNER',
'ATHENA', 'ATHENA',
'QUICKBOOKS',
'WORKDAY',
'REST',
'S3',
'SLACK',
'METABASE', 'METABASE',
'LOOKER', 'LOOKER',
'NOTION', 'NOTION',
'MYSQL', 'MYSQL',
'CLICKHOUSE', 'CLICKHOUSE',
'PLAIN',
'BETTERSTACK',
]); ]);
export type ConnectionType = z.infer<typeof connectionTypeSchema>; export type ConnectionType = z.infer<typeof connectionTypeSchema>;

View file

@ -1,5 +1,7 @@
import { KtxAthenaDialect } from '../../connectors/athena/dialect.js';
import { KtxBigQueryDialect } from '../../connectors/bigquery/dialect.js'; import { KtxBigQueryDialect } from '../../connectors/bigquery/dialect.js';
import { KtxClickHouseDialect } from '../../connectors/clickhouse/dialect.js'; import { KtxClickHouseDialect } from '../../connectors/clickhouse/dialect.js';
import { KtxDuckDbDialect } from '../../connectors/duckdb/dialect.js';
import { KtxMongoDbDialect } from '../../connectors/mongodb/dialect.js'; import { KtxMongoDbDialect } from '../../connectors/mongodb/dialect.js';
import { KtxMysqlDialect } from '../../connectors/mysql/dialect.js'; import { KtxMysqlDialect } from '../../connectors/mysql/dialect.js';
import { KtxPostgresDialect } from '../../connectors/postgres/dialect.js'; import { KtxPostgresDialect } from '../../connectors/postgres/dialect.js';
@ -53,8 +55,10 @@ export interface KtxSqlDialect extends KtxDialect {
type KtxSqlDriver = Exclude<KtxConnectionDriver, 'mongodb'>; type KtxSqlDriver = Exclude<KtxConnectionDriver, 'mongodb'>;
const sqlDialectFactories: Record<KtxSqlDriver, () => KtxSqlDialect> = { const sqlDialectFactories: Record<KtxSqlDriver, () => KtxSqlDialect> = {
athena: () => new KtxAthenaDialect(),
bigquery: () => new KtxBigQueryDialect(), bigquery: () => new KtxBigQueryDialect(),
clickhouse: () => new KtxClickHouseDialect(), clickhouse: () => new KtxClickHouseDialect(),
duckdb: () => new KtxDuckDbDialect(),
mysql: () => new KtxMysqlDialect(), mysql: () => new KtxMysqlDialect(),
postgres: () => new KtxPostgresDialect(), postgres: () => new KtxPostgresDialect(),
sqlite: () => new KtxSqliteDialect(), sqlite: () => new KtxSqliteDialect(),

View file

@ -26,6 +26,23 @@ function invalidConnectionConfig(driver: KtxConnectionDriver): Error {
/** @internal */ /** @internal */
export const driverRegistrations: Record<KtxConnectionDriver, KtxDriverRegistration> = { export const driverRegistrations: Record<KtxConnectionDriver, KtxDriverRegistration> = {
athena: {
driver: 'athena',
scopeConfigKey: 'databases',
hasHistoricSqlReader: false,
load: async () => {
const m = await import('../../connectors/athena/connector.js');
return {
isConnectionConfig: (connection) => m.isKtxAthenaConnectionConfig(connection),
createScanConnector: ({ connectionId, connection }) => {
if (!m.isKtxAthenaConnectionConfig(connection)) {
throw invalidConnectionConfig('athena');
}
return new m.KtxAthenaScanConnector({ connectionId, connection });
},
};
},
},
bigquery: { bigquery: {
driver: 'bigquery', driver: 'bigquery',
scopeConfigKey: 'dataset_ids', scopeConfigKey: 'dataset_ids',
@ -68,6 +85,27 @@ export const driverRegistrations: Record<KtxConnectionDriver, KtxDriverRegistrat
}; };
}, },
}, },
duckdb: {
driver: 'duckdb',
scopeConfigKey: null,
hasHistoricSqlReader: false,
load: async () => {
const m = await import('../../connectors/duckdb/connector.js');
return {
isConnectionConfig: (connection) => {
const typedConnection = connection as Parameters<typeof m.isKtxDuckDbConnectionConfig>[0];
return m.isKtxDuckDbConnectionConfig(typedConnection);
},
createScanConnector: ({ connectionId, connection, projectDir }) => {
const typedConnection = connection as Parameters<typeof m.isKtxDuckDbConnectionConfig>[0];
if (!m.isKtxDuckDbConnectionConfig(typedConnection)) {
throw invalidConnectionConfig('duckdb');
}
return new m.KtxDuckDbScanConnector({ connectionId, connection: typedConnection, projectDir });
},
};
},
},
mongodb: { mongodb: {
driver: 'mongodb', driver: 'mongodb',
scopeConfigKey: 'databases', scopeConfigKey: 'databases',

View file

@ -4,18 +4,19 @@ import type { KtxProjectConnectionConfig } from '../project/config.js';
export const FEDERATED_CONNECTION_ID = '_ktx_federated'; export const FEDERATED_CONNECTION_ID = '_ktx_federated';
/** /**
* Drivers DuckDB can ATTACH for federation. The driver name doubles as the * Drivers DuckDB can ATTACH for federation. Membership is governed by this set;
* DuckDB extension/TYPE name, so this set is the single source of truth for * the attach TYPE is governed by attachTypeForDriver, which returns the driver
* both membership (a driver participates iff it appears here) and attach type. * name for extension-backed engines and null for a native DuckDB file (attached
* with no INSTALL/LOAD and no TYPE).
*/ */
const ATTACH_COMPATIBLE_DRIVERS = new Set(['postgres', 'mysql', 'sqlite']); const ATTACH_COMPATIBLE_DRIVERS = new Set(['postgres', 'mysql', 'sqlite', 'duckdb']);
export function attachTypeForDriver(driver: string): string { export function attachTypeForDriver(driver: string): string | null {
const normalized = driver.toLowerCase(); const normalized = driver.toLowerCase();
if (!ATTACH_COMPATIBLE_DRIVERS.has(normalized)) { if (!ATTACH_COMPATIBLE_DRIVERS.has(normalized)) {
throw new Error(`Driver "${driver}" cannot be attached by DuckDB federation.`); throw new Error(`Driver "${driver}" cannot be attached by DuckDB federation.`);
} }
return normalized; return normalized === 'duckdb' ? null : normalized;
} }
export interface FederatedMember { export interface FederatedMember {

View file

@ -1,5 +1,6 @@
import type { KtxProjectConnectionConfig } from '../project/config.js'; import type { KtxProjectConnectionConfig } from '../project/config.js';
import type { ConnectionType } from './connection-type.js'; import type { ConnectionType } from './connection-type.js';
import { connectionQueryPolicy } from './query-policy.js';
export interface LocalWarehouseDescriptor { export interface LocalWarehouseDescriptor {
id: string; id: string;
@ -18,16 +19,19 @@ export interface LocalConnectionInfo {
connectionType: string; connectionType: string;
members?: string[]; members?: string[];
hint?: string; hint?: string;
queryPolicy?: 'semantic-layer-only';
} }
const DRIVER_TO_CONNECTION_TYPE: Record<string, ConnectionType> = { const DRIVER_TO_CONNECTION_TYPE: Record<string, ConnectionType> = {
postgres: 'POSTGRESQL', postgres: 'POSTGRESQL',
sqlite: 'SQLITE', sqlite: 'SQLITE',
duckdb: 'DUCKDB',
sqlserver: 'SQLSERVER', sqlserver: 'SQLSERVER',
mysql: 'MYSQL', mysql: 'MYSQL',
clickhouse: 'CLICKHOUSE', clickhouse: 'CLICKHOUSE',
snowflake: 'SNOWFLAKE', snowflake: 'SNOWFLAKE',
bigquery: 'BIGQUERY', bigquery: 'BIGQUERY',
athena: 'ATHENA',
}; };
export function localConnectionToWarehouseDescriptor( export function localConnectionToWarehouseDescriptor(
@ -90,10 +94,17 @@ export function localConnectionInfoFromConfig(
if (!connection) { if (!connection) {
return null; return null;
} }
const restricted = connectionQueryPolicy(connection) === 'semantic-layer-only';
return { return {
id, id,
name: id, name: id,
connectionType: localConnectionTypeForConfig(id, connection), connectionType: localConnectionTypeForConfig(id, connection),
...(restricted
? {
queryPolicy: 'semantic-layer-only' as const,
hint: 'Raw SQL is disabled on this connection; query it with sl_query using predefined measures.',
}
: {}),
}; };
} }

View file

@ -1,8 +1,15 @@
import { executeFederatedQuery } from '../../connectors/duckdb/federated-executor.js'; import { executeFederatedQuery } from '../../connectors/duckdb/federated-executor.js';
import { KtxExpectedError, KtxQueryError, isNativeProgrammingFault } from '../../errors.js';
import { sqlAnalysisDialectForDriver } from '../sql-analysis/dialect.js';
import type { SqlAnalysisPort } from '../sql-analysis/ports.js';
import { assertSafeConnectionId } from '../sl/source-files.js';
import type { KtxLocalProject } from '../project/project.js'; import type { KtxLocalProject } from '../project/project.js';
import type { KtxScanConnector, KtxScanContext } from '../scan/types.js'; import type { KtxScanConnector, KtxScanContext } from '../scan/types.js';
import { assertSqlQueryableConnection } from './dialects.js';
import { deriveFederatedConnection, FEDERATED_CONNECTION_ID } from './federation.js'; import { deriveFederatedConnection, FEDERATED_CONNECTION_ID } from './federation.js';
import { assertRawSqlAllowed } from './query-policy.js';
import type { KtxSqlQueryExecutionInput, KtxSqlQueryExecutionResult } from './query-executor.js'; import type { KtxSqlQueryExecutionInput, KtxSqlQueryExecutionResult } from './query-executor.js';
import { resolveConfiguredConnection } from './resolve-connection.js';
export interface ExecuteProjectReadOnlySqlDeps { export interface ExecuteProjectReadOnlySqlDeps {
project: KtxLocalProject; project: KtxLocalProject;
@ -56,3 +63,71 @@ export async function executeProjectReadOnlySql(
await connector?.cleanup?.(); await connector?.cleanup?.();
} }
} }
type RawSqlProgressCallback = (event: { progress: number; message: string }) => void | Promise<void>;
export interface ExecuteProjectRawSqlDeps {
project: KtxLocalProject;
connectionId: string;
sql: string;
maxRows: number;
sqlAnalysis: SqlAnalysisPort;
createConnector: (connectionId: string) => Promise<KtxScanConnector> | KtxScanConnector;
executeFederated?: typeof executeFederatedQuery;
runId: string;
onProgress?: RawSqlProgressCallback;
}
/**
* Single guarded path for user-authored (raw) SQL `ktx sql` and the MCP
* sql_execution tool. Enforces the connection's query_policy and the parser
* read-only guard before executing; ktx-internal SQL (semantic-layer, ingest)
* calls executeProjectReadOnlySql directly and is not subject to query_policy.
*/
export async function executeProjectRawSql(deps: ExecuteProjectRawSqlDeps): Promise<KtxSqlQueryExecutionResult> {
const { project } = deps;
await deps.onProgress?.({ progress: 0, message: 'Validating SQL' });
const isFederated = deps.connectionId === FEDERATED_CONNECTION_ID;
const connectionId = isFederated ? deps.connectionId : assertSafeConnectionId(deps.connectionId);
const connection = isFederated ? undefined : resolveConfiguredConnection(project.config, connectionId);
if (!isFederated) {
assertSqlQueryableConnection(connectionId, connection!.driver);
}
assertRawSqlAllowed(project.config, project.projectDir, connectionId);
const dialect = sqlAnalysisDialectForDriver(isFederated ? 'duckdb' : connection!.driver);
const validation = await deps.sqlAnalysis.validateReadOnly(deps.sql, dialect);
if (!validation.ok) {
// A read-only guard rejecting the caller's SQL is an expected outcome, not a
// ktx fault: classify it so reportException keeps it out of Error Tracking.
throw new KtxQueryError(validation.error ?? 'SQL is not read-only.');
}
await deps.onProgress?.({ progress: 0.3, message: 'Executing' });
const result = await executeProjectReadOnlySql({
project,
input: {
connectionId,
projectDir: project.projectDir,
connection,
sql: deps.sql,
maxRows: deps.maxRows,
},
createConnector: deps.createConnector,
executeFederated: deps.executeFederated,
runId: deps.runId,
}).catch((error: unknown) => {
// A warehouse/driver rejection (e.g. the caller's SQL failed to compile) is a
// surfaced operational outcome, not a ktx fault: mark it expected while
// preserving the warehouse's own diagnostics. A native JS error (TypeError,
// etc.) signals a bug in connector code — let it propagate unchanged so Error
// Tracking still sees it.
if (isNativeProgrammingFault(error) || error instanceof KtxExpectedError) {
throw error;
}
throw new KtxQueryError(error instanceof Error ? error.message : String(error), { cause: error });
});
await deps.onProgress?.({ progress: 1, message: `Fetched ${result.rowCount ?? result.rows.length} rows` });
return result;
}

View file

@ -0,0 +1,61 @@
import { KtxQueryError } from '../../errors.js';
import type { KtxProjectConfig, KtxProjectConnectionConfig } from '../project/config.js';
import { deriveFederatedConnection, FEDERATED_CONNECTION_ID } from './federation.js';
import { isSqlQueryableDriver } from './dialects.js';
export type KtxConnectionQueryPolicy = 'read-only-sql' | 'semantic-layer-only';
export function connectionQueryPolicy(
connection: KtxProjectConnectionConfig | undefined,
): KtxConnectionQueryPolicy {
return connection !== undefined && connection.query_policy === 'semantic-layer-only'
? 'semantic-layer-only'
: 'read-only-sql';
}
/** Member ids whose policy blocks raw SQL through the federated connection. */
export function restrictedFederatedMemberIds(config: KtxProjectConfig, projectDir: string): string[] {
const descriptor = deriveFederatedConnection(config.connections, projectDir);
if (!descriptor) {
return [];
}
return descriptor.members
.filter((member) => connectionQueryPolicy(member.connection) === 'semantic-layer-only')
.map((member) => member.connectionId);
}
export function assertRawSqlAllowed(config: KtxProjectConfig, projectDir: string, connectionId: string): void {
if (connectionId === FEDERATED_CONNECTION_ID) {
const restricted = restrictedFederatedMemberIds(config, projectDir);
if (restricted.length > 0) {
throw new KtxQueryError(
`Federated SQL execution is disabled: member connection(s) ${restricted
.map((id) => `"${id}"`)
.join(', ')} are restricted to semantic-layer queries (query_policy: semantic-layer-only in ktx.yaml).`,
);
}
return;
}
if (connectionQueryPolicy(config.connections[connectionId]) === 'semantic-layer-only') {
throw new KtxQueryError(
`Connection "${connectionId}" is restricted to semantic-layer queries (query_policy: semantic-layer-only in ktx.yaml); ` +
'raw SQL execution is disabled. Query it through the semantic layer with predefined measures instead ' +
'(the sl_query tool or `ktx sl query`).',
);
}
}
/**
* False only when the project has SQL-queryable connections and every one of
* them is restricted then no raw-SQL surface can succeed and the
* sql_execution tool should not be offered at all.
*/
export function projectAllowsRawSql(config: KtxProjectConfig): boolean {
const sqlConnections = Object.values(config.connections).filter((connection) =>
isSqlQueryableDriver(connection.driver),
);
if (sqlConnections.length === 0) {
return true;
}
return sqlConnections.some((connection) => connectionQueryPolicy(connection) === 'read-only-sql');
}

View file

@ -1,6 +1,7 @@
import { readFileSync } from 'node:fs'; import { readFileSync } from 'node:fs';
import { homedir } from 'node:os'; import { homedir } from 'node:os';
import { resolve } from 'node:path'; import { resolve } from 'node:path';
import { KtxExpectedError } from '../../errors.js';
/** @internal */ /** @internal */
export function resolveKtxHomePath(path: string): string { export function resolveKtxHomePath(path: string): string {
@ -28,7 +29,15 @@ export function resolveKtxConfigReference(value: string | undefined, env: NodeJS
if (value.startsWith('file:')) { if (value.startsWith('file:')) {
const filePath = resolveKtxHomePath(value.slice('file:'.length).trim()); const filePath = resolveKtxHomePath(value.slice('file:'.length).trim());
const fileValue = readFileSync(filePath, 'utf8').trim(); let fileValue: string;
try {
fileValue = readFileSync(filePath, 'utf8').trim();
} catch (error) {
// A `file:` secret whose target is missing or unreadable is a config
// problem (secrets/ is gitignored, so absent after a clone), not a ktx
// fault — classify it so it stays out of Error Tracking.
throw new KtxExpectedError(`Could not read the secret file referenced in ktx.yaml: ${filePath}`, { cause: error });
}
return fileValue.length > 0 ? fileValue : undefined; return fileValue.length > 0 ? fileValue : undefined;
} }

View file

@ -4,6 +4,44 @@ import { URL } from 'node:url';
import { spawn } from 'node:child_process'; import { spawn } from 'node:child_process';
import type { ResolvedSemanticLayerSource, SemanticLayerQueryInput } from '../sl/types.js'; import type { ResolvedSemanticLayerSource, SemanticLayerQueryInput } from '../sl/types.js';
// The daemon signals "the caller sent a well-formed request I refused as
// invalid" with a distinct subprocess exit code (INPUT_REJECTED_EXIT_CODE in
// python/ktx-daemon/__main__.py) or HTTP 400. Kept in sync across the boundary.
const DAEMON_INPUT_REJECTED_EXIT_CODE = 3;
const DAEMON_INPUT_REJECTED_HTTP_STATUS = 400;
/**
* A ktx-daemon compute call failed. `inputRejected` is true when the daemon
* refused the request as invalid (a caller-driven outcome) rather than faulting;
* callers can promote those to an expected KtxQueryError while letting genuine
* daemon crashes reach Error Tracking. `detail` is the daemon's own message,
* unwrapped for surfacing to the caller.
*/
export class KtxDaemonComputeError extends Error {
readonly inputRejected: boolean;
readonly detail: string;
constructor(message: string, options: { inputRejected: boolean; detail: string; cause?: unknown }) {
super(message, options.cause !== undefined ? { cause: options.cause } : undefined);
this.name = 'KtxDaemonComputeError';
this.inputRejected = options.inputRejected;
this.detail = options.detail;
}
}
/** The daemon reports HTTP failures as `{ "detail": "…" }`; fall back to the raw body. */
function daemonHttpErrorDetail(raw: string): string {
try {
const parsed = JSON.parse(raw) as unknown;
if (parsed && typeof parsed === 'object' && typeof (parsed as { detail?: unknown }).detail === 'string') {
return (parsed as { detail: string }).detail;
}
} catch {
// Non-JSON body (e.g. a proxy error page) — surface it verbatim.
}
return raw;
}
interface KtxSemanticLayerComputeQueryResult { interface KtxSemanticLayerComputeQueryResult {
sql: string; sql: string;
dialect: string; dialect: string;
@ -128,7 +166,13 @@ function runProcessJson(
const stdoutText = Buffer.concat(stdout).toString('utf8').trim(); const stdoutText = Buffer.concat(stdout).toString('utf8').trim();
const stderrText = Buffer.concat(stderr).toString('utf8').trim(); const stderrText = Buffer.concat(stderr).toString('utf8').trim();
if (code !== 0) { if (code !== 0) {
reject(new Error(`ktx-daemon ${subcommand} failed: ${stderrText || `exit code ${code}`}`)); const detail = stderrText || `exit code ${code}`;
reject(
new KtxDaemonComputeError(`ktx-daemon ${subcommand} failed: ${detail}`, {
inputRejected: code === DAEMON_INPUT_REJECTED_EXIT_CODE,
detail,
}),
);
return; return;
} }
try { try {
@ -168,7 +212,12 @@ function postJson(baseUrl: string): KtxDaemonHttpJsonRunner {
const text = Buffer.concat(chunks).toString('utf8'); const text = Buffer.concat(chunks).toString('utf8');
const statusCode = response.statusCode ?? 0; const statusCode = response.statusCode ?? 0;
if (statusCode < 200 || statusCode >= 300) { if (statusCode < 200 || statusCode >= 300) {
reject(new Error(`ktx-daemon HTTP ${path} failed with ${statusCode}: ${text}`)); reject(
new KtxDaemonComputeError(`ktx-daemon HTTP ${path} failed with ${statusCode}: ${text}`, {
inputRejected: statusCode === DAEMON_INPUT_REJECTED_HTTP_STATUS,
detail: daemonHttpErrorDetail(text),
}),
);
return; return;
} }
try { try {

View file

@ -1,4 +1,5 @@
import type { ParsedTargetTable } from '../../parsed-target-table.js'; import type { ParsedTargetTable } from '../../parsed-target-table.js';
import { warehouseTargetDialect } from '../../../connections/connection-type-dialect.js';
import type { LookerWarehouseConnectionInfo } from './client.js'; import type { LookerWarehouseConnectionInfo } from './client.js';
import type { LookerPullConfig, LookerRuntimeCursors, StagedExploreFile, StagedLookmlModelsFile } from './types.js'; import type { LookerPullConfig, LookerRuntimeCursors, StagedExploreFile, StagedLookmlModelsFile } from './types.js';
@ -11,6 +12,7 @@ const LOOKER_DIALECT_TO_CONNECTION_TYPE = {
sqlite: 'SQLITE', sqlite: 'SQLITE',
sqlserver: 'SQLSERVER', sqlserver: 'SQLSERVER',
clickhouse: 'CLICKHOUSE', clickhouse: 'CLICKHOUSE',
awsathena: 'ATHENA',
} as const; } as const;
/** @internal */ /** @internal */
@ -72,16 +74,6 @@ export interface LookerMappingClient {
getExplore(modelName: string, exploreName: string): Promise<StagedExploreFile>; getExplore(modelName: string, exploreName: string): Promise<StagedExploreFile>;
} }
const SQLGLOT_DIALECT_BY_CONNECTION_TYPE: Partial<Record<LookerWarehouseTargetConnectionType, string>> = {
BIGQUERY: 'bigquery',
SNOWFLAKE: 'snowflake',
POSTGRESQL: 'postgres',
MYSQL: 'mysql',
SQLITE: 'sqlite',
SQLSERVER: 'tsql',
CLICKHOUSE: 'clickhouse',
};
export async function discoverLookerConnections( export async function discoverLookerConnections(
client: Pick<LookerMappingClient, 'listLookerConnections'>, client: Pick<LookerMappingClient, 'listLookerConnections'>,
): Promise<LookerWarehouseConnectionInfo[]> { ): Promise<LookerWarehouseConnectionInfo[]> {
@ -100,7 +92,7 @@ export function lookerDialectToConnectionType(dialect: string | null): LookerWar
/** @internal */ /** @internal */
export function sqlglotDialectForConnectionType(connectionType: string): string | null { export function sqlglotDialectForConnectionType(connectionType: string): string | null {
return SQLGLOT_DIALECT_BY_CONNECTION_TYPE[connectionType as LookerWarehouseTargetConnectionType] ?? null; return warehouseTargetDialect(connectionType);
} }
/** @internal */ /** @internal */

View file

@ -8,9 +8,9 @@ export const METABASE_ENGINE_TO_CONNECTION_TYPE = {
snowflake: 'SNOWFLAKE', snowflake: 'SNOWFLAKE',
sqlserver: 'SQLSERVER', sqlserver: 'SQLSERVER',
mysql: 'MYSQL', mysql: 'MYSQL',
athena: 'ATHENA',
} as const; } as const;
export interface DiscoveredMetabaseDatabase { export interface DiscoveredMetabaseDatabase {
id: number; id: number;
name: string; name: string;

View file

@ -243,6 +243,7 @@ const connectionListOutputSchema = z.object({
connectionType: z.string(), connectionType: z.string(),
members: z.array(z.string()).optional(), members: z.array(z.string()).optional(),
hint: z.string().optional(), hint: z.string().optional(),
queryPolicy: z.literal('semantic-layer-only').optional(),
}), }),
), ),
}); });

View file

@ -3,16 +3,15 @@ import { KtxExpectedError, KtxQueryError, isNativeProgrammingFault } from '../..
import { isDatabaseDriver, normalizeConnectionDriver } from '../../connection-drivers.js'; import { isDatabaseDriver, normalizeConnectionDriver } from '../../connection-drivers.js';
import { sqlDialectNotes } from '../../context/sql-analysis/dialect-notes.js'; import { sqlDialectNotes } from '../../context/sql-analysis/dialect-notes.js';
import type { KtxProjectConnectionConfig } from '../../context/project/config.js'; import type { KtxProjectConnectionConfig } from '../../context/project/config.js';
import { executeProjectReadOnlySql } from '../../context/connections/project-sql-executor.js'; import { executeProjectRawSql } from '../../context/connections/project-sql-executor.js';
import { FEDERATED_CONNECTION_ID, federatedConnectionListing } from '../../context/connections/federation.js'; import { federatedConnectionListing } from '../../context/connections/federation.js';
import { assertSqlQueryableConnection } from '../../context/connections/dialects.js'; import { projectAllowsRawSql, restrictedFederatedMemberIds } from '../../context/connections/query-policy.js';
import { resolveConfiguredConnection } from '../../context/connections/resolve-connection.js';
import { import {
type LocalConnectionInfo, type LocalConnectionInfo,
localConnectionInfoFromConfig, localConnectionInfoFromConfig,
} from '../../context/connections/local-warehouse-descriptor.js'; } from '../../context/connections/local-warehouse-descriptor.js';
import type { KtxEmbeddingPort } from '../../context/core/embedding.js'; import type { KtxEmbeddingPort } from '../../context/core/embedding.js';
import type { KtxSemanticLayerComputePort } from '../../context/daemon/semantic-layer-compute.js'; import { KtxDaemonComputeError, type KtxSemanticLayerComputePort } from '../../context/daemon/semantic-layer-compute.js';
import type { KtxLocalProject } from '../../context/project/project.js'; import type { KtxLocalProject } from '../../context/project/project.js';
import { createKtxEntityDetailsService } from '../../context/scan/entity-details.js'; import { createKtxEntityDetailsService } from '../../context/scan/entity-details.js';
import type { LocalScanMcpOptions } from '../../context/scan/local-scan.js'; import type { LocalScanMcpOptions } from '../../context/scan/local-scan.js';
@ -35,13 +34,32 @@ interface CreateLocalProjectMcpContextPortsOptions {
embeddingService: KtxEmbeddingPort | null; embeddingService: KtxEmbeddingPort | null;
} }
/**
* Reclassify a query-path failure. Warehouse/driver rejections and daemon
* input-rejections are caller-driven outcomes (KtxQueryError, kept out of Error
* Tracking while preserving the underlying diagnostics); native JS faults, daemon
* crashes, and already-expected errors propagate unchanged so genuine ktx bugs
* still reach Error Tracking.
*/
function throwClassifiedQueryError(error: unknown): never {
if (error instanceof KtxDaemonComputeError) {
if (error.inputRejected) {
throw new KtxQueryError(error.detail, { cause: error });
}
throw error;
}
if (isNativeProgrammingFault(error) || error instanceof KtxExpectedError) {
throw error;
}
throw new KtxQueryError(error instanceof Error ? error.message : String(error), { cause: error });
}
async function executeValidatedReadOnlySql( async function executeValidatedReadOnlySql(
project: KtxLocalProject, project: KtxLocalProject,
options: CreateLocalProjectMcpContextPortsOptions, options: CreateLocalProjectMcpContextPortsOptions,
input: { connectionId: string; sql: string; maxRows: number }, input: { connectionId: string; sql: string; maxRows: number },
onProgress?: KtxMcpProgressCallback, onProgress?: KtxMcpProgressCallback,
): Promise<KtxSqlExecutionResponse> { ): Promise<KtxSqlExecutionResponse> {
await onProgress?.({ progress: 0, message: 'Validating SQL' });
if (!options.sqlAnalysis) { if (!options.sqlAnalysis) {
throw new Error('sql_execution requires parser-backed SQL validation.'); throw new Error('sql_execution requires parser-backed SQL validation.');
} }
@ -50,52 +68,22 @@ async function executeValidatedReadOnlySql(
throw new Error('sql_execution requires a local scan connector factory.'); throw new Error('sql_execution requires a local scan connector factory.');
} }
const isFederated = input.connectionId === FEDERATED_CONNECTION_ID; const result = await executeProjectRawSql({
const connectionId = isFederated ? input.connectionId : assertSafeConnectionId(input.connectionId);
const connection = isFederated ? undefined : resolveConfiguredConnection(project.config, connectionId);
if (!isFederated) {
assertSqlQueryableConnection(connectionId, connection!.driver);
}
const dialect = sqlAnalysisDialectForDriver(isFederated ? 'duckdb' : connection!.driver);
const validation = await options.sqlAnalysis.validateReadOnly(input.sql, dialect);
if (!validation.ok) {
// A read-only guard rejecting the agent's SQL is an expected outcome, not a
// ktx fault: classify it so reportException keeps it out of Error Tracking.
throw new KtxQueryError(validation.error ?? 'SQL is not read-only.');
}
await onProgress?.({ progress: 0.3, message: 'Executing' });
const result = await executeProjectReadOnlySql({
project, project,
input: { connectionId: input.connectionId,
connectionId,
projectDir: project.projectDir,
connection,
sql: input.sql, sql: input.sql,
maxRows: input.maxRows, maxRows: input.maxRows,
}, sqlAnalysis: options.sqlAnalysis,
createConnector, createConnector,
runId: 'mcp-sql-execution', runId: 'mcp-sql-execution',
}).catch((error: unknown) => { onProgress,
// A warehouse/driver rejection (e.g. the agent's SQL failed to compile) is a
// surfaced operational outcome, not a ktx fault: mark it expected while
// preserving the warehouse's own diagnostics. A native JS error (TypeError,
// etc.) signals a bug in connector code — let it propagate unchanged so Error
// Tracking still sees it.
if (isNativeProgrammingFault(error) || error instanceof KtxExpectedError) {
throw error;
}
throw new KtxQueryError(error instanceof Error ? error.message : String(error), { cause: error });
}); });
const response = { return {
headers: result.headers, headers: result.headers,
...(result.headerTypes ? { headerTypes: result.headerTypes } : {}), ...(result.headerTypes ? { headerTypes: result.headerTypes } : {}),
rows: result.rows, rows: result.rows,
rowCount: result.rowCount ?? result.rows.length, rowCount: result.rowCount ?? result.rows.length,
}; };
await onProgress?.({ progress: 1, message: `Fetched ${response.rowCount} rows` });
return response;
} }
/** @internal Resolves a connection's dialect SQL notes; throws KtxExpectedError for an unknown or non-SQL-warehouse connection. */ /** @internal Resolves a connection's dialect SQL notes; throws KtxExpectedError for an unknown or non-SQL-warehouse connection. */
@ -130,12 +118,17 @@ export function createLocalProjectMcpContextPorts(
.sort((a, b) => a.id.localeCompare(b.id)); .sort((a, b) => a.id.localeCompare(b.id));
const federated = federatedConnectionListing(project.config.connections, project.projectDir); const federated = federatedConnectionListing(project.config.connections, project.projectDir);
if (federated) { if (federated) {
const restricted = restrictedFederatedMemberIds(project.config, project.projectDir);
configured.push({ configured.push({
id: federated.id, id: federated.id,
name: federated.id, name: federated.id,
connectionType: 'DUCKDB', connectionType: 'DUCKDB',
members: federated.members, members: federated.members,
hint: federated.hint, hint:
restricted.length > 0
? `Federated SQL is disabled: member connection(s) ${restricted.join(', ')} have query_policy: semantic-layer-only.`
: federated.hint,
...(restricted.length > 0 ? { queryPolicy: 'semantic-layer-only' as const } : {}),
}); });
} }
return configured; return configured;
@ -197,7 +190,8 @@ export function createLocalProjectMcpContextPorts(
if (!options.semanticLayerCompute) { if (!options.semanticLayerCompute) {
throw new Error('sl_query requires a semantic-layer query adapter.'); throw new Error('sl_query requires a semantic-layer query adapter.');
} }
return compileLocalSlQuery(project, { try {
return await compileLocalSlQuery(project, {
connectionId: input.connectionId, connectionId: input.connectionId,
query: input.query, query: input.query,
compute: options.semanticLayerCompute, compute: options.semanticLayerCompute,
@ -206,6 +200,9 @@ export function createLocalProjectMcpContextPorts(
queryExecutor: options.queryExecutor, queryExecutor: options.queryExecutor,
onProgress: executionOptions?.onProgress, onProgress: executionOptions?.onProgress,
}); });
} catch (error) {
throwClassifiedQueryError(error);
}
}, },
}, },
entityDetails: { entityDetails: {
@ -231,7 +228,10 @@ export function createLocalProjectMcpContextPorts(
}, },
}; };
if (options.sqlAnalysis && options.localScan?.createConnector) { // Register sql_execution only when some connection can accept raw SQL; in
// mixed projects the tool stays and executeProjectRawSql rejects restricted
// connection ids at request time.
if (options.sqlAnalysis && options.localScan?.createConnector && projectAllowsRawSql(project.config)) {
ports.sqlExecution = { ports.sqlExecution = {
async execute(input, executionOptions) { async execute(input, executionOptions) {
return executeValidatedReadOnlySql(project, options, input, executionOptions?.onProgress); return executeValidatedReadOnlySql(project, options, input, executionOptions?.onProgress);

View file

@ -11,8 +11,10 @@ const warehouseDrivers = [
'snowflake', 'snowflake',
'bigquery', 'bigquery',
'sqlite', 'sqlite',
'duckdb',
'clickhouse', 'clickhouse',
'sqlserver', 'sqlserver',
'athena',
] as const; ] as const;
type WarehouseDriver = (typeof warehouseDrivers)[number]; type WarehouseDriver = (typeof warehouseDrivers)[number];
@ -40,6 +42,12 @@ function warehouseConnectionSchema<const Driver extends WarehouseDriver>(driver:
.describe( .describe(
'Maximum execution time for a single read-only query, in milliseconds (default 30000). Enforced as a server-side statement timeout for remote engines and by SIGKILL-ing a forked query subprocess for in-process SQLite. A query exceeding it is cancelled and returns a "query exceeded Ns" error so the agent can revise.', 'Maximum execution time for a single read-only query, in milliseconds (default 30000). Enforced as a server-side statement timeout for remote engines and by SIGKILL-ing a forked query subprocess for in-process SQLite. A query exceeding it is cancelled and returns a "query exceeded Ns" error so the agent can revise.',
), ),
query_policy: z
.enum(['read-only-sql', 'semantic-layer-only'])
.optional()
.describe(
'Agent-facing query authorship policy (default "read-only-sql"). "read-only-sql" allows parser-validated read-only SQL plus semantic-layer queries. "semantic-layer-only" rejects raw SQL on this connection (`ktx sql`, the sql_execution tool, and federated queries that include it) and restricts semantic-layer queries to measures predefined in the semantic-layer sources. ktx-internal scan and ingest queries are unaffected.',
),
}) })
.describe( .describe(
`${driver} warehouse connection. Additional driver-tunable fields (e.g. context.queryHistory) are accepted and passed through.`, `${driver} warehouse connection. Additional driver-tunable fields (e.g. context.queryHistory) are accepted and passed through.`,
@ -52,8 +60,10 @@ const warehouseConnectionSchemas = [
warehouseConnectionSchema('snowflake'), warehouseConnectionSchema('snowflake'),
warehouseConnectionSchema('bigquery'), warehouseConnectionSchema('bigquery'),
warehouseConnectionSchema('sqlite'), warehouseConnectionSchema('sqlite'),
warehouseConnectionSchema('duckdb'),
warehouseConnectionSchema('clickhouse'), warehouseConnectionSchema('clickhouse'),
warehouseConnectionSchema('sqlserver'), warehouseConnectionSchema('sqlserver'),
warehouseConnectionSchema('athena'),
] as const; ] as const;
const mongodbConnectionSchema = z const mongodbConnectionSchema = z

View file

@ -141,17 +141,19 @@ function normalizeDriver(driver: string | undefined): KtxConnectionDriver {
if ( if (
normalized === 'postgres' || normalized === 'postgres' ||
normalized === 'sqlite' || normalized === 'sqlite' ||
normalized === 'duckdb' ||
normalized === 'mysql' || normalized === 'mysql' ||
normalized === 'clickhouse' || normalized === 'clickhouse' ||
normalized === 'sqlserver' || normalized === 'sqlserver' ||
normalized === 'bigquery' || normalized === 'bigquery' ||
normalized === 'snowflake' || normalized === 'snowflake' ||
normalized === 'athena' ||
normalized === 'mongodb' normalized === 'mongodb'
) { ) {
return normalized; return normalized;
} }
throw new Error( throw new Error(
`Standalone ktx scan supports postgres/sqlite/mysql/clickhouse/sqlserver/bigquery/snowflake/mongodb in this phase, received "${driver ?? 'unknown'}"`, `Standalone ktx scan supports postgres/sqlite/duckdb/mysql/clickhouse/sqlserver/bigquery/snowflake/athena/mongodb in this phase, received "${driver ?? 'unknown'}"`,
); );
} }

View file

@ -2,12 +2,14 @@ import type { KtxTableRefKey } from './table-ref.js';
export type KtxConnectionDriver = export type KtxConnectionDriver =
| 'sqlite' | 'sqlite'
| 'duckdb'
| 'postgres' | 'postgres'
| 'sqlserver' | 'sqlserver'
| 'bigquery' | 'bigquery'
| 'snowflake' | 'snowflake'
| 'mysql' | 'mysql'
| 'clickhouse' | 'clickhouse'
| 'athena'
| 'mongodb'; | 'mongodb';
/** Canonical scan-mode registry. Runtime validation derives its allowlist here. */ /** Canonical scan-mode registry. Runtime validation derives its allowlist here. */

View file

@ -4,6 +4,7 @@ import type { KtxMcpProgressCallback } from '../mcp/types.js';
import type { KtxLocalProject } from '../../context/project/project.js'; import type { KtxLocalProject } from '../../context/project/project.js';
import { isSqlQueryableDriver } from '../connections/dialects.js'; import { isSqlQueryableDriver } from '../connections/dialects.js';
import { FEDERATED_CONNECTION_ID } from '../connections/federation.js'; import { FEDERATED_CONNECTION_ID } from '../connections/federation.js';
import { connectionQueryPolicy, restrictedFederatedMemberIds } from '../connections/query-policy.js';
import { resolveRequiredConnectionId } from '../connections/resolve-connection.js'; import { resolveRequiredConnectionId } from '../connections/resolve-connection.js';
import { sqlAnalysisDialectForDriver } from '../sql-analysis/dialect.js'; import { sqlAnalysisDialectForDriver } from '../sql-analysis/dialect.js';
import { loadLocalSlSourceRecords } from './local-sl.js'; import { loadLocalSlSourceRecords } from './local-sl.js';
@ -14,11 +15,31 @@ import type { SemanticLayerQueryExecutionResult, SemanticLayerQueryInput, Semant
const COMPILE_ONLY_REASON = const COMPILE_ONLY_REASON =
'Local semantic-layer query compiled SQL but no data-source execution adapter is configured.'; 'Local semantic-layer query compiled SQL but no data-source execution adapter is configured.';
const FEDERATED_SL_QUERY_UNSUPPORTED = const FEDERATED_SL_QUERY_PREFIX =
`Semantic-layer queries are per-connection and cannot target the federated connection '${FEDERATED_CONNECTION_ID}'. ` + `Semantic-layer queries are per-connection and cannot target the federated connection '${FEDERATED_CONNECTION_ID}'. `;
// The raw-SQL fallback is only valid when federated raw SQL is allowed; when a
// member is restricted (query_policy: semantic-layer-only), assertRawSqlAllowed
// rejects the same path, so directing the agent there would burn a guaranteed
// failure. Derive the message from the restricted-member set instead.
function federatedSlQueryUnsupportedMessage(project: KtxLocalProject): string {
const restricted = restrictedFederatedMemberIds(project.config, project.projectDir);
if (restricted.length > 0) {
return (
FEDERATED_SL_QUERY_PREFIX +
`Cross-database SQL through '${FEDERATED_CONNECTION_ID}' is also disabled because member connection(s) ` +
`${restricted.map((id) => `'${id}'`).join(', ')} are restricted to semantic-layer queries ` +
'(query_policy: semantic-layer-only). Query each connection on its own through the semantic layer ' +
'(the sl_query tool or `ktx sl query` with its connection id).'
);
}
return (
FEDERATED_SL_QUERY_PREFIX +
`Run a cross-database query as read-only SQL instead — ktx sql -c ${FEDERATED_CONNECTION_ID} "SELECT ..." or the sql_execution tool — ` + `Run a cross-database query as read-only SQL instead — ktx sql -c ${FEDERATED_CONNECTION_ID} "SELECT ..." or the sql_execution tool — ` +
'using catalog-qualified table names (connectionId.schema.table, or connectionId.table for sqlite; ' + 'using catalog-qualified table names (connectionId.schema.table, or connectionId.table for sqlite; ' +
'double-quote ids that are not bare identifiers, e.g. "books-db".public.books).'; 'double-quote ids that are not bare identifiers, e.g. "books-db".public.books).'
);
}
export interface CompileLocalSlQueryOptions { export interface CompileLocalSlQueryOptions {
connectionId?: string; connectionId?: string;
@ -79,7 +100,7 @@ export async function compileLocalSlQuery(
options: CompileLocalSlQueryOptions, options: CompileLocalSlQueryOptions,
): Promise<CompileLocalSlQueryResult> { ): Promise<CompileLocalSlQueryResult> {
if (options.connectionId === FEDERATED_CONNECTION_ID) { if (options.connectionId === FEDERATED_CONNECTION_ID) {
throw new Error(FEDERATED_SL_QUERY_UNSUPPORTED); throw new Error(federatedSlQueryUnsupportedMessage(project));
} }
await options.onProgress?.({ progress: 0, message: 'Compiling query' }); await options.onProgress?.({ progress: 0, message: 'Compiling query' });
const connectionId = resolveLocalConnectionId(project, options.connectionId); const connectionId = resolveLocalConnectionId(project, options.connectionId);
@ -93,11 +114,13 @@ export async function compileLocalSlQuery(
const dialect = sqlAnalysisDialectForDriver(driver); const dialect = sqlAnalysisDialectForDriver(driver);
const sources = await loadComputableSources(project, connectionId); const sources = await loadComputableSources(project, connectionId);
const predefinedMeasuresOnly =
connectionQueryPolicy(project.config.connections[connectionId]) === 'semantic-layer-only';
await options.onProgress?.({ progress: 0.3, message: 'Generating SQL' }); await options.onProgress?.({ progress: 0.3, message: 'Generating SQL' });
const response = await options.compute.query({ const response = await options.compute.query({
sources, sources,
dialect, dialect,
query: options.query, query: { ...options.query, predefined_measures_only: predefinedMeasuresOnly },
}); });
if (!options.execute) { if (!options.execute) {

View file

@ -2,6 +2,7 @@ import YAML from 'yaml';
import type { KtxFileStorePort } from '../../context/core/file-store.js'; import type { KtxFileStorePort } from '../../context/core/file-store.js';
import type { KtxLogger } from '../../context/core/config.js'; import type { KtxLogger } from '../../context/core/config.js';
import { noopLogger } from '../../context/core/config.js'; import { noopLogger } from '../../context/core/config.js';
import { dialectForConnectionType } from '../connections/connection-type-dialect.js';
import type { TableUsageOutput } from '../ingest/adapters/historic-sql/skill-schemas.js'; import type { TableUsageOutput } from '../ingest/adapters/historic-sql/skill-schemas.js';
import type { SlConnectionCatalogPort, SlPythonPort } from './ports.js'; import type { SlConnectionCatalogPort, SlPythonPort } from './ports.js';
import { normalizeSemanticLayerDescriptions } from './description-normalization.js'; import { normalizeSemanticLayerDescriptions } from './description-normalization.js';
@ -674,7 +675,7 @@ export class SemanticLayerService {
if (!connection) { if (!connection) {
throw new Error(`Data source not found: ${connectionId}`); throw new Error(`Data source not found: ${connectionId}`);
} }
return SemanticLayerService.mapDialect(connection.connectionType); return dialectForConnectionType(connection.connectionType);
} }
async listFilesForConnection(connectionId: string): Promise<string[]> { async listFilesForConnection(connectionId: string): Promise<string[]> {
@ -1107,25 +1108,6 @@ export class SemanticLayerService {
return { columns, tables }; return { columns, tables };
} }
/**
* All callers should use this instead of maintaining their own dialect maps.
*/
static mapDialect(connectionType: string): string {
const normalized = connectionType.toUpperCase();
const map: Record<string, string> = {
POSTGRES: 'postgres',
BIGQUERY: 'bigquery',
SNOWFLAKE: 'snowflake',
MYSQL: 'mysql',
SQLSERVER: 'tsql',
SQLITE: 'sqlite',
DUCKDB: 'duckdb',
CLICKHOUSE: 'clickhouse',
DATABRICKS: 'databricks',
};
return map[normalized] ?? 'postgres';
}
/** /**
* Execute a semantic layer query: load composed sources, generate SQL via * Execute a semantic layer query: load composed sources, generate SQL via
* the python SL engine, and execute the generated SQL against the data source. * the python SL engine, and execute the generated SQL against the data source.
@ -1153,7 +1135,7 @@ export class SemanticLayerService {
if (!connection) { if (!connection) {
throw new Error(`Data source not found: ${connectionId}`); throw new Error(`Data source not found: ${connectionId}`);
} }
const dialect = SemanticLayerService.mapDialect(connection.connectionType); const dialect = dialectForConnectionType(connection.connectionType);
// 3. Generate SQL via python SL engine // 3. Generate SQL via python SL engine
const { data: slResult, error: slError } = await this.python.query({ const { data: slResult, error: slError } = await this.python.query({

View file

@ -2,9 +2,10 @@ import YAML from 'yaml';
import type { GitService } from '../../../context/core/git.service.js'; import type { GitService } from '../../../context/core/git.service.js';
import type { KtxFileListResult, KtxFileReadResult, KtxFileStorePort } from '../../../context/core/file-store.js'; import type { KtxFileListResult, KtxFileReadResult, KtxFileStorePort } from '../../../context/core/file-store.js';
import { SYSTEM_GIT_AUTHOR } from '../../../context/tools/authors.js'; import { SYSTEM_GIT_AUTHOR } from '../../../context/tools/authors.js';
import { dialectForConnectionType } from '../../connections/connection-type-dialect.js';
import type { SlConnectionCatalogPort, SlSourcesIndexPort } from '../ports.js'; import type { SlConnectionCatalogPort, SlSourcesIndexPort } from '../ports.js';
import { sourceOverlaySchema } from '../schemas.js'; import { sourceOverlaySchema } from '../schemas.js';
import { SemanticLayerService } from '../semantic-layer.service.js'; import type { SemanticLayerService } from '../semantic-layer.service.js';
import { resolveSlSourceFile, slSourceFilePath } from '../source-files.js'; import { resolveSlSourceFile, slSourceFilePath } from '../source-files.js';
import type { SemanticLayerSource } from '../types.js'; import type { SemanticLayerSource } from '../types.js';
import { sourceDefinitionSchema } from './base-semantic-layer.tool.js'; import { sourceDefinitionSchema } from './base-semantic-layer.tool.js';
@ -28,7 +29,7 @@ function resolveDialect(warehouse: string | null): string | null {
if (!warehouse) { if (!warehouse) {
return null; return null;
} }
return SemanticLayerService.mapDialect(warehouse); return dialectForConnectionType(warehouse);
} }
function wrapWithZeroRowQuery(sql: string, dialect: string): string { function wrapWithZeroRowQuery(sql: string, dialect: string): string {

View file

@ -81,6 +81,9 @@ export interface SemanticLayerQueryInput {
order_by?: Array<string | { field: string; direction?: string }>; order_by?: Array<string | { field: string; direction?: string }>;
limit?: number; limit?: number;
include_empty?: boolean; include_empty?: boolean;
// Set by compileLocalSlQuery from the connection's query_policy, never from
// caller input: the planner rejects runtime-composed measures when true.
predefined_measures_only?: boolean;
} }
export interface SemanticLayerQueryExecutionResult { export interface SemanticLayerQueryExecutionResult {

View file

@ -6,8 +6,7 @@ import type { SqlAnalysisDialect } from './ports.js';
// dialect), served by the sql_dialect_notes MCP tool. They are package-internal: // dialect), served by the sql_dialect_notes MCP tool. They are package-internal:
// copy-runtime-assets.mjs ships them to dist, and they are never installed onto an // copy-runtime-assets.mjs ships them to dist, and they are never installed onto an
// agent target. The set covers every dialect reachable from a configured warehouse // agent target. The set covers every dialect reachable from a configured warehouse
// driver; duckdb/databricks are intentionally absent because no connector produces // driver; databricks is intentionally absent because no connector produces it.
// them.
/** @internal Dialects with an authored ./dialects/<dialect>.md file. */ /** @internal Dialects with an authored ./dialects/<dialect>.md file. */
export const DIALECTS_WITH_NOTES = [ export const DIALECTS_WITH_NOTES = [
@ -16,8 +15,10 @@ export const DIALECTS_WITH_NOTES = [
'snowflake', 'snowflake',
'bigquery', 'bigquery',
'sqlite', 'sqlite',
'duckdb',
'clickhouse', 'clickhouse',
'tsql', 'tsql',
'athena',
] as const; ] as const;
type DialectWithNotes = (typeof DIALECTS_WITH_NOTES)[number]; type DialectWithNotes = (typeof DIALECTS_WITH_NOTES)[number];

View file

@ -16,6 +16,7 @@ const SQLGLOT_DIALECTS: Record<string, SqlAnalysisDialect> = {
duckdb: 'duckdb', duckdb: 'duckdb',
clickhouse: 'clickhouse', clickhouse: 'clickhouse',
databricks: 'databricks', databricks: 'databricks',
athena: 'athena',
}; };
export function sqlAnalysisDialectForDriver(driver: string | undefined): SqlAnalysisDialect { export function sqlAnalysisDialectForDriver(driver: string | undefined): SqlAnalysisDialect {

View file

@ -0,0 +1,12 @@
**athena** SQL conventions (Trino engine over the Glue Data Catalog):
- **FQTN:** `database.table` (e.g. `analytics.orders`); a bare `table` resolves against the query's default database. Cross-catalog is `catalog.database.table` (e.g. `awsdatacatalog.analytics.orders`).
- **Identifiers:** case-insensitive and folded to lowercase; double-quote (`"Name"`) to keep case, spaces, or a reserved word. String literals use single quotes only.
- **Date/time:** native `DATE`/`TIMESTAMP`. Bucket with `date_trunc('month', ts)`, pull parts with `EXTRACT(YEAR FROM ts)`, shift with `date_add('day', -30, current_date)`, difference with `date_diff('day', a, b)`, and format with `date_format(ts, '%Y-%m')`. Parse text with `date_parse(str, '%Y-%m-%d')` or `from_iso8601_timestamp(str)`; `current_date` / `now()` are available.
- **Top-N / windows:** Athena has no `QUALIFY` — wrap the window in a subquery and filter it: `SELECT * FROM (SELECT ..., ROW_NUMBER() OVER (PARTITION BY key ORDER BY x DESC) AS rn FROM t) WHERE rn = 1` returns one row per key. Use `ORDER BY ... LIMIT n` for a global top-N, and paginate with `OFFSET m LIMIT n` (offset first — `LIMIT n OFFSET m` is a syntax error).
- **Series:** `CROSS JOIN UNNEST(sequence(DATE '2023-01-01', DATE '2023-12-01', INTERVAL '1' MONTH)) AS s(d)` expands a generated array into a date spine (use `sequence(1, 12)` for integers), then `LEFT JOIN` the aggregated facts onto it so empty periods still appear.
- **Rolling window over time:** a native `RANGE` frame spans real dates and tolerates gaps — `AVG(amount) OVER (ORDER BY day RANGE BETWEEN INTERVAL '29' DAY PRECEDING AND CURRENT ROW)` is a trailing 30-day average without a spine; guard minimum periods with `COUNT(*) OVER (<same frame>)`.
- **Approximate aggregates:** `approx_distinct(x)` for cardinality and `approx_percentile(x, 0.5)` for quantiles are far cheaper than exact `COUNT(DISTINCT ...)` on large scans.
- **Arrays & maps:** explode with `CROSS JOIN UNNEST(arr) AS t(x)` (add `WITH ORDINALITY` for an index); build with `array_agg(x)`, join with `array_join(arr, ',')`, index 1-based (`arr[1]`), and read a map with `element_at(m, key)`.
- **Safe cast:** `TRY_CAST(x AS DOUBLE)` yields `NULL` for a value that does not parse instead of raising, so counting residual `NULL`s catches an encoding the sample missed; `TRY(expr)` swallows other runtime errors.
- **Integer division:** `/` between integers truncates (`5 / 2``2`); cast an operand (`x / CAST(y AS DOUBLE)`) to keep the fraction, and round only in the final projection.
- **JSON:** `json_extract_scalar(col, '$.a.b')` returns varchar, `json_extract(col, '$.a')` returns json; cast a JSON string with `CAST(json_parse(col) AS ...)`.

View file

@ -0,0 +1,10 @@
**duckdb** SQL conventions:
- **FQTN:** `schema.table` within one database (e.g. `main.orders`); a bare `table` resolves against `main`. A second, attached database is `db.schema.table` (e.g. `sales.main.orders`).
- **Identifiers:** case-insensitive; double-quote (`"Name"`) to keep a name with spaces or a reserved word.
- **Date/time:** native `DATE`/`TIMESTAMP`. Bucket with `date_trunc('month', ts)`, pull parts with `EXTRACT(YEAR FROM ts)`, format with `strftime(ts, '%Y-%m')`, and use `CURRENT_DATE`; cast text with `col::DATE` (or `TRY_CAST(col AS DATE)` to null bad values).
- **Series:** `FROM generate_series(DATE '2023-01-01', DATE '2023-12-01', INTERVAL 1 MONTH) AS s(d)` builds a date spine (use `range(...)` for integers), then `LEFT JOIN` the aggregated facts onto it so empty periods still appear.
- **Rolling window over time:** a native calendar-range frame spans real dates and tolerates gaps — `AVG(amount) OVER (ORDER BY day RANGE BETWEEN INTERVAL 29 DAYS PRECEDING AND CURRENT ROW)` is a trailing 30-day average without a spine; guard minimum periods with `COUNT(*) OVER (<same frame>)`.
- **Integer division:** unlike postgres, `/` is true division (`5 / 2``2.5`), so a ratio keeps its fraction; use `//` for floor division (`5 // 2``2`) when you want the integer quotient, and round only in the final projection.
- **Safe cast:** duckdb has `TRY_CAST``TRY_CAST(x AS DOUBLE)` yields `NULL` for a value that does not parse instead of raising, so counting residual `NULL`s among non-sentinel rows catches an encoding the sample missed without a regex guard.
- **Top-N / windows:** filter a window inline with `QUALIFY``SELECT ... QUALIFY ROW_NUMBER() OVER (PARTITION BY key ORDER BY x DESC) = 1` returns one row per key without a wrapping CTE; use `ORDER BY ... LIMIT n` for a global top-N.
- **JSON / semi-structured:** `col->'k'` returns JSON, `col->>'k'` returns text, deep path `json_extract(col, '$.a.b')`; duckdb also has native `STRUCT`, `LIST`, and `MAP` — read a struct field with `col.field` and a list element with `col[1]` (1-indexed).

View file

@ -1,3 +1,5 @@
import { createAthenaLiveDatabaseIntrospection } from './connectors/athena/live-database-introspection.js';
import { isKtxAthenaConnectionConfig } from './connectors/athena/connector.js';
import { createBigQueryLiveDatabaseIntrospection } from './connectors/bigquery/live-database-introspection.js'; import { createBigQueryLiveDatabaseIntrospection } from './connectors/bigquery/live-database-introspection.js';
import { isKtxBigQueryConnectionConfig, KtxBigQueryScanConnector, type KtxBigQueryConnectionConfig } from './connectors/bigquery/connector.js'; import { isKtxBigQueryConnectionConfig, KtxBigQueryScanConnector, type KtxBigQueryConnectionConfig } from './connectors/bigquery/connector.js';
import { createClickHouseLiveDatabaseIntrospection } from './connectors/clickhouse/live-database-introspection.js'; import { createClickHouseLiveDatabaseIntrospection } from './connectors/clickhouse/live-database-introspection.js';
@ -9,6 +11,8 @@ import { isKtxPostgresConnectionConfig, type KtxPostgresConnectionConfig } from
import { KtxPostgresHistoricSqlQueryClient } from './connectors/postgres/historic-sql-query-client.js'; import { KtxPostgresHistoricSqlQueryClient } from './connectors/postgres/historic-sql-query-client.js';
import { createSqliteLiveDatabaseIntrospection } from './connectors/sqlite/live-database-introspection.js'; import { createSqliteLiveDatabaseIntrospection } from './connectors/sqlite/live-database-introspection.js';
import { isKtxSqliteConnectionConfig } from './connectors/sqlite/connector.js'; import { isKtxSqliteConnectionConfig } from './connectors/sqlite/connector.js';
import { createDuckDbLiveDatabaseIntrospection } from './connectors/duckdb/live-database-introspection.js';
import { isKtxDuckDbConnectionConfig } from './connectors/duckdb/connector.js';
import { createSqlServerLiveDatabaseIntrospection } from './connectors/sqlserver/live-database-introspection.js'; import { createSqlServerLiveDatabaseIntrospection } from './connectors/sqlserver/live-database-introspection.js';
import { isKtxSqlServerConnectionConfig } from './connectors/sqlserver/connector.js'; import { isKtxSqlServerConnectionConfig } from './connectors/sqlserver/connector.js';
import { BigQueryHistoricSqlQueryHistoryReader } from './context/ingest/adapters/historic-sql/bigquery-query-history-reader.js'; import { BigQueryHistoricSqlQueryHistoryReader } from './context/ingest/adapters/historic-sql/bigquery-query-history-reader.js';
@ -104,6 +108,10 @@ function createKtxCliLiveDatabaseIntrospection(
projectDir: project.projectDir, projectDir: project.projectDir,
connections: project.config.connections, connections: project.config.connections,
}); });
const duckdb = createDuckDbLiveDatabaseIntrospection({
projectDir: project.projectDir,
connections: project.config.connections,
});
const mysql = createMysqlLiveDatabaseIntrospection({ const mysql = createMysqlLiveDatabaseIntrospection({
connections: project.config.connections, connections: project.config.connections,
}); });
@ -119,6 +127,9 @@ function createKtxCliLiveDatabaseIntrospection(
const bigquery = createBigQueryLiveDatabaseIntrospection({ const bigquery = createBigQueryLiveDatabaseIntrospection({
connections: project.config.connections, connections: project.config.connections,
}); });
const athena = createAthenaLiveDatabaseIntrospection({
connections: project.config.connections,
});
return { return {
async extractSchema(connectionId: string, options?: LiveDatabaseIntrospectionOptions) { async extractSchema(connectionId: string, options?: LiveDatabaseIntrospectionOptions) {
const connection = project.config.connections[connectionId]; const connection = project.config.connections[connectionId];
@ -139,6 +150,9 @@ function createKtxCliLiveDatabaseIntrospection(
if (isKtxSqliteConnectionConfig(connection)) { if (isKtxSqliteConnectionConfig(connection)) {
return sqlite.extractSchema(connectionId, options); return sqlite.extractSchema(connectionId, options);
} }
if (isKtxDuckDbConnectionConfig(connection)) {
return duckdb.extractSchema(connectionId, options);
}
if (isKtxMysqlConnectionConfig(connection)) { if (isKtxMysqlConnectionConfig(connection)) {
return mysql.extractSchema(connectionId, options); return mysql.extractSchema(connectionId, options);
} }
@ -151,6 +165,9 @@ function createKtxCliLiveDatabaseIntrospection(
if (isKtxBigQueryConnectionConfig(connection)) { if (isKtxBigQueryConnectionConfig(connection)) {
return bigquery.extractSchema(connectionId, options); return bigquery.extractSchema(connectionId, options);
} }
if (isKtxAthenaConnectionConfig(connection)) {
return athena.extractSchema(connectionId, options);
}
if (hasSnowflakeDriver(connection)) { if (hasSnowflakeDriver(connection)) {
const { createSnowflakeLiveDatabaseIntrospection } = await import('./connectors/snowflake/live-database-introspection.js'); const { createSnowflakeLiveDatabaseIntrospection } = await import('./connectors/snowflake/live-database-introspection.js');
const { isKtxSnowflakeConnectionConfig } = await import('./connectors/snowflake/connector.js');; const { isKtxSnowflakeConnectionConfig } = await import('./connectors/snowflake/connector.js');;

View file

@ -64,12 +64,14 @@ const execFileAsync = promisify(execFileCallback);
export type KtxSetupDatabaseDriver = export type KtxSetupDatabaseDriver =
| 'sqlite' | 'sqlite'
| 'duckdb'
| 'postgres' | 'postgres'
| 'mysql' | 'mysql'
| 'clickhouse' | 'clickhouse'
| 'sqlserver' | 'sqlserver'
| 'bigquery' | 'bigquery'
| 'snowflake' | 'snowflake'
| 'athena'
| 'mongodb'; | 'mongodb';
export interface KtxSetupDatabasesArgs { export interface KtxSetupDatabasesArgs {
@ -157,8 +159,10 @@ const DRIVER_OPTIONS: Array<{ value: KtxSetupDatabaseDriver; label: string }> =
{ value: 'mysql', label: 'MySQL' }, { value: 'mysql', label: 'MySQL' },
{ value: 'clickhouse', label: 'ClickHouse' }, { value: 'clickhouse', label: 'ClickHouse' },
{ value: 'sqlserver', label: 'SQL Server' }, { value: 'sqlserver', label: 'SQL Server' },
{ value: 'athena', label: 'Amazon Athena' },
{ value: 'mongodb', label: 'MongoDB' }, { value: 'mongodb', label: 'MongoDB' },
{ value: 'sqlite', label: 'SQLite' }, { value: 'sqlite', label: 'SQLite' },
{ value: 'duckdb', label: 'DuckDB' },
]; ];
const DRIVER_LABELS = Object.fromEntries(DRIVER_OPTIONS.map((option) => [option.value, option.label])) as Record< const DRIVER_LABELS = Object.fromEntries(DRIVER_OPTIONS.map((option) => [option.value, option.label])) as Record<
@ -174,12 +178,14 @@ const HISTORIC_SQL_DIALECT_BY_DRIVER: Partial<Record<KtxSetupDatabaseDriver, His
const DEFAULT_CONNECTION_IDS: Record<KtxSetupDatabaseDriver, string> = { const DEFAULT_CONNECTION_IDS: Record<KtxSetupDatabaseDriver, string> = {
sqlite: 'sqlite-local', sqlite: 'sqlite-local',
duckdb: 'duckdb-local',
postgres: 'postgres-warehouse', postgres: 'postgres-warehouse',
mysql: 'mysql-warehouse', mysql: 'mysql-warehouse',
clickhouse: 'clickhouse-warehouse', clickhouse: 'clickhouse-warehouse',
sqlserver: 'sqlserver-warehouse', sqlserver: 'sqlserver-warehouse',
bigquery: 'bigquery-warehouse', bigquery: 'bigquery-warehouse',
snowflake: 'snowflake-warehouse', snowflake: 'snowflake-warehouse',
athena: 'athena-warehouse',
mongodb: 'mongodb-source', mongodb: 'mongodb-source',
}; };
@ -265,6 +271,13 @@ const SCOPE_DISCOVERY_SPECS: Partial<Record<KtxSetupDatabaseDriver, ScopeDiscove
configSingleField: 'schema_name', configSingleField: 'schema_name',
suggest: defaultSuggest, suggest: defaultSuggest,
}, },
athena: {
noun: 'database',
nounPlural: 'databases',
promptLabel: 'Glue databases',
configArrayField: 'databases',
suggest: defaultSuggest,
},
}; };
type UrlDriverType = Extract<KtxSetupDatabaseDriver, 'postgres' | 'mysql' | 'clickhouse' | 'sqlserver'>; type UrlDriverType = Extract<KtxSetupDatabaseDriver, 'postgres' | 'mysql' | 'clickhouse' | 'sqlserver'>;
@ -813,6 +826,18 @@ async function buildConnectionConfig(input: {
if (path === undefined) return 'back'; if (path === undefined) return 'back';
return path ? { driver: 'sqlite', path } : null; return path ? { driver: 'sqlite', path } : null;
} }
if (driver === 'duckdb') {
if (args.inputMode === 'disabled' && !args.databaseUrl) return null;
const path =
args.databaseUrl ??
(await promptText(
prompts,
'DuckDB database file\nEnter a relative or absolute path, for example ./warehouse.duckdb.',
stringConfigField(input.existingConnection, 'path'),
));
if (path === undefined) return 'back';
return path ? { driver: 'duckdb', path } : null;
}
if (driver === 'postgres' || driver === 'mysql' || driver === 'clickhouse' || driver === 'sqlserver') { if (driver === 'postgres' || driver === 'mysql' || driver === 'clickhouse' || driver === 'sqlserver') {
return await buildUrlConnectionConfig({ return await buildUrlConnectionConfig({
driver, driver,
@ -953,6 +978,47 @@ async function buildConnectionConfig(input: {
...(role ? { role } : {}), ...(role ? { role } : {}),
}; };
} }
if (driver === 'athena') {
if (args.inputMode === 'disabled' && !args.databaseUrl) return null;
const region = await promptText(
prompts,
'AWS region\nFor example us-east-1.',
stringConfigField(input.existingConnection, 'region'),
);
if (region === undefined) return 'back';
if (!region) return null;
const s3StagingDir = await promptText(
prompts,
'S3 staging directory\nAthena writes query results here. For example s3://my-bucket/athena-results/.',
stringConfigField(input.existingConnection, 's3_staging_dir'),
);
if (s3StagingDir === undefined) return 'back';
if (!s3StagingDir) return null;
const workgroup = await promptText(
prompts,
'Athena workgroup (optional)\nPress Enter to use the default workgroup "primary".',
stringConfigField(input.existingConnection, 'workgroup'),
);
if (workgroup === undefined) return 'back';
const catalog = await promptText(
prompts,
'Glue Data Catalog name (optional)\nPress Enter to use the default "AwsDataCatalog".',
stringConfigField(input.existingConnection, 'catalog'),
);
if (catalog === undefined) return 'back';
return {
driver: 'athena',
region,
s3_staging_dir: s3StagingDir,
...(workgroup ? { workgroup } : {}),
...(catalog ? { catalog } : {}),
...scriptedScopeConfigForDriver('athena', args.databaseSchemas),
};
}
throw new Error(`Unsupported database driver: ${driver}`); throw new Error(`Unsupported database driver: ${driver}`);
} }
@ -1417,9 +1483,11 @@ async function maybeConfigureDatabaseScope(input: {
const project = await loadKtxProject({ projectDir: input.projectDir }); const project = await loadKtxProject({ projectDir: input.projectDir });
const connection = project.config.connections[input.connectionId]; const connection = project.config.connections[input.connectionId];
const driver = normalizeDriver(connection?.driver); const driver = normalizeDriver(connection?.driver);
if (!driver || driver === 'sqlite') return okValidateResult(); const spec = driver ? SCOPE_DISCOVERY_SPECS[driver] : undefined;
// Drivers with no scope spec are single-namespace (sqlite, duckdb): there is no
// schema to choose, so skip the scope picker and ingest every table.
if (!driver || !spec) return okValidateResult();
const spec = SCOPE_DISCOVERY_SPECS[driver];
const existingTables = connection?.enabled_tables; const existingTables = connection?.enabled_tables;
const hasExistingTables = Array.isArray(existingTables) && existingTables.length > 0; const hasExistingTables = Array.isArray(existingTables) && existingTables.length > 0;
const existingScope = spec ? configuredScopeValues(connection, spec) : []; const existingScope = spec ? configuredScopeValues(connection, spec) : [];

View file

@ -1,6 +1,6 @@
import { executeFederatedQuery } from './connectors/duckdb/federated-executor.js'; import { executeFederatedQuery } from './connectors/duckdb/federated-executor.js';
import { FEDERATED_CONNECTION_ID } from './context/connections/federation.js'; import { FEDERATED_CONNECTION_ID } from './context/connections/federation.js';
import { executeProjectReadOnlySql } from './context/connections/project-sql-executor.js'; import { executeProjectRawSql } from './context/connections/project-sql-executor.js';
import type { KtxSqlQueryExecutionResult } from './context/connections/query-executor.js'; import type { KtxSqlQueryExecutionResult } from './context/connections/query-executor.js';
import { assertSqlQueryableConnection } from './context/connections/dialects.js'; import { assertSqlQueryableConnection } from './context/connections/dialects.js';
import { resolveConfiguredConnection } from './context/connections/resolve-connection.js'; import { resolveConfiguredConnection } from './context/connections/resolve-connection.js';
@ -137,6 +137,8 @@ export async function runKtxSql(args: KtxSqlArgs, io: KtxCliIo = process, deps:
const connection = isFederated ? undefined : resolveConfiguredConnection(project.config, args.connectionId); const connection = isFederated ? undefined : resolveConfiguredConnection(project.config, args.connectionId);
driver = isFederated ? 'duckdb' : String(connection?.driver ?? 'unknown').toLowerCase(); driver = isFederated ? 'duckdb' : String(connection?.driver ?? 'unknown').toLowerCase();
demoConnection = isFederated ? false : isDemoConnection(args.connectionId, connection); demoConnection = isFederated ? false : isDemoConnection(args.connectionId, connection);
// Fail fast before creating the SQL-analysis daemon port; executeProjectRawSql
// re-asserts this for every caller.
if (!isFederated) { if (!isFederated) {
assertSqlQueryableConnection(args.connectionId, connection?.driver); assertSqlQueryableConnection(args.connectionId, connection?.driver);
} }
@ -152,26 +154,19 @@ export async function runKtxSql(args: KtxSqlArgs, io: KtxCliIo = process, deps:
})); }));
const analysisPort = createSqlAnalysis(); const analysisPort = createSqlAnalysis();
const dialect: SqlAnalysisDialect = isFederated ? 'duckdb' : sqlAnalysisDialectForDriver(connection?.driver); const dialect: SqlAnalysisDialect = isFederated ? 'duckdb' : sqlAnalysisDialectForDriver(connection?.driver);
const validation = await analysisPort.validateReadOnly(args.sql, dialect);
if (!validation.ok) {
throw new Error(validation.error ?? 'SQL is not read-only.');
}
const referencedTableCount = await safeReferencedTableCount(analysisPort, args.sql, dialect);
const createScanConnector = deps.createScanConnector ?? createKtxCliScanConnector; const createScanConnector = deps.createScanConnector ?? createKtxCliScanConnector;
const result = await executeProjectReadOnlySql({ const result = await executeProjectRawSql({
project, project,
input: {
connectionId: args.connectionId, connectionId: args.connectionId,
projectDir: args.projectDir,
connection,
sql: args.sql, sql: args.sql,
maxRows: args.maxRows, maxRows: args.maxRows,
}, sqlAnalysis: analysisPort,
createConnector: (connectionId) => createScanConnector(project!, connectionId), createConnector: (connectionId) => createScanConnector(project!, connectionId),
executeFederated: deps.executeFederated, executeFederated: deps.executeFederated,
runId: 'cli-sql', runId: 'cli-sql',
}); });
const referencedTableCount = await safeReferencedTableCount(analysisPort, args.sql, dialect);
const mode = resolveOutputMode({ explicit: args.output, json: args.json, io }); const mode = resolveOutputMode({ explicit: args.output, json: args.json, io });
printSqlResult(resultOutput(args.connectionId, result), mode, io); printSqlResult(resultOutput(args.connectionId, result), mode, io);
await emitTelemetryEvent({ await emitTelemetryEvent({

View file

@ -412,6 +412,7 @@ function buildConnectionStatus(
const hint = envHint((conn as Record<string, unknown>).credentials_json); const hint = envHint((conn as Record<string, unknown>).credentials_json);
return warn(hint ? `credentials missing (env: ${hint})` : 'credentials not set', hint ? `Set ${hint}` : 'Rerun `ktx setup`'); return warn(hint ? `credentials missing (env: ${hint})` : 'credentials not set', hint ? `Set ${hint}` : 'Rerun `ktx setup`');
} }
case 'duckdb':
case 'sqlite': { case 'sqlite': {
const path = (conn as Record<string, unknown>).path; const path = (conn as Record<string, unknown>).path;
if (typeof path === 'string' && path.length > 0) return ok(`path: ${path}`); if (typeof path === 'string' && path.length > 0) return ok(`path: ${path}`);
@ -553,7 +554,7 @@ async function buildQueryHistoryStatus(
} }
const ADAPTER_DRIVER_REQUIREMENT: Record<string, string[]> = { const ADAPTER_DRIVER_REQUIREMENT: Record<string, string[]> = {
'live-database': ['postgres', 'mysql', 'snowflake', 'bigquery', 'clickhouse', 'sqlite', 'sqlserver'], 'live-database': ['postgres', 'mysql', 'snowflake', 'bigquery', 'clickhouse', 'sqlite', 'duckdb', 'sqlserver'],
dbt: ['dbt', 'dbt-core', 'dbt-cloud'], dbt: ['dbt', 'dbt-core', 'dbt-cloud'],
notion: ['notion'], notion: ['notion'],
metabase: ['metabase'], metabase: ['metabase'],

View file

@ -688,7 +688,7 @@ describe('runKtxConnection', () => {
await initKtxProject({ projectDir }); await initKtxProject({ projectDir });
await writeFile( await writeFile(
join(projectDir, 'ktx.yaml'), join(projectDir, 'ktx.yaml'),
'connections:\n mystery:\n driver: duckdb\n', 'connections:\n mystery:\n driver: nonsense\n',
'utf-8', 'utf-8',
); );
const io = makeIo(); const io = makeIo();

View file

@ -0,0 +1,630 @@
import { describe, expect, it, vi } from 'vitest';
import {
athenaConnectionConfigFromConfig,
isKtxAthenaConnectionConfig,
KtxAthenaScanConnector,
type KtxAthenaClientFactory,
type KtxAthenaClient,
type KtxGlueClient,
} from '../../../src/connectors/athena/connector.js';
import { createAthenaLiveDatabaseIntrospection } from '../../../src/connectors/athena/live-database-introspection.js';
import { tableRefSet } from '../../../src/context/scan/table-ref.js';
function fakeClientFactory(options: { queryState?: string; queryError?: string } = {}): KtxAthenaClientFactory {
const state = options.queryState ?? 'SUCCEEDED';
const queries = new Map<string, string>();
let execCounter = 0;
const fakeAthenaClient: KtxAthenaClient = {
startQueryExecution: vi.fn(async (input) => {
const id = `exec-${++execCounter}`;
queries.set(id, input.QueryString);
return { QueryExecutionId: id };
}),
getQueryExecution: vi.fn(async () => ({
QueryExecution: {
Status: {
State: state,
StateChangeReason: options.queryError,
},
},
})),
getQueryResults: vi.fn(async (input) => {
const sql = queries.get(input.QueryExecutionId) ?? '';
// Column sample query: single-column result for the queried column only.
if (sql.includes('IS NOT NULL')) {
return {
ResultSet: {
ResultSetMetadata: { ColumnInfo: [{ Name: 'status', Type: 'string' }] },
Rows: [
{ Data: [{ VarCharValue: 'status' }] }, // header row
{ Data: [{ VarCharValue: 'paid' }] },
],
},
NextToken: undefined,
};
}
return {
ResultSet: {
ResultSetMetadata: {
ColumnInfo: [
{ Name: 'id', Type: 'bigint' },
{ Name: 'status', Type: 'string' },
],
},
Rows: [
// Header row (Athena always includes it on first page)
{ Data: [{ VarCharValue: 'id' }, { VarCharValue: 'status' }] },
// Data row
{ Data: [{ VarCharValue: '1' }, { VarCharValue: 'paid' }] },
],
},
NextToken: undefined,
};
}),
};
const fakeGlueClient: KtxGlueClient = {
getDatabases: vi.fn(async () => ({
DatabaseList: [{ Name: 'analytics' }],
NextToken: undefined,
})),
getTables: vi.fn(async () => ({
TableList: [
{
Name: 'orders',
TableType: 'EXTERNAL_TABLE',
Description: 'Orders table',
StorageDescriptor: {
Columns: [
{ Name: 'id', Type: 'bigint', Comment: 'Order id' },
{ Name: 'status', Type: 'string' },
],
},
PartitionKeys: [{ Name: 'dt', Type: 'date', Comment: 'Partition date' }],
},
],
NextToken: undefined,
})),
};
return {
createAthenaClient: vi.fn(() => fakeAthenaClient),
createGlueClient: vi.fn(() => fakeGlueClient),
};
}
const connection = {
driver: 'athena',
region: 'us-east-1',
s3_staging_dir: 's3://my-bucket/athena-results/',
workgroup: 'analytics',
catalog: 'AwsDataCatalog',
database: 'analytics',
} as const;
describe('KtxAthenaScanConnector', () => {
it('identifies athena connection configs correctly', () => {
expect(isKtxAthenaConnectionConfig(connection)).toBe(true);
expect(isKtxAthenaConnectionConfig({ driver: 'bigquery' })).toBe(false);
expect(isKtxAthenaConnectionConfig(null)).toBe(false);
expect(isKtxAthenaConnectionConfig(undefined)).toBe(false);
});
it('resolves configuration and throws on missing required fields', () => {
expect(athenaConnectionConfigFromConfig({ connectionId: 'dw', connection })).toMatchObject({
region: 'us-east-1',
s3StagingDir: 's3://my-bucket/athena-results/',
workgroup: 'analytics',
catalog: 'AwsDataCatalog',
database: 'analytics',
});
expect(() =>
athenaConnectionConfigFromConfig({ connectionId: 'dw', connection: { driver: 'athena' } }),
).toThrow('connections.dw.region');
expect(() =>
athenaConnectionConfigFromConfig({
connectionId: 'dw',
connection: { driver: 'athena', region: 'us-east-1' },
}),
).toThrow('connections.dw.s3_staging_dir');
});
it('applies defaults for optional config fields', () => {
const resolved = athenaConnectionConfigFromConfig({
connectionId: 'dw',
connection: { driver: 'athena', region: 'us-east-1', s3_staging_dir: 's3://bucket/' },
});
expect(resolved.workgroup).toBe('primary');
expect(resolved.catalog).toBe('AwsDataCatalog');
expect(resolved.database).toBeUndefined();
});
it('introspects databases, tables, and columns from Glue', async () => {
const connector = new KtxAthenaScanConnector({
connectionId: 'dw',
connection,
clientFactory: fakeClientFactory(),
now: () => new Date('2026-06-21T10:00:00.000Z'),
});
const snapshot = await connector.introspect(
{ connectionId: 'dw', driver: 'athena' },
{ runId: 'scan-1' },
);
expect(snapshot).toMatchObject({
connectionId: 'dw',
driver: 'athena',
extractedAt: '2026-06-21T10:00:00.000Z',
scope: { catalogs: ['AwsDataCatalog'], datasets: ['analytics'] },
metadata: {
catalog: 'AwsDataCatalog',
databases: ['analytics'],
table_count: 1,
total_columns: 3,
},
});
expect(snapshot.tables[0]).toMatchObject({
catalog: 'AwsDataCatalog',
db: 'analytics',
name: 'orders',
kind: 'table',
comment: 'Orders table',
estimatedRows: null,
foreignKeys: [],
});
expect(snapshot.tables[0]?.columns).toEqual([
{
name: 'id',
nativeType: 'bigint',
normalizedType: 'BIGINT',
dimensionType: 'number',
nullable: true,
primaryKey: false,
comment: 'Order id',
},
{
name: 'status',
nativeType: 'string',
normalizedType: 'VARCHAR',
dimensionType: 'string',
nullable: true,
primaryKey: false,
comment: null,
},
{
name: 'dt',
nativeType: 'date',
normalizedType: 'DATE',
dimensionType: 'time',
nullable: true,
primaryKey: false,
comment: 'Partition date',
},
]);
});
it('respects tableScope and excludes tables not in scope', async () => {
const connector = new KtxAthenaScanConnector({
connectionId: 'dw',
connection,
clientFactory: fakeClientFactory(),
now: () => new Date('2026-06-21T10:00:00.000Z'),
});
const scopedSnapshot = await connector.introspect(
{
connectionId: 'dw',
driver: 'athena',
tableScope: tableRefSet([{ catalog: 'AwsDataCatalog', db: 'analytics', name: 'nonexistent' }]),
},
{ runId: 'scan-1' },
);
expect(scopedSnapshot.tables).toHaveLength(0);
const matchingSnapshot = await connector.introspect(
{
connectionId: 'dw',
driver: 'athena',
tableScope: tableRefSet([{ catalog: 'AwsDataCatalog', db: 'analytics', name: 'orders' }]),
},
{ runId: 'scan-1' },
);
expect(matchingSnapshot.tables).toHaveLength(1);
expect(matchingSnapshot.tables[0]?.name).toBe('orders');
});
it('limits introspection to the configured databases scope', async () => {
const requestedDatabases: string[] = [];
const getDatabases = vi.fn(async () => ({
DatabaseList: [{ Name: 'analytics' }, { Name: 'raw' }, { Name: 'staging' }],
NextToken: undefined,
}));
const glueClient: KtxGlueClient = {
getDatabases,
getTables: vi.fn(async (input) => {
requestedDatabases.push(input.DatabaseName);
return {
TableList: [
{
Name: `${input.DatabaseName}_orders`,
TableType: 'EXTERNAL_TABLE',
StorageDescriptor: { Columns: [{ Name: 'id', Type: 'bigint' }] },
},
],
NextToken: undefined,
};
}),
};
const clientFactory: KtxAthenaClientFactory = {
createAthenaClient: vi.fn(() => fakeClientFactory().createAthenaClient('us-east-1')),
createGlueClient: vi.fn(() => glueClient),
};
const connector = new KtxAthenaScanConnector({
connectionId: 'dw',
connection: { ...connection, databases: ['analytics', 'raw'] },
clientFactory,
now: () => new Date('2026-06-21T10:00:00.000Z'),
});
const snapshot = await connector.introspect({ connectionId: 'dw', driver: 'athena' }, { runId: 'scan-1' });
// Scope is taken from config, so the account-wide database list is never enumerated.
expect(getDatabases).not.toHaveBeenCalled();
expect(requestedDatabases).toEqual(['analytics', 'raw']);
expect(snapshot.scope).toMatchObject({ datasets: ['analytics', 'raw'] });
expect(snapshot.tables.map((t) => t.db)).toEqual(['analytics', 'raw']);
});
it('resolves optional env-referenced config to defaults when the variable is unset', () => {
const resolved = athenaConnectionConfigFromConfig({
connectionId: 'dw',
connection: {
driver: 'athena',
region: 'us-east-1',
s3_staging_dir: 's3://bucket/',
workgroup: 'env:ATHENA_WORKGROUP_UNSET',
catalog: 'env:GLUE_CATALOG_UNSET',
},
env: {},
});
expect(resolved.workgroup).toBe('primary');
expect(resolved.catalog).toBe('AwsDataCatalog');
});
it('samples a table via Athena query execution', async () => {
const connector = new KtxAthenaScanConnector({
connectionId: 'dw',
connection,
clientFactory: fakeClientFactory(),
});
const result = await connector.sampleTable(
{
connectionId: 'dw',
table: { catalog: 'AwsDataCatalog', db: 'analytics', name: 'orders' },
columns: ['id', 'status'],
limit: 10,
},
{ runId: 'scan-1' },
);
expect(result).toMatchObject({
headers: ['id', 'status'],
rows: [['1', 'paid']],
totalRows: 1,
});
});
it('samples a column via Athena query execution', async () => {
const connector = new KtxAthenaScanConnector({
connectionId: 'dw',
connection,
clientFactory: fakeClientFactory(),
});
const result = await connector.sampleColumn(
{
connectionId: 'dw',
table: { catalog: 'AwsDataCatalog', db: 'analytics', name: 'orders' },
column: 'status',
limit: 10,
},
{ runId: 'scan-1' },
);
expect(result).toMatchObject({
values: ['paid'],
nullCount: null,
distinctCount: null,
});
});
it('executes read-only SQL and rejects write statements', async () => {
const connector = new KtxAthenaScanConnector({
connectionId: 'dw',
connection,
clientFactory: fakeClientFactory(),
});
await expect(
connector.executeReadOnly(
{ connectionId: 'dw', sql: 'SELECT id, status FROM "analytics"."orders"', maxRows: 100 },
{ runId: 'scan-1' },
),
).resolves.toMatchObject({
headers: ['id', 'status'],
rows: [['1', 'paid']],
rowCount: 1,
});
await expect(
connector.executeReadOnly({ connectionId: 'dw', sql: 'DELETE FROM orders' }, { runId: 'scan-1' }),
).rejects.toThrow('Only read-only SELECT/WITH queries can be executed locally');
});
it('lists schemas (databases) from Glue', async () => {
const connector = new KtxAthenaScanConnector({
connectionId: 'dw',
connection,
clientFactory: fakeClientFactory(),
});
await expect(connector.listSchemas()).resolves.toEqual(['analytics']);
});
it('lists tables from Glue', async () => {
const connector = new KtxAthenaScanConnector({
connectionId: 'dw',
connection,
clientFactory: fakeClientFactory(),
});
await expect(connector.listTables(['analytics'])).resolves.toEqual([
{
catalog: 'AwsDataCatalog',
schema: 'analytics',
name: 'orders',
kind: 'table',
},
]);
});
it('returns null for columnStats', async () => {
const connector = new KtxAthenaScanConnector({
connectionId: 'dw',
connection,
clientFactory: fakeClientFactory(),
});
await expect(
connector.columnStats(
{ connectionId: 'dw', table: { catalog: 'AwsDataCatalog', db: 'analytics', name: 'orders' }, column: 'status' },
{ runId: 'scan-1' },
),
).resolves.toBeNull();
});
it('tests connection successfully', async () => {
const connector = new KtxAthenaScanConnector({
connectionId: 'dw',
connection,
clientFactory: fakeClientFactory(),
});
await expect(connector.testConnection()).resolves.toMatchObject({ success: true });
});
it('returns failure result when testConnection throws', async () => {
const factory = fakeClientFactory();
const glueClient = factory.createGlueClient('us-east-1');
vi.mocked(glueClient.getDatabases).mockRejectedValue(new Error('Access denied'));
const brokenFactory: KtxAthenaClientFactory = {
createAthenaClient: factory.createAthenaClient,
createGlueClient: vi.fn(() => glueClient),
};
const connector = new KtxAthenaScanConnector({
connectionId: 'dw',
connection,
clientFactory: brokenFactory,
});
await expect(connector.testConnection()).resolves.toMatchObject({
success: false,
error: 'Access denied',
});
});
it('cleans up without throwing', async () => {
const connector = new KtxAthenaScanConnector({
connectionId: 'dw',
connection,
clientFactory: fakeClientFactory(),
});
await connector.listSchemas();
await expect(connector.cleanup()).resolves.toBeUndefined();
});
it('throws when query execution fails', async () => {
const connector = new KtxAthenaScanConnector({
connectionId: 'dw',
connection,
clientFactory: fakeClientFactory({ queryState: 'FAILED', queryError: 'Syntax error in SQL' }),
});
await expect(
connector.executeReadOnly({ connectionId: 'dw', sql: 'SELECT 1' }, { runId: 'scan-1' }),
).rejects.toThrow('Athena query FAILED: Syntax error in SQL');
});
it('throws when query execution times out', async () => {
let callCount = 0;
// First now() call sets the deadline; second call simulates time past it.
const now = () => (++callCount === 1 ? new Date(0) : new Date(5 * 60 * 1000 + 1));
const connector = new KtxAthenaScanConnector({
connectionId: 'dw',
connection,
clientFactory: fakeClientFactory({ queryState: 'RUNNING' }),
now,
});
await expect(
connector.executeReadOnly({ connectionId: 'dw', sql: 'SELECT 1' }, { runId: 'scan-1' }),
).rejects.toThrow('timed out after 300s');
});
it('passes the exact column list to Athena when sampling specific columns', async () => {
const factory = fakeClientFactory();
const athenaClient = factory.createAthenaClient('us-east-1');
const connector = new KtxAthenaScanConnector({
connectionId: 'dw',
connection,
clientFactory: { createAthenaClient: vi.fn(() => athenaClient), createGlueClient: factory.createGlueClient },
});
await connector.sampleTable(
{
connectionId: 'dw',
table: { catalog: 'AwsDataCatalog', db: 'analytics', name: 'orders' },
columns: ['id', 'status'],
limit: 5,
},
{ runId: 'scan-1' },
);
expect(vi.mocked(athenaClient.startQueryExecution).mock.calls[0]?.[0].QueryString).toBe(
'SELECT "id", "status" FROM "AwsDataCatalog"."analytics"."orders" LIMIT 5',
);
});
it('paginates Glue databases and tables across multiple pages', async () => {
const glueClient: KtxGlueClient = {
getDatabases: vi.fn()
.mockResolvedValueOnce({ DatabaseList: [{ Name: 'db1' }], NextToken: 'page2' })
.mockResolvedValueOnce({ DatabaseList: [{ Name: 'db2' }], NextToken: undefined }),
getTables: vi.fn().mockImplementation(async ({ DatabaseName }: { DatabaseName: string }) => {
if (DatabaseName === 'db1') {
return {
TableList: [
{
Name: 'table_a',
TableType: 'EXTERNAL_TABLE',
StorageDescriptor: { Columns: [{ Name: 'id', Type: 'bigint' }] },
},
],
NextToken: undefined,
};
}
return {
TableList: [
{
Name: 'table_b',
TableType: 'EXTERNAL_TABLE',
StorageDescriptor: { Columns: [{ Name: 'id', Type: 'bigint' }] },
},
],
NextToken: undefined,
};
}),
};
const connector = new KtxAthenaScanConnector({
connectionId: 'dw',
connection,
clientFactory: {
createAthenaClient: vi.fn(() => fakeClientFactory().createAthenaClient('us-east-1')),
createGlueClient: vi.fn(() => glueClient),
},
now: () => new Date('2026-06-21T10:00:00.000Z'),
});
const snapshot = await connector.introspect({ connectionId: 'dw', driver: 'athena' }, { runId: 'scan-1' });
expect(vi.mocked(glueClient.getDatabases)).toHaveBeenCalledTimes(2);
expect(snapshot.metadata).toMatchObject({ databases: ['db1', 'db2'], table_count: 2 });
expect(snapshot.tables.map((t) => t.name)).toEqual(['table_a', 'table_b']);
});
it('paginates Athena query results across multiple pages', async () => {
const factory = fakeClientFactory();
const athenaClient = factory.createAthenaClient('us-east-1');
vi.mocked(athenaClient.getQueryResults)
.mockResolvedValueOnce({
ResultSet: {
ResultSetMetadata: {
ColumnInfo: [
{ Name: 'id', Type: 'bigint' },
{ Name: 'status', Type: 'string' },
],
},
Rows: [
// Header row — only present on the first page
{ Data: [{ VarCharValue: 'id' }, { VarCharValue: 'status' }] },
{ Data: [{ VarCharValue: '1' }, { VarCharValue: 'paid' }] },
{ Data: [{ VarCharValue: '2' }, { VarCharValue: 'shipped' }] },
],
},
NextToken: 'page-2',
})
.mockResolvedValueOnce({
ResultSet: {
ResultSetMetadata: { ColumnInfo: [] },
// No header row on subsequent pages
Rows: [{ Data: [{ VarCharValue: '3' }, { VarCharValue: 'pending' }] }],
},
NextToken: undefined,
});
const connector = new KtxAthenaScanConnector({
connectionId: 'dw',
connection,
clientFactory: { createAthenaClient: vi.fn(() => athenaClient), createGlueClient: factory.createGlueClient },
});
const result = await connector.executeReadOnly(
{ connectionId: 'dw', sql: 'SELECT id, status FROM "analytics"."orders"', maxRows: 100 },
{ runId: 'scan-1' },
);
expect(result.headers).toEqual(['id', 'status']);
expect(result.rows).toEqual([
['1', 'paid'],
['2', 'shipped'],
['3', 'pending'],
]);
expect(result.rowCount).toBe(3);
expect(vi.mocked(athenaClient.getQueryResults)).toHaveBeenCalledTimes(2);
expect(vi.mocked(athenaClient.getQueryResults).mock.calls[1]?.[0].NextToken).toBe('page-2');
});
it('adapts to the live-database introspection port via factory', async () => {
const introspection = createAthenaLiveDatabaseIntrospection({
connections: { dw: connection },
clientFactory: fakeClientFactory(),
now: () => new Date('2026-06-21T10:00:00.000Z'),
});
await expect(introspection.extractSchema('dw')).resolves.toMatchObject({
connectionId: 'dw',
driver: 'athena',
metadata: { catalog: 'AwsDataCatalog' },
tables: expect.arrayContaining([
expect.objectContaining({
db: 'analytics',
name: 'orders',
columns: expect.arrayContaining([
expect.objectContaining({ name: 'id', dimensionType: 'number' }),
]),
}),
]),
});
});
});

View file

@ -0,0 +1,72 @@
import { describe, expect, it } from 'vitest';
import { KtxAthenaDialect } from '../../../src/connectors/athena/dialect.js';
describe('KtxAthenaDialect', () => {
const dialect = new KtxAthenaDialect();
it('quotes identifiers and formats catalog.database.table names', () => {
expect(dialect.quoteIdentifier('my"col')).toBe('"my""col"');
expect(dialect.formatTableName({ catalog: 'AwsDataCatalog', db: 'analytics', name: 'orders' })).toBe(
'"AwsDataCatalog"."analytics"."orders"',
);
expect(dialect.formatTableName({ db: 'analytics', name: 'orders' })).toBe('"analytics"."orders"');
expect(dialect.formatTableName({ name: 'orders' })).toBe('"orders"');
});
it('maps native Athena/Glue types to normalized types and dimension types', () => {
expect(dialect.mapDataType('bigint')).toBe('BIGINT');
expect(dialect.mapDataType('string')).toBe('VARCHAR');
expect(dialect.mapDataType('array<string>')).toBe('ARRAY');
expect(dialect.mapDataType('map<string,bigint>')).toBe('MAP');
expect(dialect.mapDataType('struct<id:bigint>')).toBe('STRUCT');
expect(dialect.mapDataType('decimal(18,2)')).toBe('DECIMAL');
expect(dialect.mapDataType('UNKNOWN_TYPE')).toBe('UNKNOWN_TYPE');
expect(dialect.mapToDimensionType('timestamp')).toBe('time');
expect(dialect.mapToDimensionType('date')).toBe('time');
expect(dialect.mapToDimensionType('bigint')).toBe('number');
expect(dialect.mapToDimensionType('double')).toBe('number');
expect(dialect.mapToDimensionType('decimal(10,2)')).toBe('number');
expect(dialect.mapToDimensionType('boolean')).toBe('boolean');
expect(dialect.mapToDimensionType('string')).toBe('string');
expect(dialect.mapToDimensionType('varchar')).toBe('string');
});
it('generates correct sample and column-sample SQL', () => {
expect(dialect.generateSampleQuery('"analytics"."orders"', 10, ['id', 'status'])).toBe(
'SELECT "id", "status" FROM "analytics"."orders" LIMIT 10',
);
expect(dialect.generateSampleQuery('"analytics"."orders"', 5)).toBe(
'SELECT * FROM "analytics"."orders" LIMIT 5',
);
expect(dialect.generateColumnSampleQuery('"analytics"."orders"', 'status', 20)).toBe(
'SELECT "status" FROM "analytics"."orders" WHERE "status" IS NOT NULL LIMIT 20',
);
});
it('generates Presto-style cardinality and distinct-values SQL', () => {
expect(dialect.generateCardinalitySampleQuery('"t"', '"col"', 1000)).toContain('approx_distinct');
expect(dialect.generateRandomizedCardinalitySampleQuery('"t"', '"col"', 500)).toContain('rand()');
expect(dialect.generateDistinctValuesQuery('"t"', '"col"', 50)).toContain(
'SELECT DISTINCT CAST("col" AS VARCHAR) AS val',
);
});
it('returns null for column statistics (unsupported)', () => {
expect(dialect.generateColumnStatisticsQuery('analytics', 'orders')).toBeNull();
});
it('produces Trino-correct OFFSET-before-LIMIT ordering', () => {
expect(dialect.getLimitOffsetClause(10)).toBe('LIMIT 10');
expect(dialect.getLimitOffsetClause(10, 0)).toBe('LIMIT 10');
expect(dialect.getLimitOffsetClause(10, 20)).toBe('OFFSET 20 LIMIT 10');
});
it('uses unit-separator (U+001F) as the array_join delimiter', () => {
const sql = dialect.getSampleValueAggregation('SELECT value FROM t');
const separatorIndex =
sql.indexOf("array_join(array_agg(CAST(value AS VARCHAR)), '") +
"array_join(array_agg(CAST(value AS VARCHAR)), '".length;
expect(sql.charCodeAt(separatorIndex)).toBe(0x1f);
});
});

View file

@ -0,0 +1,280 @@
import { DuckDBInstance } from '@duckdb/node-api';
import { mkdtemp, rm } from 'node:fs/promises';
import { tmpdir } from 'node:os';
import { join } from 'node:path';
import { pathToFileURL } from 'node:url';
import { afterAll, beforeAll, describe, expect, it } from 'vitest';
import {
KtxDuckDbScanConnector,
duckDbDatabasePathFromConfig,
isKtxDuckDbConnectionConfig,
} from '../../../src/connectors/duckdb/connector.js';
import { tableRefSet } from '../../../src/context/scan/table-ref.js';
let dir: string;
let dbPath: string;
beforeAll(async () => {
dir = await mkdtemp(join(tmpdir(), 'ktx-duckdb-'));
dbPath = join(dir, 'warehouse.duckdb');
const instance = await DuckDBInstance.create(dbPath);
const connection = await instance.connect();
await connection.run('CREATE TABLE customers (id BIGINT PRIMARY KEY, name VARCHAR, big BIGINT)');
await connection.run(
`INSERT INTO customers VALUES (1, 'Ada', 9223372036854775807), (2, 'Lin', 10)`,
);
await connection.run('CREATE TABLE orders (id BIGINT, customer_id BIGINT REFERENCES customers(id))');
await connection.run('INSERT INTO orders VALUES (1, 1), (2, 2)');
// Composite primary key + composite foreign key, to exercise the parallel
// unnest() zip of constraint/referenced column names in readForeignKeys.
await connection.run('CREATE TABLE regions (country VARCHAR, code VARCHAR, PRIMARY KEY (country, code))');
await connection.run(
'CREATE TABLE stores (id BIGINT, country VARCHAR, code VARCHAR, FOREIGN KEY (country, code) REFERENCES regions(country, code))',
);
await connection.run('CREATE TABLE empty_table (id BIGINT)');
connection.closeSync();
instance.closeSync();
});
afterAll(async () => {
await rm(dir, { recursive: true, force: true });
});
function connector(connection: Record<string, unknown> = { driver: 'duckdb', path: dbPath }) {
return new KtxDuckDbScanConnector({ connectionId: 'warehouse', connection, projectDir: dir });
}
describe('isKtxDuckDbConnectionConfig', () => {
it('accepts duckdb driver, rejects others', () => {
expect(isKtxDuckDbConnectionConfig({ driver: 'duckdb' })).toBe(true);
expect(isKtxDuckDbConnectionConfig({ driver: 'sqlite' })).toBe(false);
});
});
describe('duckDbDatabasePathFromConfig', () => {
it('resolves a relative path against projectDir', () => {
const resolved = duckDbDatabasePathFromConfig({
connectionId: 'warehouse',
projectDir: dir,
connection: { driver: 'duckdb', path: 'warehouse.duckdb' },
});
expect(resolved).toBe(dbPath);
});
it('derives the path from a file: url', () => {
const resolved = duckDbDatabasePathFromConfig({
connectionId: 'warehouse',
connection: { driver: 'duckdb', url: pathToFileURL(dbPath).href },
});
expect(resolved).toBe(dbPath);
});
it('derives the path from a duckdb: url', () => {
const resolved = duckDbDatabasePathFromConfig({
connectionId: 'warehouse',
connection: { driver: 'duckdb', url: `duckdb://${dbPath}` },
});
expect(resolved).toBe(dbPath);
});
it('resolves an env: reference in path', () => {
process.env.KTX_TEST_DUCKDB_PATH = dbPath;
try {
const resolved = duckDbDatabasePathFromConfig({
connectionId: 'warehouse',
connection: { driver: 'duckdb', path: 'env:KTX_TEST_DUCKDB_PATH' },
});
expect(resolved).toBe(dbPath);
} finally {
delete process.env.KTX_TEST_DUCKDB_PATH;
}
});
it('rejects a non-duckdb driver', () => {
expect(() =>
duckDbDatabasePathFromConfig({
connectionId: 'warehouse',
connection: { driver: 'sqlite', path: dbPath },
}),
).toThrow(/cannot run driver "sqlite"/);
});
it('requires a path or url', () => {
expect(() =>
duckDbDatabasePathFromConfig({
connectionId: 'warehouse',
connection: { driver: 'duckdb' },
}),
).toThrow(/requires connections\.warehouse\.path or url/);
});
});
describe('KtxDuckDbScanConnector', () => {
it('testConnection succeeds for an existing file', async () => {
const c = connector();
expect(await c.testConnection()).toEqual({ success: true });
await c.cleanup();
});
it('testConnection fails (never creating) for a missing file', async () => {
const c = connector({ driver: 'duckdb', path: join(dir, 'absent.duckdb') });
const result = await c.testConnection();
expect(result.success).toBe(false);
await c.cleanup();
});
it('introspects main-schema tables, columns, and foreign keys', async () => {
const c = connector();
const snapshot = await c.introspect({ connectionId: 'warehouse', driver: 'duckdb' }, { runId: 't' });
const names = snapshot.tables.map((t) => t.name).sort();
expect(names).toEqual(['customers', 'empty_table', 'orders', 'regions', 'stores']);
const orders = snapshot.tables.find((t) => t.name === 'orders');
expect(orders?.foreignKeys[0]).toMatchObject({ fromColumn: 'customer_id', toTable: 'customers', toColumn: 'id' });
await c.cleanup();
});
it('maps a composite foreign key column-for-column to the referenced table', async () => {
const c = connector();
const snapshot = await c.introspect({ connectionId: 'warehouse', driver: 'duckdb' }, { runId: 't' });
const stores = snapshot.tables.find((t) => t.name === 'stores');
const fks = stores?.foreignKeys.map((fk) => ({ fromColumn: fk.fromColumn, toTable: fk.toTable, toColumn: fk.toColumn }));
expect(fks).toEqual([
{ fromColumn: 'country', toTable: 'regions', toColumn: 'country' },
{ fromColumn: 'code', toTable: 'regions', toColumn: 'code' },
]);
await c.cleanup();
});
it('lists tables', async () => {
const c = connector();
const tables = (await c.listTables()).map((t) => t.name).sort();
expect(tables).toEqual(['customers', 'empty_table', 'orders', 'regions', 'stores']);
await c.cleanup();
});
it('samples a table', async () => {
const c = connector();
const sample = await c.sampleTable(
{ connectionId: 'warehouse', table: { name: 'customers', catalog: null, db: null }, limit: 1 },
{ runId: 't' },
);
expect(sample.rows.length).toBe(1);
await c.cleanup();
});
it('stringifies BIGINT beyond 2^53 in read-only results', async () => {
const c = connector();
const result = await c.executeReadOnly(
{ connectionId: 'warehouse', sql: 'SELECT big FROM customers WHERE id = 1', maxRows: 10 },
{ runId: 't' },
);
expect(result.rows[0][0]).toBe('9223372036854775807');
await c.cleanup();
});
it('rejects non-read-only SQL', async () => {
const c = connector();
await expect(
c.executeReadOnly({ connectionId: 'warehouse', sql: 'DELETE FROM customers', maxRows: 10 }, { runId: 't' }),
).rejects.toThrow();
await c.cleanup();
});
it('returns distinct values under the cardinality cap', async () => {
const c = connector();
const distinct = await c.getColumnDistinctValues({ name: 'customers', catalog: null, db: null }, 'name', {
maxCardinality: 10,
limit: 10,
});
expect(distinct?.values?.sort()).toEqual(['Ada', 'Lin']);
await c.cleanup();
});
it('withholds values but reports the count when cardinality exceeds the cap', async () => {
const c = connector();
const distinct = await c.getColumnDistinctValues({ name: 'customers', catalog: null, db: null }, 'name', {
maxCardinality: 1,
limit: 10,
});
expect(distinct).toEqual({ values: null, cardinality: 2 });
await c.cleanup();
});
it('samples a single column, dropping null rows', async () => {
const c = connector();
const sample = await c.sampleColumn(
{ connectionId: 'warehouse', table: { name: 'customers', catalog: null, db: null }, column: 'name', limit: 10 },
{ runId: 't' },
);
expect(sample.values.sort()).toEqual(['Ada', 'Lin']);
expect(sample.nullCount).toBeNull();
await c.cleanup();
});
it('counts table rows', async () => {
const c = connector();
expect(await c.getTableRowCount('customers')).toBe(2);
await c.cleanup();
});
it('lists only the main schema and reports no column stats', async () => {
const c = connector();
expect(await c.listSchemas()).toEqual(['main']);
expect(await c.columnStats({ connectionId: 'warehouse', table: { name: 'customers', catalog: null, db: null }, column: 'id' }, { runId: 't' })).toBeNull();
await c.cleanup();
});
it('rejects operations for a mismatched connection id', async () => {
const c = connector();
await expect(
c.executeReadOnly({ connectionId: 'other', sql: 'SELECT 1', maxRows: 1 }, { runId: 't' }),
).rejects.toThrow(/cannot serve connection other/);
await c.cleanup();
});
it('exposes the dialect identifier quoting', () => {
expect(connector().quoteIdentifier('a"b')).toBe('"a""b"');
});
// Opening a connection must never create the file: the db() guard throws
// rather than letting DuckDBInstance.create() materialize a missing path.
it('refuses to open (never creating) a missing file when a query runs', async () => {
const c = connector({ driver: 'duckdb', path: join(dir, 'absent.duckdb') });
await expect(c.listTables()).rejects.toThrow(/File not found/);
await c.cleanup();
});
it('returns no tables for an empty table scope', async () => {
const c = connector();
const snapshot = await c.introspect(
{ connectionId: 'warehouse', driver: 'duckdb', tableScope: new Set() },
{ runId: 't' },
);
expect(snapshot.tables).toEqual([]);
await c.cleanup();
});
it('restricts introspection to the named tables in a non-empty scope', async () => {
const c = connector();
const snapshot = await c.introspect(
{
connectionId: 'warehouse',
driver: 'duckdb',
tableScope: tableRefSet([{ catalog: null, db: null, name: 'customers' }]),
},
{ runId: 't' },
);
expect(snapshot.tables.map((t) => t.name)).toEqual(['customers']);
await c.cleanup();
});
it('reports zero cardinality and an empty value list for an empty column', async () => {
const c = connector();
const distinct = await c.getColumnDistinctValues({ name: 'empty_table', catalog: null, db: null }, 'id', {
maxCardinality: 10,
limit: 10,
});
expect(distinct).toEqual({ values: [], cardinality: 0 });
await c.cleanup();
});
});

View file

@ -0,0 +1,108 @@
import { describe, expect, it } from 'vitest';
import { KtxDuckDbDialect } from '../../../src/connectors/duckdb/dialect.js';
describe('KtxDuckDbDialect', () => {
const dialect = new KtxDuckDbDialect();
it('quotes identifiers with double quotes and escapes embedded quotes', () => {
expect(dialect.quoteIdentifier('order"s')).toBe('"order""s"');
});
it('maps integer types to number dimension', () => {
expect(dialect.mapToDimensionType('BIGINT')).toBe('number');
expect(dialect.mapToDimensionType('DOUBLE')).toBe('number');
});
it('maps timestamp types to time dimension', () => {
expect(dialect.mapToDimensionType('TIMESTAMP')).toBe('time');
expect(dialect.mapToDimensionType('DATE')).toBe('time');
});
it('maps text types to string dimension', () => {
expect(dialect.mapToDimensionType('VARCHAR')).toBe('string');
});
it('maps boolean types to boolean dimension', () => {
expect(dialect.mapToDimensionType('BOOLEAN')).toBe('boolean');
expect(dialect.mapToDimensionType('BOOL')).toBe('boolean');
});
it('falls back to string for an empty or unknown native type', () => {
expect(dialect.mapToDimensionType('')).toBe('string');
expect(dialect.mapToDimensionType('JSON')).toBe('string');
});
// The precedence ladder strips parameters before substring rules fire, so a
// parameterized DECIMAL still resolves through the numeric branch rather than
// the string fallback.
it('strips type parameters before resolving the dimension', () => {
expect(dialect.mapToDimensionType('DECIMAL(10,2)')).toBe('number');
expect(dialect.mapToDimensionType('VARCHAR(255)')).toBe('string');
});
// Types absent from the exact-match table still resolve via substring rules:
// TIMESTAMP_NS (time), UINT128/HUGEINT-like (number), and lowercase input.
it('resolves unlisted types through substring matching, case-insensitively', () => {
expect(dialect.mapToDimensionType('timestamp_ns')).toBe('time');
expect(dialect.mapToDimensionType('INT128')).toBe('number');
expect(dialect.mapToDimensionType(' double ')).toBe('number');
});
it('generates a limited sample query', () => {
expect(dialect.generateSampleQuery('"t"', 5)).toBe('SELECT * FROM "t" LIMIT 5');
});
it('quotes selected columns in a sample query', () => {
expect(dialect.generateSampleQuery('"t"', 5, ['a', 'b'])).toBe('SELECT "a", "b" FROM "t" LIMIT 5');
});
it('builds a non-null, non-blank column sample query', () => {
expect(dialect.generateColumnSampleQuery('"t"', 'email', 3)).toBe(
`SELECT "email" FROM "t" WHERE "email" IS NOT NULL AND TRIM(CAST("email" AS VARCHAR)) != '' LIMIT 3`,
);
});
// A degenerate sample percentage (<=0 or >=1) means "no sampling", so both the
// random filter and the TABLESAMPLE clause must collapse to an empty string.
it('returns empty sample clauses outside the (0,1) range and real clauses inside it', () => {
expect(dialect.getRandomSampleFilter(0)).toBe('');
expect(dialect.getRandomSampleFilter(1)).toBe('');
expect(dialect.getRandomSampleFilter(0.25)).toBe('RANDOM() < 0.25');
expect(dialect.getTableSampleClause(0)).toBe('');
expect(dialect.getTableSampleClause(0.1)).toBe('USING SAMPLE 10 PERCENT (bernoulli)');
});
// A type missing from the exact-match table but containing BOOL still resolves
// through the substring branch rather than the string fallback.
it('resolves a BOOL-substring type to boolean', () => {
expect(dialect.mapToDimensionType('MYBOOL')).toBe('boolean');
});
it('builds limit/offset, sample-value aggregation, and randomized cardinality clauses', () => {
expect(dialect.getLimitOffsetClause(10, 5)).toContain('LIMIT 10');
expect(dialect.getSampleValueAggregation('SELECT 1')).toContain('STRING_AGG');
expect(dialect.generateRandomizedCardinalitySampleQuery('"t"', 'c', 100)).toContain('USING SAMPLE 100 ROWS');
});
it('exposes profiling expressions and a null column-statistics query', () => {
expect(dialect.getNullCountExpression('c')).toBe('SUM(CASE WHEN c IS NULL THEN 1 ELSE 0 END)');
expect(dialect.getDistinctCountExpression('c')).toBe('COUNT(DISTINCT c)');
expect(dialect.textLengthExpression('c')).toBe('LENGTH(CAST(c AS VARCHAR))');
expect(dialect.castToText('c')).toBe('CAST(c AS VARCHAR)');
expect(dialect.mapDataType('BIGINT')).toBe('BIGINT');
expect(dialect.getTopClause(5)).toBe('');
expect(dialect.generateColumnStatisticsQuery('main', 't')).toBeNull();
});
// Guards the single-namespace (db=null) display shape: v1 introspects only
// `main`, so a display ref must round-trip as a bare table name. An ANSI shape
// would emit a 1-part name it then refuses to parse, breaking column lookups.
it('round-trips a single-namespace display ref and reports a 1-part column shape', () => {
const table = { catalog: null, db: null, name: 'orders' };
const display = dialect.formatDisplayRef(table);
expect(display).toBe('orders');
expect(dialect.parseDisplayRef(display)).toMatchObject({ name: 'orders' });
expect(dialect.columnDisplayTablePartCount()).toBe(1);
expect(dialect.formatTableName(table)).toBe('"orders"');
});
});

View file

@ -135,6 +135,14 @@ describe('federatedAttachTarget', () => {
expect(target).toContain('ssl_mode=REQUIRED'); expect(target).toContain('ssl_mode=REQUIRED');
}); });
it('resolves a duckdb member to its database file path', () => {
const target = federatedAttachTarget(
{ connectionId: 'dux', driver: 'duckdb', projectDir: '/p', connection: { driver: 'duckdb', path: 'a.duckdb' } },
{},
);
expect(target).toBe('/p/a.duckdb');
});
it('throws for an unsupported driver', () => { it('throws for an unsupported driver', () => {
expect(() => federatedAttachTarget(member({ driver: 'snowflake', connection: { driver: 'snowflake' } }), {})).toThrow( expect(() => federatedAttachTarget(member({ driver: 'snowflake', connection: { driver: 'snowflake' } }), {})).toThrow(
/cannot be attached/i, /cannot be attached/i,

View file

@ -67,4 +67,27 @@ describe('buildAttachStatements', () => {
); );
expect(stmts.at(-1)).toBe('ATTACH \'postgresql://u:it\'\'s@h/db\' AS "pg" (TYPE postgres, READ_ONLY);'); expect(stmts.at(-1)).toBe('ATTACH \'postgresql://u:it\'\'s@h/db\' AS "pg" (TYPE postgres, READ_ONLY);');
}); });
it('attaches a native duckdb member with no TYPE and no INSTALL/LOAD', () => {
const statements = buildAttachStatements(
[{ connectionId: 'dux', driver: 'duckdb', projectDir: '/p', connection: { driver: 'duckdb', path: '/p/a.duckdb' } }],
{},
);
expect(statements.some((s) => s.startsWith('INSTALL'))).toBe(false);
expect(statements.find((s) => s.startsWith('ATTACH'))).toContain('(READ_ONLY)');
expect(statements.find((s) => s.startsWith('ATTACH'))).not.toContain('TYPE');
});
it('mixes a duckdb member with a postgres member, loading only postgres', () => {
const statements = buildAttachStatements(
[
{ connectionId: 'dux', driver: 'duckdb', projectDir: '/p', connection: { driver: 'duckdb', path: '/p/a.duckdb' } },
{ connectionId: 'pg', driver: 'postgres', projectDir: '/p', connection: { driver: 'postgres', url: 'postgres://h/db' } },
],
{},
);
expect(statements).toContain('INSTALL postgres; LOAD postgres;');
expect(statements.some((s) => s.includes('INSTALL duckdb'))).toBe(false);
expect(statements.filter((s) => s.startsWith('ATTACH')).length).toBe(2);
});
}); });

View file

@ -0,0 +1,45 @@
import { DuckDBInstance } from '@duckdb/node-api';
import { mkdtemp, rm } from 'node:fs/promises';
import { tmpdir } from 'node:os';
import { join } from 'node:path';
import { afterAll, beforeAll, describe, expect, it } from 'vitest';
import { createDuckDbLiveDatabaseIntrospection } from '../../../src/connectors/duckdb/live-database-introspection.js';
import { tableRefSet } from '../../../src/context/scan/table-ref.js';
let dir: string;
let dbPath: string;
beforeAll(async () => {
dir = await mkdtemp(join(tmpdir(), 'ktx-duckdb-live-'));
dbPath = join(dir, 'warehouse.duckdb');
const instance = await DuckDBInstance.create(dbPath);
const connection = await instance.connect();
await connection.run('CREATE TABLE customers (id BIGINT, name VARCHAR)');
await connection.run('CREATE TABLE orders (id BIGINT)');
connection.closeSync();
instance.closeSync();
});
afterAll(async () => {
await rm(dir, { recursive: true, force: true });
});
function port() {
return createDuckDbLiveDatabaseIntrospection({
projectDir: dir,
connections: { warehouse: { driver: 'duckdb', path: dbPath } },
});
}
describe('createDuckDbLiveDatabaseIntrospection', () => {
it('extracts the full schema for a connection', async () => {
const snapshot = await port().extractSchema('warehouse');
expect(snapshot.tables.map((t) => t.name).sort()).toEqual(['customers', 'orders']);
});
it('restricts extraction to a table scope', async () => {
const tableScope = tableRefSet([{ catalog: null, db: null, name: 'customers' }]);
const snapshot = await port().extractSchema('warehouse', { tableScope });
expect(snapshot.tables.map((t) => t.name)).toEqual(['customers']);
});
});

View file

@ -0,0 +1,17 @@
import { describe, expect, it } from 'vitest';
import { jsonSafeBigint, toJsonSafeRows } from '../../../src/connectors/shared/duckdb-json-safe.js';
describe('duckdb json-safe bigint', () => {
it('keeps safe-range bigints as numbers', () => {
expect(jsonSafeBigint(42n)).toBe(42);
});
it('stringifies bigints beyond Number.MAX_SAFE_INTEGER', () => {
const big = BigInt(Number.MAX_SAFE_INTEGER) + 10n;
expect(jsonSafeBigint(big)).toBe(big.toString());
});
it('converts only bigint cells in a row matrix', () => {
expect(toJsonSafeRows([[1n, 'a', null]])).toEqual([[1, 'a', null]]);
});
});

View file

@ -0,0 +1,59 @@
import { describe, expect, it } from 'vitest';
import { connectionTypeSchema } from '../../../src/context/connections/connection-type.js';
import {
dialectForConnectionType,
warehouseTargetDialect,
} from '../../../src/context/connections/connection-type-dialect.js';
describe('connection type dialect resolution', () => {
it('maps warehouse connection types to sqlglot dialects', () => {
const cases: Array<[string, string]> = [
['POSTGRESQL', 'postgres'],
['SQLITE', 'sqlite'],
['DUCKDB', 'duckdb'],
['SQLSERVER', 'tsql'],
['BIGQUERY', 'bigquery'],
['SNOWFLAKE', 'snowflake'],
['MYSQL', 'mysql'],
['CLICKHOUSE', 'clickhouse'],
['ATHENA', 'athena'],
];
for (const [connectionType, dialect] of cases) {
expect(dialectForConnectionType(connectionType)).toBe(dialect);
expect(warehouseTargetDialect(connectionType)).toBe(dialect);
}
});
it('normalizes case and preserves the semantic-layer postgres fallback for unknown inputs', () => {
expect(dialectForConnectionType('athena')).toBe('athena');
expect(dialectForConnectionType('postgresql')).toBe('postgres');
expect(dialectForConnectionType('not-a-real-connection-type')).toBe('postgres');
});
it('rejects non-SQL targets for BI table parsing', () => {
expect(warehouseTargetDialect('METABASE')).toBeNull();
expect(warehouseTargetDialect('LOOKER')).toBeNull();
expect(warehouseTargetDialect('NOTION')).toBeNull();
expect(warehouseTargetDialect('not-a-real-connection-type')).toBeNull();
});
it('removes inherited non-ktx connection type values while keeping AWS Athena', () => {
expect(connectionTypeSchema.safeParse('ATHENA').success).toBe(true);
for (const removed of [
'CENTRALREACH',
'EPIC',
'CERNER',
'QUICKBOOKS',
'WORKDAY',
'REST',
'S3',
'SLACK',
'PLAIN',
'BETTERSTACK',
]) {
expect(connectionTypeSchema.safeParse(removed).success).toBe(false);
}
});
});

View file

@ -305,7 +305,7 @@ describe('getDialectForDriver', () => {
it('throws with a supported-driver list for unknown drivers', () => { it('throws with a supported-driver list for unknown drivers', () => {
expect(() => getDialectForDriver('oracle')).toThrow( expect(() => getDialectForDriver('oracle')).toThrow(
'Unsupported driver "oracle". Supported drivers: bigquery, clickhouse, mongodb, mysql, postgres, snowflake, sqlite, sqlserver', 'Unsupported driver "oracle". Supported drivers: athena, bigquery, clickhouse, duckdb, mongodb, mysql, postgres, snowflake, sqlite, sqlserver',
); );
}); });

View file

@ -22,6 +22,7 @@ const connectionFixtures: Record<KtxConnectionDriver, FixtureFactory> = {
schemas: ['public'], schemas: ['public'],
}), }),
sqlite: () => ({ driver: 'sqlite', path: 'warehouse.db' }), sqlite: () => ({ driver: 'sqlite', path: 'warehouse.db' }),
duckdb: (projectDir) => ({ driver: 'duckdb', path: join(projectDir, 'warehouse.duckdb') }),
mongodb: () => ({ mongodb: () => ({
driver: 'mongodb', driver: 'mongodb',
url: 'mongodb://localhost:27017/app', url: 'mongodb://localhost:27017/app',
@ -69,6 +70,11 @@ const connectionFixtures: Record<KtxConnectionDriver, FixtureFactory> = {
database: 'ANALYTICS', database: 'ANALYTICS',
schema: 'PUBLIC', schema: 'PUBLIC',
}), }),
athena: () => ({
driver: 'athena',
region: 'us-east-1',
s3_staging_dir: 's3://my-bucket/athena-results/',
}),
}; };
const allowedScopeKeys = new Set(['dataset_ids', 'databases', 'schemas', 'schema_names']); const allowedScopeKeys = new Set(['dataset_ids', 'databases', 'schemas', 'schema_names']);
@ -99,8 +105,10 @@ describe('driverRegistrations', () => {
const registryDrivers = Object.keys(driverRegistrations).sort(); const registryDrivers = Object.keys(driverRegistrations).sort();
expect(listSupportedDrivers()).toEqual(registryDrivers); expect(listSupportedDrivers()).toEqual(registryDrivers);
expect(listSupportedDrivers()).toEqual([ expect(listSupportedDrivers()).toEqual([
'athena',
'bigquery', 'bigquery',
'clickhouse', 'clickhouse',
'duckdb',
'mongodb', 'mongodb',
'mysql', 'mysql',
'postgres', 'postgres',
@ -138,7 +146,7 @@ describe('driverRegistrations', () => {
expect(connector.listTables).toEqual(expect.any(Function)); expect(connector.listTables).toEqual(expect.any(Function));
await connector.cleanup?.(); await connector.cleanup?.();
if (registration.driver === 'sqlite') { if (registration.driver === 'sqlite' || registration.driver === 'duckdb') {
expect(registration.scopeConfigKey).toBeNull(); expect(registration.scopeConfigKey).toBeNull();
} else { } else {
expect(registration.scopeConfigKey).not.toBeNull(); expect(registration.scopeConfigKey).not.toBeNull();

View file

@ -35,6 +35,21 @@ describe('localConnectionToWarehouseDescriptor', () => {
}); });
}); });
it('maps Athena connections to canonical warehouse descriptors', () => {
expect(
localConnectionToWarehouseDescriptor('athena-warehouse', {
driver: 'athena',
region: 'us-east-1',
s3_staging_dir: 's3://my-bucket/athena-results/',
database: 'analytics',
}),
).toMatchObject({
id: 'athena-warehouse',
connection_type: 'ATHENA',
database: 'analytics',
});
});
it('returns null for non-warehouse adapters', () => { it('returns null for non-warehouse adapters', () => {
expect( expect(
localConnectionToWarehouseDescriptor('looker', { localConnectionToWarehouseDescriptor('looker', {
@ -51,6 +66,7 @@ describe('local connection info helpers', () => {
expect(localConnectionTypeForConfig('warehouse', { driver: 'postgres' })).toBe('POSTGRESQL'); expect(localConnectionTypeForConfig('warehouse', { driver: 'postgres' })).toBe('POSTGRESQL');
expect(localConnectionTypeForConfig('bq', { driver: 'bigquery', project_id: 'acme' })).toBe('BIGQUERY'); expect(localConnectionTypeForConfig('bq', { driver: 'bigquery', project_id: 'acme' })).toBe('BIGQUERY');
expect(localConnectionTypeForConfig('snowflake', { driver: 'snowflake' })).toBe('SNOWFLAKE'); expect(localConnectionTypeForConfig('snowflake', { driver: 'snowflake' })).toBe('SNOWFLAKE');
expect(localConnectionTypeForConfig('athena-warehouse', { driver: 'athena' })).toBe('ATHENA');
}); });
it('keeps removed driver aliases as display-only labels', () => { it('keeps removed driver aliases as display-only labels', () => {

View file

@ -1,10 +1,14 @@
import { describe, expect, it, vi } from 'vitest'; import { describe, expect, it, vi } from 'vitest';
import type { executeFederatedQuery } from '../../../src/connectors/duckdb/federated-executor.js'; import type { executeFederatedQuery } from '../../../src/connectors/duckdb/federated-executor.js';
import { executeProjectReadOnlySql } from '../../../src/context/connections/project-sql-executor.js'; import { executeProjectRawSql, executeProjectReadOnlySql } from '../../../src/context/connections/project-sql-executor.js';
import type { KtxLocalProject } from '../../../src/context/project/project.js'; import type { KtxLocalProject } from '../../../src/context/project/project.js';
import type { KtxScanConnector } from '../../../src/context/scan/types.js'; import type { KtxScanConnector } from '../../../src/context/scan/types.js';
import type { SqlAnalysisPort } from '../../../src/context/sql-analysis/ports.js';
import { KtxQueryError } from '../../../src/errors.js';
function fakeProject(connections: Record<string, { driver: string }>): KtxLocalProject { function fakeProject(
connections: Record<string, { driver: string; query_policy?: 'semantic-layer-only' }>,
): KtxLocalProject {
return { return {
projectDir: '/tmp/proj', projectDir: '/tmp/proj',
configPath: '/tmp/proj/ktx.yaml', configPath: '/tmp/proj/ktx.yaml',
@ -114,3 +118,97 @@ describe('executeProjectReadOnlySql headerTypes', () => {
expect(result.headerTypes).toEqual(['INTEGER']); expect(result.headerTypes).toEqual(['INTEGER']);
}); });
}); });
function fakeSqlAnalysis(validation: { ok: boolean; error: string | null }): SqlAnalysisPort {
return {
analyzeForFingerprint: vi.fn(),
analyzeBatch: vi.fn(),
validateReadOnly: vi.fn(async () => validation),
} as unknown as SqlAnalysisPort;
}
describe('executeProjectRawSql', () => {
it('validates then executes raw SQL on an unrestricted connection', async () => {
const project = fakeProject({ pg: { driver: 'postgres' } });
const sqlAnalysis = fakeSqlAnalysis({ ok: true, error: null });
const connector = connectorReturning({
headers: ['id'],
rows: [[1]],
totalRows: 1,
rowCount: 1,
});
const result = await executeProjectRawSql({
project,
connectionId: 'pg',
sql: 'SELECT id FROM orders',
maxRows: 25,
sqlAnalysis,
createConnector: () => connector,
runId: 'test-raw-sql',
});
expect(result.rows).toEqual([[1]]);
expect(sqlAnalysis.validateReadOnly).toHaveBeenCalledWith('SELECT id FROM orders', 'postgres');
});
it('rejects a restricted connection before validation or execution', async () => {
const project = fakeProject({ pg: { driver: 'postgres', query_policy: 'semantic-layer-only' } });
const sqlAnalysis = fakeSqlAnalysis({ ok: true, error: null });
const createConnector = vi.fn();
const execution = executeProjectRawSql({
project,
connectionId: 'pg',
sql: 'SELECT 1',
maxRows: 25,
sqlAnalysis,
createConnector: createConnector as never,
runId: 'test-raw-sql',
});
await expect(execution).rejects.toBeInstanceOf(KtxQueryError);
await expect(execution).rejects.toThrow(/query_policy: semantic-layer-only/);
expect(sqlAnalysis.validateReadOnly).not.toHaveBeenCalled();
expect(createConnector).not.toHaveBeenCalled();
});
it('rejects federated raw SQL when a member connection is restricted', async () => {
const project = fakeProject({
pg: { driver: 'postgres', query_policy: 'semantic-layer-only' },
lite: { driver: 'sqlite' },
});
const executeFederated = vi.fn();
await expect(
executeProjectRawSql({
project,
connectionId: '_ktx_federated',
sql: 'SELECT 1',
maxRows: 25,
sqlAnalysis: fakeSqlAnalysis({ ok: true, error: null }),
createConnector: vi.fn() as never,
executeFederated: executeFederated as never,
runId: 'test-raw-sql',
}),
).rejects.toThrow(/"pg"/);
expect(executeFederated).not.toHaveBeenCalled();
});
it('classifies a read-only validation failure as an expected query error', async () => {
const project = fakeProject({ pg: { driver: 'postgres' } });
const createConnector = vi.fn();
const execution = executeProjectRawSql({
project,
connectionId: 'pg',
sql: 'DROP TABLE orders',
maxRows: 25,
sqlAnalysis: fakeSqlAnalysis({ ok: false, error: 'SQL is not read-only: DROP.' }),
createConnector: createConnector as never,
runId: 'test-raw-sql',
});
await expect(execution).rejects.toBeInstanceOf(KtxQueryError);
await expect(execution).rejects.toThrow('SQL is not read-only: DROP.');
expect(createConnector).not.toHaveBeenCalled();
});
});

View file

@ -0,0 +1,144 @@
import { describe, expect, it } from 'vitest';
import {
assertRawSqlAllowed,
connectionQueryPolicy,
projectAllowsRawSql,
restrictedFederatedMemberIds,
} from '../../../src/context/connections/query-policy.js';
import { parseKtxProjectConfig } from '../../../src/context/project/config.js';
import { KtxQueryError } from '../../../src/errors.js';
const PROJECT_DIR = '/tmp/proj';
function config(yaml: string) {
return parseKtxProjectConfig(yaml);
}
describe('connectionQueryPolicy', () => {
it('defaults to read-only-sql when the field is absent or the connection is unknown', () => {
const parsed = config(`
connections:
warehouse:
driver: sqlite
url: file:warehouse.db
`);
expect(connectionQueryPolicy(parsed.connections.warehouse)).toBe('read-only-sql');
expect(connectionQueryPolicy(undefined)).toBe('read-only-sql');
});
it('reads semantic-layer-only from ktx.yaml', () => {
const parsed = config(`
connections:
warehouse:
driver: snowflake
url: env:SNOWFLAKE_URL
query_policy: semantic-layer-only
`);
expect(connectionQueryPolicy(parsed.connections.warehouse)).toBe('semantic-layer-only');
});
it('rejects unknown query_policy values at config parse time', () => {
expect(() =>
config(`
connections:
warehouse:
driver: sqlite
url: file:warehouse.db
query_policy: everything-goes
`),
).toThrow();
});
});
describe('assertRawSqlAllowed', () => {
it('allows raw SQL on an unrestricted connection', () => {
const parsed = config(`
connections:
warehouse:
driver: sqlite
url: file:warehouse.db
`);
expect(() => assertRawSqlAllowed(parsed, PROJECT_DIR, 'warehouse')).not.toThrow();
});
it('rejects raw SQL on a restricted connection with an expected error naming the policy', () => {
const parsed = config(`
connections:
warehouse:
driver: sqlite
url: file:warehouse.db
query_policy: semantic-layer-only
`);
expect(() => assertRawSqlAllowed(parsed, PROJECT_DIR, 'warehouse')).toThrow(KtxQueryError);
expect(() => assertRawSqlAllowed(parsed, PROJECT_DIR, 'warehouse')).toThrow(
/query_policy: semantic-layer-only/,
);
});
it('rejects federated raw SQL when any member connection is restricted', () => {
const parsed = config(`
connections:
sales:
driver: sqlite
url: file:sales.db
query_policy: semantic-layer-only
events:
driver: sqlite
url: file:events.db
`);
expect(restrictedFederatedMemberIds(parsed, PROJECT_DIR)).toEqual(['sales']);
expect(() => assertRawSqlAllowed(parsed, PROJECT_DIR, '_ktx_federated')).toThrow(/"sales"/);
});
it('allows federated raw SQL when no member is restricted', () => {
const parsed = config(`
connections:
sales:
driver: sqlite
url: file:sales.db
events:
driver: sqlite
url: file:events.db
`);
expect(restrictedFederatedMemberIds(parsed, PROJECT_DIR)).toEqual([]);
expect(() => assertRawSqlAllowed(parsed, PROJECT_DIR, '_ktx_federated')).not.toThrow();
});
});
describe('projectAllowsRawSql', () => {
it('is true when at least one SQL connection is unrestricted', () => {
const parsed = config(`
connections:
finance:
driver: postgres
url: env:FINANCE_URL
query_policy: semantic-layer-only
warehouse:
driver: sqlite
url: file:warehouse.db
`);
expect(projectAllowsRawSql(parsed)).toBe(true);
});
it('is false when every SQL connection is restricted', () => {
const parsed = config(`
connections:
finance:
driver: postgres
url: env:FINANCE_URL
query_policy: semantic-layer-only
`);
expect(projectAllowsRawSql(parsed)).toBe(false);
});
it('is true for projects with no SQL-queryable connections', () => {
const parsed = config(`
connections:
docs:
driver: mongodb
url: mongodb://localhost:27017/app
`);
expect(projectAllowsRawSql(parsed)).toBe(true);
expect(projectAllowsRawSql(config('connections: {}'))).toBe(true);
});
});

View file

@ -3,6 +3,7 @@ import { tmpdir } from 'node:os';
import { join } from 'node:path'; import { join } from 'node:path';
import { describe, expect, it } from 'vitest'; import { describe, expect, it } from 'vitest';
import { resolveKtxConfigReference, resolveKtxHomePath } from '../../../src/context/core/config-reference.js'; import { resolveKtxConfigReference, resolveKtxHomePath } from '../../../src/context/core/config-reference.js';
import { KtxExpectedError } from '../../../src/errors.js';
describe('ktx config references', () => { describe('ktx config references', () => {
it('resolves env references without returning empty values', () => { it('resolves env references without returning empty values', () => {
@ -22,6 +23,17 @@ describe('ktx config references', () => {
expect(resolveKtxConfigReference(`file:${keyPath}`, {})).toBe('file-gateway-key'); expect(resolveKtxConfigReference(`file:${keyPath}`, {})).toBe('file-gateway-key');
}); });
it('raises an expected error when a file reference is missing', () => {
const missing = join(tmpdir(), `ktx-config-reference-missing-${process.pid}`, 'absent-password');
try {
resolveKtxConfigReference(`file:${missing}`, {});
expect.unreachable('expected a thrown error for the missing secret file');
} catch (error) {
expect(error).toBeInstanceOf(KtxExpectedError);
expect((error as Error).message).toContain(missing);
}
});
it('returns literal values unchanged after trimming blank-only values', () => { it('returns literal values unchanged after trimming blank-only values', () => {
expect(resolveKtxConfigReference('provider/model', {})).toBe('provider/model'); expect(resolveKtxConfigReference('provider/model', {})).toBe('provider/model');
expect(resolveKtxConfigReference(' ', {})).toBeUndefined(); expect(resolveKtxConfigReference(' ', {})).toBeUndefined();

View file

@ -1,7 +1,11 @@
import { once } from 'node:events'; import { once } from 'node:events';
import { createServer } from 'node:http'; import { createServer } from 'node:http';
import { describe, expect, it, vi } from 'vitest'; import { describe, expect, it, vi } from 'vitest';
import { createHttpSemanticLayerComputePort, createPythonSemanticLayerComputePort } from '../../../src/context/daemon/semantic-layer-compute.js'; import {
createHttpSemanticLayerComputePort,
createPythonSemanticLayerComputePort,
KtxDaemonComputeError,
} from '../../../src/context/daemon/semantic-layer-compute.js';
const source = { const source = {
name: 'orders', name: 'orders',
@ -174,6 +178,79 @@ describe('createPythonSemanticLayerComputePort', () => {
}); });
}); });
describe('KtxDaemonComputeError classification', () => {
const query = { sources: [source], dialect: 'postgres', query: { measures: ['count(*)'], dimensions: [] } };
function exitingPort(code: number, stderr: string) {
return createPythonSemanticLayerComputePort({
command: process.execPath,
args: [
'-e',
`process.stdin.on('data',()=>{});process.stdin.on('end',()=>{process.stderr.write(${JSON.stringify(stderr)});process.exit(${code})});`,
],
});
}
async function rejection(promise: Promise<unknown>): Promise<KtxDaemonComputeError> {
const error = await promise.then(
() => null,
(thrown: unknown) => thrown,
);
expect(error).toBeInstanceOf(KtxDaemonComputeError);
return error as KtxDaemonComputeError;
}
it('marks a subprocess input-rejection (exit 3) as inputRejected', async () => {
const error = await rejection(exitingPort(3, 'Measure expr does not reference any source').query(query));
expect(error.inputRejected).toBe(true);
expect(error.detail).toContain('does not reference any source');
});
it('marks a subprocess fault (exit 1) as not inputRejected', async () => {
const error = await rejection(exitingPort(1, 'Traceback: boom').query(query));
expect(error.inputRejected).toBe(false);
expect(error.detail).toContain('boom');
});
async function statusPort(statusCode: number, body: string): Promise<{ port: ReturnType<typeof createHttpSemanticLayerComputePort>; close: () => void }> {
const server = createServer((_request, response) => {
response.writeHead(statusCode, { 'content-type': 'application/json' });
response.end(body);
});
server.listen(0, '127.0.0.1');
await once(server, 'listening');
const address = server.address();
if (!address || typeof address === 'string') {
throw new Error('expected TCP server address');
}
return {
port: createHttpSemanticLayerComputePort({ baseUrl: `http://127.0.0.1:${address.port}` }),
close: () => server.close(),
};
}
it('marks an HTTP 400 as inputRejected and unwraps the daemon detail', async () => {
const { port, close } = await statusPort(400, JSON.stringify({ detail: 'Measure expr does not reference any source' }));
try {
const error = await rejection(port.query(query));
expect(error.inputRejected).toBe(true);
expect(error.detail).toBe('Measure expr does not reference any source');
} finally {
close();
}
});
it('marks an HTTP 500 as not inputRejected', async () => {
const { port, close } = await statusPort(500, JSON.stringify({ detail: 'Daemon request failed: boom' }));
try {
const error = await rejection(port.query(query));
expect(error.inputRejected).toBe(false);
} finally {
close();
}
});
});
describe('createHttpSemanticLayerComputePort', () => { describe('createHttpSemanticLayerComputePort', () => {
it('calls semantic query and validate HTTP endpoints through an injected runner', async () => { it('calls semantic query and validate HTTP endpoints through an injected runner', async () => {
const requestJson = vi.fn(async (path: string) => { const requestJson = vi.fn(async (path: string) => {

View file

@ -72,6 +72,8 @@ describe('looker dialect and target validation helpers', () => {
it('maps Looker dialect names to ktx connection types', () => { it('maps Looker dialect names to ktx connection types', () => {
expect(lookerDialectToConnectionType('bigquery_standard_sql')).toBe('BIGQUERY'); expect(lookerDialectToConnectionType('bigquery_standard_sql')).toBe('BIGQUERY');
expect(lookerDialectToConnectionType('postgres')).toBe('POSTGRESQL'); expect(lookerDialectToConnectionType('postgres')).toBe('POSTGRESQL');
expect(lookerDialectToConnectionType('awsathena')).toBe('ATHENA');
expect(lookerDialectToConnectionType('AWSATHENA')).toBe('ATHENA');
expect(lookerDialectToConnectionType('mssql')).toBeNull(); expect(lookerDialectToConnectionType('mssql')).toBeNull();
expect(lookerDialectToConnectionType('tsql')).toBeNull(); expect(lookerDialectToConnectionType('tsql')).toBeNull();
expect(lookerDialectToConnectionType('unknown')).toBeNull(); expect(lookerDialectToConnectionType('unknown')).toBeNull();
@ -80,6 +82,7 @@ describe('looker dialect and target validation helpers', () => {
it('maps supported warehouse connection types to sqlglot dialects', () => { it('maps supported warehouse connection types to sqlglot dialects', () => {
expect(sqlglotDialectForConnectionType('BIGQUERY')).toBe('bigquery'); expect(sqlglotDialectForConnectionType('BIGQUERY')).toBe('bigquery');
expect(sqlglotDialectForConnectionType('POSTGRESQL')).toBe('postgres'); expect(sqlglotDialectForConnectionType('POSTGRESQL')).toBe('postgres');
expect(sqlglotDialectForConnectionType('ATHENA')).toBe('athena');
expect(sqlglotDialectForConnectionType('LOOKER')).toBeNull(); expect(sqlglotDialectForConnectionType('LOOKER')).toBeNull();
}); });
@ -259,6 +262,44 @@ describe('collectExploreParseItems and projectParsedIdentifier', () => {
}); });
}); });
it('collects Athena parser inputs with the athena sqlglot dialect', () => {
expect(
collectExploreParseItems({
explore: {
...mappedExplore,
rawSqlTableName: 'analytics.orders',
joins: [
{
name: 'line_items',
type: 'left_outer',
relationship: 'one_to_many',
rawSqlTableName: 'analytics.line_items',
sqlOn: null,
from: null,
targetTable: null,
},
],
},
connectionMappings: { b2b_sandbox_bq: 'athena-warehouse' },
targetConnections: new Map([['athena-warehouse', { id: 'athena-warehouse', connection_type: 'ATHENA' }]]),
}),
).toEqual({
parsedTargetTables: {},
parseItems: [
{
key: 'b2b.sales_pipeline',
sql_table_name: 'analytics.orders',
dialect: 'athena',
},
{
key: 'b2b.sales_pipeline.line_items',
sql_table_name: 'analytics.line_items',
dialect: 'athena',
},
],
});
});
it('projects successful and failed parser rows into ktx parsed target tables', () => { it('projects successful and failed parser rows into ktx parsed target tables', () => {
expect( expect(
projectParsedIdentifier({ projectParsedIdentifier({

View file

@ -121,6 +121,15 @@ describe('validateMappingPhysicalMatch', () => {
expect(reason).toContain('engine'); expect(reason).toContain('engine');
}); });
it('accepts Athena mappings when the target connection type is ATHENA', () => {
expect(
validateMappingPhysicalMatch(
{ metabaseEngine: 'athena', metabaseDbName: 'analytics', metabaseHost: null },
{ connection_type: 'ATHENA', database: 'analytics' },
),
).toBeNull();
});
it('returns null when Postgres host and database both match after normalization', () => { it('returns null when Postgres host and database both match after normalization', () => {
expect( expect(
validateMappingPhysicalMatch( validateMappingPhysicalMatch(
@ -159,6 +168,18 @@ describe('validateMappingPhysicalMatch', () => {
}); });
}); });
describe('METABASE_ENGINE_TO_CONNECTION_TYPE', () => {
it('maps Metabase Athena databases to AWS Athena warehouse connections', () => {
expect(METABASE_ENGINE_TO_CONNECTION_TYPE.athena).toBe('ATHENA');
expect(
validateMappingPhysicalMatch(
{ metabaseEngine: 'athena', metabaseDbName: 'analytics', metabaseHost: null },
{ connection_type: 'POSTGRESQL', database: 'analytics' },
),
).toBe("Metabase database engine 'athena' does not match ktx connection type 'POSTGRESQL'");
});
});
describe('computeMetabaseMappingPhysicalMismatches', () => { describe('computeMetabaseMappingPhysicalMismatches', () => {
it('returns only mismatched physical mappings', () => { it('returns only mismatched physical mappings', () => {
expect( expect(

View file

@ -33,6 +33,10 @@
}, },
"hint": { "hint": {
"type": "string" "type": "string"
},
"queryPolicy": {
"type": "string",
"const": "semantic-layer-only"
} }
}, },
"required": [ "required": [

View file

@ -33,8 +33,7 @@ describe('per-dialect SQL notes', () => {
}); });
it('does not author notes for unreachable dialects', () => { it('does not author notes for unreachable dialects', () => {
// duckdb/databricks appear in the resolver map but no connector produces them. // databricks appears in the resolver map but no connector produces it.
expect(DIALECTS_WITH_NOTES).not.toContain('duckdb');
expect(DIALECTS_WITH_NOTES).not.toContain('databricks'); expect(DIALECTS_WITH_NOTES).not.toContain('databricks');
}); });

View file

@ -4,6 +4,7 @@ import { join } from 'node:path';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import { initKtxProject } from '../../../src/context/project/project.js'; import { initKtxProject } from '../../../src/context/project/project.js';
import { KtxExpectedError, KtxQueryError } from '../../../src/errors.js'; import { KtxExpectedError, KtxQueryError } from '../../../src/errors.js';
import { KtxDaemonComputeError } from '../../../src/context/daemon/semantic-layer-compute.js';
import { createKtxConnectorCapabilities, type KtxQueryResult, type KtxScanConnector, type KtxSchemaSnapshot } from '../../../src/context/scan/types.js'; import { createKtxConnectorCapabilities, type KtxQueryResult, type KtxScanConnector, type KtxSchemaSnapshot } from '../../../src/context/scan/types.js';
import { SemanticLayerService } from '../../../src/context/sl/semantic-layer.service.js'; import { SemanticLayerService } from '../../../src/context/sl/semantic-layer.service.js';
import type { SemanticLayerSource } from '../../../src/context/sl/types.js'; import type { SemanticLayerSource } from '../../../src/context/sl/types.js';
@ -247,6 +248,80 @@ describe('createLocalProjectMcpContextPorts', () => {
expect(connector.cleanup).toHaveBeenCalled(); expect(connector.cleanup).toHaveBeenCalled();
}); });
it('omits sql_execution when every SQL connection is semantic-layer-only', async () => {
const project = await initKtxProject({ projectDir: tempDir });
project.config.connections.warehouse = {
driver: 'postgres',
url: 'env:DATABASE_URL',
query_policy: 'semantic-layer-only',
};
const ports = createLocalProjectMcpContextPorts(project, {
sqlAnalysis: {
analyzeForFingerprint: vi.fn(),
analyzeBatch: vi.fn(),
validateReadOnly: vi.fn(async () => ({ ok: true, error: null })),
},
localScan: { createConnector: vi.fn(async () => testConnector()) },
embeddingService: null,
});
expect(ports.sqlExecution).toBeUndefined();
});
it('keeps sql_execution in mixed projects but rejects restricted connections and flags them in connection_list', async () => {
const project = await initKtxProject({ projectDir: tempDir });
project.config.connections.warehouse = {
driver: 'postgres',
url: 'env:DATABASE_URL',
};
project.config.connections.finance = {
driver: 'postgres',
url: 'env:FINANCE_URL',
query_policy: 'semantic-layer-only',
};
const createConnector = vi.fn(async () => testConnector());
const sqlAnalysis = {
analyzeForFingerprint: vi.fn(),
analyzeBatch: vi.fn(),
validateReadOnly: vi.fn(async () => ({ ok: true, error: null })),
};
const ports = createLocalProjectMcpContextPorts(project, {
sqlAnalysis,
localScan: { createConnector },
embeddingService: null,
});
expect(ports.sqlExecution).toBeDefined();
const execution = ports.sqlExecution?.execute({
connectionId: 'finance',
sql: 'select 1',
maxRows: 5,
});
await expect(execution).rejects.toBeInstanceOf(KtxQueryError);
await expect(execution).rejects.toThrow(/query_policy: semantic-layer-only/);
expect(sqlAnalysis.validateReadOnly).not.toHaveBeenCalled();
expect(createConnector).not.toHaveBeenCalled();
// Both postgres members federate, so the restricted member also blocks federated raw SQL.
await expect(
ports.sqlExecution?.execute({ connectionId: '_ktx_federated', sql: 'select 1', maxRows: 5 }),
).rejects.toThrow(/"finance"/);
await expect(ports.connections?.list()).resolves.toEqual([
expect.objectContaining({
id: 'finance',
queryPolicy: 'semantic-layer-only',
hint: expect.stringContaining('sl_query'),
}),
expect.objectContaining({ id: 'warehouse' }),
expect.objectContaining({
id: '_ktx_federated',
queryPolicy: 'semantic-layer-only',
hint: expect.stringContaining('finance'),
}),
]);
});
it('rejects sql_execution against an unconfigured connection with an actionable expected error', async () => { it('rejects sql_execution against an unconfigured connection with an actionable expected error', async () => {
const project = await initKtxProject({ projectDir: tempDir }); const project = await initKtxProject({ projectDir: tempDir });
project.config.connections.warehouse = { project.config.connections.warehouse = {
@ -1103,4 +1178,86 @@ describe('createLocalProjectMcpContextPorts', () => {
}), }),
); );
}); });
async function seedOrdersWarehouse() {
const project = await initKtxProject({ projectDir: tempDir });
project.config.connections.warehouse = { driver: 'postgres', url: 'env:DATABASE_URL' };
await seedSlSourceFile(project, {
connectionId: 'warehouse',
sourceName: 'orders',
yaml: ['name: orders', 'table: public.orders', 'grain:', ' - id', 'columns:', ' - name: id', ' type: number', 'joins: []', 'measures: []', ''].join('\n'),
});
return project;
}
it('promotes a daemon input-rejection to an expected KtxQueryError carrying the daemon detail', async () => {
const project = await seedOrdersWarehouse();
const semanticLayerCompute = {
validateSources: vi.fn(),
generateSources: vi.fn(),
query: vi.fn(async () => {
throw new KtxDaemonComputeError("ktx-daemon semantic-query failed: Measure expr 'count(*)' does not reference any source", {
inputRejected: true,
detail: "Measure expr 'count(*)' does not reference any source",
});
}),
};
const ports = createLocalProjectMcpContextPorts(project, { semanticLayerCompute, embeddingService: null });
const rejection = ports.semanticLayer?.query({
connectionId: 'warehouse',
query: { measures: [{ expr: 'count(*)', name: 'n' }], dimensions: [] },
});
await expect(rejection).rejects.toBeInstanceOf(KtxQueryError);
await expect(rejection).rejects.toThrow("Measure expr 'count(*)' does not reference any source");
});
it('leaves a daemon crash as an unexpected fault', async () => {
const project = await seedOrdersWarehouse();
const semanticLayerCompute = {
validateSources: vi.fn(),
generateSources: vi.fn(),
query: vi.fn(async () => {
throw new KtxDaemonComputeError('ktx-daemon semantic-query failed: KeyError: boom', {
inputRejected: false,
detail: 'KeyError: boom',
});
}),
};
const ports = createLocalProjectMcpContextPorts(project, { semanticLayerCompute, embeddingService: null });
const rejection = ports.semanticLayer?.query({
connectionId: 'warehouse',
query: { measures: ['orders.order_count'], dimensions: [] },
});
await expect(rejection).rejects.toBeInstanceOf(KtxDaemonComputeError);
await expect(rejection).rejects.not.toBeInstanceOf(KtxQueryError);
});
it('wraps a warehouse execution rejection from sl_query as KtxQueryError', async () => {
const project = await seedOrdersWarehouse();
const semanticLayerCompute = {
validateSources: vi.fn(),
generateSources: vi.fn(),
query: vi.fn(async () => ({
sql: 'select count(*) from public.orders',
dialect: 'postgres',
columns: [{ name: 'orders.order_count' }],
plan: {},
})),
};
const queryExecutor = {
execute: vi.fn(async () => {
throw new Error("Unknown column '검사 유형' in 'SELECT'");
}),
};
const ports = createLocalProjectMcpContextPorts(project, { semanticLayerCompute, queryExecutor, embeddingService: null });
const rejection = ports.semanticLayer?.query({
connectionId: 'warehouse',
query: { measures: ['orders.order_count'], dimensions: [], limit: 5 },
});
await expect(rejection).rejects.toBeInstanceOf(KtxQueryError);
await expect(rejection).rejects.toThrow("Unknown column '검사 유형' in 'SELECT'");
});
}); });

View file

@ -2175,6 +2175,40 @@ describe('local scan', () => {
}; };
expect(manifest.tables.orders?.joins?.some((join) => join.to === 'accounts')).toBe(true); expect(manifest.tables.orders?.joins?.some((join) => join.to === 'accounts')).toBe(true);
}); });
it('accepts athena as a native standalone scan driver when the host supplies a live-database adapter', async () => {
await writeFile(
join(project.projectDir, 'ktx.yaml'),
[
'connections:',
' warehouse:',
' driver: athena',
' region: us-east-1',
' s3_staging_dir: s3://my-bucket/athena-results/',
' databases:',
' - analytics',
'ingest:',
' adapters:',
' - live-database',
'',
].join('\n'),
'utf-8',
);
project = await loadKtxProject({ projectDir: project.projectDir });
const result = await runLocalScan({
project,
adapters: [fetchOnlyAdapter()],
connectionId: 'warehouse',
jobId: 'scan-run-athena',
now: () => new Date('2026-04-29T17:00:00.000Z'),
});
expect(result.report.driver).toBe('athena');
expect(result.report.artifactPaths.reportPath).toBe(
'raw-sources/warehouse/live-database/2026-04-29-170000-scan-run-athena/scan-report.json',
);
});
}); });
describe('resolveEnabledTables', () => { describe('resolveEnabledTables', () => {

View file

@ -3,6 +3,7 @@ import { tmpdir } from 'node:os';
import { join } from 'node:path'; import { join } from 'node:path';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import type { KtxSemanticLayerComputePort } from '../../../src/context/daemon/semantic-layer-compute.js'; import type { KtxSemanticLayerComputePort } from '../../../src/context/daemon/semantic-layer-compute.js';
import { FEDERATED_CONNECTION_ID } from '../../../src/context/connections/federation.js';
import { initKtxProject, type KtxLocalProject } from '../../../src/context/project/project.js'; import { initKtxProject, type KtxLocalProject } from '../../../src/context/project/project.js';
import { compileLocalSlQuery } from '../../../src/context/sl/local-query.js'; import { compileLocalSlQuery } from '../../../src/context/sl/local-query.js';
@ -67,6 +68,59 @@ grain: []
await rm(tempDir, { recursive: true, force: true }); await rm(tempDir, { recursive: true, force: true });
}); });
it('injects predefined_measures_only when the connection query_policy is semantic-layer-only', async () => {
project.config.connections.warehouse = { driver: 'postgres', query_policy: 'semantic-layer-only' };
await compileLocalSlQuery(project, {
connectionId: 'warehouse',
query: { measures: ['orders.order_count'], dimensions: [], limit: 10 },
compute,
});
expect(compute.query).toHaveBeenCalledWith(
expect.objectContaining({
query: expect.objectContaining({ predefined_measures_only: true }),
}),
);
});
it('rejects a federated sl_query, pointing to per-connection SL when a member is restricted', async () => {
project.config.connections.warehouse = { driver: 'postgres', query_policy: 'semantic-layer-only' };
project.config.connections.analytics = { driver: 'postgres' };
let message = '';
try {
await compileLocalSlQuery(project, {
connectionId: FEDERATED_CONNECTION_ID,
query: { measures: ['orders.order_count'], dimensions: [] },
compute,
});
throw new Error('expected compileLocalSlQuery to reject');
} catch (e) {
message = (e as Error).message;
}
expect(message).toContain("member connection(s) 'warehouse'");
expect(message).toContain('query_policy: semantic-layer-only');
// Must not send the agent down the raw-SQL path that assertRawSqlAllowed rejects.
expect(message).not.toContain(`ktx sql -c ${FEDERATED_CONNECTION_ID}`);
expect(message).not.toContain('sql_execution');
expect(compute.query).not.toHaveBeenCalled();
});
it('rejects a federated sl_query, pointing to raw federated SQL when no member is restricted', async () => {
project.config.connections.analytics = { driver: 'postgres' };
await expect(
compileLocalSlQuery(project, {
connectionId: FEDERATED_CONNECTION_ID,
query: { measures: ['orders.order_count'], dimensions: [] },
compute,
}),
).rejects.toThrow(`ktx sql -c ${FEDERATED_CONNECTION_ID}`);
expect(compute.query).not.toHaveBeenCalled();
});
it('refuses a non-SQL (context-only) connection instead of compiling it as Postgres', async () => { it('refuses a non-SQL (context-only) connection instead of compiling it as Postgres', async () => {
project.config.connections['mongo-prod'] = { driver: 'mongodb', url: 'mongodb://localhost:27017/app' }; project.config.connections['mongo-prod'] = { driver: 'mongodb', url: 'mongodb://localhost:27017/app' };
await expect( await expect(
@ -109,6 +163,7 @@ grain: []
measures: ['orders.order_count'], measures: ['orders.order_count'],
dimensions: ['orders.status'], dimensions: ['orders.status'],
limit: 25, limit: 25,
predefined_measures_only: false,
}, },
}); });
expect(result).toEqual({ expect(result).toEqual({
@ -190,6 +245,7 @@ grain: []
query: { query: {
measures: ['sum(payments.amount)'], measures: ['sum(payments.amount)'],
dimensions: [], dimensions: [],
predefined_measures_only: false,
}, },
}); });
}); });

View file

@ -960,6 +960,28 @@ describe('validateWithProposedSource', () => {
); );
}); });
it('uses the Athena sqlglot dialect when validating Athena sources', async () => {
pythonPort.validateSources.mockResolvedValue({
data: { errors: [], warnings: [] },
});
service = new SemanticLayerService(configService as never, connectionCatalog('ATHENA'), pythonPort);
await service.validateWithProposedSource('conn-1', {
name: 'orders',
table: 'analytics.orders',
grain: ['id'],
columns: [{ name: 'id', type: 'number' }],
joins: [],
measures: [],
});
expect(pythonPort.validateSources).toHaveBeenCalledWith(
expect.objectContaining({
dialect: 'athena',
}),
);
});
it('composes a bare overlay with its manifest base before validating', async () => { it('composes a bare overlay with its manifest base before validating', async () => {
const schemaPath = 'semantic-layer/conn-1/_schema/core.yaml'; const schemaPath = 'semantic-layer/conn-1/_schema/core.yaml';
const listFilesImpl = (dir: string): Promise<{ files: string[] }> => { const listFilesImpl = (dir: string): Promise<{ files: string[] }> => {
@ -1326,6 +1348,44 @@ describe('validateWithProposedSource', () => {
}); });
}); });
describe('executeQuery dialect resolution', () => {
it('passes the Athena sqlglot dialect to the Python query engine for Athena connections', async () => {
pythonPort.query.mockResolvedValue({ data: { sql: 'SELECT * FROM orders' } });
const catalog = connectionCatalog('ATHENA');
catalog.executeQuery.mockResolvedValue({ headers: ['id'], rows: [[1]], totalRows: 1 });
const configService = {
listFiles: vi.fn().mockResolvedValue({
files: ['semantic-layer/conn-1/orders.yaml'],
}),
readFile: vi.fn().mockResolvedValue({
content: [
'name: orders',
'table: analytics.orders',
'grain:',
' - id',
'columns:',
' - name: id',
' type: number',
'joins: []',
'measures:',
' - name: count',
' expr: count(*)',
].join('\n'),
}),
};
const service = new SemanticLayerService(configService as never, catalog, pythonPort);
await service.executeQuery('conn-1', { measures: ['count'] } as never);
expect(pythonPort.query).toHaveBeenCalledWith(
expect.objectContaining({
dialect: 'athena',
}),
);
expect(catalog.executeQuery).toHaveBeenCalledWith('conn-1', 'SELECT * FROM orders');
});
});
describe('findDanglingSegmentRefs', () => { describe('findDanglingSegmentRefs', () => {
it('returns empty when every measure segment resolves', () => { it('returns empty when every measure segment resolves', () => {
const source = { const source = {

View file

@ -12,6 +12,7 @@ describe('sqlAnalysisDialectForDriver', () => {
expect(sqlAnalysisDialectForDriver('duckdb')).toBe('duckdb'); expect(sqlAnalysisDialectForDriver('duckdb')).toBe('duckdb');
expect(sqlAnalysisDialectForDriver('clickhouse')).toBe('clickhouse'); expect(sqlAnalysisDialectForDriver('clickhouse')).toBe('clickhouse');
expect(sqlAnalysisDialectForDriver('databricks')).toBe('databricks'); expect(sqlAnalysisDialectForDriver('databricks')).toBe('databricks');
expect(sqlAnalysisDialectForDriver('athena')).toBe('athena');
}); });
it('maps local connection-type spellings to sqlglot dialects', () => { it('maps local connection-type spellings to sqlglot dialects', () => {

View file

@ -92,7 +92,7 @@ describe('createKtxCliScanConnector', () => {
expect(bigQueryMock.constructorInputs[0]).not.toHaveProperty('maxBytesBilled'); expect(bigQueryMock.constructorInputs[0]).not.toHaveProperty('maxBytesBilled');
}); });
it('rejects daemon-only fallback driver configs at config parse time', async () => { it('resolves a duckdb connection to the DuckDB scan connector', async () => {
await initKtxProject({ projectDir: tempDir }); await initKtxProject({ projectDir: tempDir });
await writeFile( await writeFile(
join(tempDir, 'ktx.yaml'), join(tempDir, 'ktx.yaml'),
@ -105,10 +105,12 @@ describe('createKtxCliScanConnector', () => {
].join('\n'), ].join('\n'),
'utf-8', 'utf-8',
); );
const project = await loadKtxProject({ projectDir: tempDir });
await expect(loadKtxProject({ projectDir: tempDir })).rejects.toThrow( const connector = await createKtxCliScanConnector(project, 'warehouse');
/connections\.warehouse\.driver:.*Invalid discriminator value/,
); expect(connector.id).toBe('duckdb:warehouse');
expect(connector.driver).toBe('duckdb');
}); });
it('rejects connection blocks with no driver field at config parse time', async () => { it('rejects connection blocks with no driver field at config parse time', async () => {

View file

@ -243,8 +243,10 @@ describe('setup databases step', () => {
{ value: 'mysql', label: 'MySQL' }, { value: 'mysql', label: 'MySQL' },
{ value: 'clickhouse', label: 'ClickHouse' }, { value: 'clickhouse', label: 'ClickHouse' },
{ value: 'sqlserver', label: 'SQL Server' }, { value: 'sqlserver', label: 'SQL Server' },
{ value: 'athena', label: 'Amazon Athena' },
{ value: 'mongodb', label: 'MongoDB' }, { value: 'mongodb', label: 'MongoDB' },
{ value: 'sqlite', label: 'SQLite' }, { value: 'sqlite', label: 'SQLite' },
{ value: 'duckdb', label: 'DuckDB' },
], ],
required: true, required: true,
}); });
@ -617,6 +619,29 @@ describe('setup databases step', () => {
}, },
], ],
}, },
{
driver: 'athena',
textValues: ['', 'us-east-1', 's3://my-bucket/athena-results/', '', ''],
expectedTextPrompts: [
{
message: connectionNamePrompt('Amazon Athena'),
placeholder: 'athena-warehouse',
initialValue: 'athena-warehouse',
},
{
message: 'AWS region\nFor example us-east-1.',
},
{
message: 'S3 staging directory\nAthena writes query results here. For example s3://my-bucket/athena-results/.',
},
{
message: 'Athena workgroup (optional)\nPress Enter to use the default workgroup "primary".',
},
{
message: 'Glue Data Catalog name (optional)\nPress Enter to use the default "AwsDataCatalog".',
},
],
},
]; ];
for (const testCase of cases) { for (const testCase of cases) {
@ -1966,6 +1991,40 @@ describe('setup databases step', () => {
expect(project.config.connections['clickhouse-warehouse']).not.toHaveProperty('schemas'); expect(project.config.connections['clickhouse-warehouse']).not.toHaveProperty('schemas');
}); });
it('maps Athena scripted database schema input to databases field', async () => {
await writeFile(
join(tempDir, 'ktx.yaml'),
[
'connections:',
' athena-warehouse:',
' driver: athena',
' region: us-east-1',
' s3_staging_dir: s3://my-bucket/athena-results/',
'',
].join('\n'),
'utf-8',
);
await runKtxSetupDatabasesStep(
{
projectDir: tempDir,
inputMode: 'disabled',
skipDatabases: false,
databaseConnectionIds: ['athena-warehouse'],
databaseSchemas: ['analytics', 'raw'],
},
makeIo().io,
{ testConnection: vi.fn(async () => 0), scanConnection: vi.fn(async () => 0) },
);
const project = await loadKtxProject({ projectDir: tempDir });
expect(project.config.connections['athena-warehouse']).toMatchObject({
driver: 'athena',
databases: ['analytics', 'raw'],
});
expect(project.config.connections['athena-warehouse']).not.toHaveProperty('schemas');
});
it('does not prompt for a bootstrap BigQuery dataset before scope discovery', async () => { it('does not prompt for a bootstrap BigQuery dataset before scope discovery', async () => {
const prompts = makePromptAdapter({ const prompts = makePromptAdapter({
multiselectValues: [['bigquery']], multiselectValues: [['bigquery']],
@ -3515,4 +3574,72 @@ describe('setup databases step', () => {
expect(io.stdout()).toContain('ktx cannot work until you add a database.'); expect(io.stdout()).toContain('ktx cannot work until you add a database.');
expect(await readFile(join(tempDir, 'ktx.yaml'), 'utf-8')).not.toContain('completed_steps:'); expect(await readFile(join(tempDir, 'ktx.yaml'), 'utf-8')).not.toContain('completed_steps:');
}); });
it('adds one non-interactive DuckDB connection from --database-url without prompting', async () => {
const io = makeIo();
const prompts = makePromptAdapter({});
const testConnection = vi.fn(async () => 0);
const scanConnection = vi.fn(async () => 0);
const result = await runKtxSetupDatabasesStep(
{
projectDir: tempDir,
inputMode: 'disabled',
databaseDrivers: ['duckdb'],
databaseConnectionId: 'duckdb-local',
databaseUrl: './warehouse.duckdb',
databaseSchemas: [],
skipDatabases: false,
},
io.io,
{ prompts, testConnection, scanConnection },
);
expect(result.status).toBe('ready');
expect(prompts.text).not.toHaveBeenCalled();
expect(testConnection).toHaveBeenCalledWith(tempDir, 'duckdb-local', expect.anything());
expect(scanConnection).toHaveBeenCalledWith(tempDir, 'duckdb-local', expect.anything());
const config = parseKtxProjectConfig(await readFile(join(tempDir, 'ktx.yaml'), 'utf-8'));
expect(config.connections['duckdb-local']).toEqual({
driver: 'duckdb',
path: './warehouse.duckdb',
});
expect(config.setup).toEqual({
database_connection_ids: ['duckdb-local'],
});
expect((await readKtxSetupState(tempDir)).completed_steps).toContain('databases');
});
it('adds an interactive DuckDB connection without prompting for a schema', async () => {
const prompts = makePromptAdapter({
selectValues: ['no'],
textValues: ['', './warehouse.duckdb'],
});
const pickers = makePickerStubs();
const result = await runKtxSetupDatabasesStep(
{
projectDir: tempDir,
inputMode: 'auto',
databaseDrivers: ['duckdb'],
databaseSchemas: [],
skipDatabases: false,
},
makeIo().io,
{
prompts,
testConnection: vi.fn(async () => 0),
scanConnection: vi.fn(async () => 0),
pickDatabaseScope: pickers.pickDatabaseScope,
},
);
expect(result.status).toBe('ready');
expect(pickers.scopeCalls).toHaveLength(0);
const config = parseKtxProjectConfig(await readFile(join(tempDir, 'ktx.yaml'), 'utf-8'));
expect(config.connections['duckdb-local']).toEqual({
driver: 'duckdb',
path: './warehouse.duckdb',
});
});
}); });

View file

@ -717,7 +717,7 @@ joins: []
expect(query).toHaveBeenCalledWith( expect(query).toHaveBeenCalledWith(
expect.objectContaining({ expect.objectContaining({
query: { measures: ['orders.order_count'], dimensions: [] }, query: { measures: ['orders.order_count'], dimensions: [], predefined_measures_only: false },
}), }),
); );
expect(JSON.parse(String(stdout.write.mock.calls[0][0]))).toMatchObject({ expect(JSON.parse(String(stdout.write.mock.calls[0][0]))).toMatchObject({

View file

@ -141,6 +141,41 @@ describe('runKtxSql', () => {
expect(io.stderr()).toBe(''); expect(io.stderr()).toBe('');
}); });
it('refuses raw SQL when the connection query_policy is semantic-layer-only', async () => {
const projectDir = join(tempDir, 'project');
await initKtxProject({ projectDir });
await writeConnections(projectDir, {
warehouse: { driver: 'sqlite', path: 'warehouse.db', query_policy: 'semantic-layer-only' },
});
const sqlAnalysis = makeSqlAnalysis({ ok: true, error: null });
const createScanConnector = vi.fn(async () => makeConnector());
const io = makeIo();
await expect(
runKtxSql(
{
command: 'execute',
projectDir,
connectionId: 'warehouse',
sql: 'select id from orders',
maxRows: 1000,
output: 'pretty',
json: false,
cliVersion: '0.0.0-test',
},
io.io,
{
createSqlAnalysis: () => sqlAnalysis,
createScanConnector,
},
),
).resolves.toBe(1);
expect(io.stderr()).toContain('query_policy: semantic-layer-only');
expect(sqlAnalysis.validateReadOnly).not.toHaveBeenCalled();
expect(createScanConnector).not.toHaveBeenCalled();
});
it('emits debug telemetry for SQL without raw query text', async () => { it('emits debug telemetry for SQL without raw query text', async () => {
vi.stubEnv('KTX_TELEMETRY_DEBUG', '1'); vi.stubEnv('KTX_TELEMETRY_DEBUG', '1');
vi.stubEnv('CI', ''); vi.stubEnv('CI', '');

791
pnpm-lock.yaml generated

File diff suppressed because it is too large Load diff

View file

@ -5,12 +5,18 @@ packages:
overrides: overrides:
"@types/node": ^24.3.0 "@types/node": ^24.3.0
"brace-expansion@>=5.0.0 <5.0.6": 5.0.6 "brace-expansion@>=5.0.0 <5.0.6": 5.0.6
esbuild: 0.28.1
fast-uri: 3.1.2 fast-uri: 3.1.2
fast-xml-builder: 1.1.7 fast-xml-builder: 1.1.7
hono: 4.12.21 form-data: 4.0.6
hono: 4.12.25
ip-address: 10.1.1 ip-address: 10.1.1
js-yaml: 4.2.0
postcss: 8.5.10 postcss: 8.5.10
ws: 8.20.1 "undici@6": 6.27.0
"undici@7": 7.28.0
vite: 8.0.16
ws: 8.21.0
dedupePeerDependents: false dedupePeerDependents: false
preferWorkspacePackages: true preferWorkspacePackages: true

View file

@ -1,6 +1,6 @@
[project] [project]
name = "ktx-daemon" name = "ktx-daemon"
version = "0.15.0" version = "0.16.0"
description = "Portable compute package for ktx semantic-layer operations" description = "Portable compute package for ktx semantic-layer operations"
readme = "README.md" readme = "README.md"
requires-python = ">=3.13" requires-python = ">=3.13"

View file

@ -35,6 +35,12 @@ from ktx_daemon.source_generation import (
generate_sources_response, generate_sources_response,
) )
# The caller (the Node client) sent a well-formed request the compute layer
# refused as invalid — e.g. the planner rejecting an agent's query. A distinct
# exit code lets the client classify it as an expected outcome rather than a
# ktx fault. Kept in sync with DAEMON_INPUT_REJECTED_EXIT_CODE on the Node side.
INPUT_REJECTED_EXIT_CODE = 3
def build_parser() -> argparse.ArgumentParser: def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(prog="ktx-daemon") parser = argparse.ArgumentParser(prog="ktx-daemon")
@ -210,9 +216,17 @@ def main(argv: list[str] | None = None) -> int:
return 2 return 2
sys.stdout.write(response.model_dump_json() + "\n") sys.stdout.write(response.model_dump_json() + "\n")
return 0 return 0
except (json.JSONDecodeError, ValidationError, ValueError) as error: except (json.JSONDecodeError, ValidationError) as error:
# Malformed request envelope — ktx sent bad JSON or a mis-shaped payload.
# That is a ktx fault, so keep the generic non-zero code (JSONDecodeError
# subclasses ValueError, so this clause must precede the ValueError one).
sys.stderr.write(f"{error}\n") sys.stderr.write(f"{error}\n")
return 1 return 1
except ValueError as error:
# The compute layer refused a well-formed request as invalid (e.g. the
# planner rejecting the agent's measures). Expected, not a fault.
sys.stderr.write(f"{error}\n")
return INPUT_REJECTED_EXIT_CODE
except Exception as error: except Exception as error:
from ktx_daemon.telemetry import report_exception from ktx_daemon.telemetry import report_exception

View file

@ -12,6 +12,7 @@ from typing import Any
from fastapi import FastAPI, HTTPException, Request from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse, Response from fastapi.responses import JSONResponse, Response
from pydantic import ValidationError
from ktx_daemon import VERSION from ktx_daemon import VERSION
from ktx_daemon.code_execution import ( from ktx_daemon.code_execution import (
@ -268,6 +269,14 @@ def create_app(
) -> SemanticLayerQueryResponse: ) -> SemanticLayerQueryResponse:
try: try:
return query_semantic_layer(request) return query_semantic_layer(request)
except ValidationError as error:
# A malformed source or response is a ktx contract fault, not a
# caller rejection; surface it as a server fault (500) so the Node
# client does not classify it as an expected query rejection, matching
# the subprocess transport (exit 1) and query_semantic_layer's own
# report. ValidationError subclasses ValueError, so catch it first.
logger.exception("Semantic query failed on a malformed source")
raise HTTPException(status_code=500, detail=str(error)) from error
except ValueError as error: except ValueError as error:
logger.warning("Semantic query rejected: %s", error) logger.warning("Semantic query rejected: %s", error)
raise HTTPException(status_code=400, detail=str(error)) from error raise HTTPException(status_code=400, detail=str(error)) from error

View file

@ -6,7 +6,7 @@ import time
from typing import Any from typing import Any
from ktx_daemon.telemetry import error_class, report_exception, track_telemetry_event from ktx_daemon.telemetry import error_class, report_exception, track_telemetry_event
from pydantic import BaseModel, ConfigDict, Field from pydantic import BaseModel, ConfigDict, Field, ValidationError
from semantic_layer.duplicate_check import validate_measure_duplicates from semantic_layer.duplicate_check import validate_measure_duplicates
from semantic_layer.engine import SemanticEngine from semantic_layer.engine import SemanticEngine
from semantic_layer.models import QueryResult, SourceDefinition from semantic_layer.models import QueryResult, SourceDefinition
@ -150,6 +150,11 @@ def query_semantic_layer(
track_telemetry_event( track_telemetry_event(
"sql_gen_completed", sql_fields, project_id=request.project_id "sql_gen_completed", sql_fields, project_id=request.project_id
) )
# A ValueError is the engine refusing the caller's query — an expected
# rejection surfaced to the agent, not a ktx fault. Keep it out of Error
# Tracking (a pydantic ValidationError subclasses ValueError but means a
# malformed request envelope, so it stays a reported fault).
if not isinstance(error, ValueError) or isinstance(error, ValidationError):
report_exception( report_exception(
error, error,
source="semantic-query", source="semantic-query",

View file

@ -424,6 +424,24 @@ def test_semantic_query_endpoint_maps_value_error_to_400() -> None:
assert "missing.order_count" in response.json()["detail"] assert "missing.order_count" in response.json()["detail"]
def test_semantic_query_endpoint_maps_malformed_source_to_fault() -> None:
client = TestClient(create_app(), raise_server_exceptions=False)
invalid_source = {
key: value for key, value in ORDERS_SOURCE.items() if key != "table"
}
response = client.post(
"/semantic-layer/query",
json={
"sources": [invalid_source],
"dialect": "postgres",
"query": {"measures": ["orders.order_count"], "dimensions": []},
},
)
assert response.status_code == 500
def test_semantic_validate_endpoint_returns_structured_validation() -> None: def test_semantic_validate_endpoint_returns_structured_validation() -> None:
client = TestClient(create_app()) client = TestClient(create_app())
invalid_source = { invalid_source = {

View file

@ -91,6 +91,22 @@ def test_command_returns_nonzero_for_invalid_json() -> None:
assert "Expecting property name enclosed in double quotes" in result.stderr assert "Expecting property name enclosed in double quotes" in result.stderr
def test_semantic_query_rejects_invalid_query_with_input_rejected_code() -> None:
# A planner ValueError (agent's measure references no source) is an expected
# input rejection, distinguished from a fault by INPUT_REJECTED_EXIT_CODE (3).
result = run_daemon_command(
"semantic-query",
{
"sources": [ORDERS_SOURCE],
"dialect": "postgres",
"query": {"measures": ["count(*)"], "dimensions": []},
},
)
assert result.returncode == 3, result.stderr
assert "does not reference any source" in result.stderr
def test_serve_http_command_starts_uvicorn_without_reading_stdin( def test_serve_http_command_starts_uvicorn_without_reading_stdin(
monkeypatch, monkeypatch,
) -> None: ) -> None:

View file

@ -51,6 +51,34 @@ def test_query_semantic_layer_generates_sql_and_plan() -> None:
assert response.plan["sources_used"] == ["orders"] assert response.plan["sources_used"] == ["orders"]
def test_query_semantic_layer_enforces_predefined_measures_only() -> None:
with pytest.raises(ValueError, match="composed measure"):
query_semantic_layer(
SemanticLayerQueryRequest(
sources=[ORDERS_SOURCE],
dialect="postgres",
query={
"measures": ["sum(orders.amount)"],
"predefined_measures_only": True,
},
)
)
def test_query_semantic_layer_allows_predefined_measures_under_policy() -> None:
response = query_semantic_layer(
SemanticLayerQueryRequest(
sources=[ORDERS_SOURCE],
dialect="postgres",
query={
"measures": ["orders.revenue"],
"predefined_measures_only": True,
},
)
)
assert "public.orders" in response.sql
def test_query_semantic_layer_emits_plan_and_sql_debug_events( def test_query_semantic_layer_emits_plan_and_sql_debug_events(
tmp_path: Path, tmp_path: Path,
monkeypatch, monkeypatch,
@ -97,7 +125,7 @@ def test_query_semantic_layer_emits_plan_and_sql_debug_events(
assert "public.orders" not in captured.err assert "public.orders" not in captured.err
def test_query_semantic_layer_reports_exception(monkeypatch) -> None: def test_query_semantic_layer_reports_unexpected_fault(monkeypatch) -> None:
from ktx_daemon import semantic_layer as semantic_layer_module from ktx_daemon import semantic_layer as semantic_layer_module
reports: list[dict[str, object]] = [] reports: list[dict[str, object]] = []
@ -105,12 +133,16 @@ def test_query_semantic_layer_reports_exception(monkeypatch) -> None:
def fake_report(exception: BaseException, **kwargs: object) -> None: def fake_report(exception: BaseException, **kwargs: object) -> None:
reports.append({"exception": exception, **kwargs}) reports.append({"exception": exception, **kwargs})
monkeypatch.setattr(semantic_layer_module, "report_exception", fake_report) def boom(*_args: object, **_kwargs: object) -> None:
raise RuntimeError("engine construction failed")
with pytest.raises(ValueError): monkeypatch.setattr(semantic_layer_module, "report_exception", fake_report)
monkeypatch.setattr(semantic_layer_module.SemanticEngine, "from_sources", boom)
with pytest.raises(RuntimeError):
query_semantic_layer( query_semantic_layer(
SemanticLayerQueryRequest( SemanticLayerQueryRequest(
sources=[ORDERS_SOURCE, ORDERS_SOURCE], sources=[ORDERS_SOURCE],
dialect="postgres", dialect="postgres",
projectId="a" * 64, projectId="a" * 64,
query={"measures": ["orders.order_count"]}, query={"measures": ["orders.order_count"]},
@ -124,6 +156,32 @@ def test_query_semantic_layer_reports_exception(monkeypatch) -> None:
assert reports[0]["project_id"] == "a" * 64 assert reports[0]["project_id"] == "a" * 64
def test_query_semantic_layer_does_not_report_expected_query_rejection(
monkeypatch,
) -> None:
from ktx_daemon import semantic_layer as semantic_layer_module
reports: list[dict[str, object]] = []
monkeypatch.setattr(
semantic_layer_module,
"report_exception",
lambda *_args, **kwargs: reports.append(kwargs),
)
# A planner ValueError is the engine refusing the agent's query — surfaced to
# the caller and re-raised, but never filed as a ktx fault.
with pytest.raises(ValueError, match="does not reference any source"):
query_semantic_layer(
SemanticLayerQueryRequest(
sources=[ORDERS_SOURCE],
dialect="postgres",
query={"measures": ["count(*)"]},
)
)
assert reports == []
def test_semantic_layer_request_rejects_project_id_field_name() -> None: def test_semantic_layer_request_rejects_project_id_field_name() -> None:
with pytest.raises(ValueError): with pytest.raises(ValueError):
SemanticLayerQueryRequest( SemanticLayerQueryRequest(

View file

@ -1,6 +1,6 @@
[project] [project]
name = "ktx-sl" name = "ktx-sl"
version = "0.15.0" version = "0.16.0"
description = "Agent-first semantic layer engine with aggregate locality" description = "Agent-first semantic layer engine with aggregate locality"
readme = "README.md" readme = "README.md"
requires-python = ">=3.13" requires-python = ">=3.13"

View file

@ -42,7 +42,7 @@ def validate_measure_duplicates(
parsed: list[tuple[str, exp.Expression | None, str | None, frozenset[str]]] = [] parsed: list[tuple[str, exp.Expression | None, str | None, frozenset[str]]] = []
for m in source.measures: for m in source.measures:
try: try:
quoted = quote_reserved_identifiers(m.expr) quoted = quote_reserved_identifiers(m.expr, dialect)
tree = sqlglot.parse_one(f"SELECT {quoted}", read=dialect) tree = sqlglot.parse_one(f"SELECT {quoted}", read=dialect)
expr_node = tree.expressions[0] if tree.expressions else None expr_node = tree.expressions[0] if tree.expressions else None
except Exception: except Exception:

View file

@ -671,7 +671,8 @@ class SqlGenerator:
"""Apply a measure-level filter by injecting CASE WHEN into each aggregate.""" """Apply a measure-level filter by injecting CASE WHEN into each aggregate."""
try: try:
tree = sqlglot.parse_one( tree = sqlglot.parse_one(
f"SELECT {quote_reserved_identifiers(expr)}", read=self.dialect f"SELECT {quote_reserved_identifiers(expr, self.dialect)}",
read=self.dialect,
) )
select_expr = tree.expressions[0] select_expr = tree.expressions[0]
if isinstance(select_expr, exp.Alias): if isinstance(select_expr, exp.Alias):
@ -723,7 +724,8 @@ class SqlGenerator:
def _translate_custom_funcs(self, expr: str) -> str: def _translate_custom_funcs(self, expr: str) -> str:
"""Translate custom functions: median(), percentile(), count_distinct().""" """Translate custom functions: median(), percentile(), count_distinct()."""
tree = sqlglot.parse_one( tree = sqlglot.parse_one(
f"SELECT {quote_reserved_identifiers(expr)}", read=self.dialect f"SELECT {quote_reserved_identifiers(expr, self.dialect)}",
read=self.dialect,
) )
has_custom = False has_custom = False
@ -767,7 +769,8 @@ class SqlGenerator:
def _extract_outer_aggregate(self, expr: str) -> tuple[str | None, str | None]: def _extract_outer_aggregate(self, expr: str) -> tuple[str | None, str | None]:
"""Use AST to extract the outer aggregate function name and inner expression.""" """Use AST to extract the outer aggregate function name and inner expression."""
tree = sqlglot.parse_one( tree = sqlglot.parse_one(
f"SELECT {quote_reserved_identifiers(expr)}", read=self.dialect f"SELECT {quote_reserved_identifiers(expr, self.dialect)}",
read=self.dialect,
) )
select_expr = tree.expressions[0] select_expr = tree.expressions[0]
if isinstance(select_expr, exp.Alias): if isinstance(select_expr, exp.Alias):
@ -788,7 +791,8 @@ class SqlGenerator:
if not replacements: if not replacements:
return expr return expr
tree = sqlglot.parse_one( tree = sqlglot.parse_one(
f"SELECT {quote_reserved_identifiers(expr)}", read=self.dialect f"SELECT {quote_reserved_identifiers(expr, self.dialect)}",
read=self.dialect,
) )
def _replace(node): def _replace(node):
@ -857,13 +861,26 @@ class SqlGenerator:
return self._time_trunc(dim.granularity, field) return self._time_trunc(dim.granularity, field)
return field return field
_WEEKDAYS = frozenset(
{"sunday", "monday", "tuesday", "wednesday", "thursday", "friday", "saturday"}
)
def _bigquery_date_part(self, granularity: str) -> str:
# BigQuery spells a week starting on a given weekday as WEEK(MONDAY),
# not the bare WEEK_MONDAY that other systems accept.
if granularity.startswith("week_"):
weekday = granularity[len("week_") :]
if weekday in self._WEEKDAYS:
return f"WEEK({weekday.upper()})"
return granularity.upper()
def _time_trunc(self, granularity: str, field: str) -> str: def _time_trunc(self, granularity: str, field: str) -> str:
"""Generate dialect-appropriate time truncation expression.""" """Generate dialect-appropriate time truncation expression."""
g = granularity.lower() g = granularity.lower()
if self.dialect == "sqlite": if self.dialect == "sqlite":
return self._sqlite_time_trunc(g, field) return self._sqlite_time_trunc(g, field)
if self.dialect == "bigquery": if self.dialect == "bigquery":
return f"DATE_TRUNC({field}, {g.upper()})" return f"DATE_TRUNC({field}, {self._bigquery_date_part(g)})"
if self.dialect == "mysql": if self.dialect == "mysql":
return self._mysql_time_trunc(g, field) return self._mysql_time_trunc(g, field)
return f"DATE_TRUNC('{g}', {field})" return f"DATE_TRUNC('{g}', {field})"
@ -1257,7 +1274,8 @@ class SqlGenerator:
"""Qualify bare column references in a computed column expression with the source name.""" """Qualify bare column references in a computed column expression with the source name."""
try: try:
tree = sqlglot.parse_one( tree = sqlglot.parse_one(
f"SELECT {quote_reserved_identifiers(expr)}", read=self.dialect f"SELECT {quote_reserved_identifiers(expr, self.dialect)}",
read=self.dialect,
) )
def _qualify(node: exp.Expression) -> exp.Expression: def _qualify(node: exp.Expression) -> exp.Expression:
@ -1403,7 +1421,7 @@ class SqlGenerator:
try: try:
# Quote reserved-word identifiers so target dialect parsers do not # Quote reserved-word identifiers so target dialect parsers do not
# confuse them with keywords (e.g. Snowflake's SAMPLE, QUALIFY). # confuse them with keywords (e.g. Snowflake's SAMPLE, QUALIFY).
quoted_outer = quote_reserved_identifiers(outer_sql) quoted_outer = quote_reserved_identifiers(outer_sql, self.dialect)
results = sqlglot.transpile( results = sqlglot.transpile(
quoted_outer, read=self.dialect, write=self.dialect quoted_outer, read=self.dialect, write=self.dialect
) )

View file

@ -254,7 +254,7 @@ class JoinGraph:
from sqlglot import exp as _exp from sqlglot import exp as _exp
from semantic_layer.parser import quote_reserved_identifiers from semantic_layer.parser import quote_reserved_identifiers
quoted = quote_reserved_identifiers(on_clause) quoted = quote_reserved_identifiers(on_clause, self.dialect)
tree = sqlglot.parse_one( tree = sqlglot.parse_one(
f"SELECT 1 FROM _a JOIN _b ON {quoted}", read=self.dialect f"SELECT 1 FROM _a JOIN _b ON {quoted}", read=self.dialect
) )

View file

@ -169,6 +169,9 @@ class SemanticQuery(BaseModel):
order_by: list[str | dict[str, Any]] = [] order_by: list[str | dict[str, Any]] = []
limit: int = 1000 limit: int = 1000
include_empty: bool = True include_empty: bool = True
# Set by ktx from the connection's query_policy, never by agent input:
# reject runtime-composed measures so only predefined measures execute.
predefined_measures_only: bool = False
@model_validator(mode="after") @model_validator(mode="after")
def _validate_limit(self) -> SemanticQuery: def _validate_limit(self) -> SemanticQuery:

View file

@ -6,6 +6,7 @@ from dataclasses import dataclass, field
import sqlglot import sqlglot
from sqlglot import exp from sqlglot import exp
from sqlglot.dialects.dialect import Dialect
# DIALECT CONVENTION: # DIALECT CONVENTION:
# `ExpressionParser` wraps read-only AST walks over user-authored # `ExpressionParser` wraps read-only AST walks over user-authored
@ -154,12 +155,29 @@ def _strip_quotes(name: str) -> str:
return name return name
def quote_reserved_identifiers(expr: str) -> str: @functools.lru_cache(maxsize=32)
def _identifier_quote(dialect: str) -> str:
"""The character `dialect` quotes identifiers with (backtick for BigQuery/MySQL,
double-quote for ANSI dialects). A reserved-word column must be quoted with this so a
backtick dialect does not read a double-quoted identifier as a string literal."""
try:
start = Dialect.get_or_raise(dialect).IDENTIFIER_START
except Exception:
start = '"'
return start or '"'
def quote_reserved_identifiers(expr: str, dialect: str = "postgres") -> str:
"""Quote source.column references where either part is a SQL reserved word. """Quote source.column references where either part is a SQL reserved word.
String literals are masked before processing to prevent matching Quoted with `dialect`'s identifier character, because the caller parses the
dotted identifiers inside quoted strings like 'group.value'. result in that dialect: a double-quoted identifier on BigQuery/MySQL is a
string literal, not an identifier. String literals are masked before
processing to prevent matching dotted identifiers inside quoted strings like
'group.value'.
""" """
quote = _identifier_quote(dialect)
# Mask string literals to avoid matching inside them # Mask string literals to avoid matching inside them
literals: list[str] = [] literals: list[str] = []
@ -172,16 +190,16 @@ def quote_reserved_identifiers(expr: str) -> str:
def _quote_match(m: re.Match) -> str: def _quote_match(m: re.Match) -> str:
source, col = m.group(1), m.group(2) source, col = m.group(1), m.group(2)
start = m.start() start = m.start()
if start > 0 and masked[start - 1] == '"': if start > 0 and masked[start - 1] == quote:
return m.group(0) return m.group(0)
needs_quote = False needs_quote = False
source_q = source source_q = source
col_q = col col_q = col
if source.lower() in _SQL_RESERVED: if source.lower() in _SQL_RESERVED:
source_q = f'"{source}"' source_q = f"{quote}{source}{quote}"
needs_quote = True needs_quote = True
if col.lower() in _SQL_RESERVED: if col.lower() in _SQL_RESERVED:
col_q = f'"{col}"' col_q = f"{quote}{col}{quote}"
needs_quote = True needs_quote = True
if needs_quote: if needs_quote:
return f"{source_q}.{col_q}" return f"{source_q}.{col_q}"
@ -196,13 +214,13 @@ def quote_reserved_identifiers(expr: str) -> str:
return result return result
def _predicate_select(expr: str) -> str: def _predicate_select(expr: str, dialect: str = "postgres") -> str:
"""Wrap a user expression as `SELECT * WHERE …`, quoting reserved identifiers. """Wrap a user expression as `SELECT * WHERE …`, quoting reserved identifiers.
Predicate, not projection: T-SQL reads a top-level `col = 'value'` projection Predicate, not projection: T-SQL reads a top-level `col = 'value'` projection
as the `alias = expression` form and would compile the filter to `'value' AS col`. as the `alias = expression` form and would compile the filter to `'value' AS col`.
""" """
return f"SELECT * WHERE {quote_reserved_identifiers(expr)}" return f"SELECT * WHERE {quote_reserved_identifiers(expr, dialect)}"
@functools.lru_cache(maxsize=256) @functools.lru_cache(maxsize=256)
@ -220,7 +238,11 @@ def parse_predicate(expr: str, dialect: str) -> exp.Expression:
Uncached, so the result is safe to `.transform()`; raises on unparseable input. Uncached, so the result is safe to `.transform()`; raises on unparseable input.
""" """
return sqlglot.parse_one(_predicate_select(expr), read=dialect).find(exp.Where).this return (
sqlglot.parse_one(_predicate_select(expr, dialect), read=dialect)
.find(exp.Where)
.this
)
class ExpressionParser: class ExpressionParser:
@ -237,7 +259,7 @@ class ExpressionParser:
def _parse_as_select(self, expr: str) -> exp.Expression: def _parse_as_select(self, expr: str) -> exp.Expression:
"""Parse a user fragment for read-only AST walks, via the parse cache.""" """Parse a user fragment for read-only AST walks, via the parse cache."""
return _cached_parse_select(_predicate_select(expr), self.dialect) return _cached_parse_select(_predicate_select(expr, self.dialect), self.dialect)
def parse( def parse(
self, self,

View file

@ -62,6 +62,10 @@ class QueryPlanner:
# 2. Resolve measures (parse, look up pre-defined, classify) # 2. Resolve measures (parse, look up pre-defined, classify)
raw_measures = self._resolve_measures(query.measures) raw_measures = self._resolve_measures(query.measures)
if query.predefined_measures_only:
self._reject_composed_measures(raw_measures)
self._reject_composed_filter_aggregates(query.filters)
# 3. Topological sort for derived measures # 3. Topological sort for derived measures
measures = self._topological_sort_measures(raw_measures) measures = self._topological_sort_measures(raw_measures)
@ -239,6 +243,42 @@ class QueryPlanner:
measures = self._qualify_duplicate_names(measures) measures = self._qualify_duplicate_names(measures)
return measures return measures
def _reject_composed_measures(self, measures: list[ResolvedMeasure]) -> None:
composed = [m for m in measures if m.provenance is Provenance.COMPOSED]
if not composed:
return
rejected = ", ".join(f"'{m.expr}'" for m in composed)
raise ValueError(
f"Only predefined semantic-layer measures can be queried on this "
f"connection (query_policy: semantic-layer-only); composed measure "
f"expressions are not allowed: {rejected}. Use measures declared "
f"on the semantic-layer sources, referenced as 'source.measure'."
)
def _reject_composed_filter_aggregates(self, filters: list[str]) -> None:
# Predefined measures are compared in HAVING by name (e.g.
# `orders.revenue > 100`), which parses as a plain column ref. Any
# aggregate function written into a filter is therefore a runtime-composed
# aggregate the policy forbids — without this guard it would slip through
# `_classify_filters` into HAVING and evaluate an arbitrary aggregate.
composed: list[str] = []
for f in filters:
if not f or not f.strip():
continue
for clause in self._split_top_level_and(f):
if self.parser.parse(clause).is_aggregate:
composed.append(clause)
if not composed:
return
rejected = ", ".join(f"'{c}'" for c in composed)
raise ValueError(
f"Only predefined semantic-layer measures can be queried on this "
f"connection (query_policy: semantic-layer-only); filters may not "
f"compose aggregate expressions: {rejected}. Filter on declared "
f"dimensions or columns, or compare a predefined measure by name "
f"(e.g. 'orders.revenue > 100')."
)
def _collect_colliding_predefined_names(self, raw: list[str | dict]) -> set[str]: def _collect_colliding_predefined_names(self, raw: list[str | dict]) -> set[str]:
counts: Counter[str] = Counter() counts: Counter[str] = Counter()
for item in raw: for item in raw:
@ -654,7 +694,8 @@ class QueryPlanner:
all_dep_names.add(dep_name) all_dep_names.add(dep_name)
named.add(dep_name) named.add(dep_name)
tree = sqlglot.parse_one( tree = sqlglot.parse_one(
f"SELECT {quote_reserved_identifiers(expr)}", dialect=self.dialect f"SELECT {quote_reserved_identifiers(expr, self.dialect)}",
dialect=self.dialect,
) )
def _replace(node): def _replace(node):
@ -698,7 +739,8 @@ class QueryPlanner:
"""Reject expressions with nested aggregate functions (e.g., avg(sum(x))).""" """Reject expressions with nested aggregate functions (e.g., avg(sum(x)))."""
try: try:
tree = sqlglot.parse_one( tree = sqlglot.parse_one(
f"SELECT {quote_reserved_identifiers(expr)}", dialect=self.dialect f"SELECT {quote_reserved_identifiers(expr, self.dialect)}",
dialect=self.dialect,
) )
for agg_node in tree.find_all(exp.AggFunc): for agg_node in tree.find_all(exp.AggFunc):
# Check if this aggregate contains another aggregate inside # Check if this aggregate contains another aggregate inside

View file

@ -18,6 +18,8 @@ SUPPORTED_TABLE_IDENTIFIER_DIALECTS = {
"sqlite", "sqlite",
"tsql", "tsql",
"clickhouse", "clickhouse",
"athena",
"duckdb",
} }
ParseTableIdentifierReason = Literal[ ParseTableIdentifierReason = Literal[

View file

@ -0,0 +1,93 @@
"""Reserved-word identifiers and week-of-day granularity must produce valid SQL
in the connection's native dialect.
Regression: on backtick dialects (BigQuery, MySQL) a reserved-word column such
as `default` or `like` was quoted with postgres-style double quotes, which those
dialects read as a *string literal* `WHERE loans.'like' = 'x'` yielding the
production error `Syntax error: Unexpected string literal "like"`.
"""
from __future__ import annotations
import sqlglot
from conftest import make_engine
def _loans_engine(dialect: str):
source = {
"name": "loans",
"table": "mydataset.loans",
"grain": ["id"],
"columns": [
{"name": "id", "type": "number"},
{"name": "amount", "type": "number"},
{"name": "default", "type": "boolean"},
{"name": "like", "type": "string"},
{"name": "created_at", "type": "time"},
],
"measures": [{"name": "total", "expr": "sum(amount)"}],
}
return make_engine({"loans": source}, dialect=dialect)
def test_reserved_word_column_filter_valid_on_bigquery():
sql = (
_loans_engine("bigquery")
.query({"measures": ["loans.total"], "filters": ["loans.default = true"]})
.sql
)
sqlglot.parse_one(sql, read="bigquery") # must not raise
assert "`default`" in sql
assert "'default'" not in sql
def test_reserved_word_column_filter_valid_on_mysql():
sql = (
_loans_engine("mysql")
.query({"measures": ["loans.total"], "filters": ["loans.default = true"]})
.sql
)
sqlglot.parse_one(sql, read="mysql")
assert "`default`" in sql
assert "'default'" not in sql
def test_like_column_filter_valid_on_bigquery():
# Mirrors the production message: Unexpected string literal "like".
sql = (
_loans_engine("bigquery")
.query({"measures": ["loans.total"], "filters": ["loans.like = 'x'"]})
.sql
)
sqlglot.parse_one(sql, read="bigquery")
assert "`like`" in sql
assert "'like'" not in sql
def test_reserved_word_column_still_double_quoted_on_snowflake():
sql = (
_loans_engine("snowflake")
.query({"measures": ["loans.total"], "filters": ["loans.default = true"]})
.sql
)
sqlglot.parse_one(sql, read="snowflake")
assert '"default"' in sql
def test_week_weekday_granularity_translated_on_bigquery():
sql = (
_loans_engine("bigquery")
.query(
{
"measures": ["loans.total"],
"dimensions": [
{"field": "loans.created_at", "granularity": "week_monday"}
],
}
)
.sql
)
sqlglot.parse_one(sql, read="bigquery") # must not raise
assert "WEEK_MONDAY" not in sql
assert "WEEK(MONDAY)" in sql

View file

@ -0,0 +1,148 @@
"""predefined_measures_only rejects runtime-composed measures while leaving
predefined measures, predefined derived chains, dimensions, and filters usable."""
from __future__ import annotations
import pytest
from semantic_layer.engine import SemanticEngine
from semantic_layer.models import SourceDefinition
def _engine() -> SemanticEngine:
orders = SourceDefinition(
name="orders",
table="public.orders",
grain=["id"],
columns=[
{"name": "id", "type": "number"},
{"name": "amount", "type": "number"},
{"name": "status", "type": "string"},
],
measures=[
{"name": "revenue", "expr": "sum(amount)"},
{"name": "order_count", "expr": "count(*)"},
{"name": "aov", "expr": "revenue / order_count"},
],
)
return SemanticEngine.from_sources({"orders": orders})
def test_rejects_composed_string_measure() -> None:
with pytest.raises(ValueError, match="composed measure") as excinfo:
_engine().query(
{
"measures": ["sum(orders.amount)"],
"predefined_measures_only": True,
}
)
assert "sum(orders.amount)" in str(excinfo.value)
assert "query_policy: semantic-layer-only" in str(excinfo.value)
def test_rejects_composed_dict_measure() -> None:
with pytest.raises(ValueError, match="composed measure"):
_engine().query(
{
"measures": [{"expr": "avg(orders.amount)", "name": "avg_amount"}],
"predefined_measures_only": True,
}
)
def test_rejects_query_time_derivation_over_predefined_measures() -> None:
with pytest.raises(ValueError, match="composed measure"):
_engine().query(
{
"measures": [
{"expr": "orders.revenue / orders.order_count", "name": "ratio"}
],
"predefined_measures_only": True,
}
)
def test_rejects_composed_aggregate_in_filter() -> None:
# A HAVING-classified filter must not smuggle a runtime aggregate the
# measures guard would reject (threshold-probing bypass).
with pytest.raises(ValueError, match="compose aggregate expressions") as excinfo:
_engine().query(
{
"measures": ["orders.revenue"],
"dimensions": ["orders.status"],
"filters": ["avg(orders.amount) > 100"],
"predefined_measures_only": True,
}
)
assert "avg(orders.amount) > 100" in str(excinfo.value)
assert "query_policy: semantic-layer-only" in str(excinfo.value)
def test_rejects_composed_aggregate_in_compound_filter() -> None:
with pytest.raises(ValueError, match="compose aggregate expressions"):
_engine().query(
{
"measures": ["orders.revenue"],
"filters": ["orders.status = 'active' AND sum(orders.amount) > 5000"],
"predefined_measures_only": True,
}
)
def test_allows_predefined_measure_having_filter() -> None:
result = _engine().query(
{
"measures": ["orders.revenue"],
"dimensions": ["orders.status"],
"filters": ["orders.revenue > 100"],
"predefined_measures_only": True,
}
)
assert "having" in result.sql.lower()
def test_composed_aggregate_filter_allowed_when_flag_absent() -> None:
result = _engine().query(
{
"measures": ["orders.revenue"],
"filters": ["avg(orders.amount) > 100"],
}
)
assert "having" in result.sql.lower()
def test_allows_predefined_measure_with_dimensions_and_filters() -> None:
result = _engine().query(
{
"measures": ["orders.revenue"],
"dimensions": ["orders.status"],
"filters": ["orders.status != 'cancelled'"],
"predefined_measures_only": True,
}
)
assert "sum" in result.sql.lower()
def test_allows_unqualified_predefined_measure() -> None:
result = _engine().query(
{
"measures": ["revenue"],
"predefined_measures_only": True,
}
)
assert "sum" in result.sql.lower()
def test_allows_predefined_derived_measure_chain() -> None:
result = _engine().query(
{
"measures": ["orders.aov"],
"predefined_measures_only": True,
}
)
assert "sum" in result.sql.lower()
def test_composed_measures_allowed_when_flag_absent() -> None:
result = _engine().query({"measures": ["sum(orders.amount)"]})
assert "sum" in result.sql.lower()

View file

@ -23,6 +23,16 @@ def test_parse_table_identifier_supported_dialects_and_aliases() -> None:
sql_table_name="RAW.PUBLIC.ORDERS", sql_table_name="RAW.PUBLIC.ORDERS",
dialect="snowflake", dialect="snowflake",
), ),
ParseTableIdentifierItem(
key="athena",
sql_table_name="analytics.orders",
dialect="athena",
),
ParseTableIdentifierItem(
key="duckdb",
sql_table_name="analytics.orders",
dialect="duckdb",
),
] ]
) )
@ -37,6 +47,14 @@ def test_parse_table_identifier_supported_dialects_and_aliases() -> None:
assert response["sf"].catalog == "RAW" assert response["sf"].catalog == "RAW"
assert response["sf"].schema_ == "PUBLIC" assert response["sf"].schema_ == "PUBLIC"
assert response["sf"].name == "ORDERS" assert response["sf"].name == "ORDERS"
assert response["athena"].ok is True
assert response["athena"].schema_ == "analytics"
assert response["athena"].name == "orders"
assert response["athena"].canonical_table == "analytics.orders"
assert response["duckdb"].ok is True
assert response["duckdb"].schema_ == "analytics"
assert response["duckdb"].name == "orders"
assert response["duckdb"].canonical_table == "analytics.orders"
def test_parse_table_identifier_rejects_non_physical_inputs() -> None: def test_parse_table_identifier_rejects_non_physical_inputs() -> None:

Some files were not shown because too many files have changed in this diff Show more