knex-dialect-athena
Version:
A Knex dialect for AWS Athena
389 lines (381 loc) • 13.6 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/index.ts
var src_exports = {};
__export(src_exports, {
createAthenaDialect: () => createAthenaDialect
});
module.exports = __toCommonJS(src_exports);
var import_knex = __toESM(require("knex"));
// src/athena-connection.ts
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;
});
}
};
// src/athena-querycompiler.ts
var import_querycompiler = __toESM(require("knex/lib/query/querycompiler"));
var import_compact = __toESM(require("lodash/compact"));
var components = [
"comments",
"columns",
"join",
"where",
"union",
"group",
"having",
"order",
"offset",
"limit",
// TODO: remove
"lock",
"waitMode"
];
var QueryCompiler_Athena = class extends import_querycompiler.default {
// Compiles the `select` statement, or nested sub-selects by calling each of
// the component compilers, trimming out the empties, and returning a
// generated query string.
select() {
var _a;
let sql = this.with();
let unionStatement = "";
const firstStatements = [];
const endStatements = [];
components.forEach((component) => {
const statement = this[component](this);
switch (component) {
case "union":
unionStatement = statement;
break;
case "comments":
case "columns":
case "join":
case "where":
firstStatements.push(statement);
break;
default:
endStatements.push(statement);
break;
}
});
const wrapMainQuery = (_a = this.grouped.union) == null ? void 0 : _a.map((u) => u.wrap).some((u) => u);
if (this.onlyUnions()) {
const statements = (0, import_compact.default)(firstStatements.concat(endStatements)).join(
" "
);
sql += unionStatement + (statements ? " " + statements : "");
} else {
const allStatements = (wrapMainQuery ? "(" : "") + (0, import_compact.default)(firstStatements).join(" ") + (wrapMainQuery ? ")" : "");
const endStat = (0, import_compact.default)(endStatements).join(" ");
sql += allStatements + (unionStatement ? " " + unionStatement : "") + (endStat ? " " + endStat : endStat);
}
return sql;
}
};
// src/index.ts
var noOp = () => {
};
function createAthenaDialect(config) {
return class Client_Athena extends import_knex.default.Client {
constructor() {
super(...arguments);
this.dialect = "athena";
this.driverName = "athena";
// TODO: add query cancelling behavior
this.canCancelQuery = false;
this.releaseConnection = noOp;
this.queryCompiler = (builder, bindings) => new QueryCompiler_Athena(
this,
builder,
bindings
);
}
acquireConnection() {
return new AthenaConnection(config);
}
processBinding(binding) {
if (binding === null) return "null";
if (binding instanceof Date)
throw new Error("Date bindings are not (yet) supported");
if (Array.isArray(binding))
throw new Error("Array bindings are not supported");
return String(binding);
}
_query(connection, obj) {
return __async(this, null, function* () {
if (!obj.sql) throw new Error("The query is empty");
const response = yield connection.query(
obj.sql,
obj.bindings.flatMap((binding) => this.processBinding(binding))
);
obj.response = response;
return obj;
});
}
processResponse(obj) {
if (obj.method === "raw") return obj.response;
if (obj.method === "first") {
if (!obj.response[0])
throw new Error("Called `.first` but no rows were returned");
return obj.response[0];
}
if (obj.method === "pluck")
return obj.response.map((row) => {
assert(!!obj.pluck, obj);
return row[obj.pluck];
});
return obj.response;
}
};
}
// Annotate the CommonJS export names for ESM import in node:
0 && (module.exports = {
createAthenaDialect
});