Add VARIANT_JSON and VARIANT_BINARY client capabilities

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

@@ -13,6 +13,11 @@
  */
 package io.trino.client;
 
+import java.util.Set;
+
+import static java.util.Arrays.stream;
+import static java.util.stream.Collectors.toUnmodifiableSet;
+
 public enum ClientCapabilities
 {
     PATH,
@@ -33,8 +38,45 @@ public enum ClientCapabilities
      */
     NUMBER,
 
+    /**
+     * Whether client supports the `VARIANT` type encoded as JSON values on the wire.
+     * When this capability is not set, the server returns `json` for `VARIANT` columns.
+     */
+    VARIANT_JSON,
+
+    /**
+     * Whether client supports the `VARIANT` type encoded as a binary payload on the wire.
+     * This capability is opt-in, so clients continue to receive the JSON representation by default.
+     */
+    VARIANT_BINARY(false),
+
     /**
      * Whether clients support the session authorization set/reset feature
      */
     SESSION_AUTHORIZATION;
+
+    private final boolean enabledByDefault;
+
+    ClientCapabilities()
+    {
+        this(true);
+    }
+
+    ClientCapabilities(boolean enabledByDefault)
+    {
+        this.enabledByDefault = enabledByDefault;
+    }
+
+    public boolean enabledByDefault()
+    {
+        return enabledByDefault;
+    }
+
+    public static Set<String> defaultClientCapabilities()
+    {
+        return stream(values())
+                .filter(ClientCapabilities::enabledByDefault)
+                .map(Enum::name)
+                .collect(toUnmodifiableSet());
+    }
 }

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

