From 9117067ab505e08598378884c2b4d6517d5e72dc Mon Sep 17 00:00:00 2001
From: Sheen Capadngan <sheensantoscapadngan@gmail.com>
Date: Fri, 31 May 2024 21:38:16 +0800
Subject: [PATCH] feat: finalized pruning logic

---
 .../services/secret-snapshot/snapshot-dal.ts  | 21 ++++++-----
 .../secret-folder-version-dal.ts              | 36 ++++++++++++++++++-
 .../src/services/secret/secret-version-dal.ts | 22 ++++++------
 3 files changed, 56 insertions(+), 23 deletions(-)

diff --git a/backend/src/ee/services/secret-snapshot/snapshot-dal.ts b/backend/src/ee/services/secret-snapshot/snapshot-dal.ts
index 1618024bb..a35345c17 100644
--- a/backend/src/ee/services/secret-snapshot/snapshot-dal.ts
+++ b/backend/src/ee/services/secret-snapshot/snapshot-dal.ts
@@ -349,20 +349,19 @@ export const snapshotDALFactory = (db: TDbClient) => {
           )
           .as("ranked_snapshots");
 
+        const folderLimits = (tx || db)(TableName.Snapshot)
+          .join(TableName.Environment, `${TableName.Environment}.id`, `${TableName.Snapshot}.envId`)
+          .join(TableName.Project, `${TableName.Project}.id`, `${TableName.Environment}.projectId`)
+          .whereIn(`${TableName.Snapshot}.folderId`, folderBatch)
+          .groupBy(`${TableName.Snapshot}.folderId`, `${TableName.Project}.pitVersionLimit`)
+          .select("folderId", "pitVersionLimit")
+          .as("folder_limits");
+
         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)
-          );
+          .join(folderLimits, "folder_limits.folderId", "ranked_snapshots.folderId")
+          .whereRaw(`ranked_snapshots.row_num <= folder_limits."pitVersionLimit"`);
 
         await (tx || db)(TableName.Snapshot)
           .whereIn("folderId", folderBatch)
diff --git a/backend/src/services/secret-folder/secret-folder-version-dal.ts b/backend/src/services/secret-folder/secret-folder-version-dal.ts
index f133308cf..cfc370640 100644
--- a/backend/src/services/secret-folder/secret-folder-version-dal.ts
+++ b/backend/src/services/secret-folder/secret-folder-version-dal.ts
@@ -62,5 +62,39 @@ export const secretFolderVersionDALFactory = (db: TDbClient) => {
     }
   };
 
-  return { ...secretFolderVerOrm, findLatestFolderVersions, findLatestVersionByFolderId };
+  const pruneExcessVersions = async (tx?: Knex) => {
+    try {
+      const rankedFolderVersions = (tx || db)(TableName.SecretFolderVersion)
+        .select(
+          "id",
+          "folderId",
+          (tx || db).raw(
+            `ROW_NUMBER() OVER (PARTITION BY ${TableName.SecretFolderVersion}."folderId" ORDER BY ${TableName.SecretFolderVersion}."createdAt" DESC) AS row_num`
+          )
+        )
+        .as("ranked_folder_versions");
+
+      const folderLimits = (tx || db)(TableName.SecretFolderVersion)
+        .join(TableName.Environment, `${TableName.Environment}.id`, `${TableName.SecretFolderVersion}.envId`)
+        .join(TableName.Project, `${TableName.Project}.id`, `${TableName.Environment}.projectId`)
+        .groupBy(`${TableName.SecretFolderVersion}.folderId`, `${TableName.Project}.pitVersionLimit`)
+        .select("folderId", "pitVersionLimit")
+        .as("folder_limits");
+
+      const versionsToKeep = (tx || db)(rankedFolderVersions)
+        .select("id")
+        .from(rankedFolderVersions)
+        .join(folderLimits, "folder_limits.folderId", "ranked_folder_versions.folderId")
+        .whereRaw(`ranked_folder_versions.row_num <= folder_limits."pitVersionLimit"`);
+
+      await (tx || db)(TableName.SecretFolderVersion).whereNotIn("id", versionsToKeep).delete();
+    } catch (error) {
+      throw new DatabaseError({
+        error,
+        name: "Secret Version Prune"
+      });
+    }
+  };
+
+  return { ...secretFolderVerOrm, findLatestFolderVersions, findLatestVersionByFolderId, pruneExcessVersions };
 };
diff --git a/backend/src/services/secret/secret-version-dal.ts b/backend/src/services/secret/secret-version-dal.ts
index 5d52f3c07..374907c3d 100644
--- a/backend/src/services/secret/secret-version-dal.ts
+++ b/backend/src/services/secret/secret-version-dal.ts
@@ -123,19 +123,19 @@ export const secretVersionDALFactory = (db: TDbClient) => {
         )
         .as("ranked_secret_versions");
 
+      const folderLimits = (tx || db)(TableName.SecretVersion)
+        .join(TableName.SecretFolder, `${TableName.SecretFolder}.id`, `${TableName.SecretVersion}.folderId`)
+        .join(TableName.Environment, `${TableName.Environment}.id`, `${TableName.SecretFolder}.envId`)
+        .join(TableName.Project, `${TableName.Project}.id`, `${TableName.Environment}.projectId`)
+        .groupBy(`${TableName.SecretVersion}.folderId`, `${TableName.Project}.pitVersionLimit`)
+        .select("folderId", "pitVersionLimit")
+        .as("folder_limits");
+
       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)
-        );
+        .from(rankedSecretVersions)
+        .join(folderLimits, "folder_limits.folderId", "ranked_secret_versions.folderId")
+        .whereRaw(`ranked_secret_versions.row_num <= folder_limits."pitVersionLimit"`);
 
       await (tx || db)(TableName.SecretVersion).whereNotIn("id", versionsToKeep).delete();
     } catch (error) {