Skip to content

Commit 1a6af27

Browse files
daindain
committed
Add VARIANT binary support for JDBC
1 parent 30a49c6 commit 1a6af27

29 files changed

+1713
-81
lines changed

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

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,11 @@
1313
*/
1414
package io.trino.client;
1515

16+
import java.util.Set;
17+
18+
import static java.util.Arrays.stream;
19+
import static java.util.stream.Collectors.toUnmodifiableSet;
20+
1621
public enum ClientCapabilities
1722
{
1823
PATH,
@@ -39,8 +44,39 @@ public enum ClientCapabilities
3944
*/
4045
VARIANT_JSON,
4146

47+
/**
48+
* Whether client supports the `VARIANT` type encoded as a binary payload on the wire.
49+
* This capability is opt-in, so clients continue to receive the JSON representation by default.
50+
*/
51+
VARIANT_BINARY(false),
52+
4253
/**
4354
* Whether clients support the session authorization set/reset feature
4455
*/
4556
SESSION_AUTHORIZATION;
57+
58+
private final boolean enabledByDefault;
59+
60+
ClientCapabilities()
61+
{
62+
this(true);
63+
}
64+
65+
ClientCapabilities(boolean enabledByDefault)
66+
{
67+
this.enabledByDefault = enabledByDefault;
68+
}
69+
70+
public boolean enabledByDefault()
71+
{
72+
return enabledByDefault;
73+
}
74+
75+
public static Set<String> defaultClientCapabilities()
76+
{
77+
return stream(values())
78+
.filter(ClientCapabilities::enabledByDefault)
79+
.map(Enum::name)
80+
.collect(toUnmodifiableSet());
81+
}
4682
}

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

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
/*
2+
* Licensed under the Apache License, Version 2.0 (the "License");
3+
* you may not use this file except in compliance with the License.
4+
* You may obtain a copy of the License at
5+
*
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
*
8+
* Unless required by applicable law or agreed to in writing, software
9+
* distributed under the License is distributed on an "AS IS" BASIS,
10+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
* See the License for the specific language governing permissions and
12+
* limitations under the License.
13+
*/
14+
package io.trino.client;
15+
16+
import static com.google.common.base.Preconditions.checkArgument;
17+
import static java.util.Objects.requireNonNull;
18+
19+
public final class EncodedVariant
20+
{
21+
private final byte[] metadataBytes;
22+
private final byte[] valueBytes;
23+
24+
private EncodedVariant(byte[] metadataBytes, byte[] valueBytes)
25+
{
26+
this.metadataBytes = requireNonNull(metadataBytes, "metadataBytes is null");
27+
this.valueBytes = requireNonNull(valueBytes, "valueBytes is null");
28+
checkArgument(valueBytes.length > 0, "valueBytes is empty");
29+
}
30+
31+
public static EncodedVariant fromBytes(byte[] metadataBytes, byte[] valueBytes)
32+
{
33+
requireNonNull(metadataBytes, "metadataBytes is null");
34+
requireNonNull(valueBytes, "valueBytes is null");
35+
return new EncodedVariant(metadataBytes.clone(), valueBytes.clone());
36+
}
37+
38+
public byte[] getMetadataBytes()
39+
{
40+
return metadataBytes.clone();
41+
}
42+
43+
public byte[] getValueBytes()
44+
{
45+
return valueBytes.clone();
46+
}
47+
}

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

Lines changed: 62 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@
2929
import java.util.Optional;
3030

3131
import static com.fasterxml.jackson.core.JsonToken.END_ARRAY;
32+
import static com.fasterxml.jackson.core.JsonToken.END_OBJECT;
33+
import static com.fasterxml.jackson.core.JsonToken.FIELD_NAME;
3234
import static com.fasterxml.jackson.core.JsonToken.START_OBJECT;
3335
import static com.google.common.base.Preconditions.checkArgument;
3436
import static com.google.common.base.Verify.verify;
@@ -87,16 +89,22 @@ private JsonDecodingUtils() {}
8789
private static final RealDecoder REAL_DECODER = new RealDecoder();
8890
private static final BooleanDecoder BOOLEAN_DECODER = new BooleanDecoder();
8991
private static final StringDecoder STRING_DECODER = new StringDecoder();
90-
private static final VariantDecoder VARIANT_DECODER = new VariantDecoder();
92+
private static final VariantJsonDecoder VARIANT_JSON_DECODER = new VariantJsonDecoder();
93+
private static final VariantBinaryDecoder VARIANT_BINARY_DECODER = new VariantBinaryDecoder();
9194
private static final Base64Decoder BASE_64_DECODER = new Base64Decoder();
9295
private static final ObjectDecoder OBJECT_DECODER = new ObjectDecoder();
9396

