Improvement of ResultSet metadata

Author: prrvchr

client/trino-client/src/main/java/io/trino/client/Column.java

@@ -18,33 +18,80 @@
 import com.google.errorprone.annotations.Immutable;
 
 import java.util.Objects;
+import java.util.Optional;
 
 import static java.util.Objects.requireNonNull;
 
 @Immutable
 public class Column
 {
+    private final Optional<String> catalog;
+    private final Optional<String> schema;
+    private final Optional<String> table;
     private final String name;
+    private final Optional<String> label;
     private final String type;
     private final ClientTypeSignature typeSignature;
 
+    public Column(String name, String type, ClientTypeSignature typeSignature)
+    {
+        this(Optional.empty(), Optional.empty(), Optional.empty(), name, Optional.of(name), type, typeSignature);
+    }
+
+    public Column(String catalog, String schema, String table, String name, String label, String type, ClientTypeSignature typeSignature)
+    {
+        this(Optional.of(catalog), Optional.of(schema), Optional.of(table), name, Optional.of(label), type, typeSignature);
+    }
+
     @JsonCreator
     public Column(
+            @JsonProperty("catalog") Optional<String> catalog,
+            @JsonProperty("schema") Optional<String> schema,
+            @JsonProperty("table") Optional<String> table,
             @JsonProperty("name") String name,
+            @JsonProperty("label") Optional<String> label,
             @JsonProperty("type") String type,
             @JsonProperty("typeSignature") ClientTypeSignature typeSignature)
     {
+        this.catalog = requireNonNull(catalog, "catalog is null");
+        this.schema = requireNonNull(schema, "schema is null");
+        this.table = requireNonNull(table, "table is null");
         this.name = requireNonNull(name, "name is null");
+        this.label = requireNonNull(label, "label is null");
         this.type = requireNonNull(type, "type is null");
         this.typeSignature = typeSignature;
     }
 
+    @JsonProperty
+    public Optional<String> getCatalog()
+    {
+        return catalog;
+    }
+
+    @JsonProperty
+    public Optional<String> getSchema()
+    {
+        return schema;
+    }
+
+    @JsonProperty
+    public Optional<String> getTable()
+    {
+        return table;
+    }
+
     @JsonProperty
     public String getName()
     {
         return name;
     }
 
+    @JsonProperty
+    public Optional<String> getLabel()
+    {
+        return label;
+    }
+
     @JsonProperty
     public String getType()
     {

client/trino-client/src/test/java/io/trino/client/TestQueryResults.java

@@ -29,7 +29,11 @@ public class TestQueryResults
             "  \"id\" : \"20160128_214710_00012_rk68b\",\n" +
             "  \"infoUri\" : \"http://localhost:54855/query.html?20160128_214710_00012_rk68b\",\n" +
             "  \"columns\" : [ {\n" +
+            "    \"catalog\" : \"_cat0\",\n" +
+            "    \"schema\" : \"_sch0\",\n" +
+            "    \"table\" : \"_tab0\",\n" +
             "    \"name\" : \"_col0\",\n" +
+            "    \"label\" : \"_lab0\",\n" +
             "    \"type\" : \"bigint\",\n" +
             "    \"typeSignature\" : {\n" +
             "      \"rawType\" : \"varchar\",\n" +

client/trino-jdbc/src/main/java/io/trino/jdbc/AbstractTrinoResultSet.java

@@ -1970,11 +1970,11 @@ private static List<ColumnInfo> getColumnInfo(List<Column> columns)
         ImmutableList.Builder<ColumnInfo> list = ImmutableList.builderWithExpectedSize(columns.size());
         for (Column column : columns) {
             ColumnInfo.Builder builder = new ColumnInfo.Builder()
-                    .setCatalogName("") // TODO
-                    .setSchemaName("") // TODO
-                    .setTableName("") // TODO
+                    .setCatalogName(column.getCatalog().orElse(""))
+                    .setSchemaName(column.getSchema().orElse(""))
+                    .setTableName(column.getTable().orElse(""))
                     .setColumnLabel(column.getName())
-                    .setColumnName(column.getName()) // TODO
+                    .setColumnName(column.getLabel().orElse(column.getName()))
                     .setColumnTypeSignature(column.getTypeSignature())
                     .setNullable(Nullable.UNKNOWN)
                     .setCurrency(false);

client/trino-jdbc/src/test/java/io/trino/jdbc/TestTrinoDatabaseMetaData.java

@@ -0,0 +1,26 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.execution;
+
+import java.util.Optional;
+
+// XXX: ColumnInfo allows the grouping of all the naming elements necessary to deliver complete ResultSetMetaData (see issue#22306).
+public record ColumnInfo(Optional<String> catalog, Optional<String> schema, Optional<String> table, String name, Optional<String> label)
+{
+    // XXX: This constructor is used by the PlanBuilder.output() method and allows you to create ColumnInfo using only a column name.
+    public ColumnInfo(String name)
+    {
+        this(Optional.empty(), Optional.empty(), Optional.empty(), name, Optional.of(name));
+    }
+}

core/trino-main/src/main/java/io/trino/execution/ColumnInfo.java

@@ -34,6 +34,7 @@
 import java.util.Queue;
 import java.util.function.Consumer;
 
+import static com.google.common.collect.ImmutableList.toImmutableList;
 import static java.util.Objects.requireNonNull;
 
 public interface QueryExecution
@@ -99,22 +100,27 @@ interface QueryExecutionFactory<T extends QueryExecution>
      */
     class QueryOutputInfo
     {
-        private final List<String> columnNames;
+        private final List<ColumnInfo> columnInfos;
         private final List<Type> columnTypes;
         private final Queue<ExchangeInput> inputsQueue;
         private final boolean noMoreInputs;
 
-        public QueryOutputInfo(List<String> columnNames, List<Type> columnTypes, Queue<ExchangeInput> inputsQueue, boolean noMoreInputs)
+        public QueryOutputInfo(List<ColumnInfo> columnInfos, List<Type> columnTypes, Queue<ExchangeInput> inputsQueue, boolean noMoreInputs)
         {
-            this.columnNames = ImmutableList.copyOf(requireNonNull(columnNames, "columnNames is null"));
+            this.columnInfos = ImmutableList.copyOf(requireNonNull(columnInfos, "columnInfos is null"));
             this.columnTypes = ImmutableList.copyOf(requireNonNull(columnTypes, "columnTypes is null"));
             this.inputsQueue = requireNonNull(inputsQueue, "inputsQueue is null");
             this.noMoreInputs = noMoreInputs;
         }
 
+        public List<ColumnInfo> getColumnInfos()
+        {
+            return columnInfos;
+        }
+
         public List<String> getColumnNames()
         {
-            return columnNames;
+            return columnInfos.stream().map(ColumnInfo::name).collect(toImmutableList());
         }
 
         public List<Type> getColumnTypes()

core/trino-main/src/main/java/io/trino/execution/QueryExecution.java

@@ -1010,9 +1010,9 @@ public void outputTaskFailed(TaskId taskId, Throwable failure)
         outputManager.outputTaskFailed(taskId, failure);
     }
 
-    public void setColumns(List<String> columnNames, List<Type> columnTypes)
+    public void setColumns(List<ColumnInfo> columnInfos, List<Type> columnTypes)
     {
-        outputManager.setColumns(columnNames, columnTypes);
+        outputManager.setColumns(columnInfos, columnTypes);
     }
 
     public void updateInputsForQueryResults(List<ExchangeInput> inputs, boolean noMoreInputs)
@@ -1637,7 +1637,7 @@ public static class QueryOutputManager
         private Optional<Consumer<QueryOutputInfo>> listener = Optional.empty();
 
         @GuardedBy("this")
-        private List<String> columnNames;
+        private List<ColumnInfo> columnInfos;
         @GuardedBy("this")
         private List<Type> columnTypes;
         @GuardedBy("this")
@@ -1670,17 +1670,17 @@ public void setOutputInfoListener(Consumer<QueryOutputInfo> listener)
             fireStateChangedIfReady(queryOutputInfo, Optional.of(listener));
         }
 
-        public void setColumns(List<String> columnNames, List<Type> columnTypes)
+        public void setColumns(List<ColumnInfo> columnInfos, List<Type> columnTypes)
         {
-            requireNonNull(columnNames, "columnNames is null");
+            requireNonNull(columnInfos, "columnInfos is null");
             requireNonNull(columnTypes, "columnTypes is null");
-            checkArgument(columnNames.size() == columnTypes.size(), "columnNames and columnTypes must be the same size");
+            checkArgument(columnInfos.size() == columnTypes.size(), "columnInfos and columnTypes must be the same size");
 
             Optional<QueryOutputInfo> queryOutputInfo;
             Optional<Consumer<QueryOutputInfo>> listener;
             synchronized (this) {
-                checkState(this.columnNames == null && this.columnTypes == null, "output fields already set");
-                this.columnNames = ImmutableList.copyOf(columnNames);
+                checkState(this.columnInfos == null && this.columnTypes == null, "output fields already set");
+                this.columnInfos = ImmutableList.copyOf(columnInfos);
                 this.columnTypes = ImmutableList.copyOf(columnTypes);
 
                 queryOutputInfo = getQueryOutputInfo();
@@ -1755,10 +1755,10 @@ public void outputTaskFailed(TaskId taskId, Throwable failure)
 
         private synchronized Optional<QueryOutputInfo> getQueryOutputInfo()
         {
-            if (columnNames == null || columnTypes == null) {
+            if (columnInfos == null || columnTypes == null) {
                 return Optional.empty();
             }
-            return Optional.of(new QueryOutputInfo(columnNames, columnTypes, inputsQueue, noMoreInputs));
+            return Optional.of(new QueryOutputInfo(columnInfos, columnTypes, inputsQueue, noMoreInputs));
         }
 
         private void fireStateChangedIfReady(Optional<QueryOutputInfo> info, Optional<Consumer<QueryOutputInfo>> listener)

core/trino-main/src/main/java/io/trino/execution/QueryStateMachine.java

@@ -524,7 +524,7 @@ private void planDistribution(PlanRoot plan, CachingTableStatsProvider tableStat
         // record output field
         PlanFragment rootFragment = plan.getRoot().getFragment();
         stateMachine.setColumns(
-                ((OutputNode) rootFragment.getRoot()).getColumnNames(),
+                ((OutputNode) rootFragment.getRoot()).getColumnInfos(),
                 rootFragment.getTypes());
 
         RetryPolicy retryPolicy = getRetryPolicy(getSession());

core/trino-main/src/main/java/io/trino/execution/SqlQueryExecution.java

@@ -13,15 +13,16 @@
  */
 package io.trino.server.protocol;
 
+import io.trino.execution.ColumnInfo;
 import io.trino.spi.type.Type;
 
 import static java.util.Objects.requireNonNull;
 
-public record OutputColumn(int sourcePageChannel, String columnName, Type type)
+public record OutputColumn(int sourcePageChannel, ColumnInfo columnInfo, Type type)
 {
     public OutputColumn
     {
-        requireNonNull(columnName, "columnName is null");
+        requireNonNull(columnInfo, "columnInfo is null");
         requireNonNull(type, "type is null");
 
         if (sourcePageChannel < 0) {

core/trino-main/src/main/java/io/trino/server/protocol/OutputColumn.java

@@ -29,6 +29,7 @@
 import io.trino.execution.BasicStageInfo;
 import io.trino.execution.BasicStageStats;
 import io.trino.execution.BasicStagesInfo;
+import io.trino.execution.ColumnInfo;
 import io.trino.execution.ExecutionFailureInfo;
 import io.trino.execution.QueryState;
 import io.trino.execution.StageId;
@@ -53,6 +54,7 @@
 
 import java.util.HashSet;
 import java.util.List;
+import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
 
@@ -75,10 +77,20 @@ public final class ProtocolUtil
     private ProtocolUtil() {}
 
     public static Column createColumn(String name, Type type, boolean supportsParametricDateTime)
+    {
+        return createColumn(Optional.empty(), Optional.empty(), Optional.empty(), name, Optional.of(name), type, supportsParametricDateTime);
+    }
+
+    public static Column createColumn(ColumnInfo columnInfo, Type type, boolean supportsParametricDateTime)
+    {
+        return createColumn(columnInfo.catalog(), columnInfo.schema(), columnInfo.table(), columnInfo.name(), columnInfo.label(), type, supportsParametricDateTime);
+    }
+
+    public static Column createColumn(Optional<String> catalog, Optional<String> schema, Optional<String> table, String name, Optional<String> label, Type type, boolean supportsParametricDateTime)
     {
         String formatted = formatType(TypeSignatureTranslator.toSqlType(type), supportsParametricDateTime);
 
-        return new Column(name, formatted, toClientTypeSignature(type.getTypeSignature(), supportsParametricDateTime));
+        return new Column(catalog, schema, table, name, label, formatted, toClientTypeSignature(type.getTypeSignature(), supportsParametricDateTime));
     }
 
     private static String formatType(DataType type, boolean supportsParametricDateTime)