@forzalabs/remora
Version:
A powerful CLI tool for seamless data translation.
315 lines (314 loc) • 15.5 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_sqs_1 = require("@aws-sdk/client-sqs");
const Environment_1 = __importDefault(require("../Environment"));
const ProcessENVManager_1 = __importDefault(require("../ProcessENVManager"));
const SecretManager_1 = __importDefault(require("../SecretManager"));
const ExecutorOrchestrator_1 = __importDefault(require("../../executors/ExecutorOrchestrator"));
const settings_1 = require("../../settings");
class QueueManager {
constructor() {
this.queueMappings = new Map();
this.pollingIntervals = new Map();
this.isInitialized = false;
this.POLLING_INTERVAL_MS = 5000; // Poll every 5 seconds
this.MAX_MESSAGES = 10; // Maximum messages to receive in one poll
// Initialize SQS client with default configuration
// Will be reconfigured when we know the specific queue details
this.sqsClient = new client_sqs_1.SQSClient({});
}
/**
* Initialize the Queue Manager by scanning all consumers and setting up queue listeners for those with QUEUE triggers
*/
initialize() {
if (this.isInitialized) {
console.log('Queue Manager already initialized');
return;
}
console.log('Initializing Queue Manager...');
try {
const consumers = Environment_1.default.getAllConsumers();
let queueTriggerCount = 0;
for (const consumer of consumers) {
if (this.hasQueueTrigger(consumer)) {
this.setupQueueListeners(consumer);
queueTriggerCount++;
}
}
this.isInitialized = true;
console.log(`Queue Manager initialized with ${queueTriggerCount} queue triggers`);
}
catch (error) {
console.error('Failed to initialize Queue Manager:', error);
throw error;
}
}
/**
* Check if a consumer has any QUEUE triggers configured
*/
hasQueueTrigger(consumer) {
return consumer.outputs.some(output => { var _a; return ((_a = output.trigger) === null || _a === void 0 ? void 0 : _a.type) === 'QUEUE' && output.trigger.value; });
}
/**
* Setup queue listeners for a consumer with QUEUE triggers
*/
setupQueueListeners(consumer) {
consumer.outputs.forEach((output, index) => {
var _a;
if (((_a = output.trigger) === null || _a === void 0 ? void 0 : _a.type) === 'QUEUE' && output.trigger.value) {
try {
const queueConfig = this.parseQueueConfig(output.trigger.value, output.trigger.metadata);
const mapping = {
consumer,
outputIndex: index,
queueUrl: queueConfig.queueUrl,
messageType: queueConfig.messageType
};
// Add to mappings
if (!this.queueMappings.has(queueConfig.queueUrl)) {
this.queueMappings.set(queueConfig.queueUrl, []);
}
this.queueMappings.get(queueConfig.queueUrl).push(mapping);
// Start polling for this queue if not already started
if (!this.pollingIntervals.has(queueConfig.queueUrl)) {
this.startQueuePolling(queueConfig.queueUrl, queueConfig.region, queueConfig.credentials);
}
console.log(`Setup queue listener for consumer "${consumer.name}" output ${index} on queue: ${queueConfig.queueUrl}`);
}
catch (error) {
console.error(`Failed to setup queue listener for consumer ${consumer.name}:`, error);
}
}
});
}
/**
* Parse queue configuration from trigger value and metadata
*/
parseQueueConfig(triggerValue, metadata) {
// triggerValue should be the queue URL or queue name
let queueUrl = triggerValue;
// If it's not a full URL, construct it
if (!queueUrl.startsWith('https://')) {
const region = (metadata === null || metadata === void 0 ? void 0 : metadata.region) || ProcessENVManager_1.default.getEnvVariable('AWS_DEFAULT_REGION') || 'us-east-1';
const accountId = (metadata === null || metadata === void 0 ? void 0 : metadata.accountId) || ProcessENVManager_1.default.getEnvVariable('AWS_ACCOUNT_ID');
if (!accountId) {
throw new Error('AWS Account ID is required for queue trigger. Set it in metadata.accountId or AWS_ACCOUNT_ID environment variable');
}
queueUrl = `https://sqs.${region}.amazonaws.com/${accountId}/${triggerValue}`;
}
// Extract region from URL if not provided in metadata
const urlParts = queueUrl.match(/https:\/\/sqs\.([^.]+)\.amazonaws\.com\//);
const region = (metadata === null || metadata === void 0 ? void 0 : metadata.region) || (urlParts === null || urlParts === void 0 ? void 0 : urlParts[1]) || 'us-east-1';
// Get credentials from metadata or environment
let credentials;
const accessKeyId = (metadata === null || metadata === void 0 ? void 0 : metadata.accessKeyId) || ProcessENVManager_1.default.getEnvVariable('AWS_ACCESS_KEY_ID');
const secretAccessKey = (metadata === null || metadata === void 0 ? void 0 : metadata.secretAccessKey) || ProcessENVManager_1.default.getEnvVariable('AWS_SECRET_ACCESS_KEY');
const sessionToken = (metadata === null || metadata === void 0 ? void 0 : metadata.sessionToken) || ProcessENVManager_1.default.getEnvVariable('AWS_SESSION_TOKEN');
if (accessKeyId && secretAccessKey) {
credentials = {
accessKeyId: SecretManager_1.default.replaceSecret(accessKeyId),
secretAccessKey: SecretManager_1.default.replaceSecret(secretAccessKey),
sessionToken: sessionToken ? SecretManager_1.default.replaceSecret(sessionToken) : undefined
};
}
return {
queueUrl,
messageType: metadata === null || metadata === void 0 ? void 0 : metadata.messageType,
region,
credentials
};
}
/**
* Start polling a specific queue for messages
*/
startQueuePolling(queueUrl, region, credentials) {
// Create SQS client for this specific queue
const sqsClient = new client_sqs_1.SQSClient({
region,
credentials
});
const pollQueue = () => __awaiter(this, void 0, void 0, function* () {
try {
const command = new client_sqs_1.ReceiveMessageCommand({
QueueUrl: queueUrl,
MaxNumberOfMessages: this.MAX_MESSAGES,
WaitTimeSeconds: 20, // Long polling
VisibilityTimeout: 300 // 5 minutes to process the message
});
const response = yield sqsClient.send(command);
if (response.Messages && response.Messages.length > 0) {
console.log(`Received ${response.Messages.length} messages from queue: ${queueUrl}`);
for (const message of response.Messages) {
yield this.processMessage(queueUrl, message, sqsClient);
}
}
}
catch (error) {
console.error(`Error polling queue ${queueUrl}:`, error);
}
});
// Start continuous polling
const interval = setInterval(pollQueue, this.POLLING_INTERVAL_MS);
this.pollingIntervals.set(queueUrl, interval);
// Start immediately
pollQueue();
console.log(`Started polling queue: ${queueUrl}`);
}
/**
* Process a message from the queue
*/
processMessage(queueUrl, message, sqsClient) {
return __awaiter(this, void 0, void 0, function* () {
try {
const mappings = this.queueMappings.get(queueUrl);
if (!mappings || mappings.length === 0) {
console.log(`No consumer mappings found for queue: ${queueUrl}`);
return;
}
// Parse message body
let messageData;
try {
messageData = JSON.parse(message.Body || '{}');
}
catch (_a) {
console.warn(`Failed to parse message body as JSON for queue ${queueUrl}. Using raw body.`);
messageData = { body: message.Body };
}
let messageProcessedByAnyConsumer = false;
// Process message for each mapped consumer that matches the message criteria
for (const mapping of mappings) {
try {
// Check if message type matches (if specified)
if (mapping.messageType) {
const messageType = messageData.type || messageData.messageType || messageData.eventType;
if (messageType !== mapping.messageType) {
console.log(`Message type ${messageType} does not match expected ${mapping.messageType} for consumer ${mapping.consumer.name} - skipping`);
continue;
}
}
console.log(`Processing queue message for consumer "${mapping.consumer.name}" output ${mapping.outputIndex}`);
const user = settings_1.REMORA_WORKER_USER;
const result = yield ExecutorOrchestrator_1.default.launch({
consumer: mapping.consumer,
details: {
invokedBy: 'QUEUE',
user: { _id: user._id, name: user.name, type: 'actor' }
},
logProgress: false
});
console.log(`Queue trigger completed successfully for consumer "${mapping.consumer.name}" output ${mapping.outputIndex}`);
// Log execution statistics
if (result) {
console.log(`Queue trigger stats: ${result.elapsedMS}ms, size: ${result.outputCount}, cycles: ${result.cycles}`);
}
messageProcessedByAnyConsumer = true;
}
catch (error) {
console.error(`Queue trigger failed for consumer "${mapping.consumer.name}" output ${mapping.outputIndex}:`, error);
// Continue processing for other consumers even if one fails
}
}
// Only delete message from queue if it was processed by at least one consumer
// This ensures messages intended for other consumers or systems remain in the queue
if (messageProcessedByAnyConsumer && message.ReceiptHandle) {
yield sqsClient.send(new client_sqs_1.DeleteMessageCommand({
QueueUrl: queueUrl,
ReceiptHandle: message.ReceiptHandle
}));
console.log(`Deleted processed message ${message.MessageId} from queue`);
}
else if (!messageProcessedByAnyConsumer) {
console.log(`Message ${message.MessageId} was not processed by any consumer - leaving in queue for other consumers or systems`);
}
}
catch (error) {
console.error(`Error processing message from queue ${queueUrl}:`, error);
// Message will remain in queue and be retried or go to DLQ based on queue configuration
}
});
}
/**
* Add or update queue listeners for a specific consumer
*/
updateConsumerSchedule(consumer) {
// First, remove any existing listeners for this consumer
this.removeConsumerSchedule(consumer.name);
// Then, add new listeners if they have QUEUE triggers
if (this.hasQueueTrigger(consumer)) {
this.setupQueueListeners(consumer);
}
}
/**
* Remove all queue listeners for a consumer
*/
removeConsumerSchedule(consumerName) {
// Remove mappings for this consumer
for (const [queueUrl, mappings] of this.queueMappings.entries()) {
const updatedMappings = mappings.filter(mapping => mapping.consumer.name !== consumerName);
if (updatedMappings.length === 0) {
// No more consumers listening to this queue, stop polling
const interval = this.pollingIntervals.get(queueUrl);
if (interval) {
clearInterval(interval);
this.pollingIntervals.delete(queueUrl);
console.log(`Stopped polling queue: ${queueUrl}`);
}
this.queueMappings.delete(queueUrl);
}
else {
this.queueMappings.set(queueUrl, updatedMappings);
}
}
}
/**
* Stop all queue polling
*/
stopAllQueues() {
console.log('Stopping all queue polling...');
this.pollingIntervals.forEach((interval, queueUrl) => {
clearInterval(interval);
console.log(`Stopped polling queue: ${queueUrl}`);
});
this.pollingIntervals.clear();
this.queueMappings.clear();
this.isInitialized = false;
console.log('All queue polling stopped');
}
/**
* Restart the queue manager (useful for configuration reloads)
*/
restart() {
console.log('Restarting Queue Manager...');
this.stopAllQueues();
this.initialize();
}
/**
* Get the queue manager status
*/
getStatus() {
const queues = Array.from(this.queueMappings.entries()).map(([queueUrl, mappings]) => ({
queueUrl,
consumerCount: mappings.length
}));
return {
initialized: this.isInitialized,
queueCount: this.queueMappings.size,
consumerCount: Array.from(this.queueMappings.values()).flat().length,
queues
};
}
}
// Export a singleton instance
exports.default = new QueueManager();