9497
public static TypeDecoder[] createTypeDecoders(List<Column> columns)
98+
{
99+
return createTypeDecoders(columns, false);
100+
}
101+
102+
public static TypeDecoder[] createTypeDecoders(List<Column> columns, boolean supportsVariantBinary)
95103
{
96104
verify(!columns.isEmpty(), "Columns must not be empty");
97105
TypeDecoder[] decoders = new TypeDecoder[columns.size()];
98106
for (int i = 0; i < columns.size(); i++) {
99-
decoders[i] = createTypeDecoder(columns.get(i).getTypeSignature());
107+
decoders[i] = createTypeDecoder(columns.get(i).getTypeSignature(), supportsVariantBinary);
100108
}
101109
return decoders;
102110
}
@@ -107,7 +115,7 @@ Object decode(JsonParser parser)
107115
throws IOException;
108116
}
109117

110-
private static TypeDecoder createTypeDecoder(ClientTypeSignature signature)
118+
private static TypeDecoder createTypeDecoder(ClientTypeSignature signature, boolean supportsVariantBinary)
111119
{
112120
switch (signature.getRawType()) {
113121
case BIGINT:
@@ -125,13 +133,13 @@ private static TypeDecoder createTypeDecoder(ClientTypeSignature signature)
125133
case BOOLEAN:
126134
return BOOLEAN_DECODER;
127135
case VARIANT:
128-
return VARIANT_DECODER;
136+
return supportsVariantBinary ? VARIANT_BINARY_DECODER : VARIANT_JSON_DECODER;
129137
case ARRAY:
130-
return new ArrayDecoder(signature);
138+
return new ArrayDecoder(signature, supportsVariantBinary);
131139
case MAP:
132-
return new MapDecoder(signature);
140+
return new MapDecoder(signature, supportsVariantBinary);
133141
case ROW:
134-
return new RowDecoder(signature);
142+
return new RowDecoder(signature, supportsVariantBinary);
135143
case VARCHAR:
136144
case JSON:
137145
case TIME:
@@ -296,7 +304,7 @@ public Object decode(JsonParser parser)
296304
}
297305
}
298306

299-
private static class VariantDecoder
307+
private static class VariantJsonDecoder
300308
implements TypeDecoder
301309
{
302310
@Override
@@ -311,6 +319,46 @@ public Object decode(JsonParser parser)
311319
}
312320
}
313321

