Skip to content

Added jdbc fetch size config for pg/oracle/redshift #16269

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

Closed
chenjian2664 wants to merge 1 commit into from
Closed

Added jdbc fetch size config for pg/oracle/redshift #16269

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
2 changes: 1 addition & 1 deletion plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/BaseJdbcClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1064,7 +1064,7 @@ public Connection getConnection(ConnectorSession session, JdbcOutputTableHandle
}

@Override
public PreparedStatement getPreparedStatement(Connection connection, String sql)
public PreparedStatement getPreparedStatement(ConnectorSession session, Connection connection, String sql)
throws SQLException
{
return connection.prepareStatement(sql);
Expand Down
4 changes: 2 additions & 2 deletions plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/CachingJdbcClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -350,10 +350,10 @@ public Connection getConnection(ConnectorSession session, JdbcOutputTableHandle
}

@Override
public PreparedStatement getPreparedStatement(Connection connection, String sql)
public PreparedStatement getPreparedStatement(ConnectorSession session, Connection connection, String sql)
throws SQLException
{
return delegate.getPreparedStatement(connection, sql);
return delegate.getPreparedStatement(session, connection, sql);
}

@Override
Expand Down
2 changes: 1 addition & 1 deletion plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/DefaultQueryBuilder.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -176,7 +176,7 @@ public PreparedStatement prepareStatement(
{
String modifiedQuery = queryModifier.apply(session, preparedQuery.getQuery());
log.debug("Preparing query: %s", modifiedQuery);
PreparedStatement statement = client.getPreparedStatement(connection, modifiedQuery);
PreparedStatement statement = client.getPreparedStatement(session, connection, modifiedQuery);

List<QueryParameter> parameters = preparedQuery.getParameters();
for (int i = 0; i < parameters.size(); i++) {
Expand Down
4 changes: 2 additions & 2 deletions plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/ForwardingJdbcClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -240,10 +240,10 @@ public Connection getConnection(ConnectorSession session, JdbcOutputTableHandle
}

@Override
public PreparedStatement getPreparedStatement(Connection connection, String sql)
public PreparedStatement getPreparedStatement(ConnectorSession session, Connection connection, String sql)
throws SQLException
{
return delegate().getPreparedStatement(connection, sql);
return delegate().getPreparedStatement(session, connection, sql);
}

@Override
Expand Down
2 changes: 1 addition & 1 deletion plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/JdbcClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -177,7 +177,7 @@ default void setTableProperties(ConnectorSession session, JdbcTableHandle handle
Connection getConnection(ConnectorSession session, JdbcOutputTableHandle handle)
throws SQLException;

PreparedStatement getPreparedStatement(Connection connection, String sql)
PreparedStatement getPreparedStatement(ConnectorSession session, Connection connection, String sql)
throws SQLException;

/**
Expand Down
4 changes: 2 additions & 2 deletions plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/jmx/StatisticsAwareJdbcClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -320,10 +320,10 @@ public Connection getConnection(ConnectorSession session, JdbcOutputTableHandle
}

@Override
public PreparedStatement getPreparedStatement(Connection connection, String sql)
public PreparedStatement getPreparedStatement(ConnectorSession session, Connection connection, String sql)
throws SQLException
{
return stats.getGetPreparedStatement().wrap(() -> delegate().getPreparedStatement(connection, sql));
return stats.getGetPreparedStatement().wrap(() -> delegate().getPreparedStatement(session, connection, sql));
}

@Override
Expand Down
5 changes: 3 additions & 2 deletions plugin/trino-mysql/src/main/java/io/trino/plugin/mysql/MySqlClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -294,7 +294,7 @@ public void abortReadConnection(Connection connection, ResultSet resultSet)
}

@Override
public PreparedStatement getPreparedStatement(Connection connection, String sql)
public PreparedStatement getPreparedStatement(ConnectorSession session, Connection connection, String sql)
throws SQLException
{
PreparedStatement statement = connection.prepareStatement(sql);
Expand Down Expand Up @@ -472,7 +472,8 @@ public Optional<ColumnMapping> toColumnMapping(ConnectorSession session, Connect

private LongWriteFunction mySqlDateWriteFunctionUsingLocalDate()
{
return new LongWriteFunction() {
return new LongWriteFunction()
{
@Override
public String getBindExpression()
{
Expand Down
5 changes: 3 additions & 2 deletions plugin/trino-oracle/src/main/java/io/trino/plugin/oracle/OracleClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,7 @@
import static io.trino.plugin.jdbc.StandardColumnMappings.varcharWriteFunction;
import static io.trino.plugin.jdbc.TypeHandlingJdbcSessionProperties.getUnsupportedTypeHandling;
import static io.trino.plugin.jdbc.UnsupportedTypeHandling.CONVERT_TO_VARCHAR;
import static io.trino.plugin.oracle.OracleSessionProperties.getJdbcFetchSize;
import static io.trino.plugin.oracle.OracleSessionProperties.getNumberDefaultScale;
import static io.trino.plugin.oracle.OracleSessionProperties.getNumberRoundingMode;
import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED;
Expand Down Expand Up @@ -257,11 +258,11 @@ protected boolean filterSchema(String schemaName)
}

@Override
public PreparedStatement getPreparedStatement(Connection connection, String sql)
public PreparedStatement getPreparedStatement(ConnectorSession session, Connection connection, String sql)
throws SQLException
{
PreparedStatement statement = connection.prepareStatement(sql);
statement.setFetchSize(1000);
statement.setFetchSize(getJdbcFetchSize(session));
return statement;
}

Expand Down
13 changes: 13 additions & 0 deletions plugin/trino-oracle/src/main/java/io/trino/plugin/oracle/OracleConfig.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ public class OracleConfig
private int connectionPoolMinSize = 1;
private int connectionPoolMaxSize = 30;
private Duration inactiveConnectionTimeout = new Duration(20, MINUTES);
private int jdbcFetchSize = 1000;

@NotNull
public boolean isSynonymsEnabled()
Expand Down Expand Up @@ -147,4 +148,16 @@ public boolean isPoolSizedProperly()
{
return getConnectionPoolMaxSize() >= getConnectionPoolMinSize();
}

public int getJdbcFetchSize()
{
return jdbcFetchSize;
}

@Config("oracle.jdbc-fetch-size")
public OracleConfig setJdbcFetchSize(int jdbcFetchSize)
{
this.jdbcFetchSize = jdbcFetchSize;
return this;
}
}
11 changes: 11 additions & 0 deletions plugin/trino-oracle/src/main/java/io/trino/plugin/oracle/OracleSessionProperties.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ public final class OracleSessionProperties
{
public static final String NUMBER_ROUNDING_MODE = "number_rounding_mode";
public static final String NUMBER_DEFAULT_SCALE = "number_default_scale";
public static final String JDBC_FETCH_SIZE = "jdbc_fetch_size";

private final List<PropertyMetadata<?>> sessionProperties;

Expand All @@ -50,6 +51,11 @@ public OracleSessionProperties(OracleConfig config)
"Default scale for Oracle Number data type",
config.getDefaultNumberScale().orElse(null),
false))
.add(integerProperty(
JDBC_FETCH_SIZE,
"A hit number for oracle JDBC driver for fetching the result rows",
config.getJdbcFetchSize(),
false))
.build();
}

Expand All @@ -68,4 +74,9 @@ public static Optional<Integer> getNumberDefaultScale(ConnectorSession session)
{
return Optional.ofNullable(session.getProperty(NUMBER_DEFAULT_SCALE, Integer.class));
}

public static int getJdbcFetchSize(ConnectorSession session)
{
return session.getProperty(JDBC_FETCH_SIZE, Integer.class);
}
}
7 changes: 5 additions & 2 deletions plugin/trino-oracle/src/test/java/io/trino/plugin/oracle/TestOracleConfig.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,8 @@ public void testDefaults()
.setConnectionPoolEnabled(true)
.setConnectionPoolMinSize(1)
.setConnectionPoolMaxSize(30)
.setInactiveConnectionTimeout(new Duration(20, MINUTES)));
.setInactiveConnectionTimeout(new Duration(20, MINUTES))
.setJdbcFetchSize(1000));
}

@Test
Expand All @@ -58,6 +59,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.jdbc-fetch-size", "2000")
.buildOrThrow();

OracleConfig expected = new OracleConfig()
Expand All @@ -68,7 +70,8 @@ public void testExplicitPropertyMappings()
.setConnectionPoolEnabled(false)
.setConnectionPoolMinSize(10)
.setConnectionPoolMaxSize(20)
.setInactiveConnectionTimeout(new Duration(30, SECONDS));
.setInactiveConnectionTimeout(new Duration(30, SECONDS))
.setJdbcFetchSize(2000);

assertFullMapping(properties, expected);
}
Expand Down
5 changes: 3 additions & 2 deletions plugin/trino-postgresql/src/main/java/io/trino/plugin/postgresql/PostgreSqlClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -184,6 +184,7 @@
import static io.trino.plugin.postgresql.PostgreSqlConfig.ArrayMapping.AS_JSON;
import static io.trino.plugin.postgresql.PostgreSqlConfig.ArrayMapping.DISABLED;
import static io.trino.plugin.postgresql.PostgreSqlSessionProperties.getArrayMapping;
import static io.trino.plugin.postgresql.PostgreSqlSessionProperties.getJdbcFetchSize;
import static io.trino.plugin.postgresql.PostgreSqlSessionProperties.isEnableStringPushdownWithCollate;
import static io.trino.plugin.postgresql.TypeUtils.arrayDepth;
import static io.trino.plugin.postgresql.TypeUtils.getArrayElementPgTypeName;
Expand Down Expand Up @@ -378,13 +379,13 @@ protected void renameTable(ConnectorSession session, Connection connection, Stri
}

@Override
public PreparedStatement getPreparedStatement(Connection connection, String sql)
public PreparedStatement getPreparedStatement(ConnectorSession session, Connection connection, String sql)
throws SQLException
{
// fetch-size is ignored when connection is in auto-commit
connection.setAutoCommit(false);
PreparedStatement statement = connection.prepareStatement(sql);
statement.setFetchSize(1000);
statement.setFetchSize(getJdbcFetchSize(session));
return statement;
}

Expand Down
13 changes: 13 additions & 0 deletions plugin/trino-postgresql/src/main/java/io/trino/plugin/postgresql/PostgreSqlConfig.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ public class PostgreSqlConfig
private ArrayMapping arrayMapping = ArrayMapping.DISABLED;
private boolean includeSystemTables;
private boolean enableStringPushdownWithCollate;
private int jdbcFetchSize = 1000;

public enum ArrayMapping
{
Expand Down Expand Up @@ -68,4 +69,16 @@ public PostgreSqlConfig setEnableStringPushdownWithCollate(boolean enableStringP
this.enableStringPushdownWithCollate = enableStringPushdownWithCollate;
return this;
}

public int getJdbcFetchSize()
{
return jdbcFetchSize;
}

@Config("postgresql.jdbc-fetch-size")
public PostgreSqlConfig setJdbcFetchSize(int jdbcFetchSize)
{
this.jdbcFetchSize = jdbcFetchSize;
return this;
}
}
12 changes: 12 additions & 0 deletions ...rino-postgresql/src/main/java/io/trino/plugin/postgresql/PostgreSqlSessionProperties.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,12 +25,14 @@

import static io.trino.spi.session.PropertyMetadata.booleanProperty;
import static io.trino.spi.session.PropertyMetadata.enumProperty;
import static io.trino.spi.session.PropertyMetadata.integerProperty;

public final class PostgreSqlSessionProperties
implements SessionPropertiesProvider
{
public static final String ARRAY_MAPPING = "array_mapping";
public static final String ENABLE_STRING_PUSHDOWN_WITH_COLLATE = "enable_string_pushdown_with_collate";
public static final String JDBC_FETCH_SIZE = "jdbc_fetch_size";

private final List<PropertyMetadata<?>> sessionProperties;

Expand All @@ -48,6 +50,11 @@ public PostgreSqlSessionProperties(PostgreSqlConfig postgreSqlConfig)
ENABLE_STRING_PUSHDOWN_WITH_COLLATE,
"Enable string pushdown with collate (experimental)",
postgreSqlConfig.isEnableStringPushdownWithCollate(),
false),
integerProperty(
JDBC_FETCH_SIZE,
"A hit number for postgresql JDBC driver for fetching the result rows",
postgreSqlConfig.getJdbcFetchSize(),
false));
}

Expand All @@ -66,4 +73,9 @@ public static boolean isEnableStringPushdownWithCollate(ConnectorSession session
{
return session.getProperty(ENABLE_STRING_PUSHDOWN_WITH_COLLATE, Boolean.class);
}

public static int getJdbcFetchSize(ConnectorSession session)
{
return session.getProperty(JDBC_FETCH_SIZE, Integer.class);
}
}
7 changes: 5 additions & 2 deletions plugin/trino-postgresql/src/test/java/io/trino/plugin/postgresql/TestPostgreSqlConfig.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,8 @@ public void testDefaults()
assertRecordedDefaults(recordDefaults(PostgreSqlConfig.class)
.setArrayMapping(PostgreSqlConfig.ArrayMapping.DISABLED)
.setIncludeSystemTables(false)
.setEnableStringPushdownWithCollate(false));
.setEnableStringPushdownWithCollate(false)
.setJdbcFetchSize(1000));
}

@Test
Expand All @@ -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.jdbc-fetch-size", "2000")
.buildOrThrow();

PostgreSqlConfig expected = new PostgreSqlConfig()
.setArrayMapping(PostgreSqlConfig.ArrayMapping.AS_ARRAY)
.setIncludeSystemTables(true)
.setEnableStringPushdownWithCollate(true);
.setEnableStringPushdownWithCollate(true)
.setJdbcFetchSize(2000);

assertFullMapping(properties, expected);
}
Expand Down
2 changes: 1 addition & 1 deletion plugin/trino-redshift/src/main/java/io/trino/plugin/redshift/RedshiftClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -404,7 +404,7 @@ protected void renameTable(ConnectorSession session, Connection connection, Stri
}

@Override
public PreparedStatement getPreparedStatement(Connection connection, String sql)
public PreparedStatement getPreparedStatement(ConnectorSession session, Connection connection, String sql)
throws SQLException
{
// In PostgreSQL, fetch-size is ignored when connection is in auto-commit. Redshift JDBC documentation does not state this requirement
Expand Down
13 changes: 13 additions & 0 deletions plugin/trino-redshift/src/main/java/io/trino/plugin/redshift/RedshiftConfig.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
public class RedshiftConfig
{
private boolean legacyTypeMapping;
private int jdbcFetchSize = 1000;

public boolean isLegacyTypeMapping()
{
Expand All @@ -30,4 +31,16 @@ public RedshiftConfig setLegacyTypeMapping(boolean legacyTypeMapping)
this.legacyTypeMapping = legacyTypeMapping;
return this;
}

public int getJdbcFetchSize()
{
return jdbcFetchSize;
}

@Config("redshift.jdbc-fetch-size")
public RedshiftConfig setJdbcFetchSize(int jdbcFetchSize)
{
this.jdbcFetchSize = jdbcFetchSize;
return this;
}
}
7 changes: 5 additions & 2 deletions plugin/trino-redshift/src/test/java/io/trino/plugin/redshift/TestRedshiftConfig.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,18 +28,21 @@ public class TestRedshiftConfig
public void testDefaults()
{
assertRecordedDefaults(recordDefaults(RedshiftConfig.class)
.setLegacyTypeMapping(false));
.setLegacyTypeMapping(false)
.setJdbcFetchSize(1000));
}

@Test
public void testExplicitPropertyMappings()
{
Map<String, String> properties = ImmutableMap.<String, String>builder()
.put("redshift.use-legacy-type-mapping", "true")
.put("redshift.jdbc-fetch-size", "2000")
.buildOrThrow();

RedshiftConfig expected = new RedshiftConfig()
.setLegacyTypeMapping(true);
.setLegacyTypeMapping(true)
.setJdbcFetchSize(2000);

assertFullMapping(properties, expected);
}
Expand Down