feat(api): add Athena workgroups and a query editor to the cloud explorer

Create and delete Athena workgroups (with a result folder, optionally enforced), and run SQL in one: pick a database from the Glue Data Catalog, preview a table, run with Ctrl+Enter, cancel a running query, read up to 1,000 rows and download the full result CSV from S3. Recent queries in the workgroup can be reopened. TheCloud runs the SQL for real over the files a Glue table points at. The Glue page links to the editor.
2026-10-25 18:10:33 -04:00
parent 69d80c60a3
commit 5eee0b7e9b
25 changed files with 1291 additions and 121 deletions
@@ -0,0 +1,268 @@
+import {
+    type AthenaClient,
+    CreateWorkGroupCommand,
+    DeleteWorkGroupCommand,
+    GetQueryExecutionCommand,
+    GetQueryResultsCommand,
+    type GetQueryResultsCommandOutput,
+    GetWorkGroupCommand,
+    ListDatabasesCommand,
+    ListQueryExecutionsCommand,
+    ListTableMetadataCommand,
+    type ListTableMetadataCommandOutput,
+    ListWorkGroupsCommand,
+    type QueryExecution,
+    StartQueryExecutionCommand,
+    StopQueryExecutionCommand,
+    type TableMetadata,
+    type WorkGroup,
+    type WorkGroupSummary,
+} from '@aws-sdk/client-athena'
+import {NotFoundError, RuntimeError, ValidationError} from '../cloud-spi/errors'
+import {
+    awsQuerySchema,
+    PRIMARY_WORKGROUP,
+    QUERY_HISTORY_LIMIT,
+    QUERY_RESULT_ROW_LIMIT,
+    WORKGROUP_NAME_PATTERN,
+} from '../cloud-spi/querySchema'
+import type {
+    CloudResource,
+    CloudServiceAdapter,
+    CreateResourceInput,
+    QueryCatalog,
+    QueryExecutionSummary,
+    QueryInput,
+    QueryResult,
+    ResourceQuery,
+    ServiceSchema,
+} from '../cloud-spi/types'
+
+/** Athena's own catalog name for the Glue Data Catalog. */
+const DATA_CATALOG = 'AwsDataCatalog'
+/** Athena's cap on a query string. */
+const MAX_QUERY_LENGTH = 262_144
+const S3_URI = /^s3:\/\/[a-z0-9.-]{3,63}(\/.*)?$/
+
+/**
+ * Talks to Amazon Athena, verified against Floci core (community, `latest`).
+ *
+ * Notes from probing the runtime:
+ *  - Queries really run: TheCloud executes the SQL with DuckDB over the files a
+ *    Glue table points at, and writes the result to S3 as CSV, as Athena does.
+ *  - Statements that return no rows, such as SHOW TABLES, fail in the runtime,
+ *    so the editor lists databases and tables through the catalog APIs instead.
+ *  - Data scanned is always reported as 0 bytes, so the editor shows only the
+ *    run time.
+ */
+export class AwsQueryAdapter implements CloudServiceAdapter {
+    readonly cloud = 'aws' as const
+    readonly service = 'query' as const
+
+    constructor(private readonly athena: AthenaClient) {}
+
+    schema(): ServiceSchema {
+        return awsQuerySchema()
+    }
+
+    async list(query: ResourceQuery = {}): Promise<CloudResource[]> {
+        const found: WorkGroupSummary[] = []
+        let next: string | undefined
+        do {
+            const res = await this.athena.send(new ListWorkGroupsCommand({MaxResults: 50, ...(next ? {NextToken: next} : {})}))
+            found.push(...(res.WorkGroups ?? []))
+            next = res.NextToken
+        } while (next)
+
+        const search = query.search?.trim().toLowerCase()
+        return found.map(toResource).filter((workGroup) => !search
+            || workGroup.name.toLowerCase().includes(search)
+            || String(workGroup.metadata.description ?? '').toLowerCase().includes(search))
+    }
+
+    async get(id: string): Promise<CloudResource | null> {
+        try {
+            const workGroup = (await this.athena.send(new GetWorkGroupCommand({WorkGroup: id}))).WorkGroup
+            return workGroup ? toResource(workGroup) : null
+        } catch (error) {
+            if (errorName(error) === 'InvalidRequestException') return null
+            throw error
+        }
+    }
+
+    async create(input: CreateResourceInput): Promise<CloudResource> {
+        const values = input.values
+        const name = text(values.name)
+        if (!new RegExp(WORKGROUP_NAME_PATTERN).test(name)) throw new ValidationError('name must be up to 128 letters, numbers, dots, dashes or underscores')
+        if (name === PRIMARY_WORKGROUP) throw new ValidationError('Every account already has the primary workgroup')
+        const outputLocation = text(values.outputLocation)
+        if (outputLocation && !S3_URI.test(outputLocation)) throw new ValidationError('outputLocation must be an S3 URI such as s3://my-bucket/athena-results/')
+        const description = text(values.description)
+
+        await this.athena.send(new CreateWorkGroupCommand({
+            Name: name,
+            ...(description ? {Description: description} : {}),
+            Configuration: {
+                ...(outputLocation ? {ResultConfiguration: {OutputLocation: outputLocation.endsWith('/') ? outputLocation : `${outputLocation}/`}} : {}),
+                EnforceWorkGroupConfiguration: values.enforce === 'true',
+            },
+        }))
+        const created = await this.get(name)
+        if (!created) throw new RuntimeError(`Athena did not return the new workgroup ${name}`)
+        return created
+    }
+
+    async delete(id: string): Promise<void> {
+        if (id === PRIMARY_WORKGROUP) throw new ValidationError('The primary workgroup cannot be deleted')
+        await this.athena.send(new DeleteWorkGroupCommand({WorkGroup: id, RecursiveDeleteOption: true}))
+    }
+
+    async startQuery(workGroup: string, input: QueryInput): Promise<QueryExecutionSummary> {
+        const sql = typeof input.sql === 'string' ? input.sql.trim() : ''
+        if (!sql) throw new ValidationError('Enter a query to run')
+        if (sql.length > MAX_QUERY_LENGTH) throw new ValidationError(`A query can be at most ${MAX_QUERY_LENGTH} characters`)
+        const database = text(input.database)
+
+        const started = await this.athena.send(new StartQueryExecutionCommand({
+            QueryString: sql,
+            WorkGroup: workGroup,
+            ...(database ? {QueryExecutionContext: {Catalog: DATA_CATALOG, Database: database}} : {}),
+        }))
+        if (!started.QueryExecutionId) throw new RuntimeError('Athena did not return a query id')
+        return (await this.execution(workGroup, started.QueryExecutionId))
+    }
+
+    async getQueryResult(workGroup: string, queryId: string): Promise<QueryResult> {
+        const execution = await this.execution(workGroup, queryId)
+        if (execution.state !== 'SUCCEEDED' || !execution.outputLocation) return {execution, columns: [], rows: [], truncated: false}
+
+        const pages: GetQueryResultsCommandOutput[] = []
+        let next: string | undefined
+        let rowCount = 0
+        do {
+            const page = await this.athena.send(new GetQueryResultsCommand({QueryExecutionId: queryId, MaxResults: 1000, ...(next ? {NextToken: next} : {})}))
+            pages.push(page)
+            rowCount += page.ResultSet?.Rows?.length ?? 0
+            next = page.NextToken
+        } while (next && rowCount <= QUERY_RESULT_ROW_LIMIT)
+
+        const columnInfo = pages[0]?.ResultSet?.ResultSetMetadata?.ColumnInfo ?? []
+        const raw = pages.flatMap((page) => page.ResultSet?.Rows ?? []).map((row) => (row.Data ?? []).map((cell) => cell.VarCharValue ?? null))
+        const header = raw[0] ?? []
+        const columns = columnInfo.length > 0
+            ? columnInfo.map((column) => ({name: column.Label ?? column.Name ?? '', type: column.Type ?? null}))
+            : header.map((name) => ({name: name ?? '', type: null}))
+        // Athena repeats the column names as the first row of a SELECT's results.
+        const body = raw.length > 0 && header.length === columns.length && header.every((name, index) => name === columns[index].name) ? raw.slice(1) : raw
+        return {execution, columns, rows: body.slice(0, QUERY_RESULT_ROW_LIMIT), truncated: body.length > QUERY_RESULT_ROW_LIMIT || Boolean(next)}
+    }
+
+    async stopQuery(workGroup: string, queryId: string): Promise<void> {
+        await this.execution(workGroup, queryId)
+        await this.athena.send(new StopQueryExecutionCommand({QueryExecutionId: queryId}))
+    }
+
+    async listQueries(workGroup: string): Promise<QueryExecutionSummary[]> {
+        const ids = (await this.athena.send(new ListQueryExecutionsCommand({WorkGroup: workGroup, MaxResults: QUERY_HISTORY_LIMIT}))).QueryExecutionIds ?? []
+        const executions = await Promise.all(ids.slice(0, QUERY_HISTORY_LIMIT).map(async (id) => {
+            const res = await this.athena.send(new GetQueryExecutionCommand({QueryExecutionId: id}))
+            return res.QueryExecution ? toExecution(res.QueryExecution) : null
+        }))
+        return executions
+            .filter((execution): execution is QueryExecutionSummary => execution !== null && (execution.workGroup ?? PRIMARY_WORKGROUP) === workGroup)
+            .sort((a, b) => (b.submittedAt ?? '').localeCompare(a.submittedAt ?? ''))
+    }
+
+    async describeQueryCatalog(database?: string): Promise<QueryCatalog> {
+        const databases: string[] = []
+        let next: string | undefined
+        do {
+            const res = await this.athena.send(new ListDatabasesCommand({CatalogName: DATA_CATALOG, MaxResults: 50, ...(next ? {NextToken: next} : {})}))
+            databases.push(...(res.DatabaseList ?? []).map((entry) => entry.Name ?? '').filter(Boolean))
+            next = res.NextToken
+        } while (next)
+        databases.sort()
+
+        const chosen = database && databases.includes(database) ? database : databases[0] ?? null
+        if (!chosen) return {databases, database: null, tables: []}
+
+        const tables: TableMetadata[] = []
+        let tableToken: string | undefined
+        do {
+            const res: ListTableMetadataCommandOutput = await this.athena.send(new ListTableMetadataCommand({CatalogName: DATA_CATALOG, DatabaseName: chosen, MaxResults: 50, ...(tableToken ? {NextToken: tableToken} : {})}))
+            tables.push(...(res.TableMetadataList ?? []))
+            tableToken = res.NextToken
+        } while (tableToken)
+        return {
+            databases,
+            database: chosen,
+            tables: tables.map((table) => ({
+                name: table.Name ?? '',
+                columns: [...(table.Columns ?? []), ...(table.PartitionKeys ?? [])].map((column) => ({name: column.Name ?? '', type: column.Type ?? ''})),
+            })).sort((a, b) => a.name.localeCompare(b.name)),
+        }
+    }
+
+    /** A query, refused unless it ran in this workgroup, so one workgroup's page never shows another's. */
+    private async execution(workGroup: string, queryId: string): Promise<QueryExecutionSummary> {
+        let execution: QueryExecution | undefined
+        try {
+            execution = (await this.athena.send(new GetQueryExecutionCommand({QueryExecutionId: queryId}))).QueryExecution
+        } catch (error) {
+            if (errorName(error) === 'InvalidRequestException') throw new NotFoundError(`Query ${queryId} not found`)
+            throw error
+        }
+        if (!execution || (execution.WorkGroup ?? PRIMARY_WORKGROUP) !== workGroup) throw new NotFoundError(`Query ${queryId} is not in the ${workGroup} workgroup`)
+        return toExecution(execution)
+    }
+}
+
+function toResource(workGroup: WorkGroup | WorkGroupSummary): CloudResource {
+    const name = workGroup.Name ?? ''
+    const detail = 'Configuration' in workGroup ? workGroup : undefined
+    const configuration = detail?.Configuration
+    return {
+        id: name,
+        name,
+        cloud: 'aws',
+        service: 'query',
+        type: 'athenaWorkGroup',
+        region: null,
+        createdAt: workGroup.CreationTime ? workGroup.CreationTime.toISOString() : null,
+        status: workGroup.State ?? null,
+        metadata: {
+            description: workGroup.Description || null,
+            ...(detail ? {
+                outputLocation: configuration?.ResultConfiguration?.OutputLocation ?? null,
+                enforceConfiguration: configuration?.EnforceWorkGroupConfiguration === true,
+                engineVersion: configuration?.EngineVersion?.EffectiveEngineVersion ?? configuration?.EngineVersion?.SelectedEngineVersion ?? null,
+            } : {}),
+        },
+    }
+}
+
+function toExecution(execution: QueryExecution): QueryExecutionSummary {
+    const status = execution.Status
+    return {
+        id: execution.QueryExecutionId ?? '',
+        sql: execution.Query ?? '',
+        database: execution.QueryExecutionContext?.Database ?? null,
+        workGroup: execution.WorkGroup ?? null,
+        state: status?.State ?? 'QUEUED',
+        stateReason: status?.StateChangeReason ?? status?.AthenaError?.ErrorMessage ?? null,
+        statementType: execution.StatementType ?? null,
+        submittedAt: status?.SubmissionDateTime ? status.SubmissionDateTime.toISOString() : null,
+        completedAt: status?.CompletionDateTime ? status.CompletionDateTime.toISOString() : null,
+        runtimeMs: execution.Statistics?.TotalExecutionTimeInMillis ?? execution.Statistics?.EngineExecutionTimeInMillis ?? null,
+        dataScannedBytes: execution.Statistics?.DataScannedInBytes ?? null,
+        outputLocation: execution.ResultConfiguration?.OutputLocation ?? null,
+    }
+}
+
+function text(value: unknown): string {
+    return typeof value === 'string' ? value.trim() : ''
+}
+
+function errorName(error: unknown): string | undefined {
+    return typeof error === 'object' && error !== null ? (error as {name?: string}).name : undefined
+}