From f9c3608b555b2fdcf3d22c205b71e47ff13833a7 Mon Sep 17 00:00:00 2001
From: Takahiro Okumura <hfm@cookpad.com>
Date: Tue, 14 Jul 2026 16:48:06 +0900
Subject: [PATCH] Add data catalog support for Amazon Athena

Adds a catalog field to allow users to specify an Athena data catalog, supporting federated queries and nested Glue catalogs.

This change also improves the reliability of table and column metadata lookups by replacing the parsing of `show tables` and `describe` output with the ListTableMetadata API.
---
 src/lib/AthenaClient.ts                | 24 ++++++++++++++++++++++++
 src/lib/DataSourceDefinition/Athena.ts | 21 ++++++++++++++-------
 2 files changed, 38 insertions(+), 7 deletions(-)

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<athena.TableMetadata[]> {
+    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<TableSummary> {
-    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,
     };
   }