From f310391da53642a98defe19cef9584389ce03ff7 Mon Sep 17 00:00:00 2001 From: Andrey Avtomonov Date: Thu, 2 Jul 2026 15:10:29 +0200 Subject: [PATCH] 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 --- .../connections/connection-type-dialect.ts | 24 ++++++++ .../context/connections/connection-type.ts | 10 ---- .../connections/local-warehouse-descriptor.ts | 1 + .../context/ingest/adapters/looker/mapping.ts | 14 +---- .../ingest/adapters/metabase/mapping.ts | 2 +- .../src/context/sl/semantic-layer.service.ts | 24 +------- .../sl/tools/sl-warehouse-validation.ts | 5 +- .../connection-type-dialect.test.ts | 59 ++++++++++++++++++ .../local-warehouse-descriptor.test.ts | 16 +++++ .../ingest/adapters/looker/mapping.test.ts | 41 +++++++++++++ .../ingest/adapters/metabase/mapping.test.ts | 21 +++++++ .../context/sl/semantic-layer.service.test.ts | 60 +++++++++++++++++++ .../semantic_layer/table_identifier_parser.py | 2 + .../tests/test_table_identifier_parser.py | 18 ++++++ 14 files changed, 252 insertions(+), 45 deletions(-) create mode 100644 packages/cli/src/context/connections/connection-type-dialect.ts create mode 100644 packages/cli/test/context/connections/connection-type-dialect.test.ts diff --git a/packages/cli/src/context/connections/connection-type-dialect.ts b/packages/cli/src/context/connections/connection-type-dialect.ts new file mode 100644 index 00000000..f20231de --- /dev/null +++ b/packages/cli/src/context/connections/connection-type-dialect.ts @@ -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; + +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; +} diff --git a/packages/cli/src/context/connections/connection-type.ts b/packages/cli/src/context/connections/connection-type.ts index d122e950..8b12839c 100644 --- a/packages/cli/src/context/connections/connection-type.ts +++ b/packages/cli/src/context/connections/connection-type.ts @@ -7,22 +7,12 @@ export const connectionTypeSchema = z.enum([ 'SQLSERVER', 'BIGQUERY', 'SNOWFLAKE', - 'CENTRALREACH', - 'EPIC', - 'CERNER', 'ATHENA', - 'QUICKBOOKS', - 'WORKDAY', - 'REST', - 'S3', - 'SLACK', 'METABASE', 'LOOKER', 'NOTION', 'MYSQL', 'CLICKHOUSE', - 'PLAIN', - 'BETTERSTACK', ]); export type ConnectionType = z.infer; diff --git a/packages/cli/src/context/connections/local-warehouse-descriptor.ts b/packages/cli/src/context/connections/local-warehouse-descriptor.ts index b9674e09..bb4235c2 100644 --- a/packages/cli/src/context/connections/local-warehouse-descriptor.ts +++ b/packages/cli/src/context/connections/local-warehouse-descriptor.ts @@ -29,6 +29,7 @@ const DRIVER_TO_CONNECTION_TYPE: Record = { clickhouse: 'CLICKHOUSE', snowflake: 'SNOWFLAKE', bigquery: 'BIGQUERY', + athena: 'ATHENA', }; export function localConnectionToWarehouseDescriptor( diff --git a/packages/cli/src/context/ingest/adapters/looker/mapping.ts b/packages/cli/src/context/ingest/adapters/looker/mapping.ts index 3e451aca..29214a36 100644 --- a/packages/cli/src/context/ingest/adapters/looker/mapping.ts +++ b/packages/cli/src/context/ingest/adapters/looker/mapping.ts @@ -1,4 +1,5 @@ import type { ParsedTargetTable } from '../../parsed-target-table.js'; +import { warehouseTargetDialect } from '../../../connections/connection-type-dialect.js'; import type { LookerWarehouseConnectionInfo } from './client.js'; import type { LookerPullConfig, LookerRuntimeCursors, StagedExploreFile, StagedLookmlModelsFile } from './types.js'; @@ -11,6 +12,7 @@ const LOOKER_DIALECT_TO_CONNECTION_TYPE = { sqlite: 'SQLITE', sqlserver: 'SQLSERVER', clickhouse: 'CLICKHOUSE', + awsathena: 'ATHENA', } as const; /** @internal */ @@ -72,16 +74,6 @@ export interface LookerMappingClient { getExplore(modelName: string, exploreName: string): Promise; } -const SQLGLOT_DIALECT_BY_CONNECTION_TYPE: Partial> = { - BIGQUERY: 'bigquery', - SNOWFLAKE: 'snowflake', - POSTGRESQL: 'postgres', - MYSQL: 'mysql', - SQLITE: 'sqlite', - SQLSERVER: 'tsql', - CLICKHOUSE: 'clickhouse', -}; - export async function discoverLookerConnections( client: Pick, ): Promise { @@ -100,7 +92,7 @@ export function lookerDialectToConnectionType(dialect: string | null): LookerWar /** @internal */ export function sqlglotDialectForConnectionType(connectionType: string): string | null { - return SQLGLOT_DIALECT_BY_CONNECTION_TYPE[connectionType as LookerWarehouseTargetConnectionType] ?? null; + return warehouseTargetDialect(connectionType); } /** @internal */ diff --git a/packages/cli/src/context/ingest/adapters/metabase/mapping.ts b/packages/cli/src/context/ingest/adapters/metabase/mapping.ts index b5ba67a7..3123f0e2 100644 --- a/packages/cli/src/context/ingest/adapters/metabase/mapping.ts +++ b/packages/cli/src/context/ingest/adapters/metabase/mapping.ts @@ -8,9 +8,9 @@ export const METABASE_ENGINE_TO_CONNECTION_TYPE = { snowflake: 'SNOWFLAKE', sqlserver: 'SQLSERVER', mysql: 'MYSQL', + athena: 'ATHENA', } as const; - export interface DiscoveredMetabaseDatabase { id: number; name: string; diff --git a/packages/cli/src/context/sl/semantic-layer.service.ts b/packages/cli/src/context/sl/semantic-layer.service.ts index 99837247..913f62d1 100644 --- a/packages/cli/src/context/sl/semantic-layer.service.ts +++ b/packages/cli/src/context/sl/semantic-layer.service.ts @@ -2,6 +2,7 @@ import YAML from 'yaml'; import type { KtxFileStorePort } from '../../context/core/file-store.js'; import type { KtxLogger } 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 { SlConnectionCatalogPort, SlPythonPort } from './ports.js'; import { normalizeSemanticLayerDescriptions } from './description-normalization.js'; @@ -674,7 +675,7 @@ export class SemanticLayerService { if (!connection) { throw new Error(`Data source not found: ${connectionId}`); } - return SemanticLayerService.mapDialect(connection.connectionType); + return dialectForConnectionType(connection.connectionType); } async listFilesForConnection(connectionId: string): Promise { @@ -1107,25 +1108,6 @@ export class SemanticLayerService { 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 = { - 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 * the python SL engine, and execute the generated SQL against the data source. @@ -1153,7 +1135,7 @@ export class SemanticLayerService { if (!connection) { 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 const { data: slResult, error: slError } = await this.python.query({ diff --git a/packages/cli/src/context/sl/tools/sl-warehouse-validation.ts b/packages/cli/src/context/sl/tools/sl-warehouse-validation.ts index 56f80833..69acc783 100644 --- a/packages/cli/src/context/sl/tools/sl-warehouse-validation.ts +++ b/packages/cli/src/context/sl/tools/sl-warehouse-validation.ts @@ -2,9 +2,10 @@ import YAML from 'yaml'; import type { GitService } from '../../../context/core/git.service.js'; import type { KtxFileListResult, KtxFileReadResult, KtxFileStorePort } from '../../../context/core/file-store.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 { 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 type { SemanticLayerSource } from '../types.js'; import { sourceDefinitionSchema } from './base-semantic-layer.tool.js'; @@ -28,7 +29,7 @@ function resolveDialect(warehouse: string | null): string | null { if (!warehouse) { return null; } - return SemanticLayerService.mapDialect(warehouse); + return dialectForConnectionType(warehouse); } function wrapWithZeroRowQuery(sql: string, dialect: string): string { diff --git a/packages/cli/test/context/connections/connection-type-dialect.test.ts b/packages/cli/test/context/connections/connection-type-dialect.test.ts new file mode 100644 index 00000000..93bf57fa --- /dev/null +++ b/packages/cli/test/context/connections/connection-type-dialect.test.ts @@ -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); + } + }); +}); diff --git a/packages/cli/test/context/connections/local-warehouse-descriptor.test.ts b/packages/cli/test/context/connections/local-warehouse-descriptor.test.ts index e0a285a9..a368ae5d 100644 --- a/packages/cli/test/context/connections/local-warehouse-descriptor.test.ts +++ b/packages/cli/test/context/connections/local-warehouse-descriptor.test.ts @@ -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', () => { expect( localConnectionToWarehouseDescriptor('looker', { @@ -51,6 +66,7 @@ describe('local connection info helpers', () => { expect(localConnectionTypeForConfig('warehouse', { driver: 'postgres' })).toBe('POSTGRESQL'); expect(localConnectionTypeForConfig('bq', { driver: 'bigquery', project_id: 'acme' })).toBe('BIGQUERY'); expect(localConnectionTypeForConfig('snowflake', { driver: 'snowflake' })).toBe('SNOWFLAKE'); + expect(localConnectionTypeForConfig('athena-warehouse', { driver: 'athena' })).toBe('ATHENA'); }); it('keeps removed driver aliases as display-only labels', () => { diff --git a/packages/cli/test/context/ingest/adapters/looker/mapping.test.ts b/packages/cli/test/context/ingest/adapters/looker/mapping.test.ts index a5186a75..dbec700f 100644 --- a/packages/cli/test/context/ingest/adapters/looker/mapping.test.ts +++ b/packages/cli/test/context/ingest/adapters/looker/mapping.test.ts @@ -72,6 +72,8 @@ describe('looker dialect and target validation helpers', () => { it('maps Looker dialect names to ktx connection types', () => { expect(lookerDialectToConnectionType('bigquery_standard_sql')).toBe('BIGQUERY'); expect(lookerDialectToConnectionType('postgres')).toBe('POSTGRESQL'); + expect(lookerDialectToConnectionType('awsathena')).toBe('ATHENA'); + expect(lookerDialectToConnectionType('AWSATHENA')).toBe('ATHENA'); expect(lookerDialectToConnectionType('mssql')).toBeNull(); expect(lookerDialectToConnectionType('tsql')).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', () => { expect(sqlglotDialectForConnectionType('BIGQUERY')).toBe('bigquery'); expect(sqlglotDialectForConnectionType('POSTGRESQL')).toBe('postgres'); + expect(sqlglotDialectForConnectionType('ATHENA')).toBe('athena'); 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', () => { expect( projectParsedIdentifier({ diff --git a/packages/cli/test/context/ingest/adapters/metabase/mapping.test.ts b/packages/cli/test/context/ingest/adapters/metabase/mapping.test.ts index 3fb56442..dda0083e 100644 --- a/packages/cli/test/context/ingest/adapters/metabase/mapping.test.ts +++ b/packages/cli/test/context/ingest/adapters/metabase/mapping.test.ts @@ -121,6 +121,15 @@ describe('validateMappingPhysicalMatch', () => { 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', () => { expect( 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', () => { it('returns only mismatched physical mappings', () => { expect( diff --git a/packages/cli/test/context/sl/semantic-layer.service.test.ts b/packages/cli/test/context/sl/semantic-layer.service.test.ts index edc31ce1..5532986d 100644 --- a/packages/cli/test/context/sl/semantic-layer.service.test.ts +++ b/packages/cli/test/context/sl/semantic-layer.service.test.ts @@ -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 () => { const schemaPath = 'semantic-layer/conn-1/_schema/core.yaml'; 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', () => { it('returns empty when every measure segment resolves', () => { const source = { diff --git a/python/ktx-sl/semantic_layer/table_identifier_parser.py b/python/ktx-sl/semantic_layer/table_identifier_parser.py index f8df6631..a32863e1 100644 --- a/python/ktx-sl/semantic_layer/table_identifier_parser.py +++ b/python/ktx-sl/semantic_layer/table_identifier_parser.py @@ -18,6 +18,8 @@ SUPPORTED_TABLE_IDENTIFIER_DIALECTS = { "sqlite", "tsql", "clickhouse", + "athena", + "duckdb", } ParseTableIdentifierReason = Literal[ diff --git a/python/ktx-sl/tests/test_table_identifier_parser.py b/python/ktx-sl/tests/test_table_identifier_parser.py index ef18e990..6a6cdd03 100644 --- a/python/ktx-sl/tests/test_table_identifier_parser.py +++ b/python/ktx-sl/tests/test_table_identifier_parser.py @@ -23,6 +23,16 @@ def test_parse_table_identifier_supported_dialects_and_aliases() -> None: sql_table_name="RAW.PUBLIC.ORDERS", 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"].schema_ == "PUBLIC" 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: