Fix failing test cases

Author: suddendust

document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java

@@ -1091,7 +1091,7 @@ private boolean createOrReplaceWithRetry(Key key, Document document, boolean isR
       return executeUpsertReturningIsInsert(sql, parsed);
 
     } catch (PSQLException e) {
-      return handlePSQLExceptionForUpsert(e, key, document, tableName, isRetry);
+      return handlePSQLExceptionForCreateOrReplace(e, key, document, tableName, isRetry);
     } catch (SQLException e) {
       LOGGER.error("SQLException in createOrReplace. key: {} content: {}", key, document, e);
       throw new IOException(e);

document-store/src/test/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollectionTest.java

@@ -265,7 +265,6 @@ void testUpsertRetriesOnUndefinedColumn() throws Exception {
       PSQLException psqlException = createPSQLException(PSQLState.UNDEFINED_COLUMN);
       when(mockPreparedStatement.executeQuery()).thenThrow(psqlException).thenReturn(mockResultSet);
       when(mockResultSet.next()).thenReturn(true);
-      when(mockResultSet.getBoolean("is_insert")).thenReturn(true);
 
       doNothing().when(mockSchemaRegistry).invalidate(COLLECTION_NAME);
 
@@ -289,13 +288,13 @@ void testUpsertRetriesOnDatatypeMismatch() throws Exception {
       PSQLException psqlException = createPSQLException(PSQLState.DATATYPE_MISMATCH);
       when(mockPreparedStatement.executeQuery()).thenThrow(psqlException).thenReturn(mockResultSet);
       when(mockResultSet.next()).thenReturn(true);
-      when(mockResultSet.getBoolean("is_insert")).thenReturn(false);
 
       doNothing().when(mockSchemaRegistry).invalidate(COLLECTION_NAME);
 
       boolean result = flatPostgresCollection.upsert(key, document);
 
-      assertFalse(result);
+      // upsert always returns true if it succeeds
+      assertTrue(result);
       verify(mockSchemaRegistry, times(1)).invalidate(COLLECTION_NAME);
       verify(mockPreparedStatement, times(2)).executeQuery();
     }