Skip to content

Commit

Permalink
Pass JdbcClient as a first argument to QueryBuilder methods
Browse files Browse the repository at this point in the history
  • Loading branch information
@wendigo @hashhar
wendigo authored and hashhar committed Feb 25, 2022
1 parent 1f05308 commit 9143b40
Show file tree
Hide file tree
Showing 6 changed files with 35 additions and 26 deletions.
8 changes: 5 additions & 3 deletions plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/BaseJdbcClient.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -377,7 +377,7 @@ public PreparedStatement buildSql(ConnectorSession session, Connection connectio
throws SQLException
{
PreparedQuery preparedQuery = prepareQuery(session, connection, table, Optional.empty(), columns, ImmutableMap.of(), Optional.of(split));
return new QueryBuilder(this).prepareStatement(session, connection, preparedQuery);
return new QueryBuilder(this).prepareStatement(this, session, connection, preparedQuery);
}

protected PreparedQuery prepareQuery(
Expand All @@ -390,6 +390,7 @@ protected PreparedQuery prepareQuery(
Optional<JdbcSplit> split)
{
return applyQueryTransformations(table, new QueryBuilder(this).prepareQuery(
this,
session,
connection,
table.getRelationHandle(),
Expand Down Expand Up @@ -419,6 +420,7 @@ public Optional<PreparedQuery> implementJoin(

QueryBuilder queryBuilder = new QueryBuilder(this);
return Optional.of(queryBuilder.prepareJoinQuery(
this,
session,
joinType,
leftSource,
Expand Down Expand Up @@ -971,8 +973,8 @@ public OptionalLong delete(ConnectorSession session, JdbcTableHandle handle)
try (Connection connection = connectionFactory.openConnection(session)) {
verify(connection.getAutoCommit());
QueryBuilder queryBuilder = new QueryBuilder(this);
PreparedQuery preparedQuery = queryBuilder.prepareDelete(session, connection, handle.getRequiredNamedRelation(), handle.getConstraint());
try (PreparedStatement preparedStatement = queryBuilder.prepareStatement(session, connection, preparedQuery)) {
PreparedQuery preparedQuery = queryBuilder.prepareDelete(this, session, connection, handle.getRequiredNamedRelation(), handle.getConstraint());
try (PreparedStatement preparedStatement = queryBuilder.prepareStatement(this, session, connection, preparedQuery)) {
return OptionalLong.of(preparedStatement.executeUpdate());
}
}
Expand Down
4 changes: 4 additions & 0 deletions plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/QueryBuilder.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,7 @@ public QueryBuilder(JdbcClient client)
}

public PreparedQuery prepareQuery(
JdbcClient client,
ConnectorSession session,
Connection connection,
JdbcRelationHandle baseRelation,
Expand Down Expand Up @@ -101,6 +102,7 @@ public PreparedQuery prepareQuery(
}

public PreparedQuery prepareJoinQuery(
JdbcClient client,
ConnectorSession session,
JoinType joinType,
PreparedQuery leftSource,
Expand Down Expand Up @@ -159,6 +161,7 @@ protected static String formatJoinType(JoinType joinType)
}

public PreparedQuery prepareDelete(
JdbcClient client,
ConnectorSession session,
Connection connection,
JdbcNamedRelationHandle baseRelation,
Expand All @@ -176,6 +179,7 @@ public PreparedQuery prepareDelete(
}

public PreparedStatement prepareStatement(
JdbcClient client,
ConnectorSession session,
Connection connection,
PreparedQuery preparedQuery)
Expand Down
41 changes: 22 additions & 19 deletions plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/TestJdbcQueryBuilder.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -224,8 +224,8 @@ public void testNormalBuildSql()

Connection connection = database.getConnection();
QueryBuilder queryBuilder = new QueryBuilder(jdbcClient);
PreparedQuery preparedQuery = queryBuilder.prepareQuery(SESSION, connection, TEST_TABLE, Optional.empty(), columns, Map.of(), tupleDomain, Optional.empty());
try (PreparedStatement preparedStatement = queryBuilder.prepareStatement(SESSION, connection, preparedQuery)) {
PreparedQuery preparedQuery = queryBuilder.prepareQuery(jdbcClient, SESSION, connection, TEST_TABLE, Optional.empty(), columns, Map.of(), tupleDomain, Optional.empty());
try (PreparedStatement preparedStatement = queryBuilder.prepareStatement(jdbcClient, SESSION, connection, preparedQuery)) {
assertThat(preparedQuery.getQuery()).isEqualTo("" +
"SELECT \"col_0\", \"col_1\", \"col_2\", \"col_3\", \"col_4\", \"col_5\", " +
"\"col_6\", \"col_7\", \"col_8\", \"col_9\", \"col_10\", \"col_11\" " +
Expand Down Expand Up @@ -267,6 +267,7 @@ public void testBuildSqlWithDomainComplement()
Connection connection = database.getConnection();
QueryBuilder queryBuilder = new QueryBuilder(jdbcClient);
PreparedQuery preparedQuery = queryBuilder.prepareQuery(
jdbcClient,
SESSION,
connection,
TEST_TABLE,
Expand All @@ -275,7 +276,7 @@ public void testBuildSqlWithDomainComplement()
Map.of(),
tupleDomain,
Optional.empty());
try (PreparedStatement preparedStatement = queryBuilder.prepareStatement(SESSION, connection, preparedQuery)) {
try (PreparedStatement preparedStatement = queryBuilder.prepareStatement(jdbcClient, SESSION, connection, preparedQuery)) {
assertThat(preparedQuery.getQuery()).isEqualTo("" +
"SELECT \"col_0\", \"col_3\", \"col_9\" " +
"FROM \"test_table\" " +
Expand Down Expand Up @@ -306,8 +307,8 @@ public void testBuildSqlWithFloat()

Connection connection = database.getConnection();
QueryBuilder queryBuilder = new QueryBuilder(jdbcClient);
PreparedQuery preparedQuery = queryBuilder.prepareQuery(SESSION, connection, TEST_TABLE, Optional.empty(), columns, Map.of(), tupleDomain, Optional.empty());
try (PreparedStatement preparedStatement = queryBuilder.prepareStatement(SESSION, connection, preparedQuery)) {
PreparedQuery preparedQuery = queryBuilder.prepareQuery(jdbcClient, SESSION, connection, TEST_TABLE, Optional.empty(), columns, Map.of(), tupleDomain, Optional.empty());
try (PreparedStatement preparedStatement = queryBuilder.prepareStatement(jdbcClient, SESSION, connection, preparedQuery)) {
assertThat(preparedQuery.getQuery()).isEqualTo("" +
"SELECT \"col_0\", \"col_1\", \"col_2\", \"col_3\", \"col_4\", \"col_5\", " +
"\"col_6\", \"col_7\", \"col_8\", \"col_9\", \"col_10\", \"col_11\" " +
Expand Down Expand Up @@ -340,8 +341,8 @@ public void testBuildSqlWithVarchar()

Connection connection = database.getConnection();
QueryBuilder queryBuilder = new QueryBuilder(jdbcClient);
PreparedQuery preparedQuery = queryBuilder.prepareQuery(SESSION, connection, TEST_TABLE, Optional.empty(), columns, Map.of(), tupleDomain, Optional.empty());
try (PreparedStatement preparedStatement = queryBuilder.prepareStatement(SESSION, connection, preparedQuery)) {
PreparedQuery preparedQuery = queryBuilder.prepareQuery(jdbcClient, SESSION, connection, TEST_TABLE, Optional.empty(), columns, Map.of(), tupleDomain, Optional.empty());
try (PreparedStatement preparedStatement = queryBuilder.prepareStatement(jdbcClient, SESSION, connection, preparedQuery)) {
assertThat(preparedQuery.getQuery()).isEqualTo("" +
"SELECT \"col_0\", \"col_1\", \"col_2\", \"col_3\", \"col_4\", \"col_5\", " +
"\"col_6\", \"col_7\", \"col_8\", \"col_9\", \"col_10\", \"col_11\" " +
Expand Down Expand Up @@ -376,8 +377,8 @@ public void testBuildSqlWithChar()

Connection connection = database.getConnection();
QueryBuilder queryBuilder = new QueryBuilder(jdbcClient);
PreparedQuery preparedQuery = queryBuilder.prepareQuery(SESSION, connection, TEST_TABLE, Optional.empty(), columns, Map.of(), tupleDomain, Optional.empty());
try (PreparedStatement preparedStatement = queryBuilder.prepareStatement(SESSION, connection, preparedQuery)) {
PreparedQuery preparedQuery = queryBuilder.prepareQuery(jdbcClient, SESSION, connection, TEST_TABLE, Optional.empty(), columns, Map.of(), tupleDomain, Optional.empty());
try (PreparedStatement preparedStatement = queryBuilder.prepareStatement(jdbcClient, SESSION, connection, preparedQuery)) {
assertThat(preparedQuery.getQuery()).isEqualTo("" +
"SELECT \"col_0\", \"col_1\", \"col_2\", \"col_3\", \"col_4\", \"col_5\", " +
"\"col_6\", \"col_7\", \"col_8\", \"col_9\", \"col_10\", \"col_11\" " +
Expand Down Expand Up @@ -417,8 +418,8 @@ public void testBuildSqlWithDateTime()

Connection connection = database.getConnection();
QueryBuilder queryBuilder = new QueryBuilder(jdbcClient);
PreparedQuery preparedQuery = queryBuilder.prepareQuery(SESSION, connection, TEST_TABLE, Optional.empty(), columns, Map.of(), tupleDomain, Optional.empty());
try (PreparedStatement preparedStatement = queryBuilder.prepareStatement(SESSION, connection, preparedQuery)) {
PreparedQuery preparedQuery = queryBuilder.prepareQuery(jdbcClient, SESSION, connection, TEST_TABLE, Optional.empty(), columns, Map.of(), tupleDomain, Optional.empty());
try (PreparedStatement preparedStatement = queryBuilder.prepareStatement(jdbcClient, SESSION, connection, preparedQuery)) {
assertThat(preparedQuery.getQuery()).isEqualTo("" +
"SELECT \"col_0\", \"col_1\", \"col_2\", \"col_3\", \"col_4\", \"col_5\", " +
"\"col_6\", \"col_7\", \"col_8\", \"col_9\", \"col_10\", \"col_11\" " +
Expand Down Expand Up @@ -458,8 +459,8 @@ public void testBuildSqlWithTimestamp()

Connection connection = database.getConnection();
QueryBuilder queryBuilder = new QueryBuilder(jdbcClient);
PreparedQuery preparedQuery = queryBuilder.prepareQuery(SESSION, connection, TEST_TABLE, Optional.empty(), columns, Map.of(), tupleDomain, Optional.empty());
try (PreparedStatement preparedStatement = queryBuilder.prepareStatement(SESSION, connection, preparedQuery)) {
PreparedQuery preparedQuery = queryBuilder.prepareQuery(jdbcClient, SESSION, connection, TEST_TABLE, Optional.empty(), columns, Map.of(), tupleDomain, Optional.empty());
try (PreparedStatement preparedStatement = queryBuilder.prepareStatement(jdbcClient, SESSION, connection, preparedQuery)) {
assertThat(preparedQuery.getQuery()).isEqualTo("" +
"SELECT \"col_0\", \"col_1\", \"col_2\", \"col_3\", \"col_4\", \"col_5\", " +
"\"col_6\", \"col_7\", \"col_8\", \"col_9\", \"col_10\", \"col_11\" " +
Expand Down Expand Up @@ -490,9 +491,9 @@ public void testBuildSqlWithLimit()
Connection connection = database.getConnection();
Function<String, String> function = sql -> sql + " LIMIT 10";
QueryBuilder queryBuilder = new QueryBuilder(jdbcClient);
PreparedQuery preparedQuery = queryBuilder.prepareQuery(SESSION, connection, TEST_TABLE, Optional.empty(), columns, Map.of(), TupleDomain.all(), Optional.empty());
PreparedQuery preparedQuery = queryBuilder.prepareQuery(jdbcClient, SESSION, connection, TEST_TABLE, Optional.empty(), columns, Map.of(), TupleDomain.all(), Optional.empty());
preparedQuery = preparedQuery.transformQuery(function);
try (PreparedStatement preparedStatement = queryBuilder.prepareStatement(SESSION, connection, preparedQuery)) {
try (PreparedStatement preparedStatement = queryBuilder.prepareStatement(jdbcClient, SESSION, connection, preparedQuery)) {
assertThat(preparedQuery.getQuery()).isEqualTo("" +
"SELECT \"col_0\", \"col_1\", \"col_2\", \"col_3\", \"col_4\", \"col_5\", " +
"\"col_6\", \"col_7\", \"col_8\", \"col_9\", \"col_10\", \"col_11\" " +
Expand All @@ -518,8 +519,8 @@ public void testEmptyBuildSql()

Connection connection = database.getConnection();
QueryBuilder queryBuilder = new QueryBuilder(jdbcClient);
PreparedQuery preparedQuery = queryBuilder.prepareQuery(SESSION, connection, TEST_TABLE, Optional.empty(), columns, Map.of(), tupleDomain, Optional.empty());
try (PreparedStatement preparedStatement = queryBuilder.prepareStatement(SESSION, connection, preparedQuery)) {
PreparedQuery preparedQuery = queryBuilder.prepareQuery(jdbcClient, SESSION, connection, TEST_TABLE, Optional.empty(), columns, Map.of(), tupleDomain, Optional.empty());
try (PreparedStatement preparedStatement = queryBuilder.prepareStatement(jdbcClient, SESSION, connection, preparedQuery)) {
assertThat(preparedQuery.getQuery()).isEqualTo("" +
"SELECT \"col_0\", \"col_1\", \"col_2\", \"col_3\", \"col_4\", \"col_5\", " +
"\"col_6\", \"col_7\", \"col_8\", \"col_9\", \"col_10\", \"col_11\" " +
Expand Down Expand Up @@ -547,6 +548,7 @@ public void testAggregation()
Connection connection = database.getConnection();
QueryBuilder queryBuilder = new QueryBuilder(jdbcClient);
PreparedQuery preparedQuery = queryBuilder.prepareQuery(
jdbcClient,
SESSION,
connection,
TEST_TABLE,
Expand All @@ -555,7 +557,7 @@ public void testAggregation()
Map.of("s", "sum(\"col_0\")"),
TupleDomain.all(),
Optional.empty());
try (PreparedStatement preparedStatement = queryBuilder.prepareStatement(SESSION, connection, preparedQuery)) {
try (PreparedStatement preparedStatement = queryBuilder.prepareStatement(jdbcClient, SESSION, connection, preparedQuery)) {
assertThat(preparedQuery.getQuery()).isEqualTo("" +
"SELECT \"col_2\", sum(\"col_0\") AS \"s\" " +
"FROM \"test_table\" " +
Expand Down Expand Up @@ -590,6 +592,7 @@ public void testAggregationWithFilter()
Connection connection = database.getConnection();
QueryBuilder queryBuilder = new QueryBuilder(jdbcClient);
PreparedQuery preparedQuery = queryBuilder.prepareQuery(
jdbcClient,
SESSION,
connection,
TEST_TABLE,
Expand All @@ -598,7 +601,7 @@ public void testAggregationWithFilter()
Map.of("s", "sum(\"col_0\")"),
tupleDomain,
Optional.empty());
try (PreparedStatement preparedStatement = queryBuilder.prepareStatement(SESSION, connection, preparedQuery)) {
try (PreparedStatement preparedStatement = queryBuilder.prepareStatement(jdbcClient, SESSION, connection, preparedQuery)) {
assertThat(preparedQuery.getQuery()).isEqualTo("" +
"SELECT \"col_2\", sum(\"col_0\") AS \"s\" " +
"FROM \"test_table\" " +
Expand Down
2 changes: 1 addition & 1 deletion plugin/trino-phoenix/src/main/java/io/trino/plugin/phoenix/PhoenixClient.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -294,7 +294,7 @@ public PreparedStatement prepareStatement(
columns,
ImmutableMap.of(),
split);
return new QueryBuilder(this).prepareStatement(session, connection, preparedQuery);
return new QueryBuilder(this).prepareStatement(this, session, connection, preparedQuery);
}

@Override
Expand Down
2 changes: 1 addition & 1 deletion plugin/trino-phoenix5/src/main/java/io/trino/plugin/phoenix5/PhoenixClient.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -294,7 +294,7 @@ public PreparedStatement prepareStatement(
columns,
ImmutableMap.of(),
split);
return new QueryBuilder(this).prepareStatement(session, connection, preparedQuery);
return new QueryBuilder(this).prepareStatement(this, session, connection, preparedQuery);
}

@Override
Expand Down
4 changes: 2 additions & 2 deletions plugin/trino-postgresql/src/main/java/io/trino/plugin/postgresql/PostgreSqlClient.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -803,8 +803,8 @@ public OptionalLong delete(ConnectorSession session, JdbcTableHandle handle)
try (Connection connection = connectionFactory.openConnection(session)) {
verify(connection.getAutoCommit());
QueryBuilder queryBuilder = new QueryBuilder(this);
PreparedQuery preparedQuery = queryBuilder.prepareDelete(session, connection, handle.getRequiredNamedRelation(), handle.getConstraint());
try (PreparedStatement preparedStatement = queryBuilder.prepareStatement(session, connection, preparedQuery)) {
PreparedQuery preparedQuery = queryBuilder.prepareDelete(this, session, connection, handle.getRequiredNamedRelation(), handle.getConstraint());
try (PreparedStatement preparedStatement = queryBuilder.prepareStatement(this, session, connection, preparedQuery)) {
int affectedRowsCount = preparedStatement.executeUpdate();
// In getPreparedStatement we set autocommit to false so here we need an explicit commit
connection.commit();
Expand Down

0 comments on commit 9143b40

Please sign in to comment.