@forzalabs/remora
Version:
A powerful CLI tool for seamless data translation.
198 lines (197 loc) • 11.9 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 __asyncValues = (this && this.__asyncValues) || function (o) {
if (!Symbol.asyncIterator) throw new TypeError("Symbol.asyncIterator is not defined.");
var m = o[Symbol.asyncIterator], i;
return m ? m.call(o) : (o = typeof __values === "function" ? __values(o) : o[Symbol.iterator](), i = {}, verb("next"), verb("throw"), verb("return"), i[Symbol.asyncIterator] = function () { return this; }, i);
function verb(n) { i[n] = o[n] && function (v) { return new Promise(function (resolve, reject) { v = o[n](v), settle(resolve, reject, v.done, v.value); }); }; }
function settle(resolve, reject, d, v) { Promise.resolve(v).then(function(v) { resolve({ value: v, done: d }); }, reject); }
};
var __await = (this && this.__await) || function (v) { return this instanceof __await ? (this.v = v, this) : new __await(v); }
var __asyncGenerator = (this && this.__asyncGenerator) || function (thisArg, _arguments, generator) {
if (!Symbol.asyncIterator) throw new TypeError("Symbol.asyncIterator is not defined.");
var g = generator.apply(thisArg, _arguments || []), i, q = [];
return i = Object.create((typeof AsyncIterator === "function" ? AsyncIterator : Object).prototype), verb("next"), verb("throw"), verb("return", awaitReturn), i[Symbol.asyncIterator] = function () { return this; }, i;
function awaitReturn(f) { return function (v) { return Promise.resolve(v).then(f, reject); }; }
function verb(n, f) { if (g[n]) { i[n] = function (v) { return new Promise(function (a, b) { q.push([n, v, a, b]) > 1 || resume(n, v); }); }; if (f) i[n] = f(i[n]); } }
function resume(n, v) { try { step(g[n](v)); } catch (e) { settle(q[0][3], e); } }
function step(r) { r.value instanceof __await ? Promise.resolve(r.value.v).then(fulfill, reject) : settle(q[0][2], r); }
function fulfill(value) { resume("next", value); }
function reject(value) { resume("throw", value); }
function settle(f, v) { if (f(v), q.shift(), q.length) resume(q[0][0], q[0][1]); }
};
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
const client_s3_1 = require("@aws-sdk/client-s3");
const Affirm_1 = __importDefault(require("../../core/Affirm"));
const SecretManager_1 = __importDefault(require("../../engines/SecretManager"));
const path_1 = __importDefault(require("path"));
const fs_1 = __importDefault(require("fs"));
const readline_1 = __importDefault(require("readline"));
class S3DestinationDriver {
constructor() {
this.init = (source) => __awaiter(this, void 0, void 0, function* () {
this._bucketName = source.authentication['bucket'];
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
}
};
this._client = new client_s3_1.S3Client(config);
// TODO: is there a way to test if the connection was successful? like a query or scan that I can do?
return this;
});
this.copyFromS3 = (sourceBucket, sourceFileKey, destinationFileKey) => __awaiter(this, void 0, void 0, function* () {
(0, Affirm_1.default)(this._client, 'S3 client not yet initialized, call "init()" first');
(0, Affirm_1.default)(sourceBucket, 'Invalid source bucket');
(0, Affirm_1.default)(sourceFileKey, 'Invalid source file key');
(0, Affirm_1.default)(destinationFileKey, 'Invalid destination file key');
yield this._client.send(new client_s3_1.CopyObjectCommand({
CopySource: `${sourceBucket}/${sourceFileKey}`,
Bucket: this._bucketName,
Key: destinationFileKey
}));
});
this.saveFile = (fileKey, content) => __awaiter(this, void 0, void 0, function* () {
(0, Affirm_1.default)(this._client, 'S3 client not yet initialized, call "init()" first');
(0, Affirm_1.default)(fileKey, 'Invalid file key');
(0, Affirm_1.default)(content, 'Invalid content');
yield this._client.send(new client_s3_1.PutObjectCommand({
Bucket: this._bucketName,
Key: fileKey,
Body: content
}));
});
this.move = (fromPath, toName) => __awaiter(this, void 0, void 0, function* () {
(0, Affirm_1.default)(fromPath, 'Invalid source path');
(0, Affirm_1.default)(toName, 'Invalid destination name');
(0, Affirm_1.default)(fs_1.default.existsSync(fromPath), `Source file does not exist: ${fromPath}`);
const readStream = fs_1.default.createReadStream(fromPath);
return this._multipartUpload(toName, readStream);
});
this.transformAndMove = (fromPath, transform, toName) => __awaiter(this, void 0, void 0, function* () {
(0, Affirm_1.default)(fromPath, 'Invalid source path');
(0, Affirm_1.default)(transform, 'Invalid transform function');
(0, Affirm_1.default)(toName, 'Invalid destination name');
(0, Affirm_1.default)(fs_1.default.existsSync(fromPath), `Source file does not exist: ${fromPath}`);
const reader = fs_1.default.createReadStream(fromPath);
const lineReader = readline_1.default.createInterface({ input: reader, crlfDelay: Infinity });
const transformedLines = function () {
return __asyncGenerator(this, arguments, function* () {
var _a, e_1, _b, _c;
try {
for (var _d = true, lineReader_1 = __asyncValues(lineReader), lineReader_1_1; lineReader_1_1 = yield __await(lineReader_1.next()), _a = lineReader_1_1.done, !_a; _d = true) {
_c = lineReader_1_1.value;
_d = false;
const line = _c;
yield yield __await(transform(line) + '\n');
}
}
catch (e_1_1) { e_1 = { error: e_1_1 }; }
finally {
try {
if (!_d && !_a && (_b = lineReader_1.return)) yield __await(_b.call(lineReader_1));
}
finally { if (e_1) throw e_1.error; }
}
});
};
return this._multipartUpload(toName, transformedLines());
});
this._multipartUpload = (toName, dataSource) => __awaiter(this, void 0, void 0, function* () {
var _a, dataSource_1, dataSource_1_1;
var _b, e_2, _c, _d;
let uploadId;
try {
// Create the multipart upload
const createMultipartUploadRes = yield this._client.send(new client_s3_1.CreateMultipartUploadCommand({
Bucket: this._bucketName,
Key: toName
}));
uploadId = createMultipartUploadRes.UploadId;
(0, Affirm_1.default)(uploadId, 'Failed to initiate multipart upload');
const uploadedParts = [];
let partNumber = 1;
const MIN_PART_SIZE = 5 * 1024 * 1024; // 5MB
let accumulatedBuffer = Buffer.alloc(0);
const uploadPart = (buffer) => __awaiter(this, void 0, void 0, function* () {
const uploadPartRes = yield this._client.send(new client_s3_1.UploadPartCommand({
Bucket: this._bucketName,
Key: toName,
UploadId: uploadId,
PartNumber: partNumber,
Body: buffer
}));
uploadedParts.push({
PartNumber: partNumber,
ETag: uploadPartRes.ETag
});
partNumber++;
});
try {
for (_a = true, dataSource_1 = __asyncValues(dataSource); dataSource_1_1 = yield dataSource_1.next(), _b = dataSource_1_1.done, !_b; _a = true) {
_d = dataSource_1_1.value;
_a = false;
const chunk = _d;
const chunkBuffer = Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk);
accumulatedBuffer = Buffer.concat([accumulatedBuffer, chunkBuffer]);
// If accumulated buffer is at least 5MB, upload it as a part
while (accumulatedBuffer.length >= MIN_PART_SIZE) {
const partBuffer = accumulatedBuffer.subarray(0, MIN_PART_SIZE);
accumulatedBuffer = accumulatedBuffer.subarray(MIN_PART_SIZE);
yield uploadPart(partBuffer);
}
}
}
catch (e_2_1) { e_2 = { error: e_2_1 }; }
finally {
try {
if (!_a && !_b && (_c = dataSource_1.return)) yield _c.call(dataSource_1);
}
finally { if (e_2) throw e_2.error; }
}
// Upload any remaining data as the final part (even if smaller than 5MB)
if (accumulatedBuffer.length > 0) {
yield uploadPart(accumulatedBuffer);
}
// Complete the multipart upload
const completeRes = yield this._client.send(new client_s3_1.CompleteMultipartUploadCommand({
Bucket: this._bucketName,
Key: toName,
UploadId: uploadId,
MultipartUpload: {
Parts: uploadedParts
}
}));
(0, Affirm_1.default)(completeRes.$metadata.httpStatusCode === 200, `Failed to complete multipart upload for "${toName}": status code ${completeRes.$metadata.httpStatusCode}`);
return { res: true, key: path_1.default.join(this._bucketName, toName), bucket: this._bucketName };
}
catch (error) {
// If anything fails, make sure to abort the multipart upload
if (uploadId) {
yield this._client.send(new client_s3_1.AbortMultipartUploadCommand({
Bucket: this._bucketName,
Key: toName,
UploadId: uploadId
}));
}
throw error;
}
});
}
}
exports.default = S3DestinationDriver;