Clean up + tested

dust-tt · Nov 13, 2023 · 6f8fbf7 · 6f8fbf7
1 parent c6436c8
commit 6f8fbf7
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 12 deletions.
diff --git a/front/production_checks/checks/managed_data_source_gdrive_gc.ts b/front/production_checks/checks/managed_data_source_gdrive_gc.ts
@@ -1,10 +1,12 @@
 import { QueryTypes, Sequelize } from "sequelize";
 
-import { DataSource } from "@app/lib/models/data_source";
 import { getCoreDocuments } from "@app/production_checks/lib/managed_ds";
 import { CheckFunction } from "@app/production_checks/types/check";
 
-const { CONNECTORS_DATABASE_READ_REPLICA_URI } = process.env;
+const {
+  CONNECTORS_DATABASE_READ_REPLICA_URI,
+  FRONT_DATABASE_READ_REPLICA_URI,
+} = process.env;
 
 export const managedDataSourceGCGdriveCheck: CheckFunction = async (
   checkName,
@@ -17,11 +19,18 @@ export const managedDataSourceGCGdriveCheck: CheckFunction = async (
       logging: false,
     }
   );
-  const GdriveDataSources = await DataSource.findAll({
-    where: {
-      connectorProvider: "google_drive",
-    },
-  });
+  const frontSequelize = new Sequelize(
+    FRONT_DATABASE_READ_REPLICA_URI as string,
+    {
+      logging: false,
+    }
+  );
+  const GdriveDataSources: { id: number; connectorId: string }[] =
+    await frontSequelize.query(
+      `SELECT id, "connectorId" FROM data_sources WHERE "connectorProvider" = 'google_drive'`,
+      { type: QueryTypes.SELECT }
+    );
+
   for (const ds of GdriveDataSources) {
     const coreDocumentsRes = await getCoreDocuments(ds.id);
     if (coreDocumentsRes.isErr()) {
@@ -48,12 +57,12 @@ export const managedDataSourceGCGdriveCheck: CheckFunction = async (
     const connectorDocumentIds = new Set(
       connectorDocuments.map((d) => d.coreDocumentId)
     );
-    const missingDocuments = coreDocumentIds.filter(
+    const notDeleted = coreDocumentIds.filter(
       (coreId) => !connectorDocumentIds.has(coreId)
     );
-    if (missingDocuments.length > 0) {
+    if (notDeleted.length > 0) {
       reportFailure(
-        { missingDocuments },
+        { notDeleted },
         "Google Drive documents not properly Garbage collected"
       );
     } else {

diff --git a/front/production_checks/lib/managed_ds.ts b/front/production_checks/lib/managed_ds.ts
@@ -27,9 +27,14 @@ export async function getCoreDocuments(
   );
 
   const managedDsData = await front_sequelize.query(
-    'SELECT id, "connectorId", "connectorProvider", "dustAPIProjectId"\
+    'SELECT id, "connectorId", "connectorProvider", "dustAPIProjectId" \
          FROM data_sources WHERE id = :frontDataSourceId',
-    { type: QueryTypes.SELECT }
+    {
+      type: QueryTypes.SELECT,
+      replacements: {
+        frontDataSourceId: frontDataSourceId,
+      },
+    }
   );
   const managedDs = managedDsData as {
     id: number;