Add oracle, redshift, postgresql jdbc fetch size

plugin/trino-oracle/src/main/java/io/trino/plugin/oracle/OracleClient.java

@@ -216,6 +216,7 @@ public class OracleClient
     private final boolean synonymsEnabled;
     private final ConnectorExpressionRewriter<ParameterizedExpression> connectorExpressionRewriter;
     private final AggregateFunctionRewriter<JdbcExpression, ?> aggregateFunctionRewriter;
+    private final Optional<Integer> fetchSize;
 
     @Inject
     public OracleClient(
@@ -260,6 +261,8 @@ public OracleClient(
                         .add(new ImplementCovarianceSamp())
                         .add(new ImplementCovariancePop())
                         .build());
+
+        this.fetchSize = oracleConfig.getFetchSize();
     }
 
     @Override
@@ -287,8 +290,11 @@ public PreparedStatement getPreparedStatement(Connection connection, String sql,
         PreparedStatement statement = connection.prepareStatement(sql);
         // This is a heuristic, not exact science. A better formula can perhaps be found with measurements.
         // Column count is not known for non-SELECT queries. Not setting fetch size for these.
-        if (columnCount.isPresent()) {
-            statement.setFetchSize(max(100_000 / columnCount.get(), 1_000));
+        Optional<Integer> fetchSize = Optional.ofNullable(this.fetchSize.orElseGet(() ->
+                columnCount.map(count -> max(100_000 / count, 1_000))
+                        .orElse(null)));
+        if (fetchSize.isPresent()) {
+            statement.setFetchSize(fetchSize.get());
         }
         return statement;
     }

plugin/trino-oracle/src/main/java/io/trino/plugin/oracle/OracleConfig.java

@@ -38,6 +38,7 @@ public class OracleConfig
     private int connectionPoolMinSize = 1;
     private int connectionPoolMaxSize = 30;
     private Duration inactiveConnectionTimeout = new Duration(20, MINUTES);
+    private Integer fetchSize;
 
     public boolean isSynonymsEnabled()
     {
@@ -141,6 +142,19 @@ public OracleConfig setInactiveConnectionTimeout(Duration inactiveConnectionTime
         return this;
     }
 
+    public Optional<@Min(0) Integer> getFetchSize()
+    {
+        return Optional.ofNullable(fetchSize);
+    }
+
+    @Config("oracle.fetch-size")
+    @ConfigDescription("Oracle fetch size, trino specific heuristic is applied if empty")
+    public OracleConfig setFetchSize(Integer fetchSize)
+    {
+        this.fetchSize = fetchSize;
+        return this;
+    }
+
     @AssertTrue(message = "Pool min size cannot be larger than max size")
     public boolean isPoolSizedProperly()
     {

plugin/trino-oracle/src/test/java/io/trino/plugin/oracle/TestOracleConfig.java

@@ -42,7 +42,8 @@ public void testDefaults()
                 .setConnectionPoolEnabled(true)
                 .setConnectionPoolMinSize(1)
                 .setConnectionPoolMaxSize(30)
-                .setInactiveConnectionTimeout(new Duration(20, MINUTES)));
+                .setInactiveConnectionTimeout(new Duration(20, MINUTES))
+                .setFetchSize(null));
     }
 
     @Test
@@ -57,6 +58,7 @@ public void testExplicitPropertyMappings()
                 .put("oracle.connection-pool.min-size", "10")
                 .put("oracle.connection-pool.max-size", "20")
                 .put("oracle.connection-pool.inactive-timeout", "30s")
+                .put("oracle.fetch-size", "2000")
                 .buildOrThrow();
 
         OracleConfig expected = new OracleConfig()
@@ -67,7 +69,8 @@ public void testExplicitPropertyMappings()
                 .setConnectionPoolEnabled(false)
                 .setConnectionPoolMinSize(10)
                 .setConnectionPoolMaxSize(20)
-                .setInactiveConnectionTimeout(new Duration(30, SECONDS));
+                .setInactiveConnectionTimeout(new Duration(30, SECONDS))
+                .setFetchSize(2000);
 
         assertFullMapping(properties, expected);
     }

plugin/trino-postgresql/src/main/java/io/trino/plugin/postgresql/PostgreSqlClient.java

