From c685c4982e70a6f91d29ecb3b73ee10637bec08d Mon Sep 17 00:00:00 2001
From: Isaac Pedisich <isaac@pipe.com>
Date: Fri, 21 Feb 2025 15:04:35 -0500
Subject: [PATCH] Add a "resync" endpoint to pull in existing tags

When you ping api/resync, it runs a scan of all columns that don't have
an existing policy tag, checks if they match, and marks them "appropriately"
in the database.

Right now it's marking all columns that it _couldn't_ tag as benign -
might want it to mark them as "unknown" rather than explcitly having
them be benign.

Also, may want a UI for this endpoint rather than having to hit it
manually, but it is serving our purposes for the moment!
---
 src/app/api/resync/route.ts        | 36 +++++++++++++++++++++
 src/dataplatforms/BigQuery.tsx     | 51 +++++++++++++++++++++++++++++-
 src/dataplatforms/DataPlatform.tsx | 10 ++++++
 3 files changed, 96 insertions(+), 1 deletion(-)
 create mode 100644 src/app/api/resync/route.ts

diff --git a/src/app/api/resync/route.ts b/src/app/api/resync/route.ts
new file mode 100644
index 0000000..4e2aee6
--- /dev/null
+++ b/src/app/api/resync/route.ts
@@ -0,0 +1,36 @@
+import { DataPlatform } from "@/dataplatforms/DataPlatform";
+import { prisma } from "@/lib/utils";
+
+export async function GET() {
+  // Re-checks policyTagDecisions:
+  // If it has not been accepted but the policy tag is alredy applied, mark it as if it was accepted
+  // TODO: If it was accepted and the policy tag is not applied, removes that decision
+  
+  const columnsWithoutPolicyTag = await prisma.columnClassification.findMany({
+    where: {
+      PolicyTagDecision: {
+        none: {},
+      },
+    },
+    include: {
+      column: true,
+    },
+  });
+
+  const DATA_PLATFORM = DataPlatform.getInstance();
+  for (const untagged_column of columnsWithoutPolicyTag) {
+    const column = untagged_column.column;
+    if (! await DATA_PLATFORM.needsPolicyTag(column.datasetId, column.tableName, column.name, process.env.PII_POLICY_TAG_ID!)) {
+        
+        console.log(`Marking column ${column.datasetId}.${column.tableName}.${column.name} as if it has been accepted`)
+          await prisma.policyTagDecision.create({
+            data: {
+              columnId: column.id,
+              decision: true
+            }
+          })
+        }
+    
+    }
+    return new Response(JSON.stringify({status: 200}));
+}
diff --git a/src/dataplatforms/BigQuery.tsx b/src/dataplatforms/BigQuery.tsx
index 1d43045..6ba67b7 100644
--- a/src/dataplatforms/BigQuery.tsx
+++ b/src/dataplatforms/BigQuery.tsx
@@ -1,5 +1,5 @@
 import { DataPlatform, TableDataType } from "./DataPlatform";
-import { BigQuery } from "@google-cloud/bigquery";
+import { BigQuery, Dataset } from "@google-cloud/bigquery";
 import {
   GetTablesResponse,
   TableMetadata,
@@ -97,4 +97,53 @@ export class BigQueryPlatform extends DataPlatform {
 
     return rows;
   }
+
+  async needsPolicyTag(
+    datasetId: string,
+    tableName: string,
+    columnName: string,
+    policyTagId: string,
+  ): Promise<boolean> {
+    const dataset: Dataset = bigQueryClient.dataset(datasetId);
+    let tableRef : any
+    try {
+      [tableRef] = await dataset.table(tableName).get();
+    } catch (error) {
+      const errorMessage = (error as Error).message
+      if (errorMessage.includes('Not found: Dataset')) {
+        console.warn(`Dataset ${datasetId} does not exist`)
+        return false;
+      }
+      console.error(errorMessage)
+
+    }
+    if (!tableRef?.id) {
+      console.warn(`TableRef ${datasetId}.${tableName} not found`)
+      return false;
+    }
+    const [metadata]: TableMetadata[] = await tableRef.getMetadata();
+    if (metadata.type !== "TABLE") {
+      // NOTE: This will also mark "external" tables checked -- those _can_ get policy tags but it is more complex
+      console.log(`Non-table entry found of type ${metadata.type} in table ${datasetId}.${tableName}`)
+      return false;
+    }
+
+    if (!metadata.schema?.fields) {
+      console.warn(`Schema for ${datasetId}.${tableName} not found`)
+      return false;
+    }
+
+    const column = metadata.schema.fields.find((field: TableField) => field.name === columnName);
+    if (column === undefined) {
+    console.warn(`Column ${datasetId}.${tableName}.${columnName} not found`)
+      return false;
+    }
+    if (column?.policyTags?.names?.includes(policyTagId)) {
+      console.log(`Column ${datasetId}.${tableName}.${columnName} already has an appropriate policy tag`)
+      return false;
+    }
+    return true;
+
+  }
+
 }
diff --git a/src/dataplatforms/DataPlatform.tsx b/src/dataplatforms/DataPlatform.tsx
index e0e1dcd..e41926c 100644
--- a/src/dataplatforms/DataPlatform.tsx
+++ b/src/dataplatforms/DataPlatform.tsx
@@ -39,5 +39,15 @@ export abstract class DataPlatform {
     columnName: string,
     policyTagId: string,
   ): Promise<void>;
+
+  // Returns "True" if the column does not yet have the policy tag and is able to be tagged
+  abstract needsPolicyTag(
+    datasetId: string,
+    tableName: string,
+    columnName: string,
+    policyTagId: string,
+  ): Promise<boolean>;
+
+
   abstract getSampleData(datasetId: string, tableName: string): Promise<any[]>;
 }