322+
private static class VariantBinaryDecoder
323+
implements TypeDecoder
324+
{
325+
@Override
326+
public Object decode(JsonParser parser)
327+
throws IOException
328+
{
329+
if (requireNonNull(parser.currentToken()) != START_OBJECT) {
330+
throw illegalToken(parser);
331+
}
332+
333+
byte[] metadataBytes = null;
334+
byte[] valueBytes = null;
335+
while (parser.nextToken() != END_OBJECT) {
336+
if (requireNonNull(parser.currentToken()) != FIELD_NAME) {
337+
throw illegalToken(parser);
338+
}
339+
340+
String fieldName = parser.currentName();
341+
if (parser.nextToken() != JsonToken.VALUE_STRING) {
342+
throw illegalToken(parser);
343+
}
344+
345+
switch (fieldName) {
346+
case "metadata":
347+
metadataBytes = Base64.getDecoder().decode(parser.getValueAsString());
348+
break;
349+
case "value":
350+
valueBytes = Base64.getDecoder().decode(parser.getValueAsString());
351+
break;
352+
}
353+
}
354+
355+
if (metadataBytes == null || valueBytes == null) {
356+
throw illegalToken(parser);
357+
}
358+
return EncodedVariant.fromBytes(metadataBytes, valueBytes);
359+
}
360+
}
361+
314362
private static class Base64Decoder
315363
implements TypeDecoder
316364
{
@@ -327,11 +375,11 @@ private static class ArrayDecoder
327375
{
328376
private final TypeDecoder typeDecoder;
329377

330-
public ArrayDecoder(ClientTypeSignature signature)
378+
public ArrayDecoder(ClientTypeSignature signature, boolean supportsVariantBinary)
331379
{
332380
requireNonNull(signature, "signature is null");
333381
checkArgument(signature.getRawType().equals(ARRAY), "not an array type signature: %s", signature);
334-
this.typeDecoder = createTypeDecoder(signature.getArgumentsAsTypeSignatures().get(0));
382+
this.typeDecoder = createTypeDecoder(signature.getArgumentsAsTypeSignatures().get(0), supportsVariantBinary);
335383
}
336384

337385
@Override
@@ -363,12 +411,12 @@ private static class MapDecoder
363411
private final String keyType;
364412
private final TypeDecoder valueDecoder;
365413

366-
public MapDecoder(ClientTypeSignature signature)
414+
public MapDecoder(ClientTypeSignature signature, boolean supportsVariantBinary)
367415
{
368416
requireNonNull(signature, "signature is null");
369417
checkArgument(signature.getRawType().equals(MAP), "not a map type signature: %s", signature);
370418
this.keyType = signature.getArgumentsAsTypeSignatures().get(0).getRawType();
371-
this.valueDecoder = createTypeDecoder(signature.getArgumentsAsTypeSignatures().get(1));
419+
this.valueDecoder = createTypeDecoder(signature.getArgumentsAsTypeSignatures().get(1), supportsVariantBinary);
372420
}
373421

374422
@Override
@@ -446,7 +494,7 @@ private static class RowDecoder
446494
private final TypeDecoder[] fieldDecoders;
447495
private final List<Optional<String>> fieldNames;
448496

449-
private RowDecoder(ClientTypeSignature signature)
497+
private RowDecoder(ClientTypeSignature signature, boolean supportsVariantBinary)
450498
{
451499
requireNonNull(signature, "signature is null");
452500
checkArgument(signature.getRawType().equals(ROW), "not a row type signature: %s", signature);
@@ -455,7 +503,7 @@ private RowDecoder(ClientTypeSignature signature)
455503

456504
int index = 0;
457505
for (ClientTypeSignatureParameter parameter : signature.getArguments()) {
458-
fieldDecoders[index] = createTypeDecoder(parameter.getTypeSignature());
506+
fieldDecoders[index] = createTypeDecoder(parameter.getTypeSignature(), supportsVariantBinary);
459507
fieldNames.add(parameter.getName());
460508
index++;
461509
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -138,10 +138,10 @@ public void close()
138138
}
139139
}
140140

141-
public static CloseableIterator<List<Object>> forJsonParser(JsonParser parser, List<Column> columns)
141+
public static CloseableIterator<List<Object>> forJsonParser(JsonParser parser, List<Column> columns, boolean supportsVariantBinary)
142142
throws IOException
143143
{
144-
return new JsonIterator(parser, createTypeDecoders(columns));
144+
return new JsonIterator(parser, createTypeDecoders(columns, supportsVariantBinary));
145145
}
146146

147147
public static CloseableIterator<List<Object>> forInputStream(InputStream stream, TypeDecoder[] decoders)

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@
5151
import static io.trino.client.ProtocolHeaders.Headers.RESPONSE_SET_SCHEMA;
5252
import static io.trino.client.ProtocolHeaders.Headers.RESPONSE_SET_SESSION;
5353
import static io.trino.client.ProtocolHeaders.Headers.RESPONSE_STARTED_TRANSACTION_ID;
54+
import static io.trino.client.ProtocolHeaders.Headers.RESPONSE_VARIANT_ENCODING;
5455
import static java.util.Locale.ENGLISH;
5556
import static java.util.Objects.requireNonNull;
5657

@@ -88,6 +89,7 @@ enum Headers
8889
RESPONSE_SET_ROLE("Set-Role"),
8990
RESPONSE_SET_ORIGINAL_ROLES("Set-Original-Roles"),
9091
RESPONSE_QUERY_DATA_ENCODING("Query-Data-Encoding"),
92+
RESPONSE_VARIANT_ENCODING("Variant-Encoding"),
9193
RESPONSE_ADDED_PREPARE("Added-Prepare"),
9294
RESPONSE_DEALLOCATED_PREPARE("Deallocated-Prepare"),
9395
RESPONSE_STARTED_TRANSACTION_ID("Started-Transaction-Id"),
@@ -136,6 +138,7 @@ public String withProtocolName(String protocolName)
136138
private final String responseClearSession;
137139
private final String responseSetRole;
138140
private final String responseQueryDataEncoding;
141+
private final String responseVariantEncoding;
139142
private final String responseAddedPrepare;
140143
private final String responseDeallocatedPrepare;
141144
private final String responseStartedTransactionId;
@@ -185,6 +188,7 @@ private ProtocolHeaders(String name)
185188
responseClearSession = RESPONSE_CLEAR_SESSION.withProtocolName(name);
186189
responseSetRole = RESPONSE_SET_ROLE.withProtocolName(name);
187190
responseQueryDataEncoding = RESPONSE_QUERY_DATA_ENCODING.withProtocolName(name);
191+
responseVariantEncoding = RESPONSE_VARIANT_ENCODING.withProtocolName(name);
188192
responseAddedPrepare = RESPONSE_ADDED_PREPARE.withProtocolName(name);
189193
responseDeallocatedPrepare = RESPONSE_DEALLOCATED_PREPARE.withProtocolName(name);
190194
responseStartedTransactionId = RESPONSE_STARTED_TRANSACTION_ID.withProtocolName(name);
@@ -344,6 +348,11 @@ public String responseQueryDataEncoding()
344348
return responseQueryDataEncoding;
345349
}
346350

351+
public String responseVariantEncoding()
352+
{
353+
return responseVariantEncoding;
354+
}
355+
347356
public String responseAddedPrepare()
348357
{
349358
return responseAddedPrepare;

0 commit comments

Comments
 (0)