From 3a1168c7e82c7ee79938b8c53aadcdda51f41487 Mon Sep 17 00:00:00 2001 From: Sheen Capadngan Date: Fri, 31 May 2024 19:12:55 +0800 Subject: [PATCH] feat: added initial version pruning and result limiting --- .../20240531042507_add-pit-version-limit.ts | 21 ++++++++ backend/src/db/schemas/projects.ts | 3 +- .../secret-snapshot-service.ts | 17 +++++- .../services/secret-snapshot/snapshot-dal.ts | 52 ++++++++++++++++++- backend/src/server/routes/index.ts | 1 + .../src/services/project/project-service.ts | 3 +- backend/src/services/secret/secret-service.ts | 13 ++++- .../src/services/secret/secret-version-dal.ts | 37 +++++++++++++ 8 files changed, 140 insertions(+), 7 deletions(-) create mode 100644 backend/src/db/migrations/20240531042507_add-pit-version-limit.ts diff --git a/backend/src/db/migrations/20240531042507_add-pit-version-limit.ts b/backend/src/db/migrations/20240531042507_add-pit-version-limit.ts new file mode 100644 index 000000000..e37c24e2c --- /dev/null +++ b/backend/src/db/migrations/20240531042507_add-pit-version-limit.ts @@ -0,0 +1,21 @@ +import { Knex } from "knex"; + +import { TableName } from "../schemas"; + +export async function up(knex: Knex): Promise { + const hasPitVersionLimitColumn = await knex.schema.hasColumn(TableName.Project, "pitVersionLimit"); + await knex.schema.alterTable(TableName.Project, (tb) => { + if (!hasPitVersionLimitColumn) { + tb.integer("pitVersionLimit").notNullable().defaultTo(10); + } + }); +} + +export async function down(knex: Knex): Promise { + const hasPitVersionLimitColumn = await knex.schema.hasColumn(TableName.Project, "pitVersionLimit"); + await knex.schema.alterTable(TableName.Project, (tb) => { + if (hasPitVersionLimitColumn) { + tb.dropColumn("pitVersionLimit"); + } + }); +} diff --git a/backend/src/db/schemas/projects.ts b/backend/src/db/schemas/projects.ts index 3965e24c0..9197a355d 100644 --- a/backend/src/db/schemas/projects.ts +++ b/backend/src/db/schemas/projects.ts @@ -16,7 +16,8 @@ export const ProjectsSchema = z.object({ createdAt: z.date(), updatedAt: z.date(), version: z.number().default(1), - upgradeStatus: z.string().nullable().optional() + upgradeStatus: z.string().nullable().optional(), + pitVersionLimit: z.number() }); export type TProjects = z.infer; diff --git a/backend/src/ee/services/secret-snapshot/secret-snapshot-service.ts b/backend/src/ee/services/secret-snapshot/secret-snapshot-service.ts index 0e71ad126..6197c0825 100644 --- a/backend/src/ee/services/secret-snapshot/secret-snapshot-service.ts +++ b/backend/src/ee/services/secret-snapshot/secret-snapshot-service.ts @@ -4,6 +4,7 @@ import { TableName, TSecretTagJunctionInsert } from "@app/db/schemas"; import { BadRequestError, InternalServerError } from "@app/lib/errors"; import { groupBy } from "@app/lib/fn"; import { logger } from "@app/lib/logger"; +import { TProjectDALFactory } from "@app/services/project/project-dal"; import { TSecretDALFactory } from "@app/services/secret/secret-dal"; import { TSecretVersionDALFactory } from "@app/services/secret/secret-version-dal"; import { TSecretVersionTagDALFactory } from "@app/services/secret/secret-version-tag-dal"; @@ -37,6 +38,7 @@ type TSecretSnapshotServiceFactoryDep = { folderDAL: Pick; permissionService: Pick; licenseService: Pick; + projectDAL: Pick; }; export type TSecretSnapshotServiceFactory = ReturnType; @@ -48,6 +50,7 @@ export const secretSnapshotServiceFactory = ({ snapshotSecretDAL, snapshotFolderDAL, folderDAL, + projectDAL, secretDAL, permissionService, licenseService, @@ -81,8 +84,9 @@ export const secretSnapshotServiceFactory = ({ const folder = await folderDAL.findBySecretPath(projectId, environment, path); if (!folder) throw new BadRequestError({ message: "Folder not found" }); + const project = await projectDAL.findById(projectId); const count = await snapshotDAL.countOfSnapshotsByFolderId(folder.id); - return count; + return Math.min(count, project.pitVersionLimit); }; const listSnapshots = async ({ @@ -114,7 +118,16 @@ export const secretSnapshotServiceFactory = ({ const folder = await folderDAL.findBySecretPath(projectId, environment, path); if (!folder) throw new BadRequestError({ message: "Folder not found" }); - const snapshots = await snapshotDAL.find({ folderId: folder.id }, { limit, offset, sort: [["createdAt", "desc"]] }); + const { pitVersionLimit } = await projectDAL.findById(projectId); + const computedQueryLimit = Math.min(pitVersionLimit - offset, limit); + if (offset > pitVersionLimit || computedQueryLimit <= 0) { + return []; + } + + const snapshots = await snapshotDAL.find( + { folderId: folder.id }, + { limit: computedQueryLimit, offset, sort: [["createdAt", "desc"]] } + ); return snapshots; }; diff --git a/backend/src/ee/services/secret-snapshot/snapshot-dal.ts b/backend/src/ee/services/secret-snapshot/snapshot-dal.ts index cdd5a999b..1618024bb 100644 --- a/backend/src/ee/services/secret-snapshot/snapshot-dal.ts +++ b/backend/src/ee/services/secret-snapshot/snapshot-dal.ts @@ -325,12 +325,62 @@ export const snapshotDALFactory = (db: TDbClient) => { } }; + const pruneExcessSnapshots = async (tx?: Knex) => { + try { + const folders = await (tx || db)(TableName.SecretFolder).select("id"); + const folderIds = folders.map((folder) => folder.id); + const PRUNE_FOLDER_BATCH_SIZE = 500; + + const pruneBatches = []; + for (let x = 0; x < folderIds.length; x += PRUNE_FOLDER_BATCH_SIZE) { + const batch = folderIds.slice(x, x + PRUNE_FOLDER_BATCH_SIZE); + pruneBatches.push(batch); + } + + for await (const folderBatch of pruneBatches) { + const rankedSnapshots = (tx || db)(TableName.Snapshot) + .whereIn(`${TableName.Snapshot}.folderId`, folderBatch) + .select( + "folderId", + "id", + (tx || db).raw( + `ROW_NUMBER() OVER (PARTITION BY ${TableName.Snapshot}."folderId" ORDER BY ${TableName.Snapshot}."createdAt" DESC) AS row_num` + ) + ) + .as("ranked_snapshots"); + + const snapshotsToKeep = (tx || db) + .select("id") + .from(rankedSnapshots) + .where( + "row_num", + "<=", + (tx || db) + .select(`${TableName.Project}.pitVersionLimit`) + .from(TableName.Project) + .join(TableName.Environment, `${TableName.Environment}.projectId`, `${TableName.Project}.id`) + .join(TableName.Snapshot, `${TableName.Snapshot}.envId`, `${TableName.Environment}.id`) + .join(rankedSnapshots, "ranked_snapshots.folderId", `${TableName.Snapshot}.folderId`) + .limit(1) + ); + + await (tx || db)(TableName.Snapshot) + .whereIn("folderId", folderBatch) + .whereNotIn("id", snapshotsToKeep) + .delete(); + } + } catch (error) { + throw new DatabaseError({ error, name: "SnapshotPrune" }); + } + }; + return { ...secretSnapshotOrm, findById, findLatestSnapshotByFolderId, findRecursivelySnapshots, countOfSnapshotsByFolderId, - findSecretSnapshotDataById + findSecretSnapshotDataById, + pruneExcessSnapshots }; }; diff --git a/backend/src/server/routes/index.ts b/backend/src/server/routes/index.ts index 1593515ec..b0b2db9b1 100644 --- a/backend/src/server/routes/index.ts +++ b/backend/src/server/routes/index.ts @@ -535,6 +535,7 @@ export const registerRoutes = async ( licenseService, folderDAL, secretDAL, + projectDAL, snapshotDAL, snapshotFolderDAL, snapshotSecretDAL, diff --git a/backend/src/services/project/project-service.ts b/backend/src/services/project/project-service.ts index f58fd7788..1e588dd08 100644 --- a/backend/src/services/project/project-service.ts +++ b/backend/src/services/project/project-service.ts @@ -133,7 +133,8 @@ export const projectServiceFactory = ({ name: workspaceName, orgId: organization.id, slug: projectSlug || slugify(`${workspaceName}-${alphaNumericNanoId(4)}`), - version: ProjectVersion.V2 + version: ProjectVersion.V2, + pitVersionLimit: 10 }, tx ); diff --git a/backend/src/services/secret/secret-service.ts b/backend/src/services/secret/secret-service.ts index 39e47a28e..77c4024ac 100644 --- a/backend/src/services/secret/secret-service.ts +++ b/backend/src/services/secret/secret-service.ts @@ -72,7 +72,7 @@ type TSecretServiceFactoryDep = { secretDAL: TSecretDALFactory; secretTagDAL: TSecretTagDALFactory; secretVersionDAL: TSecretVersionDALFactory; - projectDAL: Pick; + projectDAL: Pick; projectEnvDAL: Pick; folderDAL: Pick< TSecretFolderDALFactory, @@ -1354,7 +1354,16 @@ export const secretServiceFactory = ({ ); ForbiddenError.from(permission).throwUnlessCan(ProjectPermissionActions.Read, ProjectPermissionSub.SecretRollback); - const secretVersions = await secretVersionDAL.find({ secretId }, { offset, limit, sort: [["createdAt", "desc"]] }); + const { pitVersionLimit } = await projectDAL.findById(folder.projectId); + const computedQueryLimit = Math.min(pitVersionLimit - offset, limit); + if (offset > pitVersionLimit || computedQueryLimit <= 0) { + return []; + } + + const secretVersions = await secretVersionDAL.find( + { secretId }, + { offset, limit: computedQueryLimit, sort: [["createdAt", "desc"]] } + ); return secretVersions; }; diff --git a/backend/src/services/secret/secret-version-dal.ts b/backend/src/services/secret/secret-version-dal.ts index 758352ed2..5d52f3c07 100644 --- a/backend/src/services/secret/secret-version-dal.ts +++ b/backend/src/services/secret/secret-version-dal.ts @@ -110,8 +110,45 @@ export const secretVersionDALFactory = (db: TDbClient) => { } }; + const pruneExcessVersions = async (tx?: Knex) => { + try { + const rankedSecretVersions = (tx || db)(TableName.SecretVersion) + .select( + "id", + "secretId", + "folderId", + (tx || db).raw( + `ROW_NUMBER() OVER (PARTITION BY ${TableName.SecretVersion}."secretId" ORDER BY ${TableName.SecretVersion}."createdAt" DESC) AS row_num` + ) + ) + .as("ranked_secret_versions"); + + const versionsToKeep = (tx || db)(rankedSecretVersions) + .select("id") + .where( + "row_num", + "<=", + (tx || db) + .select(`${TableName.Project}.pitVersionLimit`) + .from(TableName.Project) + .join(TableName.Environment, `${TableName.Environment}.projectId`, `${TableName.Project}.id`) + .join(TableName.SecretFolder, `${TableName.SecretFolder}.envId`, `${TableName.Environment}.id`) + .join(rankedSecretVersions, "ranked_secret_versions.folderId", `${TableName.SecretFolder}.id`) + .limit(1) + ); + + await (tx || db)(TableName.SecretVersion).whereNotIn("id", versionsToKeep).delete(); + } catch (error) { + throw new DatabaseError({ + error, + name: "Secret Version Prune" + }); + } + }; + return { ...secretVersionOrm, + pruneExcessVersions, findLatestVersionMany, bulkUpdate, findLatestVersionByFolderId,