From 128b2062f3c5e651027a54bf4e8470aad031e711 Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Thu, 13 Aug 2026 10:18:45 +0530 Subject: [PATCH] PG write query timeouts --- .../FlatCollectionWriteTest.java | 81 +++++++++++++++++++ .../postgres/FlatPostgresCollection.java | 26 +++--- .../postgres/PostgresQueryExecutor.java | 9 +++ .../postgres/PostgresQueryExecutorTest.java | 26 ++++++ 4 files changed, 130 insertions(+), 12 deletions(-) diff --git a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java index 5011182f7..d1721a837 100644 --- a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java +++ b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java @@ -3,6 +3,7 @@ import static org.hypertrace.core.documentstore.utils.Utils.readFileFromResource; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; @@ -14,6 +15,7 @@ import com.typesafe.config.ConfigFactory; import java.io.IOException; import java.sql.Connection; +import java.sql.DriverManager; import java.sql.PreparedStatement; import java.sql.ResultSet; import java.util.ArrayList; @@ -47,6 +49,8 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.ArgumentsSource; +import org.postgresql.util.PSQLException; +import org.postgresql.util.PSQLState; import org.testcontainers.junit.jupiter.Testcontainers; @Testcontainers @@ -4023,6 +4027,83 @@ private long getLastUpdatedEpoch(Key key) throws Exception { } } + @Test + @DisplayName("upsert on a row locked by another txn is aborted after query timeout") + void upsertHonorsQueryTimeoutOnRowLock() throws Exception { + // setup + String docId = generateDocId("qt-upsert"); + Key key = new SingleValueKey(DEFAULT_TENANT, docId); + ObjectNode initial = OBJECT_MAPPER.createObjectNode(); + initial.put("id", docId); + initial.put("item", "Seed"); + initial.put("price", 1); + flatCollection.upsert(key, new JSONDocument(initial)); + + // 2) Build a collection whose datastore has a 1s query timeout. + Collection timeoutCollection = getFlatCollectionWithQueryTimeout("1 second"); + + // 3) Hold a FOR UPDATE lock on the seeded row from a separate raw connection. + String url = + String.format( + "jdbc:postgresql://localhost:%s/postgres", postgresContainer.getMappedPort(5432)); + try (Connection lockConn = DriverManager.getConnection(url, "postgres", "postgres")) { + lockConn.setAutoCommit(false); + try (PreparedStatement lockPs = + lockConn.prepareStatement( + String.format( + "SELECT id FROM \"%s\" WHERE id = ? FOR UPDATE", FLAT_COLLECTION_NAME))) { + lockPs.setString(1, key.toString()); + try (ResultSet rs = lockPs.executeQuery()) { + assertTrue(rs.next(), "Precondition failure: Could not acquire lock on the seed row"); + } + } + + // 4) The upsert must block on the row lock and then be cancelled by + // JDBC's setQueryTimeout(1). FlatPostgresCollection wraps SQLException as IOException. + ObjectNode updateNode = OBJECT_MAPPER.createObjectNode(); + updateNode.put("id", docId); + updateNode.put("item", "Updated"); + updateNode.put("price", 2); + + long startNs = System.nanoTime(); + IOException thrown = + assertThrows( + IOException.class, () -> timeoutCollection.upsert(key, new JSONDocument(updateNode))); + long elapsedMs = (System.nanoTime() - startNs) / 1_000_000L; + + // Sanity: cancellation should fire quickly - well before the row lock would ever + // be released (this txn is held open until the try-with-resources closes). + assertTrue(elapsedMs < 15_000); + + Throwable cause = thrown.getCause(); + assertNotNull(cause); + assertInstanceOf(PSQLException.class, cause); + assertEquals(PSQLState.QUERY_CANCELED.getState(), ((PSQLException) cause).getSQLState()); + + // 5) Verify the seeded row was not modified (lock still held, timed-out UPDATE + // never committed). + queryAndAssert( + key, + rs -> { + assertTrue(rs.next()); + assertEquals("Seed", rs.getString("item")); + assertEquals(1, rs.getInt("price")); + }); + } + } + + private Collection getFlatCollectionWithQueryTimeout(String queryTimeout) { + String postgresConnectionUrl = + String.format("jdbc:postgresql://localhost:%s/", postgresContainer.getMappedPort(5432)); + Map cfg = new HashMap<>(); + cfg.put("url", postgresConnectionUrl); + cfg.put("user", "postgres"); + cfg.put("password", "postgres"); + cfg.put("queryTimeout", queryTimeout); + Datastore ds = DatastoreProvider.getDatastore("Postgres", ConfigFactory.parseMap(cfg)); + return ds.getCollectionForType(FLAT_COLLECTION_NAME, DocumentType.FLAT); + } + private static void executeInsertStatements() { PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; try { diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java index 19927abb0..100449189 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java @@ -253,7 +253,8 @@ public boolean delete(Key key) { String.format( "DELETE FROM %s WHERE %s = ?", tableIdentifier, PostgresUtils.wrapFieldNamesWithDoubleQuotes(pkForTable)); - try (PreparedStatement preparedStatement = client.getConnection().prepareStatement(deleteSQL)) { + try (PreparedStatement preparedStatement = + queryExecutor.prepareStatementWithTimeout(client.getConnection(), deleteSQL)) { preparedStatement.setString(1, key.toString()); int rowsDeleted = preparedStatement.executeUpdate(); return rowsDeleted > 0; @@ -320,7 +321,7 @@ public BulkDeleteResult delete(Set keys) { LOGGER.debug("Bulk delete SQL: {}", deleteSQL); try (Connection conn = client.getPooledConnection(); - PreparedStatement ps = conn.prepareStatement(deleteSQL)) { + PreparedStatement ps = queryExecutor.prepareStatementWithTimeout(conn, deleteSQL)) { int deletedCount = ps.executeUpdate(); LOGGER.debug("Bulk deleted {} rows", deletedCount); return new BulkDeleteResult(deletedCount); @@ -336,7 +337,7 @@ public boolean deleteAll() { LOGGER.debug("Delete all SQL: {}", deleteSQL); try (Connection conn = client.getPooledConnection(); - PreparedStatement ps = conn.prepareStatement(deleteSQL)) { + PreparedStatement ps = queryExecutor.prepareStatementWithTimeout(conn, deleteSQL)) { int deletedCount = ps.executeUpdate(); LOGGER.debug("Deleted all {} rows", deletedCount); return true; @@ -402,7 +403,7 @@ public boolean bulkUpsert(Map documents) { LOGGER.debug("Bulk upsert SQL: {}", sql); try (Connection conn = client.getPooledConnection(); - PreparedStatement ps = conn.prepareStatement(sql)) { + PreparedStatement ps = queryExecutor.prepareStatementWithTimeout(conn, sql)) { for (Map.Entry entry : parsedDocuments.entrySet()) { TypedDocument parsed = entry.getValue(); @@ -499,7 +500,7 @@ public boolean bulkCreateOrReplace(Map documents) { LOGGER.debug("Bulk createOrReplace SQL: {}", sql); try (Connection conn = client.getPooledConnection(); - PreparedStatement ps = conn.prepareStatement(sql)) { + PreparedStatement ps = queryExecutor.prepareStatementWithTimeout(conn, sql)) { for (Map.Entry entry : parsedDocuments.entrySet()) { TypedDocument parsed = entry.getValue(); @@ -680,7 +681,8 @@ private PreparedStatement getPreparedStatementForQuery( throws SQLException { String selectQuery = String.format("SELECT * FROM %s WHERE %s = ANY(?)", tableIdentifier, quotedPkColumn); - PreparedStatement preparedStatement = connection.prepareStatement(selectQuery); + PreparedStatement preparedStatement = + queryExecutor.prepareStatementWithTimeout(connection, selectQuery); String[] keyArray = documents.keySet().stream().map(Key::toString).toArray(String[]::new); Array sqlArray = connection.createArrayOf(pkType.getSqlType(), keyArray); @@ -959,7 +961,7 @@ private boolean executeKeyUpdate( LOGGER.debug("Executing key update SQL: {}", sql); - try (PreparedStatement ps = connection.prepareStatement(sql)) { + try (PreparedStatement ps = queryExecutor.prepareStatementWithTimeout(connection, sql)) { int idx = 1; for (Object param : params) { ps.setObject(idx++, param); @@ -1058,7 +1060,7 @@ private int executeBatchUpdate( LOGGER.debug("Executing batch update SQL: {} for {} keys", sql, keys.size()); - try (PreparedStatement ps = connection.prepareStatement(sql)) { + try (PreparedStatement ps = queryExecutor.prepareStatementWithTimeout(connection, sql)) { for (int i : order) { int idx = 1; for (Object param : allKeyParams.get(i)) { @@ -1264,7 +1266,7 @@ private void executeUpdate( LOGGER.debug("Executing update SQL: {}", sql); - try (PreparedStatement ps = connection.prepareStatement(sql)) { + try (PreparedStatement ps = queryExecutor.prepareStatementWithTimeout(connection, sql)) { int idx = 1; for (Object param : params) { ps.setObject(idx++, param); @@ -1580,7 +1582,7 @@ Object convertTimestampForType(long epochMillis, PostgresDataType type) { private int executeUpdate(String sql, TypedDocument parsed) throws SQLException { try (Connection conn = client.getPooledConnection(); - PreparedStatement ps = conn.prepareStatement(sql)) { + PreparedStatement ps = queryExecutor.prepareStatementWithTimeout(conn, sql)) { int index = 1; for (String column : parsed.getColumns()) { setParameter( @@ -1749,7 +1751,7 @@ private boolean executeUpsert(String sql, List allColumns, TypedDocument long ta = System.nanoTime(); try (Connection conn = client.getPooledConnection()) { long tb = System.nanoTime(); - try (PreparedStatement ps = conn.prepareStatement(sql)) { + try (PreparedStatement ps = queryExecutor.prepareStatementWithTimeout(conn, sql)) { int index = 1; Set parsedColumns = new HashSet<>(parsed.getColumns()); for (String column : allColumns) { @@ -1778,7 +1780,7 @@ private boolean executeUpsert(String sql, List allColumns, TypedDocument private boolean executeUpsertReturningIsInsert( String sql, List allColumns, TypedDocument parsed) throws SQLException { try (Connection conn = client.getPooledConnection(); - PreparedStatement ps = conn.prepareStatement(sql)) { + PreparedStatement ps = queryExecutor.prepareStatementWithTimeout(conn, sql)) { int index = 1; Set parsedColumns = new HashSet<>(parsed.getColumns()); for (String column : allColumns) { diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresQueryExecutor.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresQueryExecutor.java index 9f7cd1748..e44c63742 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresQueryExecutor.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresQueryExecutor.java @@ -53,6 +53,15 @@ public PreparedStatement buildPreparedStatement( return buildPreparedStatement(sqlQuery, params, connection, this.queryTimeoutSeconds); } + public PreparedStatement prepareStatementWithTimeout(Connection connection, String sqlQuery) + throws SQLException { + PreparedStatement preparedStatement = connection.prepareStatement(sqlQuery); + if (queryTimeoutSeconds > 0) { + preparedStatement.setQueryTimeout(queryTimeoutSeconds); + } + return preparedStatement; + } + public PreparedStatement buildPreparedStatement( String sqlQuery, Params params, Connection connection, int queryTimeoutSeconds) throws SQLException { diff --git a/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresQueryExecutorTest.java b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresQueryExecutorTest.java index 41ba3f4b7..e78e4438b 100644 --- a/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresQueryExecutorTest.java +++ b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresQueryExecutorTest.java @@ -1,5 +1,8 @@ package org.hypertrace.core.documentstore.postgres; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -43,4 +46,27 @@ void testPreparedStatementUsesOverridenTimeout() throws SQLException { executor.buildPreparedStatement(sqlQuery, params, mockConnection, 45); verify(mockPreparedStatement).setQueryTimeout(45); } + + @Test + void prepareStatementWithTimeoutAppliesConfiguredTimeout() throws SQLException { + String sqlQuery = "UPDATE foo SET x = ?"; + when(mockConnection.prepareStatement(sqlQuery)).thenReturn(mockPreparedStatement); + + PostgresQueryExecutor executor = new PostgresQueryExecutor(45); + PreparedStatement result = executor.prepareStatementWithTimeout(mockConnection, sqlQuery); + + assertSame(mockPreparedStatement, result); + verify(mockPreparedStatement).setQueryTimeout(45); + } + + @Test + void prepareStatementWithTimeoutSkipsSetWhenTimeoutIsZero() throws SQLException { + String sqlQuery = "DELETE FROM foo"; + when(mockConnection.prepareStatement(sqlQuery)).thenReturn(mockPreparedStatement); + + PostgresQueryExecutor executor = new PostgresQueryExecutor(0); + executor.prepareStatementWithTimeout(mockConnection, sqlQuery); + + verify(mockPreparedStatement, never()).setQueryTimeout(anyInt()); + } }