From f664999a30c1f78c1bb4039e837ec14fb622cfe2 Mon Sep 17 00:00:00 2001
From: "mergify[bot]" <37929162+mergify[bot]@users.noreply.github.com>
Date: Wed, 8 May 2024 23:48:30 +0530
Subject: [PATCH] [OTE-308] Update upsert querybuilder and compliance table
 upsert function (#1477) (#1478)

* Update upsert querybuilder and compliance data upsert function

* lint

(cherry picked from commit a699c24637d336f5336ce5da30e689a81e54e2a6)

Co-authored-by: Mohammed Affan <affanmd@nyu.edu>
---
 .../src/models/compliance-data-model.ts       |  5 ++---
 .../postgres/src/query-builders/upsert.ts     | 15 ++++++++-----
 .../postgres/src/stores/compliance-table.ts   | 22 ++++---------------
 3 files changed, 15 insertions(+), 27 deletions(-)

diff --git a/indexer/packages/postgres/src/models/compliance-data-model.ts b/indexer/packages/postgres/src/models/compliance-data-model.ts
index 1ba8fbd5e1..eee0fb71d3 100644
--- a/indexer/packages/postgres/src/models/compliance-data-model.ts
+++ b/indexer/packages/postgres/src/models/compliance-data-model.ts
@@ -1,10 +1,9 @@
-import { Model } from 'objection';
-
 import { NumericPattern } from '../lib/validators';
 import UpsertQueryBuilder from '../query-builders/upsert';
 import { ComplianceProvider, IsoString } from '../types';
+import BaseModel from './base-model';
 
-export default class ComplianceDataModel extends Model {
+export default class ComplianceDataModel extends BaseModel {
   static get tableName() {
     return 'compliance_data';
   }
diff --git a/indexer/packages/postgres/src/query-builders/upsert.ts b/indexer/packages/postgres/src/query-builders/upsert.ts
index 3b7dd2bfce..e0350dac38 100644
--- a/indexer/packages/postgres/src/query-builders/upsert.ts
+++ b/indexer/packages/postgres/src/query-builders/upsert.ts
@@ -14,9 +14,8 @@ export default class UpsertQueryBuilder<M extends Model, R = M[]> extends QueryB
   upsert(object: any) {
     const modelClass = this.modelClass();
 
-    const idColumn: string = modelClass.idColumn as string;
-
-    const tableDefinedId = `${modelClass.tableName}.${idColumn}`;
+    const idColumn = modelClass.idColumn;
+    const idColumns: string[] = Array.isArray(idColumn) ? idColumn : [idColumn];
 
     const knex = modelClass.knex();
 
@@ -26,11 +25,16 @@ export default class UpsertQueryBuilder<M extends Model, R = M[]> extends QueryB
     const colBindings = cols.map(() => '??').join(', ');
     const valBindings = cols.map(() => '?').join(', ');
     const setBindings = cols.map(() => '?? = ?').join(', ');
+    const idConditionBindings = idColumns.map(() => '?? = ?').join(' AND ');
 
     const setValues: string[] = [];
     for (let i = 0; i < cols.length; ++i) {
       setValues.push(cols[i], values[i]);
     }
+    const idValues: string[] = [];
+    for (let i = 0; i < idColumns.length; ++i) {
+      idValues.push(`${modelClass.tableName}.${idColumns[i]}`, object[idColumns[i]]);
+    }
 
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     return this.onBuildKnex((query: any) => {
@@ -40,15 +44,14 @@ export default class UpsertQueryBuilder<M extends Model, R = M[]> extends QueryB
             `(${colBindings}) VALUES (${valBindings})`,
             'ON CONFLICT (??) DO',
             `UPDATE SET ${setBindings}`,
-            'WHERE ?? = ?',
+            `WHERE ${idConditionBindings}`,
           ].join(' '),
           [
             ...cols,
             ...values,
             modelClass.idColumn,
             ...setValues,
-            tableDefinedId,
-            object[idColumn],
+            ...idValues,
           ],
         ),
       );
diff --git a/indexer/packages/postgres/src/stores/compliance-table.ts b/indexer/packages/postgres/src/stores/compliance-table.ts
index 6c1f8103e8..162edff16d 100644
--- a/indexer/packages/postgres/src/stores/compliance-table.ts
+++ b/indexer/packages/postgres/src/stores/compliance-table.ts
@@ -123,25 +123,11 @@ export async function upsert(
   complianceDataToUpsert: ComplianceDataCreateObject,
   options: Options = { txId: undefined },
 ): Promise<ComplianceDataFromDatabase> {
-  const complianceData: ComplianceDataFromDatabase | undefined = await findByAddressAndProvider(
-    complianceDataToUpsert.address,
-    complianceDataToUpsert.provider,
-  );
-  if (complianceData === undefined) {
-    return create({
-      ...complianceDataToUpsert,
-    }, options);
-  }
-
-  const updatedComplianceData: ComplianceDataFromDatabase | undefined = await update({
-    ...complianceDataToUpsert,
-  }, options);
-
-  if (updatedComplianceData === undefined) {
-    throw Error('order must exist after update');
-  }
+  const updatedComplianceData: ComplianceDataModel[] = await ComplianceDataModel.query(
+    Transaction.get(options.txId),
+  ).upsert(complianceDataToUpsert).returning('*');
 
-  return updatedComplianceData;
+  return updatedComplianceData[0];
 }
 
 export async function findByAddressAndProvider(