@@ -282,6 +282,7 @@ public class PostgreSqlClient
     private final ConnectorExpressionRewriter<ParameterizedExpression> connectorExpressionRewriter;
     private final ProjectFunctionRewriter<JdbcExpression, ParameterizedExpression> projectFunctionRewriter;
     private final AggregateFunctionRewriter<JdbcExpression, ?> aggregateFunctionRewriter;
+    private final Optional<Integer> fetchSize;
 
     @Inject
     public PostgreSqlClient(
@@ -368,6 +369,7 @@ public PostgreSqlClient(
                         .add(new ImplementRegrIntercept())
                         .add(new ImplementRegrSlope())
                         .build());
+        this.fetchSize = postgreSqlConfig.getFetchSize();
     }
 
     @Override
@@ -432,8 +434,11 @@ public PreparedStatement getPreparedStatement(Connection connection, String sql,
         PreparedStatement statement = connection.prepareStatement(sql);
         // This is a heuristic, not exact science. A better formula can perhaps be found with measurements.
         // Column count is not known for non-SELECT queries. Not setting fetch size for these.
-        if (columnCount.isPresent()) {
-            statement.setFetchSize(max(100_000 / columnCount.get(), 1_000));
+        Optional<Integer> fetchSize = Optional.ofNullable(this.fetchSize.orElseGet(() ->
+                columnCount.map(count -> max(100_000 / count, 1_000))
+                        .orElse(null)));
+        if (fetchSize.isPresent()) {
+            statement.setFetchSize(fetchSize.get());
         }
         return statement;
     }

plugin/trino-postgresql/src/main/java/io/trino/plugin/postgresql/PostgreSqlConfig.java

@@ -14,16 +14,21 @@
 package io.trino.plugin.postgresql;
 
 import io.airlift.configuration.Config;
+import io.airlift.configuration.ConfigDescription;
 import io.airlift.configuration.DefunctConfig;
 import io.airlift.configuration.LegacyConfig;
+import jakarta.validation.constraints.Min;
 import jakarta.validation.constraints.NotNull;
 
+import java.util.Optional;
+
 @DefunctConfig("postgresql.disable-automatic-fetch-size")
 public class PostgreSqlConfig
 {
     private ArrayMapping arrayMapping = ArrayMapping.DISABLED;
     private boolean includeSystemTables;
     private boolean enableStringPushdownWithCollate;
+    private Integer fetchSize;
 
     public enum ArrayMapping
     {
@@ -69,4 +74,17 @@ public PostgreSqlConfig setEnableStringPushdownWithCollate(boolean enableStringP
         this.enableStringPushdownWithCollate = enableStringPushdownWithCollate;
         return this;
     }
+
+    public Optional<@Min(0) Integer> getFetchSize()
+    {
+        return Optional.ofNullable(fetchSize);
+    }
+
+    @Config("postgresql.fetch-size")
+    @ConfigDescription("Postgresql fetch size, trino specific heuristic is applied if empty")
+    public PostgreSqlConfig setFetchSize(Integer fetchSize)
+    {
+        this.fetchSize = fetchSize;
+        return this;
+    }
 }

plugin/trino-postgresql/src/test/java/io/trino/plugin/postgresql/TestPostgreSqlConfig.java

@@ -30,7 +30,8 @@ public void testDefaults()
         assertRecordedDefaults(recordDefaults(PostgreSqlConfig.class)
                 .setArrayMapping(PostgreSqlConfig.ArrayMapping.DISABLED)
                 .setIncludeSystemTables(false)
-                .setEnableStringPushdownWithCollate(false));
+                .setEnableStringPushdownWithCollate(false)
+                .setFetchSize(null));
     }
 
     @Test
@@ -40,12 +41,14 @@ public void testExplicitPropertyMappings()
                 .put("postgresql.array-mapping", "AS_ARRAY")
                 .put("postgresql.include-system-tables", "true")
                 .put("postgresql.experimental.enable-string-pushdown-with-collate", "true")
+                .put("postgresql.fetch-size", "2000")
                 .buildOrThrow();
 
         PostgreSqlConfig expected = new PostgreSqlConfig()
                 .setArrayMapping(PostgreSqlConfig.ArrayMapping.AS_ARRAY)
                 .setIncludeSystemTables(true)
-                .setEnableStringPushdownWithCollate(true);
+                .setEnableStringPushdownWithCollate(true)
+                .setFetchSize(2000);
 
         assertFullMapping(properties, expected);
     }

