@forzalabs/remora
Version:
A powerful CLI tool for seamless data translation.
181 lines (180 loc) • 10.1 kB
JavaScript
;
var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
return new (P || (P = Promise))(function (resolve, reject) {
function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
step((generator = generator.apply(thisArg, _arguments || [])).next());
});
};
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
const client_redshift_data_1 = require("@aws-sdk/client-redshift-data");
const Affirm_1 = __importDefault(require("../core/Affirm"));
const Algo_1 = __importDefault(require("../core/Algo"));
const Environment_1 = __importDefault(require("../engines/Environment"));
const SecretManager_1 = __importDefault(require("../engines/SecretManager"));
class RedshiftDriver {
constructor() {
this.init = (source) => __awaiter(this, void 0, void 0, function* () {
(0, Affirm_1.default)(source, `Invalid source`);
this._SQL_MAX_QUERY_ROWS = parseInt(Environment_1.default.get('SQL_MAX_QUERY_ROWS'));
this._dbName = source.authentication['database'];
this._workgroup = source.authentication['workgroup'];
const sessionToken = SecretManager_1.default.replaceSecret(source.authentication['sessionToken']);
const config = {
region: source.authentication['region'],
credentials: {
accessKeyId: SecretManager_1.default.replaceSecret(source.authentication['accessKey']),
secretAccessKey: SecretManager_1.default.replaceSecret(source.authentication['secretKey']),
sessionToken: sessionToken ? sessionToken : undefined
}
};
try {
this._client = new client_redshift_data_1.RedshiftData(config);
// This is just to check that the connection was successful and warm up the connection
yield this.query('SELECT :diss AS number', [{ name: ':diss', value: '1' }]);
}
catch (error) {
const myError = error;
if (myError.__type === 'UnrecognizedClientException')
throw new Error(`Invalid AWS credentials for source "${source.name}": ${myError.message}`);
else
throw error;
}
return this;
});
this.execute = (sql) => __awaiter(this, void 0, void 0, function* () {
var _a;
(0, Affirm_1.default)(sql, `Invalid SQL`);
const input = {
Sql: sql,
Database: this._dbName,
WorkgroupName: this._workgroup
};
const executeStatementRes = yield this._client.executeStatement(input);
const queryId = (_a = executeStatementRes.Id) !== null && _a !== void 0 ? _a : '';
const res = (yield this.executeStatement(queryId, { retry: 5 }))[0];
return {
rows: res.result
};
});
this.query = (sql, values) => __awaiter(this, void 0, void 0, function* () {
var _a, _b;
(0, Affirm_1.default)(sql, `Invalid SQL`);
const wrapperSql = `SELECT * FROM (${sql}) LIMIT ${this._SQL_MAX_QUERY_ROWS}`;
const input = {
Sql: wrapperSql,
Database: this._dbName,
WorkgroupName: this._workgroup,
Parameters: (_a = values === null || values === void 0 ? void 0 : values.map(x => ({ name: Algo_1.default.replaceAll(x.name, ':', ''), value: x.value }))) !== null && _a !== void 0 ? _a : undefined
};
const executeStatementRes = yield this._client.executeStatement(input);
const queryId = (_b = executeStatementRes.Id) !== null && _b !== void 0 ? _b : '';
// TODO: handle errors (permission, wrong query, ...)
const res = (yield this.executeStatement(queryId, { retry: 5 }))[0];
return {
rows: res.result
};
});
this.exist = (producer) => __awaiter(this, void 0, void 0, function* () {
var _a;
(0, Affirm_1.default)(producer, `Invalid producer`);
const sql = `SELECT * FROM "${producer.settings.sqlTable}" LIMIT 1`;
const input = {
Sql: sql,
Database: this._dbName,
WorkgroupName: this._workgroup
};
const executeStatementRes = yield this._client.executeStatement(input);
const queryId = (_a = executeStatementRes.Id) !== null && _a !== void 0 ? _a : '';
try {
yield this.executeStatement(queryId, { retry: 5 });
return true;
}
catch (e) {
console.error(e);
return false;
}
});
this.executeStatement = (queryId, options) => __awaiter(this, void 0, void 0, function* () {
const MAX_WAIT_CYCLES = 20;
for (let i = 0; i < MAX_WAIT_CYCLES; i++) {
const statement = yield this._client.describeStatement({ Id: queryId });
(0, Affirm_1.default)(statement, 'SQL query failed: describe statement call failed.');
(0, Affirm_1.default)(statement.Status !== 'FAILED', `SQL query failed: ${queryId}:\n\tSQL: ${statement.QueryString}\n\tError: ${statement.Error}.`);
const { Status: status, SubStatements: subStatements } = statement;
if (status === 'FINISHED') {
if (subStatements && subStatements.length > 0) {
const result = [];
for (let i = 0; i < subStatements.length; i++) {
const subStatement = subStatements[i];
result.push({
statement: { hasResultSet: statement.HasResultSet, queryString: statement.QueryString, resultRows: statement.ResultRows, resultSize: statement.ResultSize },
subStatement: { hasResultSet: subStatement.HasResultSet, queryString: subStatement.QueryString, resultRows: subStatement.ResultRows, resultSize: subStatement.ResultSize },
result: yield this.getStatementResult(subStatement)
});
}
return result;
}
else {
return [{
statement: { hasResultSet: statement.HasResultSet, queryString: statement.QueryString, resultRows: statement.ResultRows, resultSize: statement.ResultSize },
result: yield this.getStatementResult(statement)
}];
}
}
yield Algo_1.default.sleep(1);
}
if (options && options.retry > 0)
return yield this.executeStatement(queryId, Object.assign(Object.assign({}, options), { retry: options.retry - 1 }));
else
throw new Error('SQL timeout: query failed because timeout exceeded max wait cycles');
});
this.getStatementResult = (statement) => __awaiter(this, void 0, void 0, function* () {
(0, Affirm_1.default)(statement, `SQL failed: can't extract result from invalid statement`);
if (statement.HasResultSet) {
const result = yield this._client.getStatementResult({ Id: statement.Id });
(0, Affirm_1.default)(result, 'SQL failed: invalid result statement');
(0, Affirm_1.default)(result.Records, 'SQL failed: no records in result statement');
return this.parseQueryResult(result);
}
else {
return [];
}
});
this.parseQueryResult = (queryResult) => {
(0, Affirm_1.default)(queryResult, `Parse SQL failed: invalid query result`);
(0, Affirm_1.default)(queryResult.ColumnMetadata, `Parse SQL failed: invalid column metadata in query result`);
(0, Affirm_1.default)(queryResult.Records, `Parse SQL failed: invalid records in query result`);
const columnNames = queryResult.ColumnMetadata.map(column => column.name);
const records = [];
for (const record of queryResult.Records) {
const recordData = {};
for (let i = 0; i < columnNames.length; i++) {
const columnName = columnNames[i];
(0, Affirm_1.default)(columnName, `Parse SQL failed: invalid column name`);
const value = record[i];
if (value.longValue !== undefined)
recordData[columnName] = value.longValue;
else if (value.stringValue !== undefined)
recordData[columnName] = value.stringValue;
else if (value.booleanValue !== undefined)
recordData[columnName] = value.booleanValue;
else if (value.doubleValue !== undefined)
recordData[columnName] = value.doubleValue;
else if (value.isNull === true)
recordData[columnName] = null;
else
recordData[columnName] = null;
}
records.push(recordData);
}
return records;
};
}
}
exports.default = RedshiftDriver;