@kaelio/ktx 0.15.0 → 0.16.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (72) hide show
  1. package/assets/python/kaelio_ktx-0.16.0-py3-none-any.whl +0 -0
  2. package/assets/python/manifest.json +4 -4
  3. package/dist/.tsbuildinfo +1 -1
  4. package/dist/commands/setup-commands.js +1 -0
  5. package/dist/connection-drivers.d.ts +1 -1
  6. package/dist/connection-drivers.js +2 -0
  7. package/dist/connection.js +1 -0
  8. package/dist/connectors/athena/connector.d.ts +160 -0
  9. package/dist/connectors/athena/connector.js +378 -0
  10. package/dist/connectors/athena/dialect.d.ts +31 -0
  11. package/dist/connectors/athena/dialect.js +146 -0
  12. package/dist/connectors/athena/live-database-introspection.d.ts +10 -0
  13. package/dist/connectors/athena/live-database-introspection.js +24 -0
  14. package/dist/connectors/duckdb/connector.d.ts +60 -0
  15. package/dist/connectors/duckdb/connector.js +282 -0
  16. package/dist/connectors/duckdb/dialect.d.ts +31 -0
  17. package/dist/connectors/duckdb/dialect.js +155 -0
  18. package/dist/connectors/duckdb/federated-attach.js +7 -0
  19. package/dist/connectors/duckdb/federated-executor.js +7 -13
  20. package/dist/connectors/duckdb/live-database-introspection.d.ts +8 -0
  21. package/dist/connectors/duckdb/live-database-introspection.js +24 -0
  22. package/dist/connectors/shared/duckdb-json-safe.d.ts +3 -0
  23. package/dist/connectors/shared/duckdb-json-safe.js +12 -0
  24. package/dist/context/connections/connection-type-dialect.d.ts +2 -0
  25. package/dist/context/connections/connection-type-dialect.js +20 -0
  26. package/dist/context/connections/connection-type.d.ts +6 -15
  27. package/dist/context/connections/connection-type.js +1 -10
  28. package/dist/context/connections/dialects.js +4 -0
  29. package/dist/context/connections/drivers.js +38 -0
  30. package/dist/context/connections/federation.d.ts +1 -1
  31. package/dist/context/connections/federation.js +6 -5
  32. package/dist/context/connections/local-warehouse-descriptor.d.ts +1 -0
  33. package/dist/context/connections/local-warehouse-descriptor.js +10 -0
  34. package/dist/context/connections/project-sql-executor.d.ts +24 -0
  35. package/dist/context/connections/project-sql-executor.js +56 -0
  36. package/dist/context/connections/query-policy.d.ts +12 -0
  37. package/dist/context/connections/query-policy.js +46 -0
  38. package/dist/context/core/config-reference.js +11 -1
  39. package/dist/context/daemon/semantic-layer-compute.d.ts +16 -0
  40. package/dist/context/daemon/semantic-layer-compute.js +44 -2
  41. package/dist/context/ingest/adapters/looker/mapping.d.ts +1 -0
  42. package/dist/context/ingest/adapters/looker/mapping.js +3 -10
  43. package/dist/context/ingest/adapters/looker/types.d.ts +6 -15
  44. package/dist/context/ingest/adapters/metabase/mapping.d.ts +1 -0
  45. package/dist/context/ingest/adapters/metabase/mapping.js +1 -0
  46. package/dist/context/mcp/context-tools.js +1 -0
  47. package/dist/context/mcp/local-project-ports.d.ts +1 -1
  48. package/dist/context/mcp/local-project-ports.js +53 -51
  49. package/dist/context/project/config.d.ts +92 -0
  50. package/dist/context/project/driver-schemas.d.ts +46 -0
  51. package/dist/context/project/driver-schemas.js +8 -0
  52. package/dist/context/scan/local-scan.js +3 -1
  53. package/dist/context/scan/types.d.ts +1 -1
  54. package/dist/context/sl/local-query.js +23 -6
  55. package/dist/context/sl/semantic-layer.service.d.ts +0 -4
  56. package/dist/context/sl/semantic-layer.service.js +3 -20
  57. package/dist/context/sl/tools/sl-warehouse-validation.d.ts +1 -1
  58. package/dist/context/sl/tools/sl-warehouse-validation.js +2 -2
  59. package/dist/context/sl/types.d.ts +1 -0
  60. package/dist/context/sql-analysis/dialect-notes.d.ts +1 -1
  61. package/dist/context/sql-analysis/dialect-notes.js +3 -2
  62. package/dist/context/sql-analysis/dialect.js +1 -0
  63. package/dist/context/sql-analysis/dialects/athena.md +12 -0
  64. package/dist/context/sql-analysis/dialects/duckdb.md +10 -0
  65. package/dist/local-adapters.js +17 -0
  66. package/dist/setup-databases.d.ts +1 -1
  67. package/dist/setup-databases.js +52 -2
  68. package/dist/sql.js +9 -14
  69. package/dist/status-project.js +2 -1
  70. package/dist/telemetry/events.d.ts +1 -1
  71. package/package.json +3 -1
  72. package/assets/python/kaelio_ktx-0.15.0-py3-none-any.whl +0 -0
