1import type * as athena from "@distilled.cloud/aws/athena";
2import * as Data from "effect/Data";
3import type * as Effect from "effect/Effect";
4import * as Binding from "../../Binding.ts";
5import type { Bucket } from "../S3/Bucket.ts";
6import type { WorkGroup } from "./WorkGroup.ts";
7
8/**
9 * The query request — the workgroup is injected by the binding (which also
10 * supplies the enforced result location), so callers pass the SQL text and any
11 * execution context/parameters as the raw distilled shape.
12 */
13export interface RunQueryRequest extends Omit<
14 athena.StartQueryExecutionInput,
15 "WorkGroup" | "ResultConfiguration"
16> {}
17
18/**
19 * The fully-resolved result of a query: the terminal state plus the decoded
20 * result set (column names + rows of string cells, exactly as Athena returns
21 * them — the first row of a `SELECT` is the header).
22 */
23export interface QueryResult {
24 readonly queryExecutionId: string;
25 readonly state: athena.QueryExecutionState;
26 readonly stateChangeReason: string | undefined;
27 readonly columns: string[];
28 readonly rows: string[][];
29}
30
31/**
32 * Raised when a query reaches a terminal non-`SUCCEEDED` state (FAILED or
33 * CANCELLED), or does not settle within the bounded poll window.
34 */
35export class AthenaQueryFailed extends Data.TaggedError("AthenaQueryFailed")<{
36 readonly queryExecutionId: string;
37 readonly state: string;
38 readonly reason: string | undefined;
39}> {}
40
41/**
42 * Run an Athena query end-to-end from a Lambda/Task: start the execution
43 * against a bound workgroup, poll `GetQueryExecution` until it reaches a
44 * terminal state (bounded), then read and decode the result set. Results are
45 * written to the workgroup's enforced S3 output location in `resultsBucket`.
46 *
47 * ### Running Queries
48 * **Example:** Query and read the result rows
49 * ```typescript
50 * const runQuery = yield* AWS.Athena.Query(workGroup, resultsBucket);
51 * const result = yield* runQuery({ QueryString: "SELECT 1" });
52 * // result.rows[0] === ["_col0"] (header), result.rows[1] === ["1"]
53 * ```
54 *
55 * **Example:** Bind a WorkGroup to a Lambda Function
56 * ```typescript
57 * import * as Athena from "alchemy/AWS/Athena";
58 * import * as Lambda from "alchemy/AWS/Lambda";
59 * import * as S3 from "alchemy/AWS/S3";
60 * import * as Output from "alchemy/Output";
61 *
62 * export class QueryFunction extends Lambda.Function<Lambda.Function>()(
63 * "QueryFunction",
64 * ) {}
65 *
66 * export default QueryFunction.make(
67 * { main: import.meta.url, functionUrl: true, timeout: Duration.seconds(60) },
68 * Effect.gen(function* () {
69 * const bucket = yield* S3.Bucket("Results", { forceDestroy: true });
70 * const workGroup = yield* Athena.WorkGroup("Analytics", {
71 * outputLocation: Output.interpolate`s3://${bucket.bucketName}/results/`,
72 * enforceWorkGroupConfiguration: true,
73 * });
74 *
75 * // grants athena:StartQueryExecution/GetQueryExecution/GetQueryResults
76 * // on the workgroup, S3 access on the results bucket, and Glue catalog
77 * // reads for table-backed queries
78 * const runQuery = yield* Athena.Query(workGroup, bucket);
79 *
80 * return {
81 * fetch: Effect.gen(function* () {
82 * const result = yield* runQuery({
83 * QueryString: "SELECT COUNT(*) AS c FROM my_db.people",
84 * });
85 * return yield* HttpServerResponse.json({ rows: result.rows });
86 * }).pipe(Effect.orDie),
87 * };
88 * }).pipe(Effect.provide(Athena.QueryHttp)),
89 * );
90 * ```
91 *
92 * @binding
93 */
94export interface Query extends Binding.Service<
95 Query,
96 "AWS.Athena.Query",
97 (
98 workGroup: WorkGroup,
99 resultsBucket: Bucket,
100 ) => Effect.Effect<
101 (
102 request: RunQueryRequest,
103 ) => Effect.Effect<
104 QueryResult,
105 | athena.StartQueryExecutionError
106 | athena.GetQueryExecutionError
107 | athena.GetQueryResultsError
108 | AthenaQueryFailed
109 >
110 >
111> {}
112export const Query = Binding.Service<Query>("AWS.Athena.Query");