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
This commit is contained in:
Andrey Avtomonov 2026-07-02 15:10:29 +02:00 committed by GitHub
parent fe7e6bd1fa
commit f310391da5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 252 additions and 45 deletions

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

@ -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', () => {

View file

@ -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({

View file

@ -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(

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 () => {
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 = {