aws-athena-node-client
Version:
402 lines (401 loc) • 17.4 kB
JavaScript
"use strict";
var __webpack_require__ = {};
(()=>{
__webpack_require__.n = function(module) {
var getter = module && module.__esModule ? function() {
return module['default'];
} : function() {
return module;
};
__webpack_require__.d(getter, {
a: getter
});
return getter;
};
})();
(()=>{
__webpack_require__.d = function(exports1, definition) {
for(var key in definition)if (__webpack_require__.o(definition, key) && !__webpack_require__.o(exports1, key)) Object.defineProperty(exports1, key, {
enumerable: true,
get: definition[key]
});
};
})();
(()=>{
__webpack_require__.o = function(obj, prop) {
return Object.prototype.hasOwnProperty.call(obj, prop);
};
})();
(()=>{
__webpack_require__.r = function(exports1) {
if ('undefined' != typeof Symbol && Symbol.toStringTag) Object.defineProperty(exports1, Symbol.toStringTag, {
value: 'Module'
});
Object.defineProperty(exports1, '__esModule', {
value: true
});
};
})();
var __webpack_exports__ = {};
__webpack_require__.r(__webpack_exports__);
__webpack_require__.d(__webpack_exports__, {
AthenaNodeClient: ()=>AthenaNodeClient
});
const external_lodash_es_namespaceObject = require("lodash-es");
function _define_property(obj, key, value) {
if (key in obj) Object.defineProperty(obj, key, {
value: value,
enumerable: true,
configurable: true,
writable: true
});
else obj[key] = value;
return obj;
}
class Queue {
addQuery(query) {
if ((0, external_lodash_es_namespaceObject.isEmpty)(query.id)) query.id = query.athenaId;
this.queries.push(query);
}
removeQuery(query) {
const index = this.queries.indexOf(query);
if (index > -1) this.queries.splice(index, 1);
}
getQueryById(id) {
const query = this.queries.find((query)=>query.id === id);
if (null != query) return query;
throw new Error('Query ID not found');
}
constructor(){
_define_property(this, "queries", []);
}
}
const formatting_js_namespaceObject = require("pg-promise/lib/formatting.js");
function Query_define_property(obj, key, value) {
if (key in obj) Object.defineProperty(obj, key, {
value: value,
enumerable: true,
configurable: true,
writable: true
});
else obj[key] = value;
return obj;
}
class Query {
hasColumns() {
return this.columns.length > 0;
}
hasResults() {
var _this_results;
return (null === (_this_results = this.results) || void 0 === _this_results ? void 0 : _this_results.length) > 0;
}
constructor(sql, waitTime = 0.5, parameters, id){
Query_define_property(this, "id", void 0);
Query_define_property(this, "athenaId", void 0);
Query_define_property(this, "originalSql", void 0);
Query_define_property(this, "parameters", void 0);
Query_define_property(this, "status", void 0);
Query_define_property(this, "sql", void 0);
Query_define_property(this, "waitTime", void 0);
Query_define_property(this, "results", []);
Query_define_property(this, "columns", void 0);
Query_define_property(this, "cacheInMinutes", void 0);
Query_define_property(this, "s3Location", void 0);
this.originalSql = sql;
this.parameters = parameters;
this.id = id;
this.sql = (0, formatting_js_namespaceObject.formatQuery)(sql, parameters);
this.waitTime = waitTime;
this.columns = [];
}
}
const external_destr_namespaceObject = require("destr");
var external_destr_default = /*#__PURE__*/ __webpack_require__.n(external_destr_namespaceObject);
function Column_define_property(obj, key, value) {
if (key in obj) Object.defineProperty(obj, key, {
value: value,
enumerable: true,
configurable: true,
writable: true
});
else obj[key] = value;
return obj;
}
class Column {
static parseNumber(value) {
const result = Number(value);
if (Number.isNaN(result)) throw new Error(`The value '${value} 'is not a number`);
return result;
}
static parseString(value) {
return value;
}
static parseBoolean(value) {
return 'true' === value || 'TRUE' === value || 't' === value || 'T' === value || 'yes' === value || 'YES' === value || '1' === value;
}
static parseDate(value) {
return new Date(value);
}
static parseArray(arrayInString) {
arrayInString = arrayInString.replace(/[\[\]]/gi, '');
if (null == arrayInString || '' === arrayInString) return [];
const values = arrayInString.split(', ');
const result = [];
for (const value of values){
let numberValue = Number(value);
if (Number.isNaN(numberValue)) result.push(value);
else result.push(numberValue);
}
return result;
}
static parseJson(value) {
return external_destr_default()(value);
}
constructor(name, parse){
Column_define_property(this, "name", void 0);
Column_define_property(this, "parse", void 0);
this.name = name;
this.parse = parse;
}
}
const client_athena_namespaceObject = require("@aws-sdk/client-athena");
const promises_namespaceObject = require("timers/promises");
class AthenaClientException extends Error {
}
class QueryCanceledException extends AthenaClientException {
constructor(){
super('Query cancelled');
this.name = 'AthenaQueryCancelled';
}
}
const client_s3_namespaceObject = require("@aws-sdk/client-s3");
const s3_request_presigner_namespaceObject = require("@aws-sdk/s3-request-presigner");
function src_define_property(obj, key, value) {
if (key in obj) Object.defineProperty(obj, key, {
value: value,
enumerable: true,
configurable: true,
writable: true
});
else obj[key] = value;
return obj;
}
const expiration1Day = 86400;
class AthenaNodeClient {
async executeQuery(sql, config) {
const query = await this.executeQueryCommon(sql, config);
const results = await this.getQueryResults(query);
const response = {
results,
id: query.athenaId
};
if (null == config ? void 0 : config.stats) {
const statistics = await this.getQueryStatistics(query.athenaId);
Reflect.set(response, 'statistics', statistics);
}
return response;
}
async executeQueryAndGetS3Key(sql, config) {
const query = await this.executeQueryCommon(sql, config);
if (null == query.s3Location) throw new Error('Athena no returns results S3 url');
const [bucket, key] = query.s3Location.replace('s3://', '').split('/', 1);
const response = {
bucket,
key: `${key ?? ''}${query.athenaId}.csv`,
id: query.athenaId
};
if (null == config ? void 0 : config.stats) {
const statistics = await this.getQueryStatistics(query.athenaId);
Reflect.set(response, 'statistics', statistics);
}
return response;
}
async executeQueryAndGetDownloadSignedUrl(sql, config) {
if (null == this._config.s3Client) throw new Error('[AthenaNodeClient] S3 Client is missing, you must install @aws-sdk/client-s3 and @aws-sdk/s3-request-presigner dependencies and fill s3Client field in configuration');
const { bucket, key, statistics, id } = await this.executeQueryAndGetS3Key(sql, config);
const command = new client_s3_namespaceObject.GetObjectCommand({
Bucket: bucket,
Key: key
});
const response = {
url: await (0, s3_request_presigner_namespaceObject.getSignedUrl)(this._config.s3Client, command, {
expiresIn: (null == config ? void 0 : config.s3LinkExpirationInSeconds) ?? expiration1Day
}),
id
};
if (null != statistics) Reflect.set(response, 'statistics', statistics);
return response;
}
async cancelQuery(id) {
const query = this._queue.getQueryById(id);
const input = {
QueryExecutionId: query.athenaId
};
await this._client.send(new client_athena_namespaceObject.StopQueryExecutionCommand(input));
}
async executeQueryCommon(sql, config) {
const query = new Query(sql, this._config.waitTimeInSeconds, null == config ? void 0 : config.parameters, null == config ? void 0 : config.id);
this._queue.addQuery(query);
query.athenaId = await this.startQueryExecution(query);
try {
await this.waitUntilSucceedQuery(query);
} catch (exception) {
this._queue.removeQuery(query);
throw exception;
}
return query;
}
async startQueryExecution(query) {
const input = {
QueryExecutionContext: {
Database: this._config.database
},
QueryString: query.sql,
ResultConfiguration: {
OutputLocation: this._config.bucketUri
}
};
if (null != this._config.workGroup) input.WorkGroup = this._config.workGroup;
if (query.cacheInMinutes ?? false) input.ResultReuseConfiguration = {
ResultReuseByAgeConfiguration: {
Enabled: true,
MaxAgeInMinutes: query.cacheInMinutes
}
};
const response = await this._client.send(new client_athena_namespaceObject.StartQueryExecutionCommand(input));
if (null != response.QueryExecutionId) return response.QueryExecutionId;
throw new Error('[AthenaNodeClient] Athena no returns query execution id. This may be a problem in AWS side.');
}
async getQueryResults(query, nextToken) {
const input = {
NextToken: nextToken,
QueryExecutionId: query.athenaId
};
const response = await this._client.send(new client_athena_namespaceObject.GetQueryResultsCommand(input));
const results = response.ResultSet;
if (null == results) throw new Error('[AthenaNodeClient] No query results. This may be a problem in AWS side.');
if (!query.hasColumns()) query.columns = this.setColumnParsers(results);
const isFirstPage = !query.hasResults() && null == nextToken;
query.results.push(...this.parseRows(results.Rows ?? [], query.columns, isFirstPage));
if (null != response.NextToken) query.results = await this.getQueryResults(query, response.NextToken);
return query.results;
}
async getQueryStatistics(executionId) {
var _response_QueryExecution_Statistics, _response_QueryExecution, _response_QueryExecution_Statistics1, _response_QueryExecution1;
const input = new client_athena_namespaceObject.GetQueryExecutionCommand({
QueryExecutionId: executionId
});
const response = await this._client.send(input);
const bytes = null == response ? void 0 : null === (_response_QueryExecution = response.QueryExecution) || void 0 === _response_QueryExecution ? void 0 : null === (_response_QueryExecution_Statistics = _response_QueryExecution.Statistics) || void 0 === _response_QueryExecution_Statistics ? void 0 : _response_QueryExecution_Statistics.DataScannedInBytes;
const timeInSeconds = ((null == response ? void 0 : null === (_response_QueryExecution1 = response.QueryExecution) || void 0 === _response_QueryExecution1 ? void 0 : null === (_response_QueryExecution_Statistics1 = _response_QueryExecution1.Statistics) || void 0 === _response_QueryExecution_Statistics1 ? void 0 : _response_QueryExecution_Statistics1.EngineExecutionTimeInMillis) ?? 0) / 1000;
return {
dataScannedInBytes: bytes,
executionTimeInSeconds: timeInSeconds
};
}
parseRow(row, columns) {
var _row_Data;
const result = {};
const dataLength = (null === (_row_Data = row.Data) || void 0 === _row_Data ? void 0 : _row_Data.length) ?? 0;
for(let rowDataIndex = 0; rowDataIndex < dataLength; rowDataIndex++){
var _row_Data1;
const rowData = null === (_row_Data1 = row.Data) || void 0 === _row_Data1 ? void 0 : _row_Data1[rowDataIndex];
const column = columns[rowDataIndex];
const columnName = column.name;
let value = null;
if ((null == rowData ? void 0 : rowData.VarCharValue) != null) value = column.parse(rowData.VarCharValue);
Reflect.set(result, columnName, value);
}
return result;
}
parseRows(rows, columns, isFirstPage = false) {
const results = [];
for(let rowIndex = isFirstPage ? 1 : 0, len = rows.length; rowIndex < len; rowIndex++){
const row = rows[rowIndex];
results.push(this.parseRow(row, columns));
}
return results;
}
setColumnParsers(results) {
var _results_ResultSetMetadata_ColumnInfo, _results_ResultSetMetadata;
return (null === (_results_ResultSetMetadata = results.ResultSetMetadata) || void 0 === _results_ResultSetMetadata ? void 0 : null === (_results_ResultSetMetadata_ColumnInfo = _results_ResultSetMetadata.ColumnInfo) || void 0 === _results_ResultSetMetadata_ColumnInfo ? void 0 : _results_ResultSetMetadata_ColumnInfo.map((columnInfo, index)=>{
const name = columnInfo.Name ?? `column_${index}`;
let parse;
switch(columnInfo.Type){
case "integer":
case "tinyint":
case "smallint":
case "bigint":
case "float":
case "double":
case "decimal":
parse = Column.parseNumber;
break;
case "char":
case "varchar":
case "string":
parse = Column.parseString;
break;
case "boolean":
parse = Column.parseBoolean;
break;
case "date":
case "timestamp":
case "timestamp with time zone":
parse = Column.parseDate;
break;
case "array":
parse = Column.parseArray;
break;
case "json":
parse = Column.parseJson;
break;
case "binary":
case "map":
case "struct":
default:
throw new Error(`Column type '${columnInfo.Type}' not supported`);
}
return new Column(name, parse);
})) ?? [];
}
async waitUntilSucceedQuery(query) {
const input = {
QueryExecutionId: query.athenaId
};
const waitTime = 1000 * this._config.waitTimeInSeconds;
for await (const _ of (0, promises_namespaceObject.setInterval)(waitTime)){
var _response_QueryExecution_Status, _response_QueryExecution;
const response = await this._client.send(new client_athena_namespaceObject.GetQueryExecutionCommand(input));
query.status = null === (_response_QueryExecution = response.QueryExecution) || void 0 === _response_QueryExecution ? void 0 : null === (_response_QueryExecution_Status = _response_QueryExecution.Status) || void 0 === _response_QueryExecution_Status ? void 0 : _response_QueryExecution_Status.State;
switch(query.status){
case 'SUCCEEDED':
var _response_QueryExecution_ResultConfiguration, _response_QueryExecution1;
query.s3Location = null == response ? void 0 : null === (_response_QueryExecution1 = response.QueryExecution) || void 0 === _response_QueryExecution1 ? void 0 : null === (_response_QueryExecution_ResultConfiguration = _response_QueryExecution1.ResultConfiguration) || void 0 === _response_QueryExecution_ResultConfiguration ? void 0 : _response_QueryExecution_ResultConfiguration.OutputLocation;
return;
case 'QUEUED':
case 'RUNNING':
break;
case 'CANCELLED':
throw new QueryCanceledException();
case 'FAILED':
throw new AthenaClientException('Query failed');
default:
throw new AthenaClientException(`Query Status '${query.status}' not supported`);
}
}
}
constructor(client, config){
src_define_property(this, "_client", void 0);
src_define_property(this, "_config", void 0);
src_define_property(this, "_queue", void 0);
this._client = client;
this._config = config;
this._queue = new Queue();
}
}
var __webpack_export_target__ = exports;
for(var __webpack_i__ in __webpack_exports__)__webpack_export_target__[__webpack_i__] = __webpack_exports__[__webpack_i__];
if (__webpack_exports__.__esModule) Object.defineProperty(__webpack_export_target__, '__esModule', {
value: true
});