refactor(core): refactor findPayloadByPayloadField query (#7650)

* refactor(core): refactor findPayloadByPayloadField query

refactor findPayloadByPayloadField query

* chore: update comments

Co-authored-by: Gao Sun <[email protected]>

---------

Co-authored-by: Gao Sun <[email protected]>

Author: simeng-li

packages/core/src/queries/oidc-model-instance.test.ts

@@ -40,6 +40,10 @@ describe('oidc-model-instance query', () => {
     payload: JSON.stringify(instance.payload),
   };
 
+  beforeEach(() => {
+    jest.clearAllMocks();
+  });
+
   it('upsertInstance', async () => {
     const expectSql = sql`
       insert into ${table} ("model_name", "id", "payload", "expires_at")
@@ -83,21 +87,13 @@ describe('oidc-model-instance query', () => {
     });
   });
 
-  it('findPayloadByPayloadField', async () => {
+  it('findPayloadByPayloadField - single result', async () => {
     const uid_key = 'uid';
     const uid_value = 'foo';
 
-    const expectSql = sql`
-      select ${fields.payload}, ${fields.consumedAt}
-      from ${table}
-      where ${fields.modelName}=$1
-      and ${fields.payload}->>$2=$3
-    `;
-
+    // Mock a single result
     mockQuery.mockImplementationOnce(async (sql, values) => {
-      expectSqlAssert(sql, expectSql.sql);
-      expect(values).toEqual([instance.modelName, uid_key, uid_value]);
-
+      // Simulate pool.any
       return createMockQueryResult([{ consumedAt: 10 }]);
     });
 
@@ -108,6 +104,44 @@ describe('oidc-model-instance query', () => {
     });
   });
 
+  it('findPayloadByPayloadField - multiple results (should delete and return null)', async () => {
+    const uid_key = 'uid';
+    const uid_value = 'foo';
+
+    // Mock multiple results
+    mockQuery
+      .mockImplementationOnce(async () => {
+        // Simulate pool.any
+        return createMockQueryResult([{ uid: uid_value }, { uid: uid_value }]);
+      })
+      // @ts-expect-error - mock delete query
+      .mockImplementationOnce(async () => {
+        // Simulate delete query
+        return { rowCount: 2 };
+      });
+
+    await expect(
+      findPayloadByPayloadField(instance.modelName, uid_key, uid_value)
+    ).resolves.toBeUndefined();
+
+    expect(mockQuery).toHaveBeenCalledTimes(2);
+  });
+
+  it('findPayloadByPayloadField - no result', async () => {
+    const uid_key = 'uid';
+    const uid_value = 'foo';
+
+    // Mock no result
+    mockQuery.mockImplementationOnce(async () => {
+      // Simulate pool.any
+      return createMockQueryResult([]);
+    });
+
+    await expect(
+      findPayloadByPayloadField(instance.modelName, uid_key, uid_value)
+    ).resolves.toBeUndefined();
+  });
+
   it('consumeInstanceById', async () => {
     jest.useFakeTimers().setSystemTime(100_000);
 

packages/core/src/queries/oidc-model-instance.ts

@@ -79,12 +79,26 @@ export const createOidcModelInstanceQueries = (pool: CommonQueryMethods) => {
     field: Field,
     value: T
   ) => {
-    const result = await pool.maybeOne<QueryResult>(sql`
-      ${findByModel(modelName)}
-      and ${fields.payload}->>${field}=${value}
+    // Fetch all matching records
+    const results = await pool.any<QueryResult>(sql`
+    ${findByModel(modelName)}
+    and ${fields.payload}->>${field}=${value}
+  `);
+
+    // Rarely, duplicate UIDs can exist for different sessions.
+    // This query may throw `DataIntegrityError`.
+    // If that happens, delete all duplicates and return `null`.
+    if (results.length > 1) {
+      // Delete all duplicates
+      await pool.query(sql`
+      delete from ${sql.identifier([modelName])}
+      where ${fields.payload}->>${field}=${value}
     `);
+      return;
+    }
 
-    return convertResult(result, modelName);
+    // If there is only one record, return the result.
+    return results[0] ? convertResult(results[0], modelName) : undefined;
   };
 
   const consumeInstanceById = async (modelName: string, id: string) => {