@@ -0,0 +1,47 @@
+/*
+ * 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.client;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import static java.util.Objects.requireNonNull;
+
+public final class EncodedVariant
+{
+    private final byte[] metadataBytes;
+    private final byte[] valueBytes;
+
+    private EncodedVariant(byte[] metadataBytes, byte[] valueBytes)
+    {
+        this.metadataBytes = requireNonNull(metadataBytes, "metadataBytes is null");
+        this.valueBytes = requireNonNull(valueBytes, "valueBytes is null");
+        checkArgument(valueBytes.length > 0, "valueBytes is empty");
+    }
+
+    public static EncodedVariant fromBytes(byte[] metadataBytes, byte[] valueBytes)
+    {
+        requireNonNull(metadataBytes, "metadataBytes is null");
+        requireNonNull(valueBytes, "valueBytes is null");
+        return new EncodedVariant(metadataBytes.clone(), valueBytes.clone());
+    }
+
+    public byte[] getMetadataBytes()
+    {
+        return metadataBytes.clone();
+    }
+
+    public byte[] getValueBytes()
+    {
+        return valueBytes.clone();
+    }
+}

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

@@ -29,6 +29,8 @@
 import java.util.Optional;
 
 import static com.fasterxml.jackson.core.JsonToken.END_ARRAY;
+import static com.fasterxml.jackson.core.JsonToken.END_OBJECT;
+import static com.fasterxml.jackson.core.JsonToken.FIELD_NAME;
 import static com.fasterxml.jackson.core.JsonToken.START_OBJECT;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Verify.verify;
@@ -87,16 +89,22 @@ private JsonDecodingUtils() {}
     private static final RealDecoder REAL_DECODER = new RealDecoder();
     private static final BooleanDecoder BOOLEAN_DECODER = new BooleanDecoder();
     private static final StringDecoder STRING_DECODER = new StringDecoder();
-    private static final VariantDecoder VARIANT_DECODER = new VariantDecoder();
+    private static final VariantJsonDecoder VARIANT_JSON_DECODER = new VariantJsonDecoder();
+    private static final VariantBinaryDecoder VARIANT_BINARY_DECODER = new VariantBinaryDecoder();
     private static final Base64Decoder BASE_64_DECODER = new Base64Decoder();
     private static final ObjectDecoder OBJECT_DECODER = new ObjectDecoder();
 
     public static TypeDecoder[] createTypeDecoders(List<Column> columns)
+    {
+        return createTypeDecoders(columns, false);
+    }
+
+    public static TypeDecoder[] createTypeDecoders(List<Column> columns, boolean supportsVariantBinary)
     {
         verify(!columns.isEmpty(), "Columns must not be empty");
         TypeDecoder[] decoders = new TypeDecoder[columns.size()];
         for (int i = 0; i < columns.size(); i++) {
-            decoders[i] = createTypeDecoder(columns.get(i).getTypeSignature());
+            decoders[i] = createTypeDecoder(columns.get(i).getTypeSignature(), supportsVariantBinary);
         }
         return decoders;
     }
@@ -107,7 +115,7 @@ Object decode(JsonParser parser)
                 throws IOException;
     }
 
-    private static TypeDecoder createTypeDecoder(ClientTypeSignature signature)
+    private static TypeDecoder createTypeDecoder(ClientTypeSignature signature, boolean supportsVariantBinary)
     {
         switch (signature.getRawType()) {
             case BIGINT:
@@ -125,13 +133,13 @@ private static TypeDecoder createTypeDecoder(ClientTypeSignature signature)
             case BOOLEAN:
                 return BOOLEAN_DECODER;
             case VARIANT:
-                return VARIANT_DECODER;
+                return supportsVariantBinary ? VARIANT_BINARY_DECODER : VARIANT_JSON_DECODER;
             case ARRAY:
-                return new ArrayDecoder(signature);
+                return new ArrayDecoder(signature, supportsVariantBinary);
             case MAP:
-                return new MapDecoder(signature);
+                return new MapDecoder(signature, supportsVariantBinary);
             case ROW:
-                return new RowDecoder(signature);
+                return new RowDecoder(signature, supportsVariantBinary);
             case VARCHAR:
             case JSON:
             case TIME:
@@ -296,7 +304,7 @@ public Object decode(JsonParser parser)
         }
     }
 
-    private static class VariantDecoder
+    private static class VariantJsonDecoder
             implements TypeDecoder
     {
         @Override
@@ -311,6 +319,46 @@ public Object decode(JsonParser parser)
         }
     }
 
+    private static class VariantBinaryDecoder
+            implements TypeDecoder
+    {
+        @Override
+        public Object decode(JsonParser parser)
+                throws IOException
+        {
+            if (requireNonNull(parser.currentToken()) != START_OBJECT) {
+                throw illegalToken(parser);
+            }
+
+            byte[] metadataBytes = null;
+            byte[] valueBytes = null;
+            while (parser.nextToken() != END_OBJECT) {
+                if (requireNonNull(parser.currentToken()) != FIELD_NAME) {
+                    throw illegalToken(parser);
+                }
+
+                String fieldName = parser.currentName();
+                if (parser.nextToken() != JsonToken.VALUE_STRING) {
+                    throw illegalToken(parser);
+                }
+
+                switch (fieldName) {
+                    case "metadata":
+                        metadataBytes = Base64.getDecoder().decode(parser.getValueAsString());
+                        break;
+                    case "value":
+                        valueBytes = Base64.getDecoder().decode(parser.getValueAsString());
+                        break;
+                }
+            }
+
+            if (metadataBytes == null || valueBytes == null) {
+                throw illegalToken(parser);
+            }
+            return EncodedVariant.fromBytes(metadataBytes, valueBytes);
+        }
+    }
+
     private static class Base64Decoder
             implements TypeDecoder
     {
@@ -327,11 +375,11 @@ private static class ArrayDecoder
     {
         private final TypeDecoder typeDecoder;
 
-        public ArrayDecoder(ClientTypeSignature signature)
+        public ArrayDecoder(ClientTypeSignature signature, boolean supportsVariantBinary)
         {
             requireNonNull(signature, "signature is null");
             checkArgument(signature.getRawType().equals(ARRAY), "not an array type signature: %s", signature);
-            this.typeDecoder = createTypeDecoder(signature.getArgumentsAsTypeSignatures().get(0));
+            this.typeDecoder = createTypeDecoder(signature.getArgumentsAsTypeSignatures().get(0), supportsVariantBinary);
         }
 
         @Override
@@ -363,12 +411,12 @@ private static class MapDecoder
         private final String keyType;
         private final TypeDecoder valueDecoder;
 
-        public MapDecoder(ClientTypeSignature signature)
+        public MapDecoder(ClientTypeSignature signature, boolean supportsVariantBinary)
         {
             requireNonNull(signature, "signature is null");
             checkArgument(signature.getRawType().equals(MAP), "not a map type signature: %s", signature);
             this.keyType = signature.getArgumentsAsTypeSignatures().get(0).getRawType();
-            this.valueDecoder = createTypeDecoder(signature.getArgumentsAsTypeSignatures().get(1));
+            this.valueDecoder = createTypeDecoder(signature.getArgumentsAsTypeSignatures().get(1), supportsVariantBinary);
         }
 
         @Override
@@ -446,7 +494,7 @@ private static class RowDecoder
         private final TypeDecoder[] fieldDecoders;
         private final List<Optional<String>> fieldNames;
 
-        private RowDecoder(ClientTypeSignature signature)
+        private RowDecoder(ClientTypeSignature signature, boolean supportsVariantBinary)
         {
             requireNonNull(signature, "signature is null");
             checkArgument(signature.getRawType().equals(ROW), "not a row type signature: %s", signature);
@@ -455,7 +503,7 @@ private RowDecoder(ClientTypeSignature signature)
 
             int index = 0;
             for (ClientTypeSignatureParameter parameter : signature.getArguments()) {
-                fieldDecoders[index] = createTypeDecoder(parameter.getTypeSignature());
+                fieldDecoders[index] = createTypeDecoder(parameter.getTypeSignature(), supportsVariantBinary);
                 fieldNames.add(parameter.getName());
                 index++;
             }

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

@@ -138,10 +138,10 @@ public void close()
         }
     }
 
-    public static CloseableIterator<List<Object>> forJsonParser(JsonParser parser, List<Column> columns)
+    public static CloseableIterator<List<Object>> forJsonParser(JsonParser parser, List<Column> columns, boolean supportsVariantBinary)
             throws IOException
     {
-        return new JsonIterator(parser, createTypeDecoders(columns));
+        return new JsonIterator(parser, createTypeDecoders(columns, supportsVariantBinary));
     }
 
     public static CloseableIterator<List<Object>> forInputStream(InputStream stream, TypeDecoder[] decoders)

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

@@ -51,6 +51,7 @@
 import static io.trino.client.ProtocolHeaders.Headers.RESPONSE_SET_SCHEMA;
 import static io.trino.client.ProtocolHeaders.Headers.RESPONSE_SET_SESSION;
 import static io.trino.client.ProtocolHeaders.Headers.RESPONSE_STARTED_TRANSACTION_ID;
+import static io.trino.client.ProtocolHeaders.Headers.RESPONSE_VARIANT_ENCODING;
 import static java.util.Locale.ENGLISH;
 import static java.util.Objects.requireNonNull;
 
@@ -88,6 +89,7 @@ enum Headers
         RESPONSE_SET_ROLE("Set-Role"),
         RESPONSE_SET_ORIGINAL_ROLES("Set-Original-Roles"),
         RESPONSE_QUERY_DATA_ENCODING("Query-Data-Encoding"),
+        RESPONSE_VARIANT_ENCODING("Variant-Encoding"),
         RESPONSE_ADDED_PREPARE("Added-Prepare"),
         RESPONSE_DEALLOCATED_PREPARE("Deallocated-Prepare"),
         RESPONSE_STARTED_TRANSACTION_ID("Started-Transaction-Id"),
@@ -136,6 +138,7 @@ public String withProtocolName(String protocolName)
     private final String responseClearSession;
     private final String responseSetRole;
     private final String responseQueryDataEncoding;
+    private final String responseVariantEncoding;
     private final String responseAddedPrepare;
     private final String responseDeallocatedPrepare;
     private final String responseStartedTransactionId;
@@ -185,6 +188,7 @@ private ProtocolHeaders(String name)
         responseClearSession = RESPONSE_CLEAR_SESSION.withProtocolName(name);
         responseSetRole = RESPONSE_SET_ROLE.withProtocolName(name);
         responseQueryDataEncoding = RESPONSE_QUERY_DATA_ENCODING.withProtocolName(name);
+        responseVariantEncoding = RESPONSE_VARIANT_ENCODING.withProtocolName(name);
         responseAddedPrepare = RESPONSE_ADDED_PREPARE.withProtocolName(name);
         responseDeallocatedPrepare = RESPONSE_DEALLOCATED_PREPARE.withProtocolName(name);
         responseStartedTransactionId = RESPONSE_STARTED_TRANSACTION_ID.withProtocolName(name);
@@ -344,6 +348,11 @@ public String responseQueryDataEncoding()
         return responseQueryDataEncoding;
     }
 
+    public String responseVariantEncoding()
+    {
+        return responseVariantEncoding;
+    }
+
     public String responseAddedPrepare()
     {
         return responseAddedPrepare;