@forzalabs/remora/drivers/S3SourceDriver.js

A powerful CLI tool for seamless data translation.

"use strict";
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 __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 readline_1 = __importDefault(require("readline"));
const Affirm_1 = __importDefault(require("../core/Affirm"));
const SecretManager_1 = __importDefault(require("../engines/SecretManager"));
class S3SourceDriver {
    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.download = (request) => __awaiter(this, void 0, void 0, function* () {
            var _a, e_1, _b, _c;
            (0, Affirm_1.default)(this._client, 'S3 client not yet initialized, call "connect()" first');
            (0, Affirm_1.default)(request, `Invalid download request`);
            (0, Affirm_1.default)(request.fileKey, `Invalid file key for download request`);
            const { fileKey } = request;
            const bucket = this._bucketName;
            const response = yield this._client.send(new client_s3_1.GetObjectCommand({
                Bucket: bucket,
                Key: fileKey
            }));
            (0, Affirm_1.default)(response.Body, 'Failed to fetch object from S3');
            const stream = response.Body;
            const reader = readline_1.default.createInterface({ input: stream, crlfDelay: Infinity });
            const lines = [];
            try {
                for (var _d = true, reader_1 = __asyncValues(reader), reader_1_1; reader_1_1 = yield reader_1.next(), _a = reader_1_1.done, !_a; _d = true) {
                    _c = reader_1_1.value;
                    _d = false;
                    const line = _c;
                    lines.push(line);
                }
            }
            catch (e_1_1) { e_1 = { error: e_1_1 }; }
            finally {
                try {
                    if (!_d && !_a && (_b = reader_1.return)) yield _b.call(reader_1);
                }
                finally { if (e_1) throw e_1.error; }
            }
            reader.close();
            return lines;
        });
        this.readLinesInRange = (readOptions) => __awaiter(this, void 0, void 0, function* () {
            var _a, e_2, _b, _c;
            (0, Affirm_1.default)(this._client, 'S3 client not yet initialized, call "connect()" first');
            (0, Affirm_1.default)(readOptions, 'Invalid read options');
            const { fileKey, lineFrom, lineTo } = readOptions;
            const bucket = this._bucketName;
            const response = yield this._client.send(new client_s3_1.GetObjectCommand({
                Bucket: bucket,
                Key: fileKey
            }));
            (0, Affirm_1.default)(response.Body, 'Failed to fetch object from S3');
            const stream = response.Body;
            const reader = readline_1.default.createInterface({ input: stream, crlfDelay: Infinity });
            const lines = [];
            let lineCounter = 0;
            try {
                for (var _d = true, reader_2 = __asyncValues(reader), reader_2_1; reader_2_1 = yield reader_2.next(), _a = reader_2_1.done, !_a; _d = true) {
                    _c = reader_2_1.value;
                    _d = false;
                    const line = _c;
                    if (lineCounter >= lineFrom && lineCounter < lineTo)
                        lines.push(line);
                    if (lineCounter > lineTo)
                        break;
                    lineCounter++;
                }
            }
            catch (e_2_1) { e_2 = { error: e_2_1 }; }
            finally {
                try {
                    if (!_d && !_a && (_b = reader_2.return)) yield _b.call(reader_2);
                }
                finally { if (e_2) throw e_2.error; }
            }
            reader.close();
            return lines;
        });
        this.exist = (producer) => __awaiter(this, void 0, void 0, function* () {
            var _a;
            (0, Affirm_1.default)(this._client, 'S3 client not yet initialized, call "connect()" first');
            (0, Affirm_1.default)(producer, 'Invalid read producer');
            const bucket = this._bucketName;
            const fileKey = producer.settings.fileKey;
            (0, Affirm_1.default)(fileKey, `Invalid file key for download request`);
            try {
                yield this._client.send(new client_s3_1.HeadObjectCommand({ Bucket: bucket, Key: fileKey }));
                return true;
            }
            catch (error) {
                if (((_a = error.$metadata) === null || _a === void 0 ? void 0 : _a.httpStatusCode) === 404 || error.name === 'NotFound')
                    return false;
                throw error;
            }
        });
    }
}
exports.default = S3SourceDriver;