Merge pull request #2837 from Infisical/misc/use-pg-queue-for-audit-logs-with-flag

misc: pg-queue for audit logs

backend/src/ee/services/audit-log/audit-log-queue.ts

@@ -1,6 +1,7 @@
 import { RawAxiosRequestHeaders } from "axios";
 
 import { SecretKeyEncoding } from "@app/db/schemas";
+import { getConfig } from "@app/lib/config/env";
 import { request } from "@app/lib/config/request";
 import { infisicalSymmetricDecrypt } from "@app/lib/crypto/encryption";
 import { QueueJobs, QueueName, TQueueServiceFactory } from "@app/queue";
@@ -20,27 +21,128 @@ type TAuditLogQueueServiceFactoryDep = {
   licenseService: Pick<TLicenseServiceFactory, "getPlan">;
 };
 
-export type TAuditLogQueueServiceFactory = ReturnType<typeof auditLogQueueServiceFactory>;
+export type TAuditLogQueueServiceFactory = Awaited<ReturnType<typeof auditLogQueueServiceFactory>>;
 
 // keep this timeout 5s it must be fast because else the queue will take time to finish
 // audit log is a crowded queue thus needs to be fast
 export const AUDIT_LOG_STREAM_TIMEOUT = 5 * 1000;
-export const auditLogQueueServiceFactory = ({
+
+export const auditLogQueueServiceFactory = async ({
   auditLogDAL,
   queueService,
   projectDAL,
   licenseService,
   auditLogStreamDAL
 }: TAuditLogQueueServiceFactoryDep) => {
+  const appCfg = getConfig();
+
   const pushToLog = async (data: TCreateAuditLogDTO) => {
-    await queueService.queue(QueueName.AuditLog, QueueJobs.AuditLog, data, {
-      removeOnFail: {
-        count: 3
-      },
-      removeOnComplete: true
-    });
+    if (appCfg.USE_PG_QUEUE) {
+      await queueService.queuePg<QueueName.AuditLog>(QueueJobs.AuditLog, data, {
+        retryLimit: 10,
+        retryBackoff: true
+      });
+    } else {
+      await queueService.queue<QueueName.AuditLog>(QueueName.AuditLog, QueueJobs.AuditLog, data, {
+        removeOnFail: {
+          count: 3
+        },
+        removeOnComplete: true
+      });
+    }
   };
 
+  await queueService.startPg<QueueName.AuditLog>(
+    QueueJobs.AuditLog,
+    async ([job]) => {
+      const { actor, event, ipAddress, projectId, userAgent, userAgentType } = job.data;
+      let { orgId } = job.data;
+      const MS_IN_DAY = 24 * 60 * 60 * 1000;
+      let project;
+
+      if (!orgId) {
+        // it will never be undefined for both org and project id
+        // TODO(akhilmhdh): use caching here in dal to avoid db calls
+        project = await projectDAL.findById(projectId as string);
+        orgId = project.orgId;
+      }
+
+      const plan = await licenseService.getPlan(orgId);
+      if (plan.auditLogsRetentionDays === 0) {
+        // skip inserting if audit log retention is 0 meaning its not supported
+        return;
+      }
+
+      // For project actions, set TTL to project-level audit log retention config
+      // This condition ensures that the plan's audit log retention days cannot be bypassed
+      const ttlInDays =
+        project?.auditLogsRetentionDays && project.auditLogsRetentionDays < plan.auditLogsRetentionDays
+          ? project.auditLogsRetentionDays
+          : plan.auditLogsRetentionDays;
+
+      const ttl = ttlInDays * MS_IN_DAY;
+
+      const auditLog = await auditLogDAL.create({
+        actor: actor.type,
+        actorMetadata: actor.metadata,
+        userAgent,
+        projectId,
+        projectName: project?.name,
+        ipAddress,
+        orgId,
+        eventType: event.type,
+        expiresAt: new Date(Date.now() + ttl),
+        eventMetadata: event.metadata,
+        userAgentType
+      });
+
+      const logStreams = orgId ? await auditLogStreamDAL.find({ orgId }) : [];
+      await Promise.allSettled(
+        logStreams.map(
+          async ({
+            url,
+            encryptedHeadersTag,
+            encryptedHeadersIV,
+            encryptedHeadersKeyEncoding,
+            encryptedHeadersCiphertext
+          }) => {
+            const streamHeaders =
+              encryptedHeadersIV && encryptedHeadersCiphertext && encryptedHeadersTag
+                ? (JSON.parse(
+                    infisicalSymmetricDecrypt({
+                      keyEncoding: encryptedHeadersKeyEncoding as SecretKeyEncoding,
+                      iv: encryptedHeadersIV,
+                      tag: encryptedHeadersTag,
+                      ciphertext: encryptedHeadersCiphertext
+                    })
+                  ) as LogStreamHeaders[])
+                : [];
+
+            const headers: RawAxiosRequestHeaders = { "Content-Type": "application/json" };
+
+            if (streamHeaders.length)
+              streamHeaders.forEach(({ key, value }) => {
+                headers[key] = value;
+              });
+
+            return request.post(url, auditLog, {
+              headers,
+              // request timeout
+              timeout: AUDIT_LOG_STREAM_TIMEOUT,
+              // connection timeout
+              signal: AbortSignal.timeout(AUDIT_LOG_STREAM_TIMEOUT)
+            });
+          }
+        )
+      );
+    },
+    {
+      batchSize: 1,
+      workerCount: 30,
+      pollingIntervalSeconds: 0.5
+    }
+  );
+
   queueService.start(QueueName.AuditLog, async (job) => {
     const { actor, event, ipAddress, projectId, userAgent, userAgentType } = job.data;
     let { orgId } = job.data;

backend/src/lib/config/env.ts

@@ -178,7 +178,9 @@ const envSchema = z
     HSM_LIB_PATH: zpStr(z.string().optional()),
     HSM_PIN: zpStr(z.string().optional()),
     HSM_KEY_LABEL: zpStr(z.string().optional()),
-    HSM_SLOT: z.coerce.number().optional().default(0)
+    HSM_SLOT: z.coerce.number().optional().default(0),
+
+    USE_PG_QUEUE: zodStrBool.default("false")
   })
   // To ensure that basic encryption is always possible.
   .refine(

backend/src/main.ts

@@ -55,7 +55,10 @@ const run = async () => {
   }
 
   const smtp = smtpServiceFactory(formatSmtpConfig());
-  const queue = queueServiceFactory(appCfg.REDIS_URL);
+
+  const queue = queueServiceFactory(appCfg.REDIS_URL, appCfg.DB_CONNECTION_URI);
+  await queue.initialize();
+
   const keyStore = keyStoreFactory(appCfg.REDIS_URL);
 
   const hsmModule = initializeHsmModule();

backend/src/queue/queue-service.ts

@@ -1,5 +1,6 @@
 import { Job, JobsOptions, Queue, QueueOptions, RepeatOptions, Worker, WorkerListener } from "bullmq";
 import Redis from "ioredis";
+import PgBoss, { WorkOptions } from "pg-boss";
 
 import { SecretEncryptionAlgo, SecretKeyEncoding } from "@app/db/schemas";
 import { TCreateAuditLogDTO } from "@app/ee/services/audit-log/audit-log-types";
@@ -7,6 +8,7 @@ import {
   TScanFullRepoEventPayload,
   TScanPushEventPayload
 } from "@app/ee/services/secret-scanning/secret-scanning-queue/secret-scanning-queue-types";
+import { logger } from "@app/lib/logger";
 import {
   TFailedIntegrationSyncEmailsPayload,
   TIntegrationSyncPayload,
@@ -184,17 +186,35 @@ export type TQueueJobTypes = {
 };
 
 export type TQueueServiceFactory = ReturnType<typeof queueServiceFactory>;
-export const queueServiceFactory = (redisUrl: string) => {
+export const queueServiceFactory = (redisUrl: string, dbConnectionUrl: string) => {
   const connection = new Redis(redisUrl, { maxRetriesPerRequest: null });
   const queueContainer = {} as Record<
     QueueName,
     Queue<TQueueJobTypes[QueueName]["payload"], void, TQueueJobTypes[QueueName]["name"]>
   >;
+
+  const pgBoss = new PgBoss({
+    connectionString: dbConnectionUrl,
+    archiveCompletedAfterSeconds: 60,
+    archiveFailedAfterSeconds: 1000, // we want to keep failed jobs for a longer time so that it can be retried
+    deleteAfterSeconds: 30
+  });
+
+  const queueContainerPg = {} as Record<QueueJobs, boolean>;
+
   const workerContainer = {} as Record<
     QueueName,
     Worker<TQueueJobTypes[QueueName]["payload"], void, TQueueJobTypes[QueueName]["name"]>
   >;
 
+  const initialize = async () => {
+    await pgBoss.start();
+
+    pgBoss.on("error", (error) => {
+      logger.error(error, "pg-queue error");
+    });
+  };
+
   const start = <T extends QueueName>(
     name: T,
     jobFn: (job: Job<TQueueJobTypes[T]["payload"], void, TQueueJobTypes[T]["name"]>, token?: string) => Promise<void>,
@@ -215,6 +235,27 @@ export const queueServiceFactory = (redisUrl: string) => {
     });
   };
 
+  const startPg = async <T extends QueueName>(
+    jobName: QueueJobs,
+    jobsFn: (jobs: PgBoss.Job<TQueueJobTypes[T]["payload"]>[]) => Promise<void>,
+    options: WorkOptions & {
+      workerCount: number;
+    }
+  ) => {
+    if (queueContainerPg[jobName]) {
+      throw new Error(`${jobName} queue is already initialized`);
+    }
+
+    await pgBoss.createQueue(jobName);
+    queueContainerPg[jobName] = true;
+
+    await Promise.all(
+      Array.from({ length: options.workerCount }).map(() =>
+        pgBoss.work<TQueueJobTypes[T]["payload"]>(jobName, options, jobsFn)
+      )
+    );
+  };
+
   const listen = <
     T extends QueueName,
     U extends keyof WorkerListener<TQueueJobTypes[T]["payload"], void, TQueueJobTypes[T]["name"]>
@@ -238,6 +279,18 @@ export const queueServiceFactory = (redisUrl: string) => {
     await q.add(job, data, opts);
   };
 
+  const queuePg = async <T extends QueueName>(
+    job: TQueueJobTypes[T]["name"],
+    data: TQueueJobTypes[T]["payload"],
+    opts?: PgBoss.SendOptions & { jobId?: string }
+  ) => {
+    await pgBoss.send({
+      name: job,
+      data,
+      options: opts
+    });
+  };
+
   const stopRepeatableJob = async <T extends QueueName>(
     name: T,
     job: TQueueJobTypes[T]["name"],
@@ -274,5 +327,17 @@ export const queueServiceFactory = (redisUrl: string) => {
     await Promise.all(Object.values(workerContainer).map((worker) => worker.close()));
   };
 
-  return { start, listen, queue, shutdown, stopRepeatableJob, stopRepeatableJobByJobId, clearQueue, stopJobById };
+  return {
+    initialize,
+    start,
+    listen,
+    queue,
+    shutdown,
+    stopRepeatableJob,
+    stopRepeatableJobByJobId,
+    clearQueue,
+    stopJobById,
+    startPg,
+    queuePg
+  };
 };

backend/src/server/routes/index.ts

@@ -394,13 +394,14 @@ export const registerRoutes = async (
     permissionService
   });
 
-  const auditLogQueue = auditLogQueueServiceFactory({
+  const auditLogQueue = await auditLogQueueServiceFactory({
     auditLogDAL,
     queueService,
     projectDAL,
     licenseService,
     auditLogStreamDAL
   });
+
   const auditLogService = auditLogServiceFactory({ auditLogDAL, permissionService, auditLogQueue });
   const auditLogStreamService = auditLogStreamServiceFactory({
     licenseService,