diff --git a/src/lib/AthenaClient.ts b/src/lib/AthenaClient.ts index 30cd7fc..db6dc3e 100644 --- a/src/lib/AthenaClient.ts +++ b/src/lib/AthenaClient.ts @@ -6,6 +6,7 @@ interface AthenaClientConfig { profile?: string; accessKeyId?: string; secretAccessKey?: string; + catalog?: string; database: string; outputLocation: string; } @@ -42,6 +43,7 @@ export default class AthenaClient { OutputLocation: this.config.outputLocation, }, QueryExecutionContext: { + Catalog: this.config.catalog || undefined, Database: this.config.database, }, }; @@ -88,6 +90,28 @@ export default class AthenaClient { return rows; } + async listTableMetadata(): Promise { + let tables: athena.TableMetadata[] = []; + let nextToken: string | undefined = undefined; + + for (;;) { + const result = await this.client.send( + new athena.ListTableMetadataCommand({ + CatalogName: this.config.catalog || "AwsDataCatalog", + DatabaseName: this.config.database, + NextToken: nextToken, + }) + ); + tables = tables.concat(result.TableMetadataList || []); + nextToken = result.NextToken; + if (!nextToken) { + break; + } + } + + return tables; + } + cancel(): void { this.client.send(new athena.StopQueryExecutionCommand({ QueryExecutionId: this.executionId })); } diff --git a/src/lib/DataSourceDefinition/Athena.ts b/src/lib/DataSourceDefinition/Athena.ts index a17b5b6..fd1a730 100644 --- a/src/lib/DataSourceDefinition/Athena.ts +++ b/src/lib/DataSourceDefinition/Athena.ts @@ -38,6 +38,12 @@ export default class Athena extends Base { type: "string", placeholder: "Optional - uses AWS environment/profile if empty", }, + { + name: "catalog", + label: "Catalog", + type: "string", + placeholder: "Optional - uses AwsDataCatalog if empty", + }, { name: "database", label: "Database", @@ -74,18 +80,18 @@ export default class Athena extends Base { } async fetchTables(): Promise<{ name: string; type: string; schema?: string }[]> { - const rows = await this.client.execute("show tables"); - return rows.map((row) => ({ name: row[0]!, type: "table" })); + const tables = await this.client.listTableMetadata(); + return tables.map((table) => ({ name: table.Name!, type: "table" })); } async fetchTableSummary({ name }: { name: string }): Promise { - const rows = await this.client.execute(`describe ${name}`); + const tables = await this.client.listTableMetadata(); + const table = tables.find((t) => t.Name === name); + const columns = [...(table?.Columns || []), ...(table?.PartitionKeys || [])]; + const defs = { fields: ["name", "type"], - rows: rows - .map((row) => row[0]) - .filter((v) => v !== null && v[0] !== "#" && v.trim() !== "") - .map((v) => (v || "").split("\t").map((c) => c.trim())), + rows: columns.map((column) => [column.Name!, column.Type || ""]), }; return { name, defs }; @@ -95,6 +101,7 @@ export default class Athena extends Base { return { type: Athena.label, region: this.config.region, + catalog: this.config.catalog || "AwsDataCatalog", database: this.config.database, }; }