From 590a5d491c8d611892ca093465b62340766c8dd3 Mon Sep 17 00:00:00 2001 From: VC Date: Thu, 7 May 2026 16:11:44 -0400 Subject: [PATCH 1/2] feat: add native Trino driver --- README.md | 2 +- bun.lock | 5 + docs/docs/configure/warehouses.md | 48 +++- .../data-engineering/tools/warehouse-tools.md | 5 +- docs/docs/drivers.md | 16 +- docs/docs/getting-started/index.md | 2 +- packages/drivers/package.json | 3 +- packages/drivers/src/index.ts | 1 + packages/drivers/src/normalize.ts | 10 + packages/drivers/src/trino.ts | 240 ++++++++++++++++++ packages/drivers/test/trino-unit.test.ts | 154 +++++++++++ packages/opencode/script/publish.ts | 1 + .../native/connections/dbt-profiles.ts | 2 +- .../native/connections/docker-discovery.ts | 8 + .../altimate/native/connections/register.ts | 5 + .../altimate/native/connections/registry.ts | 5 + .../altimate/native/finops/query-history.ts | 3 + .../src/altimate/tools/project-scan.ts | 18 ++ .../src/altimate/tools/warehouse-add.ts | 3 +- .../test/altimate/connections.test.ts | 9 +- .../test/altimate/driver-normalize.test.ts | 42 +++ .../test/altimate/tools/sql-explain.test.ts | 11 + .../opencode/test/tool/project-scan.test.ts | 40 +++ 23 files changed, 619 insertions(+), 14 deletions(-) create mode 100644 packages/drivers/src/trino.ts create mode 100644 packages/drivers/test/trino-unit.test.ts diff --git a/README.md b/README.md index a9c3b542e..9e9815c05 100644 --- a/README.md +++ b/README.md @@ -167,7 +167,7 @@ Each mode has scoped permissions, tool access, and SQL write-access control. ## Supported Warehouses -Snowflake · BigQuery · Databricks · PostgreSQL · Redshift · ClickHouse · DuckDB · MySQL · SQL Server (incl. Microsoft Fabric) · Oracle · SQLite · MongoDB +Snowflake · BigQuery · Databricks · PostgreSQL · Redshift · Trino · ClickHouse · DuckDB · MySQL · SQL Server (incl. Microsoft Fabric) · Oracle · SQLite · MongoDB First-class support with schema indexing, query execution, and metadata introspection. SSH tunneling available for secure connections. diff --git a/bun.lock b/bun.lock index b21e47528..2806b8527 100644 --- a/bun.lock +++ b/bun.lock @@ -53,6 +53,7 @@ "oracledb": "^6.0.0", "pg": "^8.0.0", "snowflake-sdk": "^2.0.3", + "trino-client": "^0.2.8", }, }, "packages/opencode": { @@ -2471,6 +2472,8 @@ "trim-lines": ["trim-lines@3.0.1", "", {}, "sha512-kRj8B+YHZCc9kQYdWfJB2/oUl9rA99qbowYYBtr4ui4mZyAQ2JpvVBd/6U2YloATfqBhBTSMhTpgBHtU0Mf3Rg=="], + "trino-client": ["trino-client@0.2.9", "", { "dependencies": { "axios": "1.13.2" } }, "sha512-bINcQxDH5v9DR1zqXtoNqnRJ5leYMyYzRuFZLsoqg4irWYPDh/4wBndC8c2x+QfNIOr6c35BtHKOwotelqSATg=="], + "triple-beam": ["triple-beam@1.4.1", "", {}, "sha512-aZbgViZrg1QNcG+LULa7nhZpJTZSLm/mXnHXnbAbjmN5aSa0y7V+wvv6+4WaBtpISJzThKy+PIPxc1Nq1EJ9mg=="], "ts-algebra": ["ts-algebra@2.0.0", "", {}, "sha512-FPAhNPFMrkwz76P7cdjdmiShwMynZYN6SgOujD1urY4oNm80Ou9oMdmbR45LotcKOXoy7wSmHkRFE6Mxbrhefw=="], @@ -3189,6 +3192,8 @@ "tree-sitter-bash/node-addon-api": ["node-addon-api@8.7.0", "", {}, "sha512-9MdFxmkKaOYVTV+XVRG8ArDwwQ77XIgIPyKASB1k3JPq3M8fGQQQE3YpMOrKm6g//Ktx8ivZr8xo1Qmtqub+GA=="], + "trino-client/axios": ["axios@1.13.2", "", { "dependencies": { "follow-redirects": "^1.15.6", "form-data": "^4.0.4", "proxy-from-env": "^1.1.0" } }, "sha512-VPk9ebNqPcy5lRGuSlKx752IlDatOjT9paPlm8A7yOuW2Fbvp4X3JznJtT4f0GzGLLiWE9W8onz51SqLYwzGaA=="], + "tuf-js/make-fetch-happen": ["make-fetch-happen@15.0.5", "", { "dependencies": { "@gar/promise-retry": "^1.0.0", "@npmcli/agent": "^4.0.0", "@npmcli/redact": "^4.0.0", "cacache": "^20.0.1", "http-cache-semantics": "^4.1.1", "minipass": "^7.0.2", "minipass-fetch": "^5.0.0", "minipass-flush": "^1.0.5", "minipass-pipeline": "^1.2.4", "negotiator": "^1.0.0", "proc-log": "^6.0.0", "ssri": "^13.0.0" } }, "sha512-uCbIa8jWWmQZt4dSnEStkVC6gdakiinAm4PiGsywIkguF0eWMdcjDz0ECYhUolFU3pFLOev9VNPCEygydXnddg=="], "wide-align/string-width": ["string-width@4.2.3", "", { "dependencies": { "emoji-regex": "^8.0.0", "is-fullwidth-code-point": "^3.0.0", "strip-ansi": "^6.0.1" } }, "sha512-wKyQRQpjJ0sIp62ErSZdGsjMJWsap5oRNihHhu6G7JVO/9jIB6UyevL+tXuOqrng8j/cxKTWyWUwvSTriiZz/g=="], diff --git a/docs/docs/configure/warehouses.md b/docs/docs/configure/warehouses.md index d5f32cd84..6d7604842 100644 --- a/docs/docs/configure/warehouses.md +++ b/docs/docs/configure/warehouses.md @@ -1,6 +1,6 @@ # Warehouses -Altimate Code connects to 12 warehouse types. Configure them in `.altimate-code/connections.json` (project-local) or `~/.altimate-code/connections.json` (global). +Altimate Code connects to 13 warehouse types. Configure them in `.altimate-code/connections.json` (project-local) or `~/.altimate-code/connections.json` (global). ## Configuration @@ -118,6 +118,50 @@ If you're already authenticated via `gcloud`, omit `credentials_path`: | `catalog` | No | Unity Catalog name | | `schema` | No | Schema/database name | +## Trino + +```json +{ + "trino-prod": { + "type": "trino", + "host": "trino.example.com", + "port": 8443, + "protocol": "https", + "catalog": "iceberg", + "schema": "analytics", + "user": "analyst", + "password": "{env:TRINO_PASSWORD}" + } +} +``` + +| Field | Required | Description | +|-------|----------|-------------| +| `connection_string` | No | Full server URL (alternative to host/port, e.g. `https://trino.example.com:8443`) | +| `host` | No | Hostname (default: `localhost`) | +| `port` | No | Port (default: `8080`, or `8443` when `protocol` is `https`) | +| `protocol` | No | `http` or `https` (default: `http`) | +| `catalog` | Recommended | Default catalog for query execution and required for schema/table introspection | +| `schema` | No | Default schema | +| `user` | No | Trino user (default sent by the client if omitted) | +| `password` | Auth | Basic authentication password | +| `access_token` | Auth | Bearer/JWT token | +| `extra_headers` | No | Additional HTTP headers to send to Trino | + +### Using a bearer token + +```json +{ + "trino-prod": { + "type": "trino", + "connection_string": "https://trino.example.com:8443", + "catalog": "hive", + "schema": "default", + "access_token": "{env:TRINO_TOKEN}" + } +} +``` + ## PostgreSQL ```json @@ -471,7 +515,7 @@ The `/discover` command can automatically detect warehouse connections from: | Source | Detection | |--------|-----------| | dbt profiles | Searches for `profiles.yml` (see resolution order below) | -| Docker containers | Finds running PostgreSQL, MySQL, SQL Server, and ClickHouse containers | +| Docker containers | Finds running PostgreSQL, MySQL, SQL Server, Trino, and ClickHouse containers | | Environment variables | Scans for `SNOWFLAKE_ACCOUNT`, `PGHOST`, `DATABRICKS_HOST`, etc. | ### dbt profiles.yml resolution order diff --git a/docs/docs/data-engineering/tools/warehouse-tools.md b/docs/docs/data-engineering/tools/warehouse-tools.md index 89a1c2c46..85fd6fdf6 100644 --- a/docs/docs/data-engineering/tools/warehouse-tools.md +++ b/docs/docs/data-engineering/tools/warehouse-tools.md @@ -54,7 +54,7 @@ env_bigquery | bigquery | GOOGLE_APPLICATION_CREDENTIALS | **dbt project** | Walks up directories for `dbt_project.yml`, reads name/profile | | **dbt manifest** | Parses `target/manifest.json` for model/source/test counts | | **dbt profiles** | Searches for `profiles.yml`: `DBT_PROFILES_DIR` env var → project root → `/.dbt/profiles.yml` | -| **Docker DBs** | Bridge call to discover running PostgreSQL/MySQL/MSSQL containers | +| **Docker DBs** | Bridge call to discover running PostgreSQL/MySQL/MSSQL/Trino containers | | **Existing connections** | Bridge call to list already-configured warehouses | | **Environment variables** | Scans `process.env` for warehouse signals (see table below) | | **Schema cache** | Bridge call for indexed warehouse status | @@ -73,6 +73,7 @@ env_bigquery | bigquery | GOOGLE_APPLICATION_CREDENTIALS | MongoDB | `MONGODB_URI`, `MONGO_URL` | | Redshift | `REDSHIFT_HOST` | | ClickHouse | `CLICKHOUSE_HOST`, `CLICKHOUSE_URL` | +| Trino | `TRINO_HOST`, `TRINO_SERVER`, `TRINO_URL`, `DATABASE_URL` | ### Parameters @@ -166,7 +167,7 @@ Remove an existing warehouse connection. ## warehouse_discover -Discover database containers running in Docker. Detects PostgreSQL, MySQL/MariaDB, SQL Server, ClickHouse, and MongoDB containers with their connection details. +Discover database containers running in Docker. Detects PostgreSQL, MySQL/MariaDB, SQL Server, Trino, ClickHouse, and MongoDB containers with their connection details. ``` > warehouse_discover diff --git a/docs/docs/drivers.md b/docs/docs/drivers.md index 2b5c29e95..bcd1e0009 100644 --- a/docs/docs/drivers.md +++ b/docs/docs/drivers.md @@ -2,7 +2,7 @@ ## Overview -Altimate Code connects to 12 databases natively via TypeScript drivers. No Python dependency required. Drivers are loaded lazily, so only the driver you need is imported at runtime. +Altimate Code connects to 13 databases natively via TypeScript drivers. No Python dependency required. Drivers are loaded lazily, so only the driver you need is imported at runtime. ## Support Matrix @@ -19,6 +19,7 @@ Altimate Code connects to 12 databases natively via TypeScript drivers. No Pytho | Databricks | `@databricks/sql` | PAT, OAuth | ✅ Live account | 24 E2E tests, Unity Catalog support | | MongoDB | `mongodb` | Password, Connection String | ✅ Docker | 90 E2E tests, MQL queries, aggregation pipelines | | ClickHouse | `@clickhouse/client` | Password, Connection String, TLS | ✅ Docker | HTTP(S) protocol, ClickHouse Cloud support | +| Trino | `trino-client` | None, Basic, Bearer Token | ❌ Needs Trino cluster | HTTP(S) protocol, catalog/schema support | | Oracle | `oracledb` (thin) | Password | ❌ Needs Oracle 12.1+ | Thin mode only, no Instant Client | ## Installation @@ -43,6 +44,7 @@ bun add snowflake-sdk # Snowflake bun add @google-cloud/bigquery # BigQuery bun add @databricks/sql # Databricks bun add @clickhouse/client # ClickHouse +bun add trino-client # Trino bun add oracledb # Oracle (thin mode) ``` @@ -121,6 +123,16 @@ Optional: `location` (e.g. `us`, `eu`, `us-central1`, `asia-northeast1`). Requir |--------|--------------| | PAT | `server_hostname`, `http_path`, `access_token` | +### Trino +| Method | Config Fields | +|--------|--------------| +| No auth | `host`, `port`, `catalog`, `schema`, `user` | +| Basic | `host`, `port`, `catalog`, `schema`, `user`, `password` | +| Bearer token | `host`, `port`, `catalog`, `schema`, `access_token` | +| Connection String | `connection_string: "https://trino.example.com:8443"` | + +Trino uses the official `trino-client` package over HTTP(S). Set `catalog` for schema/table introspection. Query history is not exposed as a durable history source by the native driver. + ### MySQL | Method | Config Fields | |--------|--------------| @@ -193,7 +205,7 @@ SSH auth types: `"key"` (default) or `"password"` (set `ssh_password`). The CLI auto-discovers connections from: -1. **Docker containers**: detects running PostgreSQL, MySQL, MariaDB, SQL Server, Oracle, ClickHouse, MongoDB containers +1. **Docker containers**: detects running PostgreSQL, MySQL, MariaDB, SQL Server, Oracle, Trino, ClickHouse, MongoDB containers 2. **dbt profiles**: parses `~/.dbt/profiles.yml` for all supported adapters 3. **Environment variables**: detects `SNOWFLAKE_ACCOUNT`, `PGHOST`, `MYSQL_HOST`, `MSSQL_HOST`, `ORACLE_HOST`, `DUCKDB_PATH`, `SQLITE_PATH`, etc. diff --git a/docs/docs/getting-started/index.md b/docs/docs/getting-started/index.md index 72f71a8b9..ec0bafd33 100644 --- a/docs/docs/getting-started/index.md +++ b/docs/docs/getting-started/index.md @@ -58,7 +58,7 @@ Altimate Code goes the other direction. It connects to your **entire** stack and --- - Optimize a Snowflake query in the morning. Migrate a SQL Server pipeline to BigQuery in the afternoon. Same agent, same tools. No warehouse subscription required. First-class support for :material-snowflake: Snowflake, :material-google-cloud: BigQuery, :simple-databricks: Databricks, :material-elephant: PostgreSQL, :material-aws: Redshift, :material-database: ClickHouse, :material-duck: DuckDB, :material-database: MySQL, :material-microsoft: SQL Server, and :material-leaf: MongoDB. + Optimize a Snowflake query in the morning. Migrate a SQL Server pipeline to BigQuery in the afternoon. Same agent, same tools. No warehouse subscription required. First-class support for :material-snowflake: Snowflake, :material-google-cloud: BigQuery, :simple-databricks: Databricks, :material-elephant: PostgreSQL, :material-aws: Redshift, :material-database: Trino, :material-database: ClickHouse, :material-duck: DuckDB, :material-database: MySQL, :material-microsoft: SQL Server, and :material-leaf: MongoDB. - :material-cloud-outline:{ .lg .middle } **Works with any LLM** diff --git a/packages/drivers/package.json b/packages/drivers/package.json index 361c1dd96..c2437ac6b 100644 --- a/packages/drivers/package.json +++ b/packages/drivers/package.json @@ -21,6 +21,7 @@ "oracledb": "^6.0.0", "duckdb": "^1.0.0", "mongodb": "^6.0.0", - "@clickhouse/client": "^1.0.0" + "@clickhouse/client": "^1.0.0", + "trino-client": "^0.2.8" } } diff --git a/packages/drivers/src/index.ts b/packages/drivers/src/index.ts index 4854785c6..d3c755c31 100644 --- a/packages/drivers/src/index.ts +++ b/packages/drivers/src/index.ts @@ -17,3 +17,4 @@ export { connect as connectDuckdb } from "./duckdb" export { connect as connectSqlite } from "./sqlite" export { connect as connectMongodb } from "./mongodb" export { connect as connectClickhouse } from "./clickhouse" +export { connect as connectTrino } from "./trino" diff --git a/packages/drivers/src/normalize.ts b/packages/drivers/src/normalize.ts index 2d3c36127..56668154c 100644 --- a/packages/drivers/src/normalize.ts +++ b/packages/drivers/src/normalize.ts @@ -98,6 +98,15 @@ const CLICKHOUSE_ALIASES: AliasMap = { tls_key: ["tlsKey", "ssl_key"], } +const TRINO_ALIASES: AliasMap = { + user: ["username"], + catalog: ["database", "dbname", "db"], + connection_string: ["connectionString", "server", "url", "uri"], + access_token: ["token", "accessToken", "jwt", "bearer_token", "bearerToken"], + extra_headers: ["extraHeaders", "headers"], + extra_credential: ["extraCredential"], +} + /** Map of warehouse type to its alias map. */ const DRIVER_ALIASES: Record = { snowflake: SNOWFLAKE_ALIASES, @@ -115,6 +124,7 @@ const DRIVER_ALIASES: Record = { mongodb: MONGODB_ALIASES, mongo: MONGODB_ALIASES, clickhouse: CLICKHOUSE_ALIASES, + trino: TRINO_ALIASES, // duckdb and sqlite have simple configs — no aliases needed } diff --git a/packages/drivers/src/trino.ts b/packages/drivers/src/trino.ts new file mode 100644 index 000000000..e7498d1bb --- /dev/null +++ b/packages/drivers/src/trino.ts @@ -0,0 +1,240 @@ +/** + * Trino driver using the official `trino-client` package. + * + * Supports Trino's HTTP(S) protocol via trino-js-client. Catalog should be set + * for schema/table introspection. + */ + +import type { ConnectionConfig, Connector, ConnectorResult, ExecuteOptions, SchemaColumn } from "./types" + +type QueryResult = { + columns?: Array<{ name: string; type: string }> + data?: any[][] + error?: { message?: string; errorName?: string; errorCode?: number } +} + +function cleanSql(sql: string): string { + return sql + .replace(/'(?:[^'\\]|\\.|'{2})*'/g, "") + .replace(/--[^\n]*/g, "") + .replace(/\/\*[\s\S]*?\*\//g, "") +} + +function escapeStringLiteral(value: string): string { + return value.replace(/'/g, "''") +} + +function quoteIdent(value: unknown): string { + const text = String(value ?? "").trim() + if (!text) { + throw new Error("Trino identifier cannot be empty") + } + return `"${text.replace(/"/g, '""')}"` +} + +function defaultCatalog(config: ConnectionConfig): string { + return String(config.catalog ?? "").trim() +} + +function serverUrl(config: ConnectionConfig): string { + const configured = config.connection_string ?? config.server ?? config.url + if (typeof configured === "string" && configured.trim()) { + return configured.trim() + } + + const protocol = typeof config.protocol === "string" ? config.protocol : config.ssl || config.tls ? "https" : "http" + const host = String(config.host ?? "localhost") + const port = Number(config.port ?? (protocol === "https" ? 8443 : 8080)) + return `${protocol}://${host}:${port}` +} + +function extraHeaders(config: ConnectionConfig): Record { + const headers: Record = {} + + if (config.extra_headers && typeof config.extra_headers === "object" && !Array.isArray(config.extra_headers)) { + for (const [key, value] of Object.entries(config.extra_headers as Record)) { + if (value !== undefined && value !== null) headers[key] = String(value) + } + } + + const token = config.access_token ?? config.token + if (typeof token === "string" && token.trim()) { + headers.Authorization = `Bearer ${token.trim()}` + } + + if (config.user && !headers["X-Trino-User"]) { + headers["X-Trino-User"] = String(config.user) + } + + return headers +} + +function trinoError(result: QueryResult): Error | null { + if (!result.error) return null + const name = result.error.errorName ? ` ${result.error.errorName}` : "" + const code = result.error.errorCode !== undefined ? ` (${result.error.errorCode})` : "" + return new Error(`Trino query failed${name}${code}: ${result.error.message ?? "unknown error"}`) +} + +export async function connect(config: ConnectionConfig): Promise { + let Trino: any + let BasicAuth: any + try { + const mod = await import("trino-client") + Trino = mod.Trino ?? mod.default?.Trino ?? mod.default + BasicAuth = mod.BasicAuth ?? mod.default?.BasicAuth + if (!Trino?.create) { + throw new Error("Trino.create export not found in trino-client") + } + } catch { + throw new Error("Trino driver not installed. Run: npm install trino-client") + } + + let client: any + + const connector: Connector = { + async connect() { + const headers = extraHeaders(config) + const options: Record = { + server: serverUrl(config), + source: config.source ?? "altimate-code", + catalog: config.catalog, + schema: config.schema, + extraHeaders: headers, + } + + if (config.password) { + if (typeof config.password !== "string") { + throw new Error("Trino password must be a string. Check your warehouse configuration.") + } + if (!BasicAuth) { + throw new Error("Trino BasicAuth export not found in trino-client") + } + options.auth = new BasicAuth(String(config.user ?? "trino"), config.password) + } + if (config.ssl && typeof config.ssl === "object") { + options.ssl = config.ssl + } + if (config.session && typeof config.session === "object" && !Array.isArray(config.session)) { + options.session = config.session + } + if ( + config.extra_credential && + typeof config.extra_credential === "object" && + !Array.isArray(config.extra_credential) + ) { + options.extraCredential = config.extra_credential + } + + client = Trino.create(options) + }, + + async execute(sql: string, limit?: number, _binds?: any[], options?: ExecuteOptions): Promise { + if (!client) { + throw new Error("Trino client not connected — call connect() first") + } + + const effectiveLimit = options?.noLimit ? 0 : (limit ?? 1000) + const sqlCleaned = cleanSql(sql) + const isSelectLike = /^\s*(SELECT|WITH|VALUES|TABLE)\b/i.test(sqlCleaned) + const hasDML = /\b(INSERT|UPDATE|DELETE|MERGE|CREATE|DROP|ALTER|TRUNCATE|CALL|GRANT|REVOKE)\b/i.test(sqlCleaned) + const hasLimit = /\b(LIMIT|FETCH\s+FIRST)\b/i.test(sqlCleaned) + + let query = sql + if (isSelectLike && !hasDML && effectiveLimit > 0 && !hasLimit) { + query = `${sql.replace(/;\s*$/, "")}\nLIMIT ${effectiveLimit + 1}` + } + + const iter = await client.query(query) + let columns: string[] = [] + const rows: any[][] = [] + + for await (const result of iter as AsyncIterable) { + const err = trinoError(result) + if (err) throw err + if (result.columns) { + columns = result.columns.map((c) => c.name) + } + if (result.data) { + rows.push(...result.data) + } + } + + const truncated = effectiveLimit > 0 && rows.length > effectiveLimit + const limitedRows = truncated ? rows.slice(0, effectiveLimit) : rows + return { + columns, + rows: limitedRows, + row_count: limitedRows.length, + truncated, + } + }, + + async listSchemas(): Promise { + const catalog = defaultCatalog(config) + if (!catalog) { + const result = await connector.execute("SHOW CATALOGS", 10000) + return result.rows.map((r) => String(r[0])) + } + + const result = await connector.execute( + `SELECT schema_name + FROM ${quoteIdent(catalog)}.information_schema.schemata + WHERE schema_name NOT IN ('information_schema') + ORDER BY schema_name`, + 10000, + ) + return result.rows.map((r) => String(r[0])) + }, + + async listTables(schema: string): Promise> { + const catalog = defaultCatalog(config) + if (!catalog) { + throw new Error("Trino catalog is required to list tables. Set catalog in the warehouse config.") + } + + const result = await connector.execute( + `SELECT table_name, table_type + FROM ${quoteIdent(catalog)}.information_schema.tables + WHERE table_schema = '${escapeStringLiteral(schema)}' + ORDER BY table_name`, + 10000, + ) + return result.rows.map((r) => ({ + name: String(r[0]), + type: String(r[1] ?? "") + .toUpperCase() + .includes("VIEW") + ? "view" + : "table", + })) + }, + + async describeTable(schema: string, table: string): Promise { + const catalog = defaultCatalog(config) + if (!catalog) { + throw new Error("Trino catalog is required to describe tables. Set catalog in the warehouse config.") + } + + const result = await connector.execute( + `SELECT column_name, data_type, is_nullable + FROM ${quoteIdent(catalog)}.information_schema.columns + WHERE table_schema = '${escapeStringLiteral(schema)}' + AND table_name = '${escapeStringLiteral(table)}' + ORDER BY ordinal_position`, + 10000, + ) + return result.rows.map((r) => ({ + name: String(r[0]), + data_type: String(r[1]), + nullable: String(r[2]).toUpperCase() === "YES", + })) + }, + + async close() { + client = null + }, + } + + return connector +} diff --git a/packages/drivers/test/trino-unit.test.ts b/packages/drivers/test/trino-unit.test.ts new file mode 100644 index 000000000..2ad22e785 --- /dev/null +++ b/packages/drivers/test/trino-unit.test.ts @@ -0,0 +1,154 @@ +/** + * Unit tests for Trino driver logic: + * - client configuration + * - LIMIT injection and truncation + * - Trino result iteration + * - catalog-scoped introspection SQL + */ +import { beforeEach, describe, expect, mock, test } from "bun:test" + +let mockCreateCalls: any[] = [] +let mockQueryCalls: string[] = [] +let mockResults: any[][] = [] + +function resetMocks() { + mockCreateCalls = [] + mockQueryCalls = [] + mockResults = [] +} + +function iteratorFor(results: any[]) { + let idx = 0 + let last: any + return { + [Symbol.asyncIterator]() { + return this + }, + async next() { + if (idx < results.length) { + last = results[idx++] + return { value: last, done: false } + } + return { value: last, done: true } + }, + } +} + +mock.module("trino-client", () => ({ + BasicAuth: class { + username: string + password: string + constructor(username: string, password: string) { + this.username = username + this.password = password + } + }, + Trino: { + create: (options: any) => { + mockCreateCalls.push(options) + return { + query: async (sql: string) => { + mockQueryCalls.push(sql) + return iteratorFor(mockResults.shift() ?? []) + }, + } + }, + }, +})) + +const { connect } = await import("../src/trino") + +describe("Trino driver unit tests", () => { + beforeEach(() => { + resetMocks() + }) + + test("configures server, catalog, schema, token headers, and source", async () => { + const connector = await connect({ + type: "trino", + host: "trino.example.com", + port: 8443, + protocol: "https", + catalog: "iceberg", + schema: "analytics", + user: "analyst", + access_token: "jwt-token", + source: "unit-test", + }) + await connector.connect() + + expect(mockCreateCalls[0].server).toBe("https://trino.example.com:8443") + expect(mockCreateCalls[0].catalog).toBe("iceberg") + expect(mockCreateCalls[0].schema).toBe("analytics") + expect(mockCreateCalls[0].source).toBe("unit-test") + expect(mockCreateCalls[0].extraHeaders.Authorization).toBe("Bearer jwt-token") + expect(mockCreateCalls[0].extraHeaders["X-Trino-User"]).toBe("analyst") + }) + + test("uses BasicAuth when password is configured", async () => { + const connector = await connect({ type: "trino", user: "analyst", password: "secret" }) + await connector.connect() + expect(mockCreateCalls[0].auth.username).toBe("analyst") + expect(mockCreateCalls[0].auth.password).toBe("secret") + }) + + test("appends LIMIT to SELECT and detects truncation", async () => { + mockResults = [[{ columns: [{ name: "id", type: "integer" }], data: [[1], [2], [3]] }]] + const connector = await connect({ type: "trino", catalog: "iceberg" }) + await connector.connect() + + const result = await connector.execute("SELECT id FROM orders", 2) + expect(mockQueryCalls[0]).toBe("SELECT id FROM orders\nLIMIT 3") + expect(result.columns).toEqual(["id"]) + expect(result.rows).toEqual([[1], [2]]) + expect(result.truncated).toBe(true) + }) + + test("does not double-limit or limit write statements", async () => { + mockResults = [[{ columns: [{ name: "id", type: "integer" }], data: [[1]] }], [{ data: [] }]] + const connector = await connect({ type: "trino", catalog: "iceberg" }) + await connector.connect() + + await connector.execute("SELECT id FROM orders LIMIT 5", 2) + expect(mockQueryCalls[0]).toBe("SELECT id FROM orders LIMIT 5") + + await connector.execute("INSERT INTO t VALUES (1)", 2) + expect(mockQueryCalls[1]).toBe("INSERT INTO t VALUES (1)") + }) + + test("listTables and describeTable use catalog-scoped information_schema", async () => { + mockResults = [ + [ + { + columns: [ + { name: "table_name", type: "varchar" }, + { name: "table_type", type: "varchar" }, + ], + data: [["orders", "BASE TABLE"]], + }, + ], + [ + { + columns: [ + { name: "column_name", type: "varchar" }, + { name: "data_type", type: "varchar" }, + { name: "is_nullable", type: "varchar" }, + ], + data: [["id", "integer", "NO"]], + }, + ], + ] + const connector = await connect({ type: "trino", catalog: "iceberg" }) + await connector.connect() + + await expect(connector.listTables("analytics")).resolves.toEqual([{ name: "orders", type: "table" }]) + expect(mockQueryCalls[0]).toContain('FROM "iceberg".information_schema.tables') + expect(mockQueryCalls[0]).toContain("table_schema = 'analytics'") + + await expect(connector.describeTable("analytics", "orders")).resolves.toEqual([ + { name: "id", data_type: "integer", nullable: false }, + ]) + expect(mockQueryCalls[1]).toContain('FROM "iceberg".information_schema.columns') + expect(mockQueryCalls[1]).toContain("table_name = 'orders'") + }) +}) diff --git a/packages/opencode/script/publish.ts b/packages/opencode/script/publish.ts index 516fdda86..b4f47de4e 100755 --- a/packages/opencode/script/publish.ts +++ b/packages/opencode/script/publish.ts @@ -30,6 +30,7 @@ const driverPeerDependencies: Record = { oracledb: ">=6", duckdb: ">=1", "@clickhouse/client": ">=1", + "trino-client": ">=0.2", } const driverPeerDependenciesMeta: Record = Object.fromEntries( diff --git a/packages/opencode/src/altimate/native/connections/dbt-profiles.ts b/packages/opencode/src/altimate/native/connections/dbt-profiles.ts index 679f168b3..e36b5bd0f 100644 --- a/packages/opencode/src/altimate/native/connections/dbt-profiles.ts +++ b/packages/opencode/src/altimate/native/connections/dbt-profiles.ts @@ -24,7 +24,7 @@ const ADAPTER_TYPE_MAP: Record = { oracle: "oracle", sqlite: "sqlite", spark: "databricks", - trino: "postgres", // wire-compatible + trino: "trino", clickhouse: "clickhouse", } diff --git a/packages/opencode/src/altimate/native/connections/docker-discovery.ts b/packages/opencode/src/altimate/native/connections/docker-discovery.ts index 45a174bc9..8b09491f6 100644 --- a/packages/opencode/src/altimate/native/connections/docker-discovery.ts +++ b/packages/opencode/src/altimate/native/connections/docker-discovery.ts @@ -18,6 +18,7 @@ const IMAGE_MAP: Array<{ pattern: RegExp; type: string }> = [ { pattern: /oracle/i, type: "oracle" }, { pattern: /gvenzl\/oracle/i, type: "oracle" }, { pattern: /clickhouse/i, type: "clickhouse" }, + { pattern: /trinodb\/trino|prestosql|presto/i, type: "trino" }, ] /** Map environment variable names to connection config fields by db type. */ @@ -48,6 +49,11 @@ const ENV_MAP: Record> = { CLICKHOUSE_PASSWORD: "password", CLICKHOUSE_DB: "database", }, + trino: { + TRINO_USER: "user", + TRINO_PASSWORD: "password", + TRINO_CATALOG: "database", + }, } /** Default ports by database type. */ @@ -57,6 +63,7 @@ const DEFAULT_PORTS: Record = { sqlserver: 1433, oracle: 1521, clickhouse: 8123, + trino: 8080, } /** Default users by database type. */ @@ -66,6 +73,7 @@ const DEFAULT_USERS: Record = { sqlserver: "sa", oracle: "system", clickhouse: "default", + trino: "trino", } function detectDbType(image: string): string | null { diff --git a/packages/opencode/src/altimate/native/connections/register.ts b/packages/opencode/src/altimate/native/connections/register.ts index 53d185adc..5b32cb560 100644 --- a/packages/opencode/src/altimate/native/connections/register.ts +++ b/packages/opencode/src/altimate/native/connections/register.ts @@ -210,6 +210,7 @@ export interface ExplainPlan { * - DuckDB: `EXPLAIN` or `EXPLAIN ANALYZE`. * - Databricks / Spark: `EXPLAIN` or `EXPLAIN FORMATTED`. * - ClickHouse: `EXPLAIN` (no ANALYZE form accepted via statement prefix). + * - Trino: `EXPLAIN` or `EXPLAIN ANALYZE`. * - BigQuery: uses a dry-run API instead of any EXPLAIN statement. Not * supported via this code path — return empty prefix. * - Oracle: `EXPLAIN PLAN FOR` stores the plan in PLAN_TABLE rather than @@ -249,6 +250,10 @@ export function buildExplainPlan(warehouseType: string | undefined, analyze: boo return { prefix: analyze ? "EXPLAIN FORMATTED" : "EXPLAIN", actuallyAnalyzed: false } case "clickhouse": return { prefix: "EXPLAIN", actuallyAnalyzed: false } + case "trino": + return analyze + ? { prefix: "EXPLAIN ANALYZE", actuallyAnalyzed: true } + : { prefix: "EXPLAIN", actuallyAnalyzed: false } case "bigquery": // BigQuery has no EXPLAIN statement — the correct answer is a dry-run // job via the BigQuery API, which this tool does not support today. diff --git a/packages/opencode/src/altimate/native/connections/registry.ts b/packages/opencode/src/altimate/native/connections/registry.ts index cc871682c..60a63e384 100644 --- a/packages/opencode/src/altimate/native/connections/registry.ts +++ b/packages/opencode/src/altimate/native/connections/registry.ts @@ -130,6 +130,7 @@ const DRIVER_MAP: Record = { mongodb: "@altimateai/drivers/mongodb", mongo: "@altimateai/drivers/mongodb", clickhouse: "@altimateai/drivers/clickhouse", + trino: "@altimateai/drivers/trino", } async function createConnector(name: string, config: ConnectionConfig): Promise { @@ -170,6 +171,7 @@ async function createConnector(name: string, config: ConnectionConfig): Promise< "oracle", "snowflake", "clickhouse", + "trino", ]) if ( PASSWORD_DRIVERS.has(resolvedConfig.type.toLowerCase()) && @@ -237,6 +239,9 @@ async function createConnector(name: string, config: ConnectionConfig): Promise< case "@altimateai/drivers/clickhouse": mod = await import("@altimateai/drivers/clickhouse") break + case "@altimateai/drivers/trino": + mod = await import("@altimateai/drivers/trino") + break default: throw new Error(`No static import available for driver: ${driverPath}`) } diff --git a/packages/opencode/src/altimate/native/finops/query-history.ts b/packages/opencode/src/altimate/native/finops/query-history.ts index 5bf75ab36..5fe7f8f70 100644 --- a/packages/opencode/src/altimate/native/finops/query-history.ts +++ b/packages/opencode/src/altimate/native/finops/query-history.ts @@ -192,6 +192,9 @@ function buildHistoryQuery( if (whType === "duckdb") { return null // DuckDB has no native query history } + if (whType === "trino") { + return null // Trino exposes live runtime query state, not durable query history + } return null } diff --git a/packages/opencode/src/altimate/tools/project-scan.ts b/packages/opencode/src/altimate/tools/project-scan.ts index 4d7038046..81dc0fe04 100644 --- a/packages/opencode/src/altimate/tools/project-scan.ts +++ b/packages/opencode/src/altimate/tools/project-scan.ts @@ -233,6 +233,20 @@ export async function detectEnvVars(): Promise { connection_string: "CLICKHOUSE_URL", }, }, + { + type: "trino", + signals: ["TRINO_HOST", "TRINO_SERVER", "TRINO_URL"], + configMap: { + host: "TRINO_HOST", + port: "TRINO_PORT", + user: ["TRINO_USER", "TRINO_USERNAME"], + password: "TRINO_PASSWORD", + catalog: ["TRINO_CATALOG", "TRINO_DATABASE"], + schema: "TRINO_SCHEMA", + connection_string: ["TRINO_URL", "TRINO_SERVER"], + access_token: ["TRINO_TOKEN", "TRINO_ACCESS_TOKEN"], + }, + }, ] for (const wh of warehouses) { @@ -296,6 +310,10 @@ export async function detectEnvVars(): Promise { clickhouse: "clickhouse", "clickhouse+http": "clickhouse", "clickhouse+https": "clickhouse", + trino: "trino", + "trino+http": "trino", + "trino+https": "trino", + presto: "trino", } const dbType = schemeTypeMap[scheme] ?? "postgres" // Only add if we don't already have this type detected from other env vars diff --git a/packages/opencode/src/altimate/tools/warehouse-add.ts b/packages/opencode/src/altimate/tools/warehouse-add.ts index 3758afaf1..9e9e9e8c4 100644 --- a/packages/opencode/src/altimate/tools/warehouse-add.ts +++ b/packages/opencode/src/altimate/tools/warehouse-add.ts @@ -24,6 +24,7 @@ export const WarehouseAddTool = Tool.define("warehouse_add", { - duckdb: path (file path or ":memory:") - sqlite: path (file path) - clickhouse: host, port, database, user, password, protocol (http/https), connection_string, request_timeout, tls_ca_cert, tls_cert, tls_key, clickhouse_settings +- trino: host, port, catalog, schema, user, password, protocol (http/https), connection_string, access_token, extra_headers Snowflake auth examples: (1) Password: {"type":"snowflake","account":"xy12345","user":"admin","password":"secret","warehouse":"WH","database":"db"}. (2) Key-pair: {"type":"snowflake","account":"xy12345","user":"admin","private_key_path":"/path/rsa_key.p8","warehouse":"WH","database":"db"}. (3) OAuth: {"type":"snowflake","account":"xy12345","authenticator":"oauth","token":"","warehouse":"WH","database":"db"}. (4) SSO: {"type":"snowflake","account":"xy12345","user":"admin","authenticator":"externalbrowser","warehouse":"WH","database":"db"}. IMPORTANT: For private key file paths, always use "private_key_path" (not "private_key").`, ), @@ -33,7 +34,7 @@ IMPORTANT: For private key file paths, always use "private_key_path" (not "priva return { title: `Add '${args.name}': FAILED`, metadata: { success: false, name: args.name, type: "" }, - output: `Missing required field "type" in config. Specify the database type (postgres, snowflake, bigquery, databricks, redshift, clickhouse, duckdb, mysql, sqlserver, oracle, sqlite, mongodb).`, + output: `Missing required field "type" in config. Specify the database type (postgres, snowflake, bigquery, databricks, redshift, clickhouse, trino, duckdb, mysql, sqlserver, oracle, sqlite, mongodb).`, } } diff --git a/packages/opencode/test/altimate/connections.test.ts b/packages/opencode/test/altimate/connections.test.ts index 5c9680297..5b52972ee 100644 --- a/packages/opencode/test/altimate/connections.test.ts +++ b/packages/opencode/test/altimate/connections.test.ts @@ -624,7 +624,7 @@ spark_project: } }) - test("trino adapter maps to postgres (wire-compatible)", async () => { + test("trino adapter maps to native trino connector", async () => { const fs = await import("fs") const os = await import("os") const path = await import("path") @@ -643,14 +643,17 @@ trino_project: port: 8080 user: analyst dbname: hive + schema: analytics `, ) try { const connections = await parseDbtProfiles(profilesPath) expect(connections).toHaveLength(1) - expect(connections[0].type).toBe("postgres") - expect(connections[0].config.type).toBe("postgres") + expect(connections[0].type).toBe("trino") + expect(connections[0].config.type).toBe("trino") + expect(connections[0].config.database).toBe("hive") + expect(connections[0].config.schema).toBe("analytics") } finally { fs.rmSync(tmpDir, { recursive: true }) } diff --git a/packages/opencode/test/altimate/driver-normalize.test.ts b/packages/opencode/test/altimate/driver-normalize.test.ts index 43b31c4e8..dc3eb34e2 100644 --- a/packages/opencode/test/altimate/driver-normalize.test.ts +++ b/packages/opencode/test/altimate/driver-normalize.test.ts @@ -310,6 +310,48 @@ describe("normalizeConfig — BigQuery", () => { }) }) +// --------------------------------------------------------------------------- +// normalizeConfig — Trino aliases +// --------------------------------------------------------------------------- + +describe("normalizeConfig — Trino", () => { + test("dbname → catalog for dbt-trino profiles", () => { + const result = normalizeConfig({ + type: "trino", + dbname: "hive", + }) + expect(result.catalog).toBe("hive") + expect(result.dbname).toBeUndefined() + }) + + test("database → catalog for compatibility with generic config", () => { + const result = normalizeConfig({ + type: "trino", + database: "iceberg", + }) + expect(result.catalog).toBe("iceberg") + expect(result.database).toBeUndefined() + }) + + test("token → access_token", () => { + const result = normalizeConfig({ + type: "trino", + token: "jwt-token", + }) + expect(result.access_token).toBe("jwt-token") + expect(result.token).toBeUndefined() + }) + + test("server → connection_string", () => { + const result = normalizeConfig({ + type: "trino", + server: "https://trino.example.com:8443", + }) + expect(result.connection_string).toBe("https://trino.example.com:8443") + expect(result.server).toBeUndefined() + }) +}) + // --------------------------------------------------------------------------- // normalizeConfig — Databricks aliases // --------------------------------------------------------------------------- diff --git a/packages/opencode/test/altimate/tools/sql-explain.test.ts b/packages/opencode/test/altimate/tools/sql-explain.test.ts index e9b8ae92c..b92585fa0 100644 --- a/packages/opencode/test/altimate/tools/sql-explain.test.ts +++ b/packages/opencode/test/altimate/tools/sql-explain.test.ts @@ -222,6 +222,17 @@ describe("buildExplainPlan", () => { expect(buildExplainPlan("clickhouse", false).prefix).toBe("EXPLAIN") }) + test("Trino: plain EXPLAIN or EXPLAIN ANALYZE", () => { + expect(buildExplainPlan("trino", false)).toEqual({ + prefix: "EXPLAIN", + actuallyAnalyzed: false, + }) + expect(buildExplainPlan("trino", true)).toEqual({ + prefix: "EXPLAIN ANALYZE", + actuallyAnalyzed: true, + }) + }) + test("BigQuery: not supported via statement prefix", () => { // BigQuery requires a dry-run API call — no statement EXPLAIN. expect(buildExplainPlan("bigquery", false).prefix).toBe("") diff --git a/packages/opencode/test/tool/project-scan.test.ts b/packages/opencode/test/tool/project-scan.test.ts index 2b4c031e5..2465f0bbf 100644 --- a/packages/opencode/test/tool/project-scan.test.ts +++ b/packages/opencode/test/tool/project-scan.test.ts @@ -314,6 +314,9 @@ describe("detectEnvVars", () => { "REDSHIFT_HOST", "REDSHIFT_PORT", "REDSHIFT_DATABASE", "REDSHIFT_USER", "REDSHIFT_PASSWORD", "CLICKHOUSE_HOST", "CLICKHOUSE_URL", "CLICKHOUSE_PORT", "CLICKHOUSE_DB", "CLICKHOUSE_DATABASE", "CLICKHOUSE_USER", "CLICKHOUSE_USERNAME", "CLICKHOUSE_PASSWORD", + "TRINO_HOST", "TRINO_SERVER", "TRINO_URL", "TRINO_PORT", "TRINO_CATALOG", + "TRINO_DATABASE", "TRINO_SCHEMA", "TRINO_USER", "TRINO_USERNAME", "TRINO_PASSWORD", + "TRINO_TOKEN", "TRINO_ACCESS_TOKEN", ] for (const v of vars) { delete process.env[v] @@ -561,6 +564,43 @@ describe("detectEnvVars", () => { } }) + test("detects Trino via TRINO_HOST", async () => { + clearWarehouseEnvVars() + process.env.TRINO_HOST = "trino.example.com" + process.env.TRINO_PORT = "8443" + process.env.TRINO_CATALOG = "iceberg" + process.env.TRINO_SCHEMA = "analytics" + process.env.TRINO_USER = "analyst" + process.env.TRINO_PASSWORD = "secret" + + const result = await detectEnvVars() + const trino = result.find((r) => r.type === "trino") + expect(trino).toBeDefined() + expect(trino!.name).toBe("env_trino") + expect(trino!.source).toBe("env-var") + expect(trino!.signal).toBe("TRINO_HOST") + expect(trino!.config.host).toBe("trino.example.com") + expect(trino!.config.port).toBe("8443") + expect(trino!.config.catalog).toBe("iceberg") + expect(trino!.config.schema).toBe("analytics") + expect(trino!.config.user).toBe("analyst") + expect(trino!.config.password).toBe("***") + }) + + test("detects Trino via DATABASE_URL with trino schemes", async () => { + for (const scheme of ["trino", "trino+http", "trino+https", "presto"]) { + clearWarehouseEnvVars() + process.env.DATABASE_URL = `${scheme}://trino.example.com:8080` + + const result = await detectEnvVars() + const trino = result.find((r) => r.type === "trino") + expect(trino).toBeDefined() + expect(trino!.signal).toBe("DATABASE_URL") + expect(trino!.type).toBe("trino") + expect(trino!.config.connection_string).toBe("***") + } + }) + test("detects multiple warehouses simultaneously", async () => { clearWarehouseEnvVars() process.env.SNOWFLAKE_ACCOUNT = "sf_account" From 30a0034a4aca233c644183e093c18740d9994691 Mon Sep 17 00:00:00 2001 From: VC Date: Fri, 8 May 2026 07:25:01 -0400 Subject: [PATCH 2/2] fix: harden altimate tool responses --- .../tools/altimate-core-column-lineage.ts | 48 +++++++++++++++++-- .../src/altimate/tools/lineage-check.ts | 45 +++++++++++------ .../src/altimate/tools/schema-inspect.ts | 48 ++++++++++++++----- .../src/altimate/tools/sql-analyze.ts | 44 +++++++++++------ 4 files changed, 141 insertions(+), 44 deletions(-) diff --git a/packages/opencode/src/altimate/tools/altimate-core-column-lineage.ts b/packages/opencode/src/altimate/tools/altimate-core-column-lineage.ts index 180836d12..82b54175d 100644 --- a/packages/opencode/src/altimate/tools/altimate-core-column-lineage.ts +++ b/packages/opencode/src/altimate/tools/altimate-core-column-lineage.ts @@ -47,8 +47,7 @@ function formatColumnLineage(data: Record): string { if (data.column_dict && Object.keys(data.column_dict).length > 0) { lines.push("Column Mappings:") for (const [target, sources] of Object.entries(data.column_dict)) { - const srcList = Array.isArray(sources) ? (sources as string[]).join(", ") : JSON.stringify(sources) - lines.push(` ${target} ← ${srcList}`) + lines.push(` ${target} ← ${formatLineageValue(sources)}`) } lines.push("") } @@ -56,10 +55,51 @@ function formatColumnLineage(data: Record): string { if (data.column_lineage?.length) { lines.push("Lineage Edges:") for (const edge of data.column_lineage) { - const transform = edge.lens_type ?? edge.transform_type ?? edge.transform ?? "" - lines.push(` ${edge.source ?? "?"} → ${edge.target ?? "?"}${transform ? ` (${transform})` : ""}`) + const source = formatLineageEndpoint(edge, "source") + const target = formatLineageEndpoint(edge, "target") + const transform = formatLineageValue(edge.lens_type ?? edge.transform_type ?? edge.transform ?? "") + lines.push(` ${source} → ${target}${transform ? ` (${transform})` : ""}`) } } return lines.length ? lines.join("\n") : "No column lineage edges found." } + +function formatLineageEndpoint(edge: Record, side: "source" | "target"): string { + if (edge[side] !== null && edge[side] !== undefined) return formatLineageValue(edge[side]) + + const table = edge[`${side}_table`] ?? edge[`${side}Table`] + const column = edge[`${side}_column`] ?? edge[`${side}Column`] + if (table !== null && table !== undefined && column !== null && column !== undefined) { + return `${formatLineageValue(table)}.${formatLineageValue(column)}` + } + return "?" +} + +function formatLineageValue(value: unknown): string { + if (value === null || value === undefined) return "" + if (typeof value === "string") return value + if (typeof value === "number" || typeof value === "boolean" || typeof value === "bigint") return String(value) + + if (Array.isArray(value)) { + return value.map(formatLineageValue).filter(Boolean).join(", ") + } + + if (typeof value === "object") { + const obj = value as Record + const table = obj.source_table ?? obj.sourceTable ?? obj.target_table ?? obj.targetTable ?? obj.table + const column = obj.source_column ?? obj.sourceColumn ?? obj.target_column ?? obj.targetColumn ?? obj.column ?? obj.name + if (table !== null && table !== undefined && column !== null && column !== undefined) { + return `${formatLineageValue(table)}.${formatLineageValue(column)}` + } + if (obj.source !== null && obj.source !== undefined) return formatLineageValue(obj.source) + if (obj.target !== null && obj.target !== undefined) return formatLineageValue(obj.target) + try { + return JSON.stringify(value) + } catch { + return "unserializable object" + } + } + + return String(value) +} diff --git a/packages/opencode/src/altimate/tools/lineage-check.ts b/packages/opencode/src/altimate/tools/lineage-check.ts index dbe19fc0f..04c490298 100644 --- a/packages/opencode/src/altimate/tools/lineage-check.ts +++ b/packages/opencode/src/altimate/tools/lineage-check.ts @@ -20,22 +20,24 @@ export const LineageCheckTool = Tool.define("lineage_check", { }), async execute(args, ctx) { try { - const result = await Dispatcher.call("lineage.check", { + const rawResult = (await Dispatcher.call("lineage.check", { sql: args.sql, dialect: args.dialect, schema_context: args.schema_context, - }) + })) as unknown - const data = (result.data ?? {}) as Record + if (!isRecord(rawResult)) { + return lineageError("Invalid lineage response from dispatcher.") + } + + const result = rawResult as Partial + + const data = isRecord(result.data) ? result.data : {} if (result.error) { - return { - title: "Lineage: ERROR", - metadata: { success: false, error: result.error }, - output: `Error: ${result.error}`, - } + return lineageError(result.error) } - const error = data.error + const error = normalizeError(data.error) return { title: `Lineage: ${result.success ? "OK" : "PARTIAL"}`, metadata: { success: result.success, ...(error && { error }) }, @@ -43,15 +45,30 @@ export const LineageCheckTool = Tool.define("lineage_check", { } } catch (e) { const msg = e instanceof Error ? e.message : String(e) - return { - title: "Lineage: ERROR", - metadata: { success: false, error: msg }, - output: `Failed to check lineage: ${msg}\n\nEnsure the dispatcher is running and altimate-core is initialized.`, - } + return lineageError(msg) } }, }) +function lineageError(msg: string) { + return { + title: "Lineage: ERROR", + metadata: { success: false, error: msg }, + output: `Failed to check lineage: ${msg}\n\nEnsure the dispatcher is running and altimate-core is initialized.`, + } +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value) +} + +function normalizeError(value: unknown): string | undefined { + if (value instanceof Error) return value.message + if (typeof value === "string") return value + if (value === null || value === undefined) return undefined + return String(value) +} + function formatLineage(data: Record): string { const lines: string[] = [] diff --git a/packages/opencode/src/altimate/tools/schema-inspect.ts b/packages/opencode/src/altimate/tools/schema-inspect.ts index 92f11b48f..f9475da75 100644 --- a/packages/opencode/src/altimate/tools/schema-inspect.ts +++ b/packages/opencode/src/altimate/tools/schema-inspect.ts @@ -15,14 +15,25 @@ export const SchemaInspectTool = Tool.define("schema_inspect", { }), async execute(args, ctx) { try { - const result = await Dispatcher.call("schema.inspect", { + const result = (await Dispatcher.call("schema.inspect", { table: args.table, schema_name: args.schema_name, warehouse: args.warehouse, - }) + })) as unknown + + if (!isRecord(result)) { + return schemaError("Invalid schema response from dispatcher.") + } + + const responseError = normalizeError(result.error) + if (result.success === false || responseError) { + return schemaError(responseError || "Schema inspection failed.") + } + + const schemaResult = result as Partial // altimate_change start — progressive disclosure suggestions - let output = formatSchema(result) + let output = formatSchema(schemaResult) const suggestion = PostConnectSuggestions.getProgressiveSuggestion("schema_inspect") if (suggestion) { output += "\n\n" + suggestion @@ -34,22 +45,37 @@ export const SchemaInspectTool = Tool.define("schema_inspect", { } // altimate_change end return { - title: `Schema: ${result.table ?? args.table}`, - metadata: { columnCount: (result.columns ?? []).length, rowCount: result.row_count }, + title: `Schema: ${schemaResult.table ?? args.table}`, + metadata: { columnCount: (schemaResult.columns ?? []).length, rowCount: schemaResult.row_count }, output, } } catch (e) { const msg = e instanceof Error ? e.message : String(e) - return { - title: "Schema: ERROR", - metadata: { columnCount: 0, rowCount: undefined, error: msg }, - output: `Failed to inspect schema: ${msg}\n\nEnsure the dispatcher is running and a warehouse connection is configured.`, - } + return schemaError(msg) } }, }) -function formatSchema(result: SchemaInspectResult): string { +function schemaError(msg: string) { + return { + title: "Schema: ERROR", + metadata: { columnCount: 0, rowCount: undefined, error: msg }, + output: `Failed to inspect schema: ${msg}\n\nEnsure the dispatcher is running and a warehouse connection is configured.`, + } +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value) +} + +function normalizeError(value: unknown): string | undefined { + if (value instanceof Error) return value.message + if (typeof value === "string") return value + if (value === null || value === undefined) return undefined + return String(value) +} + +function formatSchema(result: Partial): string { const lines: string[] = [] const table = result.table ?? "unknown" const qualified = result.schema_name ? `${result.schema_name}.${table}` : table diff --git a/packages/opencode/src/altimate/tools/sql-analyze.ts b/packages/opencode/src/altimate/tools/sql-analyze.ts index 87c123727..eb18fcec6 100644 --- a/packages/opencode/src/altimate/tools/sql-analyze.ts +++ b/packages/opencode/src/altimate/tools/sql-analyze.ts @@ -26,12 +26,18 @@ export const SqlAnalyzeTool = Tool.define("sql_analyze", { async execute(args, ctx) { const hasSchema = !!(args.schema_path || (args.schema_context && Object.keys(args.schema_context).length > 0)) try { - const result = await Dispatcher.call("sql.analyze", { + const rawResult = (await Dispatcher.call("sql.analyze", { sql: args.sql, dialect: args.dialect, schema_path: args.schema_path, schema_context: args.schema_context, - }) + })) as unknown + + if (!isRecord(rawResult)) { + return analysisError(args, hasSchema, "Invalid analysis response from dispatcher.") + } + + const result = rawResult as Partial // The handler returns success=true when analysis completes (issues are // reported via issues/issue_count). Only treat it as a failure when @@ -70,23 +76,31 @@ export const SqlAnalyzeTool = Tool.define("sql_analyze", { } } catch (e) { const msg = e instanceof Error ? e.message : String(e) - return { - title: "Analyze: ERROR", - metadata: { - success: false, - issueCount: 0, - confidence: "unknown", - dialect: args.dialect, - has_schema: hasSchema, - error: msg, - }, - output: `Failed to analyze SQL: ${msg}\n\nCheck your connection configuration and try again.`, - } + return analysisError(args, hasSchema, msg) } }, }) -function formatAnalysis(result: SqlAnalyzeResult): string { +function analysisError(args: { dialect?: string }, hasSchema: boolean, msg: string) { + return { + title: "Analyze: ERROR", + metadata: { + success: false, + issueCount: 0, + confidence: "unknown", + dialect: args.dialect, + has_schema: hasSchema, + error: msg, + }, + output: `Failed to analyze SQL: ${msg}\n\nCheck your connection configuration and try again.`, + } +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value) +} + +function formatAnalysis(result: Partial): string { if (result.error) { return `Analysis failed: ${result.error}` }