plugin/trino-redshift/pom.xml

@@ -50,6 +50,11 @@
             <artifactId>trino-plugin-toolkit</artifactId>
         </dependency>
 
+        <dependency>
+            <groupId>jakarta.validation</groupId>
+            <artifactId>jakarta.validation-api</artifactId>
+        </dependency>
+
         <dependency>
             <groupId>org.jdbi</groupId>
             <artifactId>jdbi3-core</artifactId>

plugin/trino-redshift/src/main/java/io/trino/plugin/redshift/RedshiftClient.java

@@ -228,10 +228,12 @@ public class RedshiftClient
     private final boolean statisticsEnabled;
     private final RedshiftTableStatisticsReader statisticsReader;
     private final ConnectorExpressionRewriter<ParameterizedExpression> connectorExpressionRewriter;
+    private final Optional<Integer> fetchSize;
 
     @Inject
     public RedshiftClient(
             BaseJdbcConfig config,
+            RedshiftConfig redshiftConfig,
             ConnectionFactory connectionFactory,
             JdbcStatisticsConfig statisticsConfig,
             QueryBuilder queryBuilder,
@@ -269,6 +271,7 @@ public RedshiftClient(
 
         this.statisticsEnabled = requireNonNull(statisticsConfig, "statisticsConfig is null").isEnabled();
         this.statisticsReader = new RedshiftTableStatisticsReader(connectionFactory);
+        this.fetchSize = redshiftConfig.getFetchSize();
     }
 
     private static Optional<JdbcTypeHandle> toTypeHandle(DecimalType decimalType)
@@ -487,8 +490,11 @@ public PreparedStatement getPreparedStatement(Connection connection, String sql,
         PreparedStatement statement = connection.prepareStatement(sql);
         // This is a heuristic, not exact science. A better formula can perhaps be found with measurements.
         // Column count is not known for non-SELECT queries. Not setting fetch size for these.
-        if (columnCount.isPresent()) {
-            statement.setFetchSize(max(100_000 / columnCount.get(), 1_000));
+        Optional<Integer> fetchSize = Optional.ofNullable(this.fetchSize.orElseGet(() ->
+                columnCount.map(count -> max(100_000 / count, 1_000))
+                        .orElse(null)));
+        if (fetchSize.isPresent()) {
+            statement.setFetchSize(fetchSize.get());
         }
         return statement;
     }

plugin/trino-redshift/src/main/java/io/trino/plugin/redshift/RedshiftConfig.java

@@ -13,12 +13,31 @@
  */
 package io.trino.plugin.redshift;
 
+import io.airlift.configuration.Config;
+import io.airlift.configuration.ConfigDescription;
 import io.airlift.configuration.DefunctConfig;
+import jakarta.validation.constraints.Min;
+
+import java.util.Optional;
 
 @DefunctConfig({
         "redshift.disable-automatic-fetch-size",
         "redshift.use-legacy-type-mapping",
 })
 public class RedshiftConfig
 {
+    private Integer fetchSize;
+
+    public Optional<@Min(0) Integer> getFetchSize()
+    {
+        return Optional.ofNullable(fetchSize);
+    }
+
+    @Config("redshift.fetch-size")
+    @ConfigDescription("Redshift fetch size, trino specific heuristic is applied if empty")
+    public RedshiftConfig setFetchSize(Integer fetchSize)
+    {
+        this.fetchSize = fetchSize;
+        return this;
+    }
 }

plugin/trino-redshift/src/test/java/io/trino/plugin/redshift/TestRedshiftConfig.java

@@ -27,16 +27,19 @@ public class TestRedshiftConfig
     @Test
     public void testDefaults()
     {
-        assertRecordedDefaults(recordDefaults(RedshiftConfig.class));
+        assertRecordedDefaults(recordDefaults(RedshiftConfig.class)
+                .setFetchSize(null));
     }
 
     @Test
     public void testExplicitPropertyMappings()
     {
         Map<String, String> properties = ImmutableMap.<String, String>builder()
+                .put("redshift.fetch-size", "2000")
                 .buildOrThrow();
 
-        RedshiftConfig expected = new RedshiftConfig();
+        RedshiftConfig expected = new RedshiftConfig()
+                .setFetchSize(2000);
 
         assertFullMapping(properties, expected);
     }