Accept PreparedQuery in QueryBuilder

trinodb · Jan 25, 2021 · 6442c96 · 6442c96
1 parent 75baaae
commit 6442c96
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 7 deletions.
diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/BaseJdbcClient.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/BaseJdbcClient.java
@@ -428,7 +428,7 @@ public PreparedStatement buildSql(ConnectorSession session, Connection connectio
         PreparedQuery preparedQuery = queryBuilder.prepareQuery(
                 session,
                 connection,
-                table.getRemoteTableName(),
+                table.getRelationHandle(),
                 table.getGroupingSets(),
                 columns,
                 table.getConstraint(),

diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/QueryBuilder.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/QueryBuilder.java
@@ -20,6 +20,7 @@
 import io.airlift.slice.Slice;
 import io.trino.spi.connector.ColumnHandle;
 import io.trino.spi.connector.ConnectorSession;
+import io.trino.spi.connector.SchemaTableName;
 import io.trino.spi.predicate.Domain;
 import io.trino.spi.predicate.Range;
 import io.trino.spi.predicate.TupleDomain;
@@ -78,7 +79,10 @@ public PreparedStatement buildSql(
         PreparedQuery preparedQuery = prepareQuery(
                 session,
                 connection,
-                remoteTableName,
+                new JdbcNamedRelationHandle(
+                        // This dummy SchemaTableName is not used for anything here. It's provided only to implement the deprecated buildSql() method
+                        new SchemaTableName(remoteTableName.getSchemaName().orElse(""), remoteTableName.getTableName()),
+                        remoteTableName),
                 groupingSets,
                 columns,
                 tupleDomain,
@@ -90,17 +94,27 @@ public PreparedStatement buildSql(
     public PreparedQuery prepareQuery(
             ConnectorSession session,
             Connection connection,
-            RemoteTableName remoteTableName,
+            JdbcRelationHandle baseRelation,
             Optional<List<List<JdbcColumnHandle>>> groupingSets,
             List<JdbcColumnHandle> columns,
             TupleDomain<ColumnHandle> tupleDomain,
             Optional<String> additionalPredicate)
     {
-        String sql = "SELECT " + getProjection(columns);
-        sql += " FROM " + getRelation(remoteTableName);
-
         ImmutableList.Builder<QueryParameter> accumulator = ImmutableList.builder();
 
+        String sql = "SELECT " + getProjection(columns);
+        if (baseRelation instanceof JdbcNamedRelationHandle) {
+            sql += " FROM " + getRelation(((JdbcNamedRelationHandle) baseRelation).getRemoteTableName());
+        }
+        else if (baseRelation instanceof JdbcQueryRelationHandle) {
+            PreparedQuery preparedQuery = ((JdbcQueryRelationHandle) baseRelation).getPreparedQuery();
+            sql += " FROM (" + preparedQuery.getQuery() + ") o";
+            accumulator.addAll(preparedQuery.getParameters());
+        }
+        else {
+            throw new IllegalArgumentException("Unsupported relation: " + baseRelation);
+        }
+
         List<String> clauses = toConjuncts(client, session, connection, tupleDomain, accumulator::add);
         if (additionalPredicate.isPresent()) {
             clauses = ImmutableList.<String>builder()

diff --git a/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/TestJdbcQueryBuilder.java b/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/TestJdbcQueryBuilder.java
@@ -20,6 +20,7 @@
 import com.google.common.collect.Multiset;
 import io.trino.spi.connector.ColumnHandle;
 import io.trino.spi.connector.ConnectorSession;
+import io.trino.spi.connector.SchemaTableName;
 import io.trino.spi.predicate.Domain;
 import io.trino.spi.predicate.Range;
 import io.trino.spi.predicate.SortedRangeSet;
@@ -87,7 +88,9 @@
 @Test(singleThreaded = true)
 public class TestJdbcQueryBuilder
 {
-    private static final RemoteTableName TEST_TABLE = new RemoteTableName(Optional.empty(), Optional.empty(), "test_table");
+    private static final JdbcNamedRelationHandle TEST_TABLE = new JdbcNamedRelationHandle(new SchemaTableName(
+            "some_test_schema", "test_table"),
+            new RemoteTableName(Optional.empty(), Optional.empty(), "test_table"));
     private static final ConnectorSession SESSION = TestingConnectorSession.builder()
             .setPropertyMetadata(new JdbcMetadataSessionProperties(new JdbcMetadataConfig(), Optional.empty()).getSessionProperties())
             .build();