knex-dialect-athena
Version:
A Knex dialect for AWS Athena
204 lines (201 loc) • 7 kB
JavaScript
import {
assert
} from "./chunk-54RCQBX4.mjs";
import {
packageDebug
} from "./chunk-X5PQ2CFY.mjs";
import {
__async,
__objRest
} from "./chunk-SGFGNEI2.mjs";
// src/athena-connection.ts
import {
AthenaClient,
GetQueryExecutionCommand,
GetQueryResultsCommand,
QueryExecutionState,
StartQueryExecutionCommand
} from "@aws-sdk/client-athena";
var debug = packageDebug.extend("connection");
var AthenaConnection = class {
constructor(_a) {
var _b = _a, {
database,
outputLocation,
workGroup,
maxTimeoutMilliseconds
} = _b, config = __objRest(_b, [
"database",
"outputLocation",
"workGroup",
"maxTimeoutMilliseconds"
]);
this.client = new AthenaClient(config);
this.database = database;
this.outputLocation = outputLocation;
this.workGroup = workGroup;
this.maxTimeoutMilliseconds = maxTimeoutMilliseconds != null ? maxTimeoutMilliseconds : 3e4;
}
// Athena Commands
startQueryExecutionCommand(queryString, parameters) {
return new StartQueryExecutionCommand({
QueryString: queryString,
ExecutionParameters: parameters.length <= 0 ? void 0 : parameters,
QueryExecutionContext: {
Database: this.database
},
WorkGroup: this.workGroup,
ResultConfiguration: {
OutputLocation: this.outputLocation
}
});
}
getQueryExecutionCommand(queryExecutionId) {
return new GetQueryExecutionCommand({
QueryExecutionId: queryExecutionId
});
}
getQueryResultsCommand(queryExecutionId) {
return new GetQueryResultsCommand({
QueryExecutionId: queryExecutionId
});
}
// Query execution waiting logic
isQueryExecutionUnfinished(queryExecution) {
var _a;
return !((_a = queryExecution.Status) == null ? void 0 : _a.State) || queryExecution.Status.State === QueryExecutionState.QUEUED || queryExecution.Status.State === QueryExecutionState.RUNNING;
}
waitForQueryExecution(queryExecutionId) {
return __async(this, null, function* () {
var _a, _b;
debug("starting wait for query execution (id %s)", queryExecutionId);
const backoffStepMilliseconds = 250;
let overallTimeElapsedMilliseconds = 0;
let retryInMilliseconds = 0;
let queryExecution;
do {
if (overallTimeElapsedMilliseconds < this.maxTimeoutMilliseconds) {
retryInMilliseconds += backoffStepMilliseconds;
overallTimeElapsedMilliseconds += retryInMilliseconds;
}
debug("waiting for %dms (id %s)", retryInMilliseconds, queryExecutionId);
yield new Promise((resolve) => setTimeout(resolve, retryInMilliseconds));
const executionResponse = yield this.client.send(
this.getQueryExecutionCommand(queryExecutionId)
);
debug(
"query execution state (id %s): %o",
queryExecutionId,
(_b = (_a = executionResponse.QueryExecution) == null ? void 0 : _a.Status) == null ? void 0 : _b.State
);
assert(
!!executionResponse.QueryExecution,
executionResponse,
"missing QueryExecution"
);
queryExecution = executionResponse.QueryExecution;
} while (this.isQueryExecutionUnfinished(queryExecution));
assert(
!this.isQueryExecutionUnfinished(queryExecution),
queryExecution,
`query did not finish in ${this.maxTimeoutMilliseconds.toString()}ms`
);
return queryExecution;
});
}
// Mapping query results
mapQueryResults(resultSet) {
var _a;
debug("mapping over result set: %o", resultSet);
assert(!!resultSet.Rows, resultSet, "missing Rows");
const columns = (_a = resultSet.ResultSetMetadata) == null ? void 0 : _a.ColumnInfo;
assert(!!columns, resultSet.ResultSetMetadata, "missing column specifiers");
debug("using columns: %o", columns);
return resultSet.Rows.slice(1).map((row) => {
var _a2, _b, _c;
const result = {};
let index = 0;
for (const column of columns) {
const datum = (_c = (_b = (_a2 = row.Data) == null ? void 0 : _a2[index++]) == null ? void 0 : _b.VarCharValue) != null ? _c : null;
assert(!!column.Name, column, "missing Name");
switch (column.Type) {
case "boolean":
result[column.Name] = datum === null ? datum : datum === "true";
break;
case "tinyint":
case "smallint":
case "integer":
case "int":
case "bigint":
result[column.Name] = datum === null ? datum : parseInt(datum);
break;
case "real":
case "double":
case "decimal":
result[column.Name] = datum === null ? datum : parseFloat(datum);
break;
case "json":
result[column.Name] = datum === null ? datum : JSON.parse(datum);
break;
case "varbinary":
result[column.Name] = datum === null ? datum : new Uint8Array(
datum.split(" ").map((byte) => parseInt(byte, 16))
);
break;
default:
result[column.Name] = datum;
break;
}
}
return result;
});
}
// Public API
query(_0) {
return __async(this, arguments, function* (queryString, parameters = []) {
var _a, _b, _c, _d;
debug("starting query execution");
debug("query: %o", queryString);
debug("parameters: %o", parameters);
const startQueryExecutionResponse = yield this.client.send(
this.startQueryExecutionCommand(queryString, parameters)
);
debug(
"got start query execution response: %o",
startQueryExecutionResponse
);
assert(
!!startQueryExecutionResponse.QueryExecutionId,
startQueryExecutionResponse,
"missing QueryExecutionId"
);
const queryExecution = yield this.waitForQueryExecution(
startQueryExecutionResponse.QueryExecutionId
);
debug(
"final query execution state (id %s): %o",
startQueryExecutionResponse.QueryExecutionId,
(_a = queryExecution.Status) == null ? void 0 : _a.State
);
assert(
((_b = queryExecution.Status) == null ? void 0 : _b.State) === QueryExecutionState.SUCCEEDED,
queryExecution,
"query failed"
);
const resultsResponse = yield this.client.send(
this.getQueryResultsCommand(startQueryExecutionResponse.QueryExecutionId)
);
debug("got full results response: %o", resultsResponse);
if (!((_d = (_c = resultsResponse.ResultSet) == null ? void 0 : _c.Rows) == null ? void 0 : _d[0])) {
debug("no rows (or column specifiers); returning update count");
return resultsResponse.UpdateCount;
}
const results = this.mapQueryResults(resultsResponse.ResultSet);
debug("mapped results: %o", results);
return results;
});
}
};
export {
AthenaConnection
};