From 2e3cc903a6cde777dc00f7d56cb200036fced64b Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Thu, 1 Jan 2026 11:26:07 +0530 Subject: [PATCH 1/2] Postgres query timeout --- .../TypesafeDatastoreConfigAdapter.java | 1 + .../model/config/ConnectionConfig.java | 1 + .../postgres/PostgresConnectionConfig.java | 4 ++++ .../postgres/PostgresClient.java | 4 ++++ .../postgres/PostgresCollection.java | 2 +- .../postgres/PostgresQueryExecutor.java | 20 ++++++++++++++++--- ...afeConfigDatastoreConfigExtractorTest.java | 2 ++ 7 files changed, 30 insertions(+), 4 deletions(-) diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/TypesafeDatastoreConfigAdapter.java b/document-store/src/main/java/org/hypertrace/core/documentstore/TypesafeDatastoreConfigAdapter.java index c1fcacde..abc00e88 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/TypesafeDatastoreConfigAdapter.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/TypesafeDatastoreConfigAdapter.java @@ -93,6 +93,7 @@ public DatastoreConfig convert(final Config config) { connectionConfig.credentials(), connectionConfig.applicationName(), connectionConfig.connectionPoolConfig(), + connectionConfig.queryTimeout(), connectionConfig.customParameters()) { @Override public String toConnectionString() { diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/model/config/ConnectionConfig.java b/document-store/src/main/java/org/hypertrace/core/documentstore/model/config/ConnectionConfig.java index 18d9975c..11127b81 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/model/config/ConnectionConfig.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/model/config/ConnectionConfig.java @@ -126,6 +126,7 @@ public ConnectionConfig build() { credentials, applicationName, connectionPoolConfig, + queryTimeout, customParameters); } diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/model/config/postgres/PostgresConnectionConfig.java b/document-store/src/main/java/org/hypertrace/core/documentstore/model/config/postgres/PostgresConnectionConfig.java index 925e4b1e..dd809587 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/model/config/postgres/PostgresConnectionConfig.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/model/config/postgres/PostgresConnectionConfig.java @@ -3,6 +3,7 @@ import static java.util.Collections.unmodifiableList; import static java.util.function.Predicate.not; +import java.time.Duration; import java.util.ArrayList; import java.util.List; import java.util.Map; @@ -36,6 +37,7 @@ public class PostgresConnectionConfig extends ConnectionConfig { @NonNull String applicationName; @NonNull ConnectionPoolConfig connectionPoolConfig; + @NonNull Duration queryTimeout; public static ConnectionConfigBuilder builder() { return ConnectionConfig.builder().type(DatabaseType.POSTGRES); @@ -47,6 +49,7 @@ public PostgresConnectionConfig( @Nullable final ConnectionCredentials credentials, @NonNull final String applicationName, @Nullable final ConnectionPoolConfig connectionPoolConfig, + @NonNull final Duration queryTimeout, @NonNull final Map customParameters) { super( ensureSingleEndpoint(endpoints), @@ -55,6 +58,7 @@ public PostgresConnectionConfig( customParameters); this.applicationName = applicationName; this.connectionPoolConfig = getConnectionPoolConfigOrDefault(connectionPoolConfig); + this.queryTimeout = queryTimeout; } public String toConnectionString() { diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresClient.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresClient.java index fa6a12cf..d0ced723 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresClient.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresClient.java @@ -69,6 +69,10 @@ public Map getCustomParameters() { return connectionConfig.customParameters(); } + public int getQueryTimeoutSeconds() { + return (int) connectionConfig.queryTimeout().toSeconds(); + } + public void close() { if (connection != null) { try { diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java index fa2519bf..ffcc283c 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java @@ -113,7 +113,7 @@ public PostgresCollection(final PostgresClient client, final String collectionNa this.tableIdentifier = tableIdentifier; this.subDocUpdater = new PostgresSubDocumentUpdater(new PostgresQueryBuilder(this.tableIdentifier)); - this.queryExecutor = new PostgresQueryExecutor(); + this.queryExecutor = new PostgresQueryExecutor(client.getQueryTimeoutSeconds()); this.updateValidator = new CommonUpdateValidator(); } 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 1a5fdade..721ce9a1 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 @@ -6,15 +6,19 @@ import java.sql.PreparedStatement; import java.sql.ResultSet; import java.sql.SQLException; -import lombok.AllArgsConstructor; import lombok.extern.slf4j.Slf4j; import org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser; import org.hypertrace.core.documentstore.postgres.query.v1.transformer.PostgresQueryTransformer; @Slf4j -@AllArgsConstructor public class PostgresQueryExecutor { + private final int queryTimeoutSeconds; + + public PostgresQueryExecutor(int queryTimeoutSeconds) { + this.queryTimeoutSeconds = queryTimeoutSeconds; + } + static org.hypertrace.core.documentstore.query.Query transformAndLog( org.hypertrace.core.documentstore.query.Query query) { log.debug("Original query before transformation: {}", query); @@ -28,7 +32,8 @@ protected ResultSet execute(final Connection connection, PostgresQueryParser que final String sqlQuery = queryParser.parse(); final Params params = queryParser.getParamsBuilder().build(); // this is closed when the corresponding ResultSet is closed in the iterators - PreparedStatement preparedStatement = buildPreparedStatement(sqlQuery, params, connection); + PreparedStatement preparedStatement = + buildPreparedStatement(sqlQuery, params, connection, queryTimeoutSeconds); try { log.debug("Executing SQL query: {}", sqlQuery); return preparedStatement.executeQuery(); @@ -45,7 +50,16 @@ protected ResultSet execute(final Connection connection, PostgresQueryParser que public PreparedStatement buildPreparedStatement( String sqlQuery, Params params, Connection connection) throws SQLException { + return buildPreparedStatement(sqlQuery, params, connection, 0); + } + + public PreparedStatement buildPreparedStatement( + String sqlQuery, Params params, Connection connection, int queryTimeoutSeconds) + throws SQLException { PreparedStatement preparedStatement = connection.prepareStatement(sqlQuery); + if (queryTimeoutSeconds > 0) { + preparedStatement.setQueryTimeout(queryTimeoutSeconds); + } enrichPreparedStatementWithParams(preparedStatement, params); return preparedStatement; } diff --git a/document-store/src/test/java/org/hypertrace/core/documentstore/model/config/TypesafeConfigDatastoreConfigExtractorTest.java b/document-store/src/test/java/org/hypertrace/core/documentstore/model/config/TypesafeConfigDatastoreConfigExtractorTest.java index cfb1cd52..01062175 100644 --- a/document-store/src/test/java/org/hypertrace/core/documentstore/model/config/TypesafeConfigDatastoreConfigExtractorTest.java +++ b/document-store/src/test/java/org/hypertrace/core/documentstore/model/config/TypesafeConfigDatastoreConfigExtractorTest.java @@ -158,6 +158,7 @@ void testBuildPostgres() { .poolMaxConnectionsKey(MAX_CONNECTIONS_KEY) .poolConnectionAccessTimeoutKey(CONNECTION_ACCESS_TIMEOUT_KEY) .poolConnectionSurrenderTimeoutKey(CONNECTION_SURRENDER_TIMEOUT_KEY) + .queryTimeoutKey(QUERY_TIMEOUT_KEY) .extract() .connectionConfig(); final ConnectionConfig expected = @@ -179,6 +180,7 @@ void testBuildPostgres() { .connectionSurrenderTimeout(surrenderTimeout) .build()) .aggregationPipelineMode(SORT_OPTIMIZED_IF_POSSIBLE) + .queryTimeout(queryTimeout) .build(); assertEquals(expected, config); From 5b36407cb50de6e89643e175bac71f21648c70f5 Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Thu, 1 Jan 2026 12:07:36 +0530 Subject: [PATCH 2/2] Added UTs --- .../postgres/PostgresQueryExecutor.java | 2 +- .../postgres/PostgresQueryExecutorTest.java | 46 +++++++++++++++++++ 2 files changed, 47 insertions(+), 1 deletion(-) create mode 100644 document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresQueryExecutorTest.java 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 721ce9a1..9f7cd174 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 @@ -50,7 +50,7 @@ protected ResultSet execute(final Connection connection, PostgresQueryParser que public PreparedStatement buildPreparedStatement( String sqlQuery, Params params, Connection connection) throws SQLException { - return buildPreparedStatement(sqlQuery, params, connection, 0); + return buildPreparedStatement(sqlQuery, params, connection, this.queryTimeoutSeconds); } public PreparedStatement buildPreparedStatement( 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 new file mode 100644 index 00000000..41ba3f4b --- /dev/null +++ b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresQueryExecutorTest.java @@ -0,0 +1,46 @@ +package org.hypertrace.core.documentstore.postgres; + +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.SQLException; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +@ExtendWith(MockitoExtension.class) +class PostgresQueryExecutorTest { + + @Mock private Connection mockConnection; + @Mock private PreparedStatement mockPreparedStatement; + + @Test + void testPreparedStatementUsesConfiguredTimeout() throws SQLException { + int configuredTimeoutSeconds = 45; + String sqlQuery = "SELECT * FROM test_table"; + Params params = Params.newBuilder().build(); + + when(mockConnection.prepareStatement(sqlQuery)).thenReturn(mockPreparedStatement); + + PostgresQueryExecutor executor = new PostgresQueryExecutor(configuredTimeoutSeconds); + executor.buildPreparedStatement(sqlQuery, params, mockConnection); + verify(mockPreparedStatement).setQueryTimeout(configuredTimeoutSeconds); + } + + @Test + void testPreparedStatementUsesOverridenTimeout() throws SQLException { + int defaultTimeoutSeconds = 30; + String sqlQuery = "SELECT * FROM test_table"; + Params params = Params.newBuilder().build(); + + when(mockConnection.prepareStatement(sqlQuery)).thenReturn(mockPreparedStatement); + + PostgresQueryExecutor executor = new PostgresQueryExecutor(defaultTimeoutSeconds); + // override timeout to 45s + executor.buildPreparedStatement(sqlQuery, params, mockConnection, 45); + verify(mockPreparedStatement).setQueryTimeout(45); + } +}