From 29258b811b0f26bcc44dd99d234c8f49463be0ce Mon Sep 17 00:00:00 2001 From: Hazmi Date: Thu, 9 Oct 2025 10:25:06 +0300 Subject: [PATCH] Add new generic thrift toolkit module for connectors --- CODEOWNERS | 1 + pom.xml | 7 + .../connector/ConnectorCodecManager.java | 7 + .../thrift/ColumnHandleThriftCodec.java | 82 ++++++++ .../server/thrift/HandleThriftModule.java | 3 + .../server/remotetask/TestHttpRemoteTask.java | 4 + .../TestHttpRemoteTaskWithEventLoop.java | 4 + .../spi/connector/ConnectorCodecProvider.java | 6 + presto-thrift-connector-toolkit/pom.xml | 37 ++++ .../thrift/codec/GenericThriftCodec.java | 73 +++++++ .../thrift/codec/ThriftCodecProvider.java | 183 ++++++++++++++++++ .../thrift/codec}/ThriftCodecUtils.java | 2 +- presto-thrift-spec/pom.xml | 1 + presto-tpcds/pom.xml | 14 +- .../presto/tpcds/TpcdsConnectorFactory.java | 8 +- .../tpcds/thrift/TpcdsCodecProvider.java | 61 ------ .../presto/tpcds/thrift/TpcdsSplitCodec.java | 60 ------ .../tpcds/thrift/TpcdsTableHandleCodec.java | 60 ------ .../thrift/TpcdsTableLayoutHandleCodec.java | 60 ------ .../thrift/TpcdsTransactionHandleCodec.java | 60 ------ 20 files changed, 422 insertions(+), 311 deletions(-) create mode 100644 presto-main-base/src/main/java/com/facebook/presto/server/thrift/ColumnHandleThriftCodec.java create mode 100644 presto-thrift-connector-toolkit/pom.xml create mode 100644 presto-thrift-connector-toolkit/src/main/java/com/facebook/presto/thrift/codec/GenericThriftCodec.java create mode 100644 presto-thrift-connector-toolkit/src/main/java/com/facebook/presto/thrift/codec/ThriftCodecProvider.java rename {presto-tpcds/src/main/java/com/facebook/presto/tpcds/thrift => presto-thrift-connector-toolkit/src/main/java/com/facebook/presto/thrift/codec}/ThriftCodecUtils.java (97%) delete mode 100644 presto-tpcds/src/main/java/com/facebook/presto/tpcds/thrift/TpcdsCodecProvider.java delete mode 100644 presto-tpcds/src/main/java/com/facebook/presto/tpcds/thrift/TpcdsSplitCodec.java delete mode 100644 presto-tpcds/src/main/java/com/facebook/presto/tpcds/thrift/TpcdsTableHandleCodec.java delete mode 100644 presto-tpcds/src/main/java/com/facebook/presto/tpcds/thrift/TpcdsTableLayoutHandleCodec.java delete mode 100644 presto-tpcds/src/main/java/com/facebook/presto/tpcds/thrift/TpcdsTransactionHandleCodec.java diff --git a/CODEOWNERS b/CODEOWNERS index e4a0e3bb56aed..fa05b630f8e5e 100644 --- a/CODEOWNERS +++ b/CODEOWNERS @@ -89,6 +89,7 @@ /presto-testng-services @prestodb/committers /presto-thrift-api @prestodb/committers /presto-thrift-connector @prestodb/committers +/presto-thrift-connector-toolkit @prestodb/committers /presto-thrift-spec @prestodb/committers /presto-thrift-testing-server @prestodb/committers /presto-thrift-testing-udf-server @prestodb/committers diff --git a/pom.xml b/pom.xml index afcddad87b223..abbbf3f210057 100644 --- a/pom.xml +++ b/pom.xml @@ -227,6 +227,7 @@ presto-plan-checker-router-plugin presto-sql-helpers/presto-sql-invoked-functions-plugin presto-sql-helpers/presto-native-sql-invoked-functions-plugin + presto-thrift-connector-toolkit @@ -1110,6 +1111,12 @@ test-jar + + com.facebook.presto + presto-thrift-connector-toolkit + ${project.version} + + com.facebook.presto presto-thrift-testing-server diff --git a/presto-main-base/src/main/java/com/facebook/presto/connector/ConnectorCodecManager.java b/presto-main-base/src/main/java/com/facebook/presto/connector/ConnectorCodecManager.java index fb5179d141507..09a3e1634fc0a 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/connector/ConnectorCodecManager.java +++ b/presto-main-base/src/main/java/com/facebook/presto/connector/ConnectorCodecManager.java @@ -14,6 +14,7 @@ package com.facebook.presto.connector; import com.facebook.drift.codec.ThriftCodecManager; +import com.facebook.presto.spi.ColumnHandle; import com.facebook.presto.spi.ConnectorCodec; import com.facebook.presto.spi.ConnectorDeleteTableHandle; import com.facebook.presto.spi.ConnectorId; @@ -96,4 +97,10 @@ public Optional> getTableHandleCodec(String requireNonNull(connectorId, "connectorId is null"); return Optional.ofNullable(connectorCodecProviders.get(connectorId)).flatMap(ConnectorCodecProvider::getConnectorTableHandleCodec); } + + public Optional> getColumnHandleCodec(String connectorId) + { + requireNonNull(connectorId, "connectorId is null"); + return Optional.ofNullable(connectorCodecProviders.get(connectorId)).flatMap(ConnectorCodecProvider::getColumnHandleCodec); + } } diff --git a/presto-main-base/src/main/java/com/facebook/presto/server/thrift/ColumnHandleThriftCodec.java b/presto-main-base/src/main/java/com/facebook/presto/server/thrift/ColumnHandleThriftCodec.java new file mode 100644 index 0000000000000..6d09f882c9e86 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/server/thrift/ColumnHandleThriftCodec.java @@ -0,0 +1,82 @@ +/* + * 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 com.facebook.presto.server.thrift; + +import com.facebook.airlift.json.JsonCodec; +import com.facebook.drift.codec.CodecThriftType; +import com.facebook.drift.codec.metadata.ThriftType; +import com.facebook.drift.protocol.TProtocolReader; +import com.facebook.drift.protocol.TProtocolWriter; +import com.facebook.presto.connector.ConnectorCodecManager; +import com.facebook.presto.metadata.HandleResolver; +import com.facebook.presto.spi.ColumnHandle; + +import javax.inject.Inject; + +import java.nio.ByteBuffer; + +import static java.util.Objects.requireNonNull; + +public class ColumnHandleThriftCodec + extends AbstractTypedThriftCodec +{ + private static final ThriftType THRIFT_TYPE = createThriftType(ColumnHandle.class); + private final ConnectorCodecManager connectorCodecManager; + + @Inject + public ColumnHandleThriftCodec(HandleResolver handleResolver, ConnectorCodecManager connectorCodecManager, JsonCodec jsonCodec) + { + super(ColumnHandle.class, + requireNonNull(jsonCodec, "jsonCodec is null"), + requireNonNull(handleResolver, "handleResolver is null")::getId, + handleResolver::getColumnHandleClass); + this.connectorCodecManager = requireNonNull(connectorCodecManager, "connectorThriftCodecManager is null"); + } + + @CodecThriftType + public static ThriftType getThriftType() + { + return THRIFT_TYPE; + } + + @Override + public ThriftType getType() + { + return THRIFT_TYPE; + } + + @Override + public ColumnHandle readConcreteValue(String connectorId, TProtocolReader reader) + throws Exception + { + ByteBuffer byteBuffer = reader.readBinary(); + assert (byteBuffer.position() == 0); + byte[] bytes = byteBuffer.array(); + return connectorCodecManager.getColumnHandleCodec(connectorId).map(codec -> codec.deserialize(bytes)).orElse(null); + } + + @Override + public void writeConcreteValue(String connectorId, ColumnHandle value, TProtocolWriter writer) + throws Exception + { + requireNonNull(value, "value is null"); + writer.writeBinary(ByteBuffer.wrap(connectorCodecManager.getColumnHandleCodec(connectorId).map(codec -> codec.serialize(value)).orElseThrow(() -> new IllegalArgumentException("Can not serialize " + value)))); + } + + @Override + public boolean isThriftCodecAvailable(String connectorId) + { + return connectorCodecManager.getColumnHandleCodec(connectorId).isPresent(); + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/server/thrift/HandleThriftModule.java b/presto-main-base/src/main/java/com/facebook/presto/server/thrift/HandleThriftModule.java index 6d4c169e82032..58c7a592c45ec 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/server/thrift/HandleThriftModule.java +++ b/presto-main-base/src/main/java/com/facebook/presto/server/thrift/HandleThriftModule.java @@ -14,6 +14,7 @@ package com.facebook.presto.server.thrift; import com.facebook.presto.metadata.HandleResolver; +import com.facebook.presto.spi.ColumnHandle; import com.facebook.presto.spi.ConnectorDeleteTableHandle; import com.facebook.presto.spi.ConnectorInsertTableHandle; import com.facebook.presto.spi.ConnectorOutputTableHandle; @@ -41,6 +42,7 @@ public void configure(Binder binder) thriftCodecBinder(binder).bindCustomThriftCodec(DeleteTableHandleThriftCodec.class); thriftCodecBinder(binder).bindCustomThriftCodec(TableLayoutHandleThriftCodec.class); thriftCodecBinder(binder).bindCustomThriftCodec(TableHandleThriftCodec.class); + thriftCodecBinder(binder).bindCustomThriftCodec(ColumnHandleThriftCodec.class); jsonCodecBinder(binder).bindJsonCodec(ConnectorSplit.class); jsonCodecBinder(binder).bindJsonCodec(ConnectorTransactionHandle.class); @@ -49,6 +51,7 @@ public void configure(Binder binder) jsonCodecBinder(binder).bindJsonCodec(ConnectorDeleteTableHandle.class); jsonCodecBinder(binder).bindJsonCodec(ConnectorTableLayoutHandle.class); jsonCodecBinder(binder).bindJsonCodec(ConnectorTableHandle.class); + jsonCodecBinder(binder).bindJsonCodec(ColumnHandle.class); binder.bind(HandleResolver.class).in(Scopes.SINGLETON); } diff --git a/presto-main/src/test/java/com/facebook/presto/server/remotetask/TestHttpRemoteTask.java b/presto-main/src/test/java/com/facebook/presto/server/remotetask/TestHttpRemoteTask.java index 2de40f4fa2173..15dbe0bdb6e0a 100644 --- a/presto-main/src/test/java/com/facebook/presto/server/remotetask/TestHttpRemoteTask.java +++ b/presto-main/src/test/java/com/facebook/presto/server/remotetask/TestHttpRemoteTask.java @@ -59,6 +59,7 @@ import com.facebook.presto.metadata.Split; import com.facebook.presto.server.InternalCommunicationConfig; import com.facebook.presto.server.TaskUpdateRequest; +import com.facebook.presto.server.thrift.ColumnHandleThriftCodec; import com.facebook.presto.server.thrift.ConnectorSplitThriftCodec; import com.facebook.presto.server.thrift.DeleteTableHandleThriftCodec; import com.facebook.presto.server.thrift.InsertTableHandleThriftCodec; @@ -66,6 +67,7 @@ import com.facebook.presto.server.thrift.TableHandleThriftCodec; import com.facebook.presto.server.thrift.TableLayoutHandleThriftCodec; import com.facebook.presto.server.thrift.TransactionHandleThriftCodec; +import com.facebook.presto.spi.ColumnHandle; import com.facebook.presto.spi.ConnectorDeleteTableHandle; import com.facebook.presto.spi.ConnectorId; import com.facebook.presto.spi.ConnectorInsertTableHandle; @@ -392,6 +394,7 @@ public void configure(Binder binder) jsonCodecBinder(binder).bindJsonCodec(ConnectorInsertTableHandle.class); jsonCodecBinder(binder).bindJsonCodec(ConnectorTableHandle.class); jsonCodecBinder(binder).bindJsonCodec(ConnectorTableLayoutHandle.class); + jsonCodecBinder(binder).bindJsonCodec(ColumnHandle.class); binder.bind(ConnectorCodecManager.class).in(Scopes.SINGLETON); @@ -402,6 +405,7 @@ public void configure(Binder binder) thriftCodecBinder(binder).bindCustomThriftCodec(DeleteTableHandleThriftCodec.class); thriftCodecBinder(binder).bindCustomThriftCodec(TableHandleThriftCodec.class); thriftCodecBinder(binder).bindCustomThriftCodec(TableLayoutHandleThriftCodec.class); + thriftCodecBinder(binder).bindCustomThriftCodec(ColumnHandleThriftCodec.class); thriftCodecBinder(binder).bindThriftCodec(TaskStatus.class); thriftCodecBinder(binder).bindThriftCodec(TaskInfo.class); thriftCodecBinder(binder).bindThriftCodec(TaskUpdateRequest.class); diff --git a/presto-main/src/test/java/com/facebook/presto/server/remotetask/TestHttpRemoteTaskWithEventLoop.java b/presto-main/src/test/java/com/facebook/presto/server/remotetask/TestHttpRemoteTaskWithEventLoop.java index 3e8cc6d11e4b4..dc828bd46b5b3 100644 --- a/presto-main/src/test/java/com/facebook/presto/server/remotetask/TestHttpRemoteTaskWithEventLoop.java +++ b/presto-main/src/test/java/com/facebook/presto/server/remotetask/TestHttpRemoteTaskWithEventLoop.java @@ -57,6 +57,7 @@ import com.facebook.presto.metadata.Split; import com.facebook.presto.server.InternalCommunicationConfig; import com.facebook.presto.server.TaskUpdateRequest; +import com.facebook.presto.server.thrift.ColumnHandleThriftCodec; import com.facebook.presto.server.thrift.ConnectorSplitThriftCodec; import com.facebook.presto.server.thrift.DeleteTableHandleThriftCodec; import com.facebook.presto.server.thrift.InsertTableHandleThriftCodec; @@ -64,6 +65,7 @@ import com.facebook.presto.server.thrift.TableHandleThriftCodec; import com.facebook.presto.server.thrift.TableLayoutHandleThriftCodec; import com.facebook.presto.server.thrift.TransactionHandleThriftCodec; +import com.facebook.presto.spi.ColumnHandle; import com.facebook.presto.spi.ConnectorDeleteTableHandle; import com.facebook.presto.spi.ConnectorId; import com.facebook.presto.spi.ConnectorInsertTableHandle; @@ -400,6 +402,7 @@ public void configure(Binder binder) jsonCodecBinder(binder).bindJsonCodec(ConnectorInsertTableHandle.class); jsonCodecBinder(binder).bindJsonCodec(ConnectorTableHandle.class); jsonCodecBinder(binder).bindJsonCodec(ConnectorTableLayoutHandle.class); + jsonCodecBinder(binder).bindJsonCodec(ColumnHandle.class); binder.bind(ConnectorCodecManager.class).in(Scopes.SINGLETON); @@ -410,6 +413,7 @@ public void configure(Binder binder) thriftCodecBinder(binder).bindCustomThriftCodec(DeleteTableHandleThriftCodec.class); thriftCodecBinder(binder).bindCustomThriftCodec(TableHandleThriftCodec.class); thriftCodecBinder(binder).bindCustomThriftCodec(TableLayoutHandleThriftCodec.class); + thriftCodecBinder(binder).bindCustomThriftCodec(ColumnHandleThriftCodec.class); thriftCodecBinder(binder).bindThriftCodec(TaskStatus.class); thriftCodecBinder(binder).bindThriftCodec(TaskInfo.class); thriftCodecBinder(binder).bindThriftCodec(TaskUpdateRequest.class); diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/connector/ConnectorCodecProvider.java b/presto-spi/src/main/java/com/facebook/presto/spi/connector/ConnectorCodecProvider.java index 4bd2d81d456b4..48c47fe9fea98 100644 --- a/presto-spi/src/main/java/com/facebook/presto/spi/connector/ConnectorCodecProvider.java +++ b/presto-spi/src/main/java/com/facebook/presto/spi/connector/ConnectorCodecProvider.java @@ -13,6 +13,7 @@ */ package com.facebook.presto.spi.connector; +import com.facebook.presto.spi.ColumnHandle; import com.facebook.presto.spi.ConnectorCodec; import com.facebook.presto.spi.ConnectorDeleteTableHandle; import com.facebook.presto.spi.ConnectorInsertTableHandle; @@ -59,4 +60,9 @@ default Optional> getConnectorTableHandleCo { return Optional.empty(); } + + default Optional> getColumnHandleCodec() + { + return Optional.empty(); + } } diff --git a/presto-thrift-connector-toolkit/pom.xml b/presto-thrift-connector-toolkit/pom.xml new file mode 100644 index 0000000000000..160de2d06648b --- /dev/null +++ b/presto-thrift-connector-toolkit/pom.xml @@ -0,0 +1,37 @@ + + + 4.0.0 + + + com.facebook.presto + presto-root + 0.296-SNAPSHOT + + + presto-thrift-connector-toolkit + presto-thrift-toolkit + jar + + + ${project.parent.basedir} + true + + + + + com.facebook.airlift.drift + drift-codec + + + + com.facebook.airlift.drift + drift-protocol + + + + com.facebook.presto + presto-spi + provided + + + diff --git a/presto-thrift-connector-toolkit/src/main/java/com/facebook/presto/thrift/codec/GenericThriftCodec.java b/presto-thrift-connector-toolkit/src/main/java/com/facebook/presto/thrift/codec/GenericThriftCodec.java new file mode 100644 index 0000000000000..9b7c8e25883e1 --- /dev/null +++ b/presto-thrift-connector-toolkit/src/main/java/com/facebook/presto/thrift/codec/GenericThriftCodec.java @@ -0,0 +1,73 @@ +/* + * 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 com.facebook.presto.thrift.codec; + +import com.facebook.drift.codec.ThriftCodec; +import com.facebook.drift.codec.ThriftCodecManager; +import com.facebook.drift.protocol.TProtocolException; +import com.facebook.presto.spi.ConnectorCodec; +import com.facebook.presto.spi.PrestoException; + +import java.lang.reflect.Type; + +import static com.facebook.presto.spi.StandardErrorCode.INVALID_ARGUMENTS; +import static com.facebook.presto.thrift.codec.ThriftCodecUtils.fromThrift; +import static com.facebook.presto.thrift.codec.ThriftCodecUtils.toThrift; +import static java.util.Objects.requireNonNull; + +public class GenericThriftCodec + implements ConnectorCodec +{ + private final ThriftCodec thriftCodec; + + public GenericThriftCodec(ThriftCodecManager codecManager, Type javaType, Class expectedType) + { + requireNonNull(codecManager, "codecManager is null"); + requireNonNull(javaType, "javaType is null"); + requireNonNull(expectedType, "expectedType is null"); + + if (!(javaType instanceof Class)) { + throw new IllegalArgumentException("Expected a Class type for javaType, but got: " + javaType.getTypeName()); + } + + Class clazz = (Class) javaType; + if (!expectedType.isAssignableFrom(clazz)) { + throw new IllegalArgumentException("javaType must be a subclass of " + expectedType.getName() + ", but got: " + clazz.getName()); + } + + this.thriftCodec = (ThriftCodec) codecManager.getCodec(clazz); + } + + @Override + public byte[] serialize(T value) + { + try { + return toThrift(value, thriftCodec); + } + catch (TProtocolException e) { + throw new PrestoException(INVALID_ARGUMENTS, "Unable to serialize object of type " + value.getClass().getSimpleName(), e); + } + } + + @Override + public T deserialize(byte[] bytes) + { + try { + return fromThrift(bytes, thriftCodec); + } + catch (TProtocolException e) { + throw new PrestoException(INVALID_ARGUMENTS, "Unable to deserialize bytes to object of expected type", e); + } + } +} diff --git a/presto-thrift-connector-toolkit/src/main/java/com/facebook/presto/thrift/codec/ThriftCodecProvider.java b/presto-thrift-connector-toolkit/src/main/java/com/facebook/presto/thrift/codec/ThriftCodecProvider.java new file mode 100644 index 0000000000000..d4dddbcf0dc89 --- /dev/null +++ b/presto-thrift-connector-toolkit/src/main/java/com/facebook/presto/thrift/codec/ThriftCodecProvider.java @@ -0,0 +1,183 @@ +/* + * 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 com.facebook.presto.thrift.codec; + +import com.facebook.drift.codec.ThriftCodecManager; +import com.facebook.presto.spi.ColumnHandle; +import com.facebook.presto.spi.ConnectorCodec; +import com.facebook.presto.spi.ConnectorDeleteTableHandle; +import com.facebook.presto.spi.ConnectorInsertTableHandle; +import com.facebook.presto.spi.ConnectorOutputTableHandle; +import com.facebook.presto.spi.ConnectorSplit; +import com.facebook.presto.spi.ConnectorTableHandle; +import com.facebook.presto.spi.ConnectorTableLayoutHandle; +import com.facebook.presto.spi.connector.ConnectorCodecProvider; +import com.facebook.presto.spi.connector.ConnectorTransactionHandle; + +import java.lang.reflect.Type; +import java.util.Optional; + +import static java.util.Objects.requireNonNull; + +public class ThriftCodecProvider + implements ConnectorCodecProvider +{ + private final ThriftCodecManager thriftCodecManager; + private final Optional connectorSplitType; + private final Optional connectorTransactionHandle; + private final Optional connectorTableLayoutHandle; + private final Optional connectorTableHandle; + private final Optional connectorOutputTableHandle; + private final Optional connectorInsertTableHandle; + private final Optional connectorDeleteTableHandle; + private final Optional connectorColumnHandle; + + private ThriftCodecProvider(Builder builder) + { + this.thriftCodecManager = requireNonNull(builder.thriftCodecManager, "thriftCodecManager is null"); + this.connectorSplitType = builder.connectorSplitType; + this.connectorTransactionHandle = builder.connectorTransactionHandle; + this.connectorTableLayoutHandle = builder.connectorTableLayoutHandle; + this.connectorTableHandle = builder.connectorTableHandle; + this.connectorOutputTableHandle = builder.connectorOutputTableHandle; + this.connectorInsertTableHandle = builder.connectorInsertTableHandle; + this.connectorDeleteTableHandle = builder.connectorDeleteTableHandle; + this.connectorColumnHandle = builder.connectorColumnHandle; + } + + @Override + public Optional> getConnectorSplitCodec() + { + return connectorSplitType.map(type -> new GenericThriftCodec<>(thriftCodecManager, type, ConnectorSplit.class)); + } + + @Override + public Optional> getConnectorTransactionHandleCodec() + { + return connectorTransactionHandle.map(type -> new GenericThriftCodec<>(thriftCodecManager, type, ConnectorTransactionHandle.class)); + } + + @Override + public Optional> getConnectorTableLayoutHandleCodec() + { + return connectorTableLayoutHandle.map(type -> new GenericThriftCodec<>(thriftCodecManager, type, ConnectorTableLayoutHandle.class)); + } + + @Override + public Optional> getConnectorTableHandleCodec() + { + return connectorTableHandle.map(type -> new GenericThriftCodec<>(thriftCodecManager, type, ConnectorTableHandle.class)); + } + + @Override + public Optional> getConnectorOutputTableHandleCodec() + { + return connectorOutputTableHandle.map(type -> new GenericThriftCodec<>(thriftCodecManager, type, ConnectorOutputTableHandle.class)); + } + + @Override + public Optional> getConnectorInsertTableHandleCodec() + { + return connectorInsertTableHandle.map(type -> new GenericThriftCodec<>(thriftCodecManager, type, ConnectorInsertTableHandle.class)); + } + + @Override + public Optional> getConnectorDeleteTableHandleCodec() + { + return connectorDeleteTableHandle.map(type -> new GenericThriftCodec<>(thriftCodecManager, type, ConnectorDeleteTableHandle.class)); + } + + @Override + public Optional> getColumnHandleCodec() + { + return connectorColumnHandle.map(type -> new GenericThriftCodec<>(thriftCodecManager, type, ColumnHandle.class)); + } + + public ThriftCodecManager getThriftCodecManager() + { + return thriftCodecManager; + } + + public static class Builder + { + private ThriftCodecManager thriftCodecManager; + private Optional connectorSplitType = Optional.empty(); + private Optional connectorTransactionHandle = Optional.empty(); + private Optional connectorTableLayoutHandle = Optional.empty(); + private Optional connectorTableHandle = Optional.empty(); + private Optional connectorOutputTableHandle = Optional.empty(); + private Optional connectorInsertTableHandle = Optional.empty(); + private Optional connectorDeleteTableHandle = Optional.empty(); + private Optional connectorColumnHandle = Optional.empty(); + + public Builder setThriftCodecManager(ThriftCodecManager thriftCodecManager) + { + this.thriftCodecManager = thriftCodecManager; + return this; + } + + public Builder setConnectorSplitType(Class type) + { + this.connectorSplitType = Optional.ofNullable(type); + return this; + } + + public Builder setConnectorTransactionHandle(Class type) + { + this.connectorTransactionHandle = Optional.ofNullable(type); + return this; + } + + public Builder setConnectorTableLayoutHandle(Class type) + { + this.connectorTableLayoutHandle = Optional.ofNullable(type); + return this; + } + + public Builder setConnectorTableHandle(Class type) + { + this.connectorTableHandle = Optional.ofNullable(type); + return this; + } + + public Builder setConnectorOutputTableHandle(Class type) + { + this.connectorOutputTableHandle = Optional.ofNullable(type); + return this; + } + + public Builder setConnectorInsertTableHandle(Class type) + { + this.connectorInsertTableHandle = Optional.ofNullable(type); + return this; + } + + public Builder setConnectorDeleteTableHandle(Class type) + { + this.connectorDeleteTableHandle = Optional.ofNullable(type); + return this; + } + + public Builder setConnectorColumnHandle(Class type) + { + this.connectorColumnHandle = Optional.ofNullable(type); + return this; + } + + public ThriftCodecProvider build() + { + return new ThriftCodecProvider(this); + } + } +} diff --git a/presto-tpcds/src/main/java/com/facebook/presto/tpcds/thrift/ThriftCodecUtils.java b/presto-thrift-connector-toolkit/src/main/java/com/facebook/presto/thrift/codec/ThriftCodecUtils.java similarity index 97% rename from presto-tpcds/src/main/java/com/facebook/presto/tpcds/thrift/ThriftCodecUtils.java rename to presto-thrift-connector-toolkit/src/main/java/com/facebook/presto/thrift/codec/ThriftCodecUtils.java index 4dbff7ef02038..18bc66e911b8e 100644 --- a/presto-tpcds/src/main/java/com/facebook/presto/tpcds/thrift/ThriftCodecUtils.java +++ b/presto-thrift-connector-toolkit/src/main/java/com/facebook/presto/thrift/codec/ThriftCodecUtils.java @@ -11,7 +11,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package com.facebook.presto.tpcds.thrift; +package com.facebook.presto.thrift.codec; import com.facebook.drift.codec.ThriftCodec; import com.facebook.drift.protocol.TBinaryProtocol; diff --git a/presto-thrift-spec/pom.xml b/presto-thrift-spec/pom.xml index e6079576f1b89..fcccfaaa4e67b 100644 --- a/presto-thrift-spec/pom.xml +++ b/presto-thrift-spec/pom.xml @@ -74,6 +74,7 @@ com.facebook.presto.server.thrift.InsertTableHandleThriftCodec com.facebook.presto.server.thrift.TableHandleThriftCodec com.facebook.presto.server.thrift.TableLayoutHandleThriftCodec + com.facebook.presto.server.thrift.ColumnHandleThriftCodec com.facebook.drift.codec.utils.DurationToMillisThriftCodec com.facebook.drift.codec.utils.DataSizeToBytesThriftCodec com.facebook.drift.codec.utils.JodaDateTimeToEpochMillisThriftCodec diff --git a/presto-tpcds/pom.xml b/presto-tpcds/pom.xml index 26d2ed201032f..4d24030a06c6a 100644 --- a/presto-tpcds/pom.xml +++ b/presto-tpcds/pom.xml @@ -44,13 +44,8 @@ - com.facebook.airlift.drift - drift-codec - - - - com.facebook.airlift.drift - drift-protocol + com.facebook.presto + presto-thrift-connector-toolkit @@ -95,6 +90,11 @@ log + + com.facebook.airlift.drift + drift-codec + + com.facebook.presto diff --git a/presto-tpcds/src/main/java/com/facebook/presto/tpcds/TpcdsConnectorFactory.java b/presto-tpcds/src/main/java/com/facebook/presto/tpcds/TpcdsConnectorFactory.java index 0e3880c20c022..6b1e0e0990c13 100644 --- a/presto-tpcds/src/main/java/com/facebook/presto/tpcds/TpcdsConnectorFactory.java +++ b/presto-tpcds/src/main/java/com/facebook/presto/tpcds/TpcdsConnectorFactory.java @@ -26,7 +26,7 @@ import com.facebook.presto.spi.connector.ConnectorSplitManager; import com.facebook.presto.spi.connector.ConnectorTransactionHandle; import com.facebook.presto.spi.transaction.IsolationLevel; -import com.facebook.presto.tpcds.thrift.TpcdsCodecProvider; +import com.facebook.presto.thrift.codec.ThriftCodecProvider; import java.util.Map; @@ -110,7 +110,11 @@ public ConnectorNodePartitioningProvider getNodePartitioningProvider() @Override public ConnectorCodecProvider getConnectorCodecProvider() { - return new TpcdsCodecProvider(new ThriftCodecManager()); + return new ThriftCodecProvider.Builder().setThriftCodecManager(new ThriftCodecManager()) + .setConnectorSplitType(TpcdsSplit.class) + .setConnectorTransactionHandle(TpcdsTransactionHandle.class) + .setConnectorTableLayoutHandle(TpcdsTableLayoutHandle.class) + .setConnectorTableHandle(TpcdsTableHandle.class).build(); } }; } diff --git a/presto-tpcds/src/main/java/com/facebook/presto/tpcds/thrift/TpcdsCodecProvider.java b/presto-tpcds/src/main/java/com/facebook/presto/tpcds/thrift/TpcdsCodecProvider.java deleted file mode 100644 index 09545c6b96fe7..0000000000000 --- a/presto-tpcds/src/main/java/com/facebook/presto/tpcds/thrift/TpcdsCodecProvider.java +++ /dev/null @@ -1,61 +0,0 @@ -/* - * 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 com.facebook.presto.tpcds.thrift; - -import com.facebook.drift.codec.ThriftCodecManager; -import com.facebook.presto.spi.ConnectorCodec; -import com.facebook.presto.spi.ConnectorSplit; -import com.facebook.presto.spi.ConnectorTableHandle; -import com.facebook.presto.spi.ConnectorTableLayoutHandle; -import com.facebook.presto.spi.connector.ConnectorCodecProvider; -import com.facebook.presto.spi.connector.ConnectorTransactionHandle; - -import java.util.Optional; - -import static java.util.Objects.requireNonNull; - -public class TpcdsCodecProvider - implements ConnectorCodecProvider -{ - private final ThriftCodecManager thriftCodecManager; - - public TpcdsCodecProvider(ThriftCodecManager thriftCodecManager) - { - this.thriftCodecManager = requireNonNull(thriftCodecManager, "thriftCodecManager is null"); - } - - @Override - public Optional> getConnectorSplitCodec() - { - return Optional.of(new TpcdsSplitCodec(thriftCodecManager)); - } - - @Override - public Optional> getConnectorTransactionHandleCodec() - { - return Optional.of(new TpcdsTransactionHandleCodec(thriftCodecManager)); - } - - @Override - public Optional> getConnectorTableLayoutHandleCodec() - { - return Optional.of(new TpcdsTableLayoutHandleCodec(thriftCodecManager)); - } - - @Override - public Optional> getConnectorTableHandleCodec() - { - return Optional.of(new TpcdsTableHandleCodec(thriftCodecManager)); - } -} diff --git a/presto-tpcds/src/main/java/com/facebook/presto/tpcds/thrift/TpcdsSplitCodec.java b/presto-tpcds/src/main/java/com/facebook/presto/tpcds/thrift/TpcdsSplitCodec.java deleted file mode 100644 index 32e45525de417..0000000000000 --- a/presto-tpcds/src/main/java/com/facebook/presto/tpcds/thrift/TpcdsSplitCodec.java +++ /dev/null @@ -1,60 +0,0 @@ -/* - * 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 com.facebook.presto.tpcds.thrift; - -import com.facebook.drift.codec.ThriftCodec; -import com.facebook.drift.codec.ThriftCodecManager; -import com.facebook.drift.protocol.TProtocolException; -import com.facebook.presto.spi.ConnectorCodec; -import com.facebook.presto.spi.ConnectorSplit; -import com.facebook.presto.spi.PrestoException; -import com.facebook.presto.tpcds.TpcdsSplit; - -import static com.facebook.presto.spi.StandardErrorCode.INVALID_ARGUMENTS; -import static com.facebook.presto.tpcds.thrift.ThriftCodecUtils.fromThrift; -import static com.facebook.presto.tpcds.thrift.ThriftCodecUtils.toThrift; -import static java.util.Objects.requireNonNull; - -public class TpcdsSplitCodec - implements ConnectorCodec -{ - private final ThriftCodec thriftCodec; - - public TpcdsSplitCodec(ThriftCodecManager thriftCodecManager) - { - this.thriftCodec = requireNonNull(thriftCodecManager, "thriftCodecManager is null").getCodec(TpcdsSplit.class); - } - - @Override - public byte[] serialize(ConnectorSplit split) - { - try { - return toThrift((TpcdsSplit) split, thriftCodec); - } - catch (TProtocolException e) { - throw new PrestoException(INVALID_ARGUMENTS, "Can not serialize tpcds split", e); - } - } - - @Override - public ConnectorSplit deserialize(byte[] bytes) - { - try { - return fromThrift(bytes, thriftCodec); - } - catch (TProtocolException e) { - throw new PrestoException(INVALID_ARGUMENTS, "Can not deserialize tpcds split", e); - } - } -} diff --git a/presto-tpcds/src/main/java/com/facebook/presto/tpcds/thrift/TpcdsTableHandleCodec.java b/presto-tpcds/src/main/java/com/facebook/presto/tpcds/thrift/TpcdsTableHandleCodec.java deleted file mode 100644 index 815d981bc1059..0000000000000 --- a/presto-tpcds/src/main/java/com/facebook/presto/tpcds/thrift/TpcdsTableHandleCodec.java +++ /dev/null @@ -1,60 +0,0 @@ -/* - * 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 com.facebook.presto.tpcds.thrift; - -import com.facebook.drift.codec.ThriftCodec; -import com.facebook.drift.codec.ThriftCodecManager; -import com.facebook.drift.protocol.TProtocolException; -import com.facebook.presto.spi.ConnectorCodec; -import com.facebook.presto.spi.ConnectorTableHandle; -import com.facebook.presto.spi.PrestoException; -import com.facebook.presto.tpcds.TpcdsTableHandle; - -import static com.facebook.presto.spi.StandardErrorCode.INVALID_ARGUMENTS; -import static com.facebook.presto.tpcds.thrift.ThriftCodecUtils.fromThrift; -import static com.facebook.presto.tpcds.thrift.ThriftCodecUtils.toThrift; -import static java.util.Objects.requireNonNull; - -public class TpcdsTableHandleCodec - implements ConnectorCodec -{ - private final ThriftCodec thriftCodec; - - public TpcdsTableHandleCodec(ThriftCodecManager thriftCodecManager) - { - this.thriftCodec = requireNonNull(thriftCodecManager, "thriftCodecManager is null").getCodec(TpcdsTableHandle.class); - } - - @Override - public byte[] serialize(ConnectorTableHandle handle) - { - try { - return toThrift((TpcdsTableHandle) handle, thriftCodec); - } - catch (TProtocolException e) { - throw new PrestoException(INVALID_ARGUMENTS, "Can not serialize tpcds table handle", e); - } - } - - @Override - public ConnectorTableHandle deserialize(byte[] bytes) - { - try { - return fromThrift(bytes, thriftCodec); - } - catch (TProtocolException e) { - throw new PrestoException(INVALID_ARGUMENTS, "Can not deserialize tpcds table handle", e); - } - } -} diff --git a/presto-tpcds/src/main/java/com/facebook/presto/tpcds/thrift/TpcdsTableLayoutHandleCodec.java b/presto-tpcds/src/main/java/com/facebook/presto/tpcds/thrift/TpcdsTableLayoutHandleCodec.java deleted file mode 100644 index c7a84e168d3ce..0000000000000 --- a/presto-tpcds/src/main/java/com/facebook/presto/tpcds/thrift/TpcdsTableLayoutHandleCodec.java +++ /dev/null @@ -1,60 +0,0 @@ -/* - * 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 com.facebook.presto.tpcds.thrift; - -import com.facebook.drift.codec.ThriftCodec; -import com.facebook.drift.codec.ThriftCodecManager; -import com.facebook.drift.protocol.TProtocolException; -import com.facebook.presto.spi.ConnectorCodec; -import com.facebook.presto.spi.ConnectorTableLayoutHandle; -import com.facebook.presto.spi.PrestoException; -import com.facebook.presto.tpcds.TpcdsTableLayoutHandle; - -import static com.facebook.presto.spi.StandardErrorCode.INVALID_ARGUMENTS; -import static com.facebook.presto.tpcds.thrift.ThriftCodecUtils.fromThrift; -import static com.facebook.presto.tpcds.thrift.ThriftCodecUtils.toThrift; -import static java.util.Objects.requireNonNull; - -public class TpcdsTableLayoutHandleCodec - implements ConnectorCodec -{ - private final ThriftCodec thriftCodec; - - public TpcdsTableLayoutHandleCodec(ThriftCodecManager thriftCodecManager) - { - this.thriftCodec = requireNonNull(thriftCodecManager, "thriftCodecManager is null").getCodec(TpcdsTableLayoutHandle.class); - } - - @Override - public byte[] serialize(ConnectorTableLayoutHandle handle) - { - try { - return toThrift((TpcdsTableLayoutHandle) handle, thriftCodec); - } - catch (TProtocolException e) { - throw new PrestoException(INVALID_ARGUMENTS, "Can not serialize tpcds table Layout handle", e); - } - } - - @Override - public ConnectorTableLayoutHandle deserialize(byte[] bytes) - { - try { - return fromThrift(bytes, thriftCodec); - } - catch (TProtocolException e) { - throw new PrestoException(INVALID_ARGUMENTS, "Can not deserialize tpcds table Layout handle", e); - } - } -} diff --git a/presto-tpcds/src/main/java/com/facebook/presto/tpcds/thrift/TpcdsTransactionHandleCodec.java b/presto-tpcds/src/main/java/com/facebook/presto/tpcds/thrift/TpcdsTransactionHandleCodec.java deleted file mode 100644 index 0f98eefb4d40d..0000000000000 --- a/presto-tpcds/src/main/java/com/facebook/presto/tpcds/thrift/TpcdsTransactionHandleCodec.java +++ /dev/null @@ -1,60 +0,0 @@ -/* - * 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 com.facebook.presto.tpcds.thrift; - -import com.facebook.drift.codec.ThriftCodec; -import com.facebook.drift.codec.ThriftCodecManager; -import com.facebook.drift.protocol.TProtocolException; -import com.facebook.presto.spi.ConnectorCodec; -import com.facebook.presto.spi.PrestoException; -import com.facebook.presto.spi.connector.ConnectorTransactionHandle; -import com.facebook.presto.tpcds.TpcdsTransactionHandle; - -import static com.facebook.presto.spi.StandardErrorCode.INVALID_ARGUMENTS; -import static com.facebook.presto.tpcds.thrift.ThriftCodecUtils.fromThrift; -import static com.facebook.presto.tpcds.thrift.ThriftCodecUtils.toThrift; -import static java.util.Objects.requireNonNull; - -public class TpcdsTransactionHandleCodec - implements ConnectorCodec -{ - private final ThriftCodec thriftCodec; - - public TpcdsTransactionHandleCodec(ThriftCodecManager thriftCodecManager) - { - this.thriftCodec = requireNonNull(thriftCodecManager, "thriftCodecManager is null").getCodec(TpcdsTransactionHandle.class); - } - - @Override - public byte[] serialize(ConnectorTransactionHandle handle) - { - try { - return toThrift((TpcdsTransactionHandle) handle, thriftCodec); - } - catch (TProtocolException e) { - throw new PrestoException(INVALID_ARGUMENTS, "Can not serialize tpcds transaction handle", e); - } - } - - @Override - public ConnectorTransactionHandle deserialize(byte[] bytes) - { - try { - return fromThrift(bytes, thriftCodec); - } - catch (TProtocolException e) { - throw new PrestoException(INVALID_ARGUMENTS, "Can not deserialize tpcds transaction handle", e); - } - } -}