knex-dialect-athena
Version:
A Knex dialect for AWS Athena
269 lines (264 loc) • 9.99 kB
JavaScript
;
var __create = Object.create;
var __defProp = Object.defineProperty;
var __getOwnPropDesc = Object.getOwnPropertyDescriptor;
var __getOwnPropNames = Object.getOwnPropertyNames;
var __getOwnPropSymbols = Object.getOwnPropertySymbols;
var __getProtoOf = Object.getPrototypeOf;
var __hasOwnProp = Object.prototype.hasOwnProperty;
var __propIsEnum = Object.prototype.propertyIsEnumerable;
var __objRest = (source, exclude) => {
var target = {};
for (var prop in source)
if (__hasOwnProp.call(source, prop) && exclude.indexOf(prop) < 0)
target[prop] = source[prop];
if (source != null && __getOwnPropSymbols)
for (var prop of __getOwnPropSymbols(source)) {
if (exclude.indexOf(prop) < 0 && __propIsEnum.call(source, prop))
target[prop] = source[prop];
}
return target;
};
var __export = (target, all) => {
for (var name in all)
__defProp(target, name, { get: all[name], enumerable: true });
};
var __copyProps = (to, from, except, desc) => {
if (from && typeof from === "object" || typeof from === "function") {
for (let key of __getOwnPropNames(from))
if (!__hasOwnProp.call(to, key) && key !== except)
__defProp(to, key, { get: () => from[key], enumerable: !(desc = __getOwnPropDesc(from, key)) || desc.enumerable });
}
return to;
};
var __toESM = (mod, isNodeMode, target) => (target = mod != null ? __create(__getProtoOf(mod)) : {}, __copyProps(
// If the importer is in node compatibility mode or this is not an ESM
// file that has been converted to a CommonJS file using a Babel-
// compatible transform (i.e. "__esModule" has not been set), then set
// "default" to the CommonJS "module.exports" for node compatibility.
isNodeMode || !mod || !mod.__esModule ? __defProp(target, "default", { value: mod, enumerable: true }) : target,
mod
));
var __toCommonJS = (mod) => __copyProps(__defProp({}, "__esModule", { value: true }), mod);
var __async = (__this, __arguments, generator) => {
return new Promise((resolve, reject) => {
var fulfilled = (value) => {
try {
step(generator.next(value));
} catch (e) {
reject(e);
}
};
var rejected = (value) => {
try {
step(generator.throw(value));
} catch (e) {
reject(e);
}
};
var step = (x) => x.done ? resolve(x.value) : Promise.resolve(x.value).then(fulfilled, rejected);
step((generator = generator.apply(__this, __arguments)).next());
});
};
// src/athena-connection.ts
var athena_connection_exports = {};
__export(athena_connection_exports, {
AthenaConnection: () => AthenaConnection
});
module.exports = __toCommonJS(athena_connection_exports);
var import_client_athena = require("@aws-sdk/client-athena");
// src/assert.ts
function assert(condition, contextObject, message = "") {
if (!condition)
throw new Error(
`Assertion failed${message ? `: ${message}` : ""}: ${JSON.stringify(contextObject)}`
);
}
// src/debug.ts
var import_debug = __toESM(require("debug"));
var packageDebug = (0, import_debug.default)("knex-dialect-athena");
// src/athena-connection.ts
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 import_client_athena.AthenaClient(config);
this.database = database;
this.outputLocation = outputLocation;
this.workGroup = workGroup;
this.maxTimeoutMilliseconds = maxTimeoutMilliseconds != null ? maxTimeoutMilliseconds : 3e4;
}
// Athena Commands
startQueryExecutionCommand(queryString, parameters) {
return new import_client_athena.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 import_client_athena.GetQueryExecutionCommand({
QueryExecutionId: queryExecutionId
});
}
getQueryResultsCommand(queryExecutionId) {
return new import_client_athena.GetQueryResultsCommand({
QueryExecutionId: queryExecutionId
});
}
// Query execution waiting logic
isQueryExecutionUnfinished(queryExecution) {
var _a;
return !((_a = queryExecution.Status) == null ? void 0 : _a.State) || queryExecution.Status.State === import_client_athena.QueryExecutionState.QUEUED || queryExecution.Status.State === import_client_athena.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) === import_client_athena.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;
});
}
};
// Annotate the CommonJS export names for ESM import in node:
0 && (module.exports = {
AthenaConnection
});