[OTE-308] Update upsert querybuilder and compliance table upsert func…

…tion (#1477) * Update upsert querybuilder and compliance data upsert function * lint
dydxprotocol · May 8, 2024 · a699c24 · a699c24
1 parent eadfc82
commit a699c24
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 27 deletions.
diff --git 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
@@ -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
@@ -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(