@@ -26,6 +26,7 @@ function llmBackend(value) {
26
26
  }
27
27
  function databaseDriver(value) {
28
28
  if (value === 'sqlite' ||
29
+ value === 'duckdb' ||
29
30
  value === 'postgres' ||
30
31
  value === 'mysql' ||
31
32
  value === 'clickhouse' ||
@@ -1,5 +1,5 @@
1
1
  import type { KtxProjectConnectionConfig } from './context/project/config.js';
2
2
  /** @internal Canonical SQL-warehouse driver ids; the dialect-notes coverage test derives its required coverage from this set. */
3
- export declare const KTX_DATABASE_DRIVER_IDS: readonly ["sqlite", "postgres", "mysql", "clickhouse", "sqlserver", "bigquery", "snowflake"];
3
+ export declare const KTX_DATABASE_DRIVER_IDS: readonly ["sqlite", "duckdb", "postgres", "mysql", "clickhouse", "sqlserver", "bigquery", "snowflake", "athena"];
4
4
  export declare function normalizeConnectionDriver(connection: KtxProjectConnectionConfig): string;
5
5
  export declare function isDatabaseDriver(driver: string): boolean;
@@ -1,12 +1,14 @@
1
1
  /** @internal Canonical SQL-warehouse driver ids; the dialect-notes coverage test derives its required coverage from this set. */
2
2
  export const KTX_DATABASE_DRIVER_IDS = [
3
3
  'sqlite',
4
+ 'duckdb',
4
5
  'postgres',
5
6
  'mysql',
6
7
  'clickhouse',
7
8
  'sqlserver',
8
9
  'bigquery',
9
10
  'snowflake',
11
+ 'athena',
10
12
  ];
11
13
  // mongodb is a database driver but has no SQL dialect, so it sits outside the
12
14
  // dialect-notes coverage set above.
@@ -22,6 +22,7 @@ import { formatErrorDetail, scrubErrorClass } from './telemetry/scrubber.js';
22
22
  profileMark('module:connection');
23
23
  const SUPPORTED_TEST_DRIVERS = [
24
24
  'sqlite',
25
+ 'duckdb',
25
26
  'postgres',
26
27
  'mysql',
27
28
  'clickhouse',
@@ -0,0 +1,160 @@
1
+ import { type KtxConnectorTestResult, type KtxColumnSampleInput, type KtxColumnSampleResult, type KtxColumnStatsInput, type KtxColumnStatsResult, type KtxQueryResult, type KtxReadOnlyQueryInput, type KtxScanConnector, type KtxScanContext, type KtxScanInput, type KtxSchemaSnapshot, type KtxTableListEntry, type KtxTableRef, type KtxTableSampleInput, type KtxTableSampleResult } from '../../context/scan/types.js';
2
+ export interface KtxAthenaConnectionConfig {
3
+ driver?: string;
4
+ region?: string;
5
+ s3_staging_dir?: string;
6
+ workgroup?: string;
7
+ catalog?: string;
8
+ database?: string;
9
+ databases?: string[];
10
+ [key: string]: unknown;
11
+ }
12
+ export interface KtxAthenaResolvedConnectionConfig {
13
+ region: string;
14
+ s3StagingDir: string;
15
+ workgroup: string;
16
+ catalog: string;
17
+ database: string | undefined;
18
+ databases: string[];
19
+ }
20
+ interface KtxAthenaQueryExecutionStatus {
21
+ State?: string;
22
+ StateChangeReason?: string;
23
+ }
24
+ interface KtxAthenaQueryExecution {
25
+ Status?: KtxAthenaQueryExecutionStatus;
26
+ }
27
+ interface KtxAthenaColumnInfo {
28
+ Name?: string;
29
+ Type?: string;
30
+ }
31
+ interface KtxAthenaDatum {
32
+ VarCharValue?: string;
33
+ }
34
+ interface KtxAthenaRow {
35
+ Data?: KtxAthenaDatum[];
36
+ }
37
+ interface KtxAthenaResultSet {
38
+ Rows?: KtxAthenaRow[];
39
+ ResultSetMetadata?: {
40
+ ColumnInfo?: KtxAthenaColumnInfo[];
41
+ };
42
+ }
43
+ /** @internal */
44
+ export interface KtxAthenaClient {
45
+ startQueryExecution(input: {
46
+ QueryString: string;
47
+ ResultConfiguration: {
48
+ OutputLocation: string;
49
+ };
50
+ WorkGroup: string;
51
+ QueryExecutionContext?: {
52
+ Database?: string;
53
+ Catalog?: string;
54
+ };
55
+ }): Promise<{
56
+ QueryExecutionId?: string;
57
+ }>;
58
+ getQueryExecution(input: {
59
+ QueryExecutionId: string;
60
+ }): Promise<{
61
+ QueryExecution?: KtxAthenaQueryExecution;
62
+ }>;
63
+ getQueryResults(input: {
64
+ QueryExecutionId: string;
65
+ NextToken?: string;
66
+ }): Promise<{
67
+ ResultSet?: KtxAthenaResultSet;
68
+ NextToken?: string;
69
+ }>;
70
+ }
71
+ interface KtxGlueColumnDef {
72
+ Name?: string;
73
+ Type?: string;
74
+ Comment?: string;
75
+ }
76
+ interface KtxGlueStorageDescriptor {
77
+ Columns?: KtxGlueColumnDef[];
78
+ }
79
+ /** @internal */
80
+ export interface KtxGlueTable {
81
+ Name?: string;
82
+ TableType?: string;
83
+ StorageDescriptor?: KtxGlueStorageDescriptor;
84
+ PartitionKeys?: KtxGlueColumnDef[];
85
+ Description?: string;
86
+ Parameters?: Record<string, string>;
87
+ }
88
+ /** @internal */
89
+ export interface KtxGlueClient {
90
+ getDatabases(input: {
91
+ CatalogId?: string;
92
+ NextToken?: string;
93
+ }): Promise<{
94
+ DatabaseList?: Array<{
95
+ Name?: string;
96
+ }>;
97
+ NextToken?: string;
98
+ }>;
99
+ getTables(input: {
100
+ DatabaseName: string;
101
+ CatalogId?: string;
102
+ NextToken?: string;
103
+ }): Promise<{
104
+ TableList?: KtxGlueTable[];
105
+ NextToken?: string;
106
+ }>;
107
+ }
108
+ export interface KtxAthenaClientFactory {
109
+ createAthenaClient(region: string): KtxAthenaClient;
110
+ createGlueClient(region: string): KtxGlueClient;
111
+ }
112
+ export declare function isKtxAthenaConnectionConfig(connection: unknown): connection is KtxAthenaConnectionConfig;
113
+ /** @internal */
114
+ export declare function athenaConnectionConfigFromConfig(input: {
115
+ connectionId: string;
116
+ connection: KtxAthenaConnectionConfig | undefined;
117
+ env?: NodeJS.ProcessEnv;
118
+ }): KtxAthenaResolvedConnectionConfig;
119
+ export interface KtxAthenaScanConnectorOptions {
120
+ connectionId: string;
121
+ connection: KtxAthenaConnectionConfig | undefined;
122
+ clientFactory?: KtxAthenaClientFactory;
123
+ env?: NodeJS.ProcessEnv;
124
+ now?: () => Date;
125
+ }
126
+ export declare class KtxAthenaScanConnector implements KtxScanConnector {
127
+ readonly id: string;
128
+ readonly driver: "athena";
129
+ readonly capabilities: import("../../context/scan/types.js").KtxConnectorCapabilities;
130
+ private readonly connectionId;
131
+ private readonly resolved;
132
+ private readonly clientFactory;
133
+ private readonly now;
134
+ private readonly dialect;
135
+ private athenaClient;
136
+ private glueClient;
137
+ constructor(options: KtxAthenaScanConnectorOptions);
138
+ testConnection(): Promise<KtxConnectorTestResult>;
139
+ introspect(input: KtxScanInput, _ctx: KtxScanContext): Promise<KtxSchemaSnapshot>;
140
+ sampleTable(input: KtxTableSampleInput, _ctx: KtxScanContext): Promise<KtxTableSampleResult & {
141
+ headerTypes?: string[];
142
+ }>;
143
+ sampleColumn(input: KtxColumnSampleInput, _ctx: KtxScanContext): Promise<KtxColumnSampleResult>;
144
+ columnStats(_input: KtxColumnStatsInput, _ctx: KtxScanContext): Promise<KtxColumnStatsResult | null>;
145
+ executeReadOnly(input: KtxReadOnlyQueryInput, _ctx: KtxScanContext): Promise<KtxQueryResult>;
146
+ listSchemas(): Promise<string[]>;
147
+ listTables(databases?: string[]): Promise<KtxTableListEntry[]>;
148
+ cleanup(): Promise<void>;
149
+ qTableName(table: Pick<KtxTableRef, 'name'> & Partial<Pick<KtxTableRef, 'catalog' | 'db'>>): string;
150
+ private getAthenaClient;
151
+ private getGlueClient;
152
+ private listDatabasesPaginated;
153
+ private listGlueTablesPaginated;
154
+ private introspectDatabase;
155
+ private toSchemaColumns;
156
+ private query;
157
+ private waitForQueryCompletion;
158
+ private assertConnection;
159
+ }
160
+ export {};
@@ -0,0 +1,378 @@
1
+ import { AthenaClient, StartQueryExecutionCommand, GetQueryExecutionCommand, GetQueryResultsCommand } from '@aws-sdk/client-athena';
2
+ import { GlueClient, GetDatabasesCommand, GetTablesCommand } from '@aws-sdk/client-glue';
3
+ import { getSqlDialectForDriver } from '../../context/connections/dialects.js';
4
+ import { assertReadOnlySql, limitSqlForExecution } from '../../context/connections/read-only-sql.js';
5
+ import { connectorTestFailure, createKtxConnectorCapabilities, } from '../../context/scan/types.js';
6
+ import { scopedTableNames } from '../../context/scan/table-ref.js';
7
+ import { resolveStringReference } from '../shared/string-reference.js';
8
+ class DefaultAthenaClientFactory {
9
+ createAthenaClient(region) {
10
+ const client = new AthenaClient({ region });
11
+ return {
12
+ startQueryExecution: async (input) => {
13
+ const result = await client.send(new StartQueryExecutionCommand({
14
+ QueryString: input.QueryString,
15
+ ResultConfiguration: { OutputLocation: input.ResultConfiguration.OutputLocation },
16
+ WorkGroup: input.WorkGroup,
17
+ QueryExecutionContext: input.QueryExecutionContext,
18
+ }));
19
+ return { QueryExecutionId: result.QueryExecutionId };
20
+ },
21
+ getQueryExecution: async (input) => {
22
+ const result = await client.send(new GetQueryExecutionCommand({ QueryExecutionId: input.QueryExecutionId }));
23
+ return {
24
+ QueryExecution: result.QueryExecution
25
+ ? {
26
+ Status: {
27
+ State: result.QueryExecution.Status?.State,
28
+ StateChangeReason: result.QueryExecution.Status?.StateChangeReason,
29
+ },
30
+ }
31
+ : undefined,
32
+ };
33
+ },
34
+ getQueryResults: async (input) => {
35
+ const result = await client.send(new GetQueryResultsCommand({ QueryExecutionId: input.QueryExecutionId, NextToken: input.NextToken }));
36
+ return {
37
+ ResultSet: result.ResultSet,
38
+ NextToken: result.NextToken,
39
+ };
40
+ },
41
+ };
42
+ }
43
+ createGlueClient(region) {
44
+ const client = new GlueClient({ region });
45
+ return {
46
+ getDatabases: async (input) => {
47
+ const result = await client.send(new GetDatabasesCommand({ CatalogId: input.CatalogId, NextToken: input.NextToken }));
48
+ return {
49
+ DatabaseList: result.DatabaseList?.map((db) => ({ Name: db.Name })),
50
+ NextToken: result.NextToken,
51
+ };
52
+ },
53
+ getTables: async (input) => {
54
+ const result = await client.send(new GetTablesCommand({ DatabaseName: input.DatabaseName, CatalogId: input.CatalogId, NextToken: input.NextToken }));
55
+ return {
56
+ TableList: result.TableList,
57
+ NextToken: result.NextToken,
58
+ };
59
+ },
60
+ };
61
+ }
62
+ }
63
+ function stringConfigValue(connection, key, env) {
64
+ const value = connection?.[key];
65
+ if (typeof value !== 'string' || value.trim().length === 0)
66
+ return undefined;
67
+ // Resolve before checking emptiness: an unset `env:` reference resolves to '',
68
+ // which must become undefined so `?? default` applies instead of keeping ''.
69
+ const resolved = resolveStringReference(value.trim(), env).trim();
70
+ return resolved.length > 0 ? resolved : undefined;
71
+ }
72
+ function configuredAthenaDatabases(connection) {
73
+ if (!Array.isArray(connection.databases))
74
+ return [];
75
+ const selected = connection.databases
76
+ .filter((database) => typeof database === 'string' && database.trim().length > 0)
77
+ .map((database) => database.trim());
78
+ return [...new Set(selected)];
79
+ }
80
+ export function isKtxAthenaConnectionConfig(connection) {
81
+ return (typeof connection === 'object' &&
82
+ connection !== null &&
83
+ String(connection.driver ?? '').toLowerCase() === 'athena');
84
+ }
85
+ /** @internal */
86
+ export function athenaConnectionConfigFromConfig(input) {
87
+ const inputDriver = input.connection?.driver ?? 'unknown';
88
+ if (!isKtxAthenaConnectionConfig(input.connection)) {
89
+ throw new Error(`Native Athena connector cannot run driver "${String(inputDriver)}"`);
90
+ }
91
+ const env = input.env ?? process.env;
92
+ const region = stringConfigValue(input.connection, 'region', env);
93
+ if (!region) {
94
+ throw new Error(`Native Athena connector requires connections.${input.connectionId}.region`);
95
+ }
96
+ const s3StagingDir = stringConfigValue(input.connection, 's3_staging_dir', env);
97
+ if (!s3StagingDir) {
98
+ throw new Error(`Native Athena connector requires connections.${input.connectionId}.s3_staging_dir`);
99
+ }
100
+ return {
101
+ region,
102
+ s3StagingDir,
103
+ workgroup: stringConfigValue(input.connection, 'workgroup', env) ?? 'primary',
104
+ catalog: stringConfigValue(input.connection, 'catalog', env) ?? 'AwsDataCatalog',
105
+ database: stringConfigValue(input.connection, 'database', env),
106
+ databases: configuredAthenaDatabases(input.connection),
107
+ };
108
+ }
109
+ function glueTableKind(tableType) {
110
+ const t = String(tableType ?? '').toUpperCase();
111
+ if (t === 'VIRTUAL_VIEW')
112
+ return 'view';
113
+ return 'table';
114
+ }
115
+ const POLL_INTERVAL_MS = 250;
116
+ const QUERY_TIMEOUT_MS = 5 * 60 * 1000;
117
+ export class KtxAthenaScanConnector {
118
+ id;
119
+ driver = 'athena';
120
+ capabilities = createKtxConnectorCapabilities({
121
+ tableSampling: true,
122
+ columnSampling: true,
123
+ columnStats: false,
124
+ readOnlySql: true,
125
+ nestedAnalysis: false,
126
+ formalForeignKeys: false,
127
+ estimatedRowCounts: false,
128
+ });
129
+ connectionId;
130
+ resolved;
131
+ clientFactory;
132
+ now;
133
+ dialect = getSqlDialectForDriver('athena');
134
+ athenaClient = null;
135
+ glueClient = null;
136
+ constructor(options) {
137
+ this.connectionId = options.connectionId;
138
+ this.resolved = athenaConnectionConfigFromConfig({
139
+ connectionId: options.connectionId,
140
+ connection: options.connection,
141
+ env: options.env,
142
+ });
143
+ this.clientFactory = options.clientFactory ?? new DefaultAthenaClientFactory();
144
+ this.now = options.now ?? (() => new Date());
145
+ this.id = `athena:${options.connectionId}`;
146
+ }
147
+ async testConnection() {
148
+ try {
149
+ await this.listDatabasesPaginated({ maxResults: 1 });
150
+ return { success: true };
151
+ }
152
+ catch (error) {
153
+ return connectorTestFailure(error);
154
+ }
155
+ }
156
+ async introspect(input, _ctx) {
157
+ this.assertConnection(input.connectionId);
158
+ // Honor the configured `databases` scope (written by `ktx setup`); fall back
159
+ // to every Glue database only when the scope is unset.
160
+ const databases = this.resolved.databases.length > 0 ? this.resolved.databases : await this.listDatabasesPaginated({});
161
+ const tables = [];
162
+ for (const database of databases) {
163
+ const scopedNames = input.tableScope
164
+ ? scopedTableNames(input.tableScope, { catalog: this.resolved.catalog, db: database })
165
+ : null;
166
+ tables.push(...(await this.introspectDatabase(database, scopedNames)));
167
+ }
168
+ return {
169
+ connectionId: this.connectionId,
170
+ driver: 'athena',
171
+ extractedAt: this.now().toISOString(),
172
+ scope: { catalogs: [this.resolved.catalog], datasets: databases },
173
+ metadata: {
174
+ catalog: this.resolved.catalog,
175
+ databases,
176
+ table_count: tables.length,
177
+ total_columns: tables.reduce((sum, t) => sum + t.columns.length, 0),
178
+ },
179
+ tables,
180
+ warnings: [],
181
+ };
182
+ }
183
+ async sampleTable(input, _ctx) {
184
+ this.assertConnection(input.connectionId);
185
+ const result = await this.query(this.dialect.generateSampleQuery(this.qTableName(input.table), input.limit, input.columns));
186
+ return { headers: result.headers, headerTypes: result.headerTypes, rows: result.rows, totalRows: result.totalRows };
187
+ }
188
+ async sampleColumn(input, _ctx) {
189
+ this.assertConnection(input.connectionId);
190
+ const result = await this.query(this.dialect.generateColumnSampleQuery(this.qTableName(input.table), input.column, input.limit));
191
+ return {
192
+ values: result.rows.filter((row) => row.length > 0 && row[0] !== null).map((row) => row[0]),
193
+ nullCount: null,
194
+ distinctCount: null,
195
+ };
196
+ }
197
+ async columnStats(_input, _ctx) {
198
+ return null;
199
+ }
200
+ async executeReadOnly(input, _ctx) {
201
+ this.assertConnection(input.connectionId);
202
+ const limitedSql = limitSqlForExecution(assertReadOnlySql(input.sql), input.maxRows);
203
+ const result = await this.query(limitedSql);
204
+ return { ...result, rowCount: result.rows.length };
205
+ }
206
+ async listSchemas() {
207
+ return this.listDatabasesPaginated({});
208
+ }
209
+ async listTables(databases) {
210
+ const targetDatabases = databases && databases.length > 0 ? databases : await this.listDatabasesPaginated({});
211
+ const entries = [];
212
+ for (const database of targetDatabases) {
213
+ const glueTables = await this.listGlueTablesPaginated(database);
214
+ for (const t of glueTables) {
215
+ if (!t.Name)
216
+ continue;
217
+ entries.push({
218
+ catalog: this.resolved.catalog,
219
+ schema: database,
220
+ name: t.Name,
221
+ kind: glueTableKind(t.TableType),
222
+ });
223
+ }
224
+ }
225
+ return entries;
226
+ }
227
+ async cleanup() {
228
+ this.athenaClient = null;
229
+ this.glueClient = null;
230
+ }
231
+ qTableName(table) {
232
+ return this.dialect.formatTableName(table);
233
+ }
234
+ getAthenaClient() {
235
+ if (!this.athenaClient) {
236
+ this.athenaClient = this.clientFactory.createAthenaClient(this.resolved.region);
237
+ }
238
+ return this.athenaClient;
239
+ }
240
+ getGlueClient() {
241
+ if (!this.glueClient) {
242
+ this.glueClient = this.clientFactory.createGlueClient(this.resolved.region);
243
+ }
244
+ return this.glueClient;
245
+ }
246
+ async listDatabasesPaginated(opts) {
247
+ const names = [];
248
+ let nextToken;
249
+ do {
250
+ const result = await this.getGlueClient().getDatabases({ NextToken: nextToken });
251
+ for (const db of result.DatabaseList ?? []) {
252
+ if (db.Name)
253
+ names.push(db.Name);
254
+ if (opts.maxResults && names.length >= opts.maxResults)
255
+ return names;
256
+ }
257
+ nextToken = result.NextToken;
258
+ } while (nextToken);
259
+ return names;
260
+ }
261
+ async listGlueTablesPaginated(database) {
262
+ const tables = [];
263
+ let nextToken;
264
+ do {
265
+ const result = await this.getGlueClient().getTables({ DatabaseName: database, NextToken: nextToken });
266
+ tables.push(...(result.TableList ?? []));
267
+ nextToken = result.NextToken;
268
+ } while (nextToken);
269
+ return tables;
270
+ }
271
+ async introspectDatabase(database, scopedNames) {
272
+ if (scopedNames && scopedNames.length === 0)
273
+ return [];
274
+ const glueTables = await this.listGlueTablesPaginated(database);
275
+ const scopeSet = scopedNames ? new Set(scopedNames) : null;
276
+ return glueTables
277
+ .filter((t) => Boolean(t.Name) && (!scopeSet || scopeSet.has(t.Name)))
278
+ .map((t) => ({
279
+ catalog: this.resolved.catalog,
280
+ db: database,
281
+ name: t.Name,
282
+ kind: glueTableKind(t.TableType),
283
+ comment: t.Description ?? null,
284
+ estimatedRows: null,
285
+ columns: this.toSchemaColumns(t),
286
+ foreignKeys: [],
287
+ }));
288
+ }
289
+ toSchemaColumns(table) {
290
+ const columns = [...(table.StorageDescriptor?.Columns ?? []), ...(table.PartitionKeys ?? [])];
291
+ return columns
292
+ .filter((col) => Boolean(col.Name))
293
+ .map((col) => {
294
+ const nativeType = String(col.Type ?? 'string').toLowerCase();
295
+ return {
296
+ name: col.Name,
297
+ nativeType,
298
+ normalizedType: this.dialect.mapDataType(nativeType),
299
+ dimensionType: this.dialect.mapToDimensionType(nativeType),
300
+ nullable: true,
301
+ primaryKey: false,
302
+ comment: col.Comment ?? null,
303
+ };
304
+ });
305
+ }
306
+ async query(sql) {
307
+ const athena = this.getAthenaClient();
308
+ const { QueryExecutionId } = await athena.startQueryExecution({
309
+ QueryString: sql,
310
+ ResultConfiguration: { OutputLocation: this.resolved.s3StagingDir },
311
+ WorkGroup: this.resolved.workgroup,
312
+ ...(this.resolved.database || this.resolved.catalog
313
+ ? {
314
+ QueryExecutionContext: {
315
+ ...(this.resolved.database ? { Database: this.resolved.database } : {}),
316
+ ...(this.resolved.catalog ? { Catalog: this.resolved.catalog } : {}),
317
+ },
318
+ }
319
+ : {}),
320
+ });
321
+ if (!QueryExecutionId) {
322
+ throw new Error('Athena did not return a QueryExecutionId');
323
+ }
324
+ await this.waitForQueryCompletion(athena, QueryExecutionId);
325
+ const rows = [];
326
+ let headers = [];
327
+ let headerTypes = [];
328
+ let nextToken;
329
+ let firstPage = true;
330
+ do {
331
+ const result = await athena.getQueryResults({ QueryExecutionId, NextToken: nextToken });
332
+ const resultSet = result.ResultSet;
333
+ if (firstPage) {
334
+ const columnInfo = resultSet?.ResultSetMetadata?.ColumnInfo ?? [];
335
+ headers = columnInfo.map((col) => col.Name ?? '');
336
+ headerTypes = columnInfo.map((col) => String(col.Type ?? 'varchar').toUpperCase());
337
+ firstPage = false;
338
+ }
339
+ const pageRows = resultSet?.Rows ?? [];
340
+ // Athena includes the header row as the first row of the first page — skip it.
341
+ const dataRows = nextToken === undefined ? pageRows.slice(1) : pageRows;
342
+ for (const row of dataRows) {
343
+ rows.push((row.Data ?? []).map((d) => d.VarCharValue ?? null));
344
+ }
345
+ nextToken = result.NextToken;
346
+ } while (nextToken);
347
+ return {
348
+ headers,
349
+ headerTypes: headerTypes.length > 0 ? headerTypes : undefined,
350
+ rows,
351
+ totalRows: rows.length,
352
+ rowCount: rows.length,
353
+ };
354
+ }
355
+ async waitForQueryCompletion(athena, queryExecutionId) {
356
+ const terminalStates = new Set(['SUCCEEDED', 'FAILED', 'CANCELLED']);
357
+ const deadline = this.now().getTime() + QUERY_TIMEOUT_MS;
358
+ for (;;) {
359
+ const { QueryExecution } = await athena.getQueryExecution({ QueryExecutionId: queryExecutionId });
360
+ const state = QueryExecution?.Status?.State ?? '';
361
+ if (state === 'SUCCEEDED')
362
+ return;
363
+ if (terminalStates.has(state)) {
364
+ const reason = QueryExecution?.Status?.StateChangeReason ?? state;
365
+ throw new Error(`Athena query ${state}: ${reason}`);
366
+ }
367
+ if (this.now().getTime() >= deadline) {
368
+ throw new Error(`Athena query ${queryExecutionId} timed out after ${QUERY_TIMEOUT_MS / 1000}s`);
369
+ }
370
+ await new Promise((resolve) => setTimeout(resolve, POLL_INTERVAL_MS));
371
+ }
372
+ }
373
+ assertConnection(connectionId) {
374
+ if (connectionId !== this.connectionId) {
375
+ throw new Error(`Athena connector ${this.connectionId} cannot scan connection ${connectionId}`);
376
+ }
377
+ }
378
+ }
@@ -0,0 +1,31 @@
1
+ import type { KtxSqlDialect } from '../../context/connections/dialects.js';
2
+ import type { KtxSchemaDimensionType, KtxTableRef } from '../../context/scan/types.js';
3
+ type AthenaTableNameRef = Pick<KtxTableRef, 'name'> & Partial<Pick<KtxTableRef, 'catalog' | 'db'>>;
4
+ /** @internal */
5
+ export declare class KtxAthenaDialect implements KtxSqlDialect {
6
+ readonly type: "athena";
7
+ private readonly dimensionTypeMappings;
8
+ quoteIdentifier(identifier: string): string;
9
+ formatTableName(table: AthenaTableNameRef): string;
10
+ formatDisplayRef(table: AthenaTableNameRef): string;
11
+ parseDisplayRef(display: string): KtxTableRef | null;
12
+ columnDisplayTablePartCount(): 1 | 2 | 3;
13
+ mapDataType(nativeType: string): string;
14
+ mapToDimensionType(nativeType: string): KtxSchemaDimensionType;
15
+ generateSampleQuery(tableName: string, limit: number, columns?: string[]): string;
16
+ generateColumnSampleQuery(tableName: string, columnName: string, limit: number): string;
17
+ generateCardinalitySampleQuery(tableName: string, columnName: string, sampleSize: number): string;
18
+ generateRandomizedCardinalitySampleQuery(tableName: string, columnName: string, sampleSize: number): string;
19
+ generateDistinctValuesQuery(tableName: string, columnName: string, limit: number): string;
20
+ generateColumnStatisticsQuery(_schemaName: string, _tableName: string): string | null;
21
+ getNullCountExpression(column: string): string;
22
+ getDistinctCountExpression(column: string): string;
23
+ textLengthExpression(columnSql: string): string;
24
+ castToText(columnSql: string): string;
25
+ getSampleValueAggregation(innerSql: string): string;
26
+ getLimitOffsetClause(limit: number, offset?: number): string;
27
+ getTopClause(_limit: number): string;
28
+ getTableSampleClause(_samplePct: number): string;
29
+ getRandomSampleFilter(samplePct: number): string;
30
+ }
31
+ export {};