Skip to content

Commit 330cbc2

Browse files
suddendustsuddendust
authored
PG write query timeouts (#321)
1 parent c534f39 commit 330cbc2

4 files changed

Lines changed: 130 additions & 12 deletions

File tree

‎document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import static org.hypertrace.core.documentstore.utils.Utils.readFileFromResource;
44
import static org.junit.jupiter.api.Assertions.assertEquals;
55
import static org.junit.jupiter.api.Assertions.assertFalse;
6+
import static org.junit.jupiter.api.Assertions.assertInstanceOf;
67
import static org.junit.jupiter.api.Assertions.assertNotNull;
78
import static org.junit.jupiter.api.Assertions.assertNull;
89
import static org.junit.jupiter.api.Assertions.assertThrows;
@@ -14,6 +15,7 @@
1415
import com.typesafe.config.ConfigFactory;
1516
import java.io.IOException;
1617
import java.sql.Connection;
18+
import java.sql.DriverManager;
1719
import java.sql.PreparedStatement;
1820
import java.sql.ResultSet;
1921
import java.util.ArrayList;
@@ -47,6 +49,8 @@
4749
import org.junit.jupiter.api.Test;
4850
import org.junit.jupiter.params.ParameterizedTest;
4951
import org.junit.jupiter.params.provider.ArgumentsSource;
52+
import org.postgresql.util.PSQLException;
53+
import org.postgresql.util.PSQLState;
5054
import org.testcontainers.junit.jupiter.Testcontainers;
5155

5256
@Testcontainers
@@ -4023,6 +4027,83 @@ private long getLastUpdatedEpoch(Key key) throws Exception {
40234027
}
40244028
}
40254029

4030+
@Test
4031+
@DisplayName("upsert on a row locked by another txn is aborted after query timeout")
4032+
void upsertHonorsQueryTimeoutOnRowLock() throws Exception {
4033+
// setup
4034+
String docId = generateDocId("qt-upsert");
4035+
Key key = new SingleValueKey(DEFAULT_TENANT, docId);
4036+
ObjectNode initial = OBJECT_MAPPER.createObjectNode();
4037+
initial.put("id", docId);
4038+
initial.put("item", "Seed");
4039+
initial.put("price", 1);
4040+
flatCollection.upsert(key, new JSONDocument(initial));
4041+
4042+
// 2) Build a collection whose datastore has a 1s query timeout.
4043+
Collection timeoutCollection = getFlatCollectionWithQueryTimeout("1 second");
4044+
4045+
// 3) Hold a FOR UPDATE lock on the seeded row from a separate raw connection.
4046+
String url =
4047+
String.format(
4048+
"jdbc:postgresql://localhost:%s/postgres", postgresContainer.getMappedPort(5432));
4049+
try (Connection lockConn = DriverManager.getConnection(url, "postgres", "postgres")) {
4050+
lockConn.setAutoCommit(false);
4051+
try (PreparedStatement lockPs =
4052+
lockConn.prepareStatement(
4053+
String.format(
4054+
"SELECT id FROM \"%s\" WHERE id = ? FOR UPDATE", FLAT_COLLECTION_NAME))) {
4055+
lockPs.setString(1, key.toString());
4056+
try (ResultSet rs = lockPs.executeQuery()) {
4057+
assertTrue(rs.next(), "Precondition failure: Could not acquire lock on the seed row");
4058+
}
4059+
}
4060+
4061+
// 4) The upsert must block on the row lock and then be cancelled by
4062+
// JDBC's setQueryTimeout(1). FlatPostgresCollection wraps SQLException as IOException.
4063+
ObjectNode updateNode = OBJECT_MAPPER.createObjectNode();
4064+
updateNode.put("id", docId);
4065+
updateNode.put("item", "Updated");
4066+
updateNode.put("price", 2);
4067+
4068+
long startNs = System.nanoTime();
4069+
IOException thrown =
4070+
assertThrows(
4071+
IOException.class, () -> timeoutCollection.upsert(key, new JSONDocument(updateNode)));
4072+
long elapsedMs = (System.nanoTime() - startNs) / 1_000_000L;
4073+
4074+
// Sanity: cancellation should fire quickly - well before the row lock would ever
4075+
// be released (this txn is held open until the try-with-resources closes).
4076+
assertTrue(elapsedMs < 15_000);
4077+
4078+
Throwable cause = thrown.getCause();
4079+
assertNotNull(cause);
4080+
assertInstanceOf(PSQLException.class, cause);
4081+
assertEquals(PSQLState.QUERY_CANCELED.getState(), ((PSQLException) cause).getSQLState());
4082+
4083+
// 5) Verify the seeded row was not modified (lock still held, timed-out UPDATE
4084+
// never committed).
4085+
queryAndAssert(
4086+
key,
4087+
rs -> {
4088+
assertTrue(rs.next());
4089+
assertEquals("Seed", rs.getString("item"));
4090+
assertEquals(1, rs.getInt("price"));
4091+
});
4092+
}
4093+
}
4094+
4095+
private Collection getFlatCollectionWithQueryTimeout(String queryTimeout) {
4096+
String postgresConnectionUrl =
4097+
String.format("jdbc:postgresql://localhost:%s/", postgresContainer.getMappedPort(5432));
4098+
Map<String, String> cfg = new HashMap<>();
4099+
cfg.put("url", postgresConnectionUrl);
4100+
cfg.put("user", "postgres");
4101+
cfg.put("password", "postgres");
4102+
cfg.put("queryTimeout", queryTimeout);
4103+
Datastore ds = DatastoreProvider.getDatastore("Postgres", ConfigFactory.parseMap(cfg));
4104+
return ds.getCollectionForType(FLAT_COLLECTION_NAME, DocumentType.FLAT);
4105+
}
4106+
40264107
private static void executeInsertStatements() {
40274108
PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore;
40284109
try {

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

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -253,7 +253,8 @@ public boolean delete(Key key) {
253253
String.format(
254254
"DELETE FROM %s WHERE %s = ?",
255255
tableIdentifier, PostgresUtils.wrapFieldNamesWithDoubleQuotes(pkForTable));
256-
try (PreparedStatement preparedStatement = client.getConnection().prepareStatement(deleteSQL)) {
256+
try (PreparedStatement preparedStatement =
257+
queryExecutor.prepareStatementWithTimeout(client.getConnection(), deleteSQL)) {
257258
preparedStatement.setString(1, key.toString());
258259
int rowsDeleted = preparedStatement.executeUpdate();
259260
return rowsDeleted > 0;
@@ -320,7 +321,7 @@ public BulkDeleteResult delete(Set<Key> keys) {
320321
LOGGER.debug("Bulk delete SQL: {}", deleteSQL);
321322

322323
try (Connection conn = client.getPooledConnection();
323-
PreparedStatement ps = conn.prepareStatement(deleteSQL)) {
324+
PreparedStatement ps = queryExecutor.prepareStatementWithTimeout(conn, deleteSQL)) {
324325
int deletedCount = ps.executeUpdate();
325326
LOGGER.debug("Bulk deleted {} rows", deletedCount);
326327
return new BulkDeleteResult(deletedCount);
@@ -336,7 +337,7 @@ public boolean deleteAll() {
336337
LOGGER.debug("Delete all SQL: {}", deleteSQL);
337338

338339
try (Connection conn = client.getPooledConnection();
339-
PreparedStatement ps = conn.prepareStatement(deleteSQL)) {
340+
PreparedStatement ps = queryExecutor.prepareStatementWithTimeout(conn, deleteSQL)) {
340341
int deletedCount = ps.executeUpdate();
341342
LOGGER.debug("Deleted all {} rows", deletedCount);
342343
return true;
@@ -402,7 +403,7 @@ public boolean bulkUpsert(Map<Key, Document> documents) {
402403
LOGGER.debug("Bulk upsert SQL: {}", sql);
403404

404405
try (Connection conn = client.getPooledConnection();
405-
PreparedStatement ps = conn.prepareStatement(sql)) {
406+
PreparedStatement ps = queryExecutor.prepareStatementWithTimeout(conn, sql)) {
406407

407408
for (Map.Entry<Key, TypedDocument> entry : parsedDocuments.entrySet()) {
408409
TypedDocument parsed = entry.getValue();
@@ -499,7 +500,7 @@ public boolean bulkCreateOrReplace(Map<Key, Document> documents) {
499500
LOGGER.debug("Bulk createOrReplace SQL: {}", sql);
500501

501502
try (Connection conn = client.getPooledConnection();
502-
PreparedStatement ps = conn.prepareStatement(sql)) {
503+
PreparedStatement ps = queryExecutor.prepareStatementWithTimeout(conn, sql)) {
503504

504505
for (Map.Entry<Key, TypedDocument> entry : parsedDocuments.entrySet()) {
505506
TypedDocument parsed = entry.getValue();
@@ -680,7 +681,8 @@ private PreparedStatement getPreparedStatementForQuery(
680681
throws SQLException {
681682
String selectQuery =
682683
String.format("SELECT * FROM %s WHERE %s = ANY(?)", tableIdentifier, quotedPkColumn);
683-
PreparedStatement preparedStatement = connection.prepareStatement(selectQuery);
684+
PreparedStatement preparedStatement =
685+
queryExecutor.prepareStatementWithTimeout(connection, selectQuery);
684686

685687
String[] keyArray = documents.keySet().stream().map(Key::toString).toArray(String[]::new);
686688
Array sqlArray = connection.createArrayOf(pkType.getSqlType(), keyArray);
@@ -959,7 +961,7 @@ private boolean executeKeyUpdate(
959961

960962
LOGGER.debug("Executing key update SQL: {}", sql);
961963

962-
try (PreparedStatement ps = connection.prepareStatement(sql)) {
964+
try (PreparedStatement ps = queryExecutor.prepareStatementWithTimeout(connection, sql)) {
963965
int idx = 1;
964966
for (Object param : params) {
965967
ps.setObject(idx++, param);
@@ -1058,7 +1060,7 @@ private int executeBatchUpdate(
10581060

10591061
LOGGER.debug("Executing batch update SQL: {} for {} keys", sql, keys.size());
10601062

1061-
try (PreparedStatement ps = connection.prepareStatement(sql)) {
1063+
try (PreparedStatement ps = queryExecutor.prepareStatementWithTimeout(connection, sql)) {
10621064
for (int i : order) {
10631065
int idx = 1;
10641066
for (Object param : allKeyParams.get(i)) {
@@ -1264,7 +1266,7 @@ private void executeUpdate(
12641266

12651267
LOGGER.debug("Executing update SQL: {}", sql);
12661268

1267-
try (PreparedStatement ps = connection.prepareStatement(sql)) {
1269+
try (PreparedStatement ps = queryExecutor.prepareStatementWithTimeout(connection, sql)) {
12681270
int idx = 1;
12691271
for (Object param : params) {
12701272
ps.setObject(idx++, param);
@@ -1580,7 +1582,7 @@ Object convertTimestampForType(long epochMillis, PostgresDataType type) {
15801582

15811583
private int executeUpdate(String sql, TypedDocument parsed) throws SQLException {
15821584
try (Connection conn = client.getPooledConnection();
1583-
PreparedStatement ps = conn.prepareStatement(sql)) {
1585+
PreparedStatement ps = queryExecutor.prepareStatementWithTimeout(conn, sql)) {
15841586
int index = 1;
15851587
for (String column : parsed.getColumns()) {
15861588
setParameter(
@@ -1749,7 +1751,7 @@ private boolean executeUpsert(String sql, List<String> allColumns, TypedDocument
17491751
long ta = System.nanoTime();
17501752
try (Connection conn = client.getPooledConnection()) {
17511753
long tb = System.nanoTime();
1752-
try (PreparedStatement ps = conn.prepareStatement(sql)) {
1754+
try (PreparedStatement ps = queryExecutor.prepareStatementWithTimeout(conn, sql)) {
17531755
int index = 1;
17541756
Set<String> parsedColumns = new HashSet<>(parsed.getColumns());
17551757
for (String column : allColumns) {
@@ -1778,7 +1780,7 @@ private boolean executeUpsert(String sql, List<String> allColumns, TypedDocument
17781780
private boolean executeUpsertReturningIsInsert(
17791781
String sql, List<String> allColumns, TypedDocument parsed) throws SQLException {
17801782
try (Connection conn = client.getPooledConnection();
1781-
PreparedStatement ps = conn.prepareStatement(sql)) {
1783+
PreparedStatement ps = queryExecutor.prepareStatementWithTimeout(conn, sql)) {
17821784
int index = 1;
17831785
Set<String> parsedColumns = new HashSet<>(parsed.getColumns());
17841786
for (String column : allColumns) {

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,15 @@ public PreparedStatement buildPreparedStatement(
5353
return buildPreparedStatement(sqlQuery, params, connection, this.queryTimeoutSeconds);
5454
}
5555

56+
public PreparedStatement prepareStatementWithTimeout(Connection connection, String sqlQuery)
57+
throws SQLException {
58+
PreparedStatement preparedStatement = connection.prepareStatement(sqlQuery);
59+
if (queryTimeoutSeconds > 0) {
60+
preparedStatement.setQueryTimeout(queryTimeoutSeconds);
61+
}
62+
return preparedStatement;
63+
}
64+
5665
public PreparedStatement buildPreparedStatement(
5766
String sqlQuery, Params params, Connection connection, int queryTimeoutSeconds)
5867
throws SQLException {

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

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,8 @@
11
package org.hypertrace.core.documentstore.postgres;
22

3+
import static org.junit.jupiter.api.Assertions.assertSame;
4+
import static org.mockito.ArgumentMatchers.anyInt;
5+
import static org.mockito.Mockito.never;
36
import static org.mockito.Mockito.verify;
47
import static org.mockito.Mockito.when;
58

@@ -43,4 +46,27 @@ void testPreparedStatementUsesOverridenTimeout() throws SQLException {
4346
executor.buildPreparedStatement(sqlQuery, params, mockConnection, 45);
4447
verify(mockPreparedStatement).setQueryTimeout(45);
4548
}
49+
50+
@Test
51+
void prepareStatementWithTimeoutAppliesConfiguredTimeout() throws SQLException {
52+
String sqlQuery = "UPDATE foo SET x = ?";
53+
when(mockConnection.prepareStatement(sqlQuery)).thenReturn(mockPreparedStatement);
54+
55+
PostgresQueryExecutor executor = new PostgresQueryExecutor(45);
56+
PreparedStatement result = executor.prepareStatementWithTimeout(mockConnection, sqlQuery);
57+
58+
assertSame(mockPreparedStatement, result);
59+
verify(mockPreparedStatement).setQueryTimeout(45);
60+
}
61+
62+
@Test
63+
void prepareStatementWithTimeoutSkipsSetWhenTimeoutIsZero() throws SQLException {
64+
String sqlQuery = "DELETE FROM foo";
65+
when(mockConnection.prepareStatement(sqlQuery)).thenReturn(mockPreparedStatement);
66+
67+
PostgresQueryExecutor executor = new PostgresQueryExecutor(0);
68+
executor.prepareStatementWithTimeout(mockConnection, sqlQuery);
69+
70+
verify(mockPreparedStatement, never()).setQueryTimeout(anyInt());
71+
}
4672
}

0 commit comments

Comments
 (0)