Skip to content

PG write query timeouts #321

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
suddendust merged 1 commit into from
Aug 13, 2026
Merged

PG write query timeouts #321

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
81 changes: 81 additions & 0 deletions ...e/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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<String, String> 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 {
Expand Down
26 changes: 14 additions & 12 deletions ...tore/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -320,7 +321,7 @@ public BulkDeleteResult delete(Set<Key> 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);
Expand All @@ -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;
Expand Down Expand Up @@ -402,7 +403,7 @@ public boolean bulkUpsert(Map<Key, Document> 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<Key, TypedDocument> entry : parsedDocuments.entrySet()) {
TypedDocument parsed = entry.getValue();
Expand Down Expand Up @@ -499,7 +500,7 @@ public boolean bulkCreateOrReplace(Map<Key, Document> 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<Key, TypedDocument> entry : parsedDocuments.entrySet()) {
TypedDocument parsed = entry.getValue();
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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)) {
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -1749,7 +1751,7 @@ private boolean executeUpsert(String sql, List<String> 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<String> parsedColumns = new HashSet<>(parsed.getColumns());
for (String column : allColumns) {
Expand Down Expand Up @@ -1778,7 +1780,7 @@ private boolean executeUpsert(String sql, List<String> allColumns, TypedDocument
private boolean executeUpsertReturningIsInsert(
String sql, List<String> 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<String> parsedColumns = new HashSet<>(parsed.getColumns());
for (String column : allColumns) {
Expand Down
9 changes: 9 additions & 0 deletions ...store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresQueryExecutor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

@puneet-traceable puneet-traceable Aug 13, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this set per prepared statement? or is it possible to set it at connection?

@suddendust suddendust Aug 13, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So it's a bit nuanced. There're three ways to do this:

  1. Set it in PS - Like what we're doing here. In this case, the driver cancels the query by sending a new cancel request to the server.
  2. Set it at a connection level - By setting SET statement_timeout = 60s in the connection pool config. However, this might pin connections at the proxy (using SET commands does but the doc is not very clear on whether SET statement_timeout will do this).
  3. Set it in the RDS parameter group. This applies this timeout to any client regardless of what they've configured. Client can still use setQueryTimeout to set a value lesser than the global value in the param group.

Given our issues with pinning, I'll go with 1 + 3.

}

public PreparedStatement buildPreparedStatement(
String sqlQuery, Params params, Connection connection, int queryTimeoutSeconds)
throws SQLException {
Expand Down
26 changes: 26 additions & 0 deletions ...e/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresQueryExecutorTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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;

Expand Down Expand Up @@ -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());
}
}
Loading