misc: initial setup for audit log partition:

backend/src/@types/knex.d.ts

@@ -170,6 +170,9 @@ import {
   TOrgRoles,
   TOrgRolesInsert,
   TOrgRolesUpdate,
+  TPartitionedAuditLogs,
+  TPartitionedAuditLogsInsert,
+  TPartitionedAuditLogsUpdate,
   TPkiAlerts,
   TPkiAlertsInsert,
   TPkiAlertsUpdate,
@@ -715,6 +718,11 @@ declare module "knex/types/tables" {
       TAuditLogStreamsInsert,
       TAuditLogStreamsUpdate
     >;
+    [TableName.PartitionedAuditLog]: KnexOriginal.CompositeTableType<
+      TPartitionedAuditLogs,
+      TPartitionedAuditLogsInsert,
+      TPartitionedAuditLogsUpdate
+    >;
     [TableName.GitAppInstallSession]: KnexOriginal.CompositeTableType<
       TGitAppInstallSessions,
       TGitAppInstallSessionsInsert,

backend/src/db/migrations/20241002092243_audit-log-drop-fk.ts

@@ -1,53 +0,0 @@
-import { Knex } from "knex";
-
-import { TableName } from "../schemas";
-
-export async function up(knex: Knex): Promise<void> {
-  const doesProjectIdExist = await knex.schema.hasColumn(TableName.AuditLog, "projectId");
-  const doesOrgIdExist = await knex.schema.hasColumn(TableName.AuditLog, "orgId");
-  const doesTableExist = await knex.schema.hasTable(TableName.AuditLog);
-
-  const doesProjectNameExist = await knex.schema.hasColumn(TableName.AuditLog, "projectName");
-
-  if (doesTableExist) {
-    await knex.schema.alterTable(TableName.AuditLog, (t) => {
-      // remove existing FKs
-      if (doesOrgIdExist) {
-        t.dropForeign("orgId");
-      }
-
-      if (doesProjectIdExist) {
-        t.dropForeign("projectId");
-      }
-
-      // add normalized fields necessary after FK removal
-      if (!doesProjectNameExist) {
-        t.string("projectName");
-      }
-    });
-  }
-}
-
-export async function down(knex: Knex): Promise<void> {
-  const doesProjectIdExist = await knex.schema.hasColumn(TableName.AuditLog, "projectId");
-  const doesOrgIdExist = await knex.schema.hasColumn(TableName.AuditLog, "orgId");
-  const doesTableExist = await knex.schema.hasTable(TableName.AuditLog);
-  const doesProjectNameExist = await knex.schema.hasColumn(TableName.AuditLog, "projectName");
-
-  if (doesTableExist) {
-    await knex.schema.alterTable(TableName.AuditLog, (t) => {
-      // add back FKs
-      if (doesOrgIdExist) {
-        t.foreign("orgId").references("id").inTable(TableName.Organization).onDelete("CASCADE");
-      }
-      if (doesProjectIdExist) {
-        t.foreign("projectId").references("id").inTable(TableName.Project).onDelete("CASCADE");
-      }
-
-      // remove normalized fields
-      if (doesProjectNameExist) {
-        t.dropColumn("projectName");
-      }
-    });
-  }
-}

backend/src/db/migrations/20241002110531_add-audit-log-metadata-index.ts

@@ -1,21 +0,0 @@
-import { Knex } from "knex";
-
-import { TableName } from "../schemas";
-
-export async function up(knex: Knex): Promise<void> {
-  if (await knex.schema.hasColumn(TableName.AuditLog, "actorMetadata")) {
-    await knex.raw(
-      `CREATE INDEX "audit_logs_actorMetadata_idx" ON ${TableName.AuditLog} USING gin("actorMetadata" jsonb_path_ops)`
-    );
-  }
-  if (await knex.schema.hasColumn(TableName.AuditLog, "eventMetadata")) {
-    await knex.raw(
-      `CREATE INDEX "audit_logs_eventMetadata_idx" ON ${TableName.AuditLog} USING gin("eventMetadata" jsonb_path_ops)`
-    );
-  }
-}
-
-export async function down(knex: Knex): Promise<void> {
-  await knex.raw(`DROP INDEX IF EXISTS "audit_logs_actorMetadata_idx"`);
-  await knex.raw(`DROP INDEX IF EXISTS "audit_logs_eventMetadata_idx"`);
-}

backend/src/db/migrations/20241003075413_partition-audit-logs.ts

@@ -0,0 +1,161 @@
+import { Knex } from "knex";
+
+import { TableName } from "../schemas";
+
+const formatDateToYYYYMMDD = (date: Date) => {
+  const year = date.getFullYear();
+  const month = String(date.getMonth() + 1).padStart(2, "0"); // getMonth() returns 0-based month, so add 1
+  const day = String(date.getDate()).padStart(2, "0");
+
+  return `${year}-${month}-${day}`;
+};
+
+const createAuditLogPartition = async (knex: Knex, startDate: Date, endDate: Date) => {
+  const startDateStr = formatDateToYYYYMMDD(startDate);
+  const endDateStr = formatDateToYYYYMMDD(endDate);
+
+  const partitionName = `${TableName.PartitionedAuditLog}_${startDateStr.replace(/-/g, "")}_${endDateStr.replace(
+    /-/g,
+    ""
+  )}`;
+
+  await knex.schema.raw(
+    `CREATE TABLE ${partitionName} PARTITION OF ${TableName.PartitionedAuditLog} FOR VALUES FROM ('${startDateStr}') TO ('${endDateStr}')`
+  );
+};
+
+export async function up(knex: Knex): Promise<void> {
+  // prepare the existing audit log table for it to become a partition
+  if (await knex.schema.hasTable(TableName.AuditLog)) {
+    const doesProjectIdExist = await knex.schema.hasColumn(TableName.AuditLog, "projectId");
+    const doesOrgIdExist = await knex.schema.hasColumn(TableName.AuditLog, "orgId");
+    const doesProjectNameExist = await knex.schema.hasColumn(TableName.AuditLog, "projectName");
+
+    await knex.schema.alterTable(TableName.AuditLog, (t) => {
+      // remove existing keys
+      t.dropPrimary();
+
+      if (doesOrgIdExist) {
+        t.dropForeign("orgId");
+      }
+
+      if (doesProjectIdExist) {
+        t.dropForeign("projectId");
+      }
+
+      // add normalized fields present in the partition table
+      if (!doesProjectNameExist) {
+        t.string("projectName");
+      }
+    });
+  }
+
+  // create a new partitioned table for audit logs
+  if (!(await knex.schema.hasTable(TableName.PartitionedAuditLog))) {
+    const createTableSql = knex.schema
+      .createTable(TableName.PartitionedAuditLog, (t) => {
+        t.uuid("id").defaultTo(knex.fn.uuid());
+        t.string("actor").notNullable();
+        t.jsonb("actorMetadata").notNullable();
+        t.string("ipAddress");
+        t.string("eventType").notNullable();
+        t.jsonb("eventMetadata");
+        t.string("userAgent");
+        t.string("userAgentType");
+        t.datetime("expiresAt");
+        t.timestamps(true, true, true);
+        t.uuid("orgId");
+        t.string("projectId");
+        t.string("projectName");
+        t.primary(["id", "createdAt"]);
+      })
+      .toString();
+
+    await knex.schema.raw(`
+        ${createTableSql} PARTITION BY RANGE ("createdAt");
+    `);
+
+    // add indices
+    await knex.raw(
+      `CREATE INDEX "audit_logs_actorMetadata_idx" ON ${TableName.PartitionedAuditLog} USING gin("actorMetadata" jsonb_path_ops)`
+    );
+
+    await knex.raw(
+      `CREATE INDEX "audit_logs_eventMetadata_idx" ON ${TableName.PartitionedAuditLog} USING gin("eventMetadata" jsonb_path_ops)`
+    );
+
+    // create default partition
+    await knex.schema.raw(
+      `CREATE TABLE ${TableName.PartitionedAuditLog}_default PARTITION OF ${TableName.PartitionedAuditLog} DEFAULT`
+    );
+
+    const nextDate = new Date();
+    nextDate.setDate(nextDate.getDate() + 1);
+    const nextDateStr = formatDateToYYYYMMDD(nextDate);
+
+    // attach existing audit log table as a partition
+    await knex.schema.raw(`
+    ALTER TABLE ${TableName.AuditLog} ADD CONSTRAINT audit_log_old
+    CHECK ( "createdAt" < DATE '${nextDateStr}' );
+
+    ALTER TABLE ${TableName.PartitionedAuditLog} ATTACH PARTITION ${TableName.AuditLog}
+    FOR VALUES FROM (MINVALUE) TO ('${nextDateStr}' );
+    `);
+
+    // create partitions 3 months ahead
+    await createAuditLogPartition(knex, nextDate, new Date(nextDate.getFullYear(), nextDate.getMonth() + 1, 1));
+
+    await createAuditLogPartition(
+      knex,
+      new Date(nextDate.getFullYear(), nextDate.getMonth() + 1, 1),
+      new Date(nextDate.getFullYear(), nextDate.getMonth() + 2, 1)
+    );
+
+    await createAuditLogPartition(
+      knex,
+      new Date(nextDate.getFullYear(), nextDate.getMonth() + 2, 1),
+      new Date(nextDate.getFullYear(), nextDate.getMonth() + 3, 1)
+    );
+  }
+}
+
+export async function down(knex: Knex): Promise<void> {
+  // detach audit log from partition
+  await knex.schema.raw(`
+    ALTER TABLE ${TableName.PartitionedAuditLog} DETACH PARTITION ${TableName.AuditLog};
+
+    ALTER TABLE ${TableName.AuditLog} DROP CONSTRAINT audit_log_old;
+  `);
+
+  // revert audit log modifications
+  const doesProjectIdExist = await knex.schema.hasColumn(TableName.AuditLog, "projectId");
+  const doesOrgIdExist = await knex.schema.hasColumn(TableName.AuditLog, "orgId");
+  const doesTableExist = await knex.schema.hasTable(TableName.AuditLog);
+  const doesProjectNameExist = await knex.schema.hasColumn(TableName.AuditLog, "projectName");
+
+  if (doesTableExist) {
+    await knex.schema.alterTable(TableName.AuditLog, (t) => {
+      // we drop this first because adding to the partition results in a new primary key
+      t.dropPrimary();
+
+      // add back the original keys of the audit logs table
+      t.primary(["id"], {
+        constraintName: "audit_logs_pkey"
+      });
+
+      if (doesOrgIdExist) {
+        t.foreign("orgId").references("id").inTable(TableName.Organization).onDelete("CASCADE");
+      }
+      if (doesProjectIdExist) {
+        t.foreign("projectId").references("id").inTable(TableName.Project).onDelete("CASCADE");
+      }
+
+      // remove normalized fields
+      if (doesProjectNameExist) {
+        t.dropColumn("projectName");
+      }
+    });
+  }
+
+  await knex.schema.dropTableIfExists(TableName.PartitionedAuditLog);
+}

backend/src/db/schemas/index.ts

@@ -55,6 +55,7 @@ export * from "./org-bots";
 export * from "./org-memberships";
 export * from "./org-roles";
 export * from "./organizations";
+export * from "./partitioned-audit-logs";
 export * from "./pki-alerts";
 export * from "./pki-collection-items";
 export * from "./pki-collections";

backend/src/db/schemas/models.ts

@@ -90,6 +90,7 @@ export enum TableName {
   OidcConfig = "oidc_configs",
   LdapGroupMap = "ldap_group_maps",
   AuditLog = "audit_logs",
+  PartitionedAuditLog = "partitioned_audit_logs",
   AuditLogStream = "audit_log_streams",
   GitAppInstallSession = "git_app_install_sessions",
   GitAppOrg = "git_app_org",

backend/src/db/schemas/partitioned-audit-logs.ts

@@ -0,0 +1,29 @@
+// Code generated by automation script, DO NOT EDIT.
+// Automated by pulling database and generating zod schema
+// To update. Just run npm run generate:schema
+// Written by akhilmhdh.
+
+import { z } from "zod";
+
+import { TImmutableDBKeys } from "./models";
+
+export const PartitionedAuditLogsSchema = z.object({
+  id: z.string().uuid(),
+  actor: z.string(),
+  actorMetadata: z.unknown(),
+  ipAddress: z.string().nullable().optional(),
+  eventType: z.string(),
+  eventMetadata: z.unknown().nullable().optional(),
+  userAgent: z.string().nullable().optional(),
+  userAgentType: z.string().nullable().optional(),
+  expiresAt: z.date().nullable().optional(),
+  createdAt: z.date(),
+  updatedAt: z.date(),
+  orgId: z.string().uuid().nullable().optional(),
+  projectId: z.string().nullable().optional(),
+  projectName: z.string().nullable().optional()
+});
+
+export type TPartitionedAuditLogs = z.infer<typeof PartitionedAuditLogsSchema>;
+export type TPartitionedAuditLogsInsert = Omit<z.input<typeof PartitionedAuditLogsSchema>, TImmutableDBKeys>;
+export type TPartitionedAuditLogsUpdate = Partial<Omit<z.input<typeof PartitionedAuditLogsSchema>, TImmutableDBKeys>>;

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

@@ -25,7 +25,7 @@ type TFindQuery = {
 };
 
 export const auditLogDALFactory = (db: TDbClient) => {
-  const auditLogOrm = ormify(db, TableName.AuditLog);
+  const auditLogOrm = ormify(db, TableName.PartitionedAuditLog);
 
   const find = async (
     {
@@ -54,13 +54,13 @@ export const auditLogDALFactory = (db: TDbClient) => {
 
     try {
       // Find statements
-      const sqlQuery = (tx || db.replicaNode())(TableName.AuditLog)
+      const sqlQuery = (tx || db.replicaNode())(TableName.PartitionedAuditLog)
         // eslint-disable-next-line func-names
         .where(function () {
           if (orgId) {
-            void this.where(`${TableName.AuditLog}.orgId`, orgId);
+            void this.where(`${TableName.PartitionedAuditLog}.orgId`, orgId);
           } else if (projectId) {
-            void this.where(`${TableName.AuditLog}.projectId`, projectId);
+            void this.where(`${TableName.PartitionedAuditLog}.projectId`, projectId);
           }
         });
 
@@ -70,10 +70,10 @@ export const auditLogDALFactory = (db: TDbClient) => {
 
       // Select statements
       void sqlQuery
-        .select(selectAllTableCols(TableName.AuditLog))
+        .select(selectAllTableCols(TableName.PartitionedAuditLog))
         .limit(limit)
         .offset(offset)
-        .orderBy(`${TableName.AuditLog}.createdAt`, "desc");
+        .orderBy(`${TableName.PartitionedAuditLog}.createdAt`, "desc");
 
       // Special case: Filter by actor ID
       if (actorId) {
@@ -99,10 +99,10 @@ export const auditLogDALFactory = (db: TDbClient) => {
 
       // Filter by date range
       if (startDate) {
-        void sqlQuery.where(`${TableName.AuditLog}.createdAt`, ">=", startDate);
+        void sqlQuery.where(`${TableName.PartitionedAuditLog}.createdAt`, ">=", startDate);
       }
       if (endDate) {
-        void sqlQuery.where(`${TableName.AuditLog}.createdAt`, "<=", endDate);
+        void sqlQuery.where(`${TableName.PartitionedAuditLog}.createdAt`, "<=", endDate);
       }
 
       const docs = await sqlQuery;
@@ -126,13 +126,13 @@ export const auditLogDALFactory = (db: TDbClient) => {
     logger.info(`${QueueName.DailyResourceCleanUp}: audit log started`);
     do {
       try {
-        const findExpiredLogSubQuery = (tx || db)(TableName.AuditLog)
+        const findExpiredLogSubQuery = (tx || db)(TableName.PartitionedAuditLog)
           .where("expiresAt", "<", today)
           .select("id")
           .limit(AUDIT_LOG_PRUNE_BATCH_SIZE);
 
         // eslint-disable-next-line no-await-in-loop
-        deletedAuditLogIds = await (tx || db)(TableName.AuditLog)
+        deletedAuditLogIds = await (tx || db)(TableName.PartitionedAuditLog)
           .whereIn("id", findExpiredLogSubQuery)
           .del()
           .returning("id");