diff --git a/CODEOWNERS b/CODEOWNERS
index e4a0e3bb56aed..ddd1a78ce1754 100644
--- a/CODEOWNERS
+++ b/CODEOWNERS
@@ -92,6 +92,7 @@
/presto-thrift-spec @prestodb/committers
/presto-thrift-testing-server @prestodb/committers
/presto-thrift-testing-udf-server @prestodb/committers
+/presto-thrift-connector-toolkit @prestodb/committers
/presto-tpcds @prestodb/committers
/presto-tpch @prestodb/committers
/presto-verifier @prestodb/committers
diff --git a/pom.xml b/pom.xml
index 69dc11f44cdc8..3d4e5e7ec3fc4 100644
--- a/pom.xml
+++ b/pom.xml
@@ -210,6 +210,7 @@
presto-druid
presto-common
presto-thrift-testing-udf-server
+ presto-thrift-connector-toolkit
presto-thrift-spec
presto-testng-services
presto-node-ttl-fetchers
@@ -1140,6 +1141,12 @@
test-jar
+
+ com.facebook.presto
+ presto-thrift-connector-toolkit
+ ${project.version}
+
+
com.facebook.presto
presto-thrift-testing-server
diff --git a/presto-thrift-connector-toolkit/pom.xml b/presto-thrift-connector-toolkit/pom.xml
new file mode 100644
index 0000000000000..e506b747f7448
--- /dev/null
+++ b/presto-thrift-connector-toolkit/pom.xml
@@ -0,0 +1,63 @@
+
+
+ 4.0.0
+
+
+ com.facebook.presto
+ presto-root
+ 0.297-SNAPSHOT
+
+
+ presto-thrift-connector-toolkit
+ presto-thrift-connector-toolkit
+ jar
+
+
+ ${project.parent.basedir}
+ true
+
+
+
+
+ com.facebook.airlift.drift
+ drift-codec
+
+
+
+ com.facebook.airlift.drift
+ drift-protocol
+
+
+
+ com.facebook.presto
+ presto-spi
+ provided
+
+
+
+ com.google.guava
+ guava
+ provided
+
+
+
+ com.facebook.airlift.drift
+ drift-api
+ ${dep.drift.version}
+ provided
+
+
+
+
+ org.testng
+ testng
+ test
+
+
+
+ com.facebook.presto
+ presto-testng-services
+ test
+
+
+
diff --git a/presto-tpcds/src/main/java/com/facebook/presto/tpcds/thrift/TpcdsTransactionHandleCodec.java b/presto-thrift-connector-toolkit/src/main/java/com/facebook/presto/thrift/codec/GenericThriftCodec.java
similarity index 50%
rename from presto-tpcds/src/main/java/com/facebook/presto/tpcds/thrift/TpcdsTransactionHandleCodec.java
rename to presto-thrift-connector-toolkit/src/main/java/com/facebook/presto/thrift/codec/GenericThriftCodec.java
index 0f98eefb4d40d..03e4b8e3c4d4a 100644
--- a/presto-tpcds/src/main/java/com/facebook/presto/tpcds/thrift/TpcdsTransactionHandleCodec.java
+++ b/presto-thrift-connector-toolkit/src/main/java/com/facebook/presto/thrift/codec/GenericThriftCodec.java
@@ -11,50 +11,59 @@
* 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.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 java.lang.reflect.Type;
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 com.facebook.presto.thrift.codec.ThriftCodecUtils.fromThrift;
+import static com.facebook.presto.thrift.codec.ThriftCodecUtils.toThrift;
import static java.util.Objects.requireNonNull;
-public class TpcdsTransactionHandleCodec
- implements ConnectorCodec
+public class GenericThriftCodec
+ implements ConnectorCodec
{
- private final ThriftCodec thriftCodec;
+ private final ThriftCodec thriftCodec;
- public TpcdsTransactionHandleCodec(ThriftCodecManager thriftCodecManager)
+ public GenericThriftCodec(ThriftCodecManager codecManager, Type javaType)
{
- this.thriftCodec = requireNonNull(thriftCodecManager, "thriftCodecManager is null").getCodec(TpcdsTransactionHandle.class);
+ requireNonNull(codecManager, "codecManager is null");
+ requireNonNull(javaType, "javaType is null");
+
+ if (!(javaType instanceof Class>)) {
+ throw new IllegalArgumentException("Expected a Class type for javaType, but got: " + javaType.getTypeName());
+ }
+
+ Class> clazz = (Class>) javaType;
+
+ this.thriftCodec = (ThriftCodec) codecManager.getCodec(clazz);
}
@Override
- public byte[] serialize(ConnectorTransactionHandle handle)
+ public byte[] serialize(T value)
{
try {
- return toThrift((TpcdsTransactionHandle) handle, thriftCodec);
+ return toThrift(value, thriftCodec);
}
catch (TProtocolException e) {
- throw new PrestoException(INVALID_ARGUMENTS, "Can not serialize tpcds transaction handle", e);
+ throw new PrestoException(INVALID_ARGUMENTS, "Unable to serialize object of type " + value.getClass().getSimpleName(), e);
}
}
@Override
- public ConnectorTransactionHandle deserialize(byte[] bytes)
+ public T deserialize(byte[] bytes)
{
try {
return fromThrift(bytes, thriftCodec);
}
catch (TProtocolException e) {
- throw new PrestoException(INVALID_ARGUMENTS, "Can not deserialize tpcds transaction handle", 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..6c57089f573b3
--- /dev/null
+++ b/presto-thrift-connector-toolkit/src/main/java/com/facebook/presto/thrift/codec/ThriftCodecProvider.java
@@ -0,0 +1,168 @@
+/*
+ * 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.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> connectorSplitCodec;
+ private final Optional> connectorTransactionHandleCodec;
+ private final Optional> connectorTableLayoutHandleCodec;
+ private final Optional> connectorTableHandleCodec;
+ private final Optional> connectorOutputTableHandleCodec;
+ private final Optional> connectorInsertTableHandleCodec;
+ private final Optional> connectorDeleteTableHandleCodec;
+
+ private ThriftCodecProvider(Builder builder)
+ {
+ this.thriftCodecManager = requireNonNull(builder.thriftCodecManager, "thriftCodecManager is null");
+ this.connectorSplitCodec = builder.connectorSplitType.map(type -> new GenericThriftCodec<>(thriftCodecManager, type));
+ this.connectorTransactionHandleCodec = builder.connectorTransactionHandle.map(type -> new GenericThriftCodec<>(thriftCodecManager, type));
+ this.connectorTableLayoutHandleCodec = builder.connectorTableLayoutHandle.map(type -> new GenericThriftCodec<>(thriftCodecManager, type));
+ this.connectorTableHandleCodec = builder.connectorTableHandle.map(type -> new GenericThriftCodec<>(thriftCodecManager, type));
+ this.connectorOutputTableHandleCodec = builder.connectorOutputTableHandle.map(type -> new GenericThriftCodec<>(thriftCodecManager, type));
+ this.connectorInsertTableHandleCodec = builder.connectorInsertTableHandle.map(type -> new GenericThriftCodec<>(thriftCodecManager, type));
+ this.connectorDeleteTableHandleCodec = builder.connectorDeleteTableHandle.map(type -> new GenericThriftCodec<>(thriftCodecManager, type));
+ }
+
+ @Override
+ public Optional> getConnectorSplitCodec()
+ {
+ return connectorSplitCodec;
+ }
+
+ @Override
+ public Optional> getConnectorTransactionHandleCodec()
+ {
+ return connectorTransactionHandleCodec;
+ }
+
+ @Override
+ public Optional> getConnectorTableLayoutHandleCodec()
+ {
+ return connectorTableLayoutHandleCodec;
+ }
+
+ @Override
+ public Optional> getConnectorTableHandleCodec()
+ {
+ return connectorTableHandleCodec;
+ }
+
+ @Override
+ public Optional> getConnectorOutputTableHandleCodec()
+ {
+ return connectorOutputTableHandleCodec;
+ }
+
+ @Override
+ public Optional> getConnectorInsertTableHandleCodec()
+ {
+ return connectorInsertTableHandleCodec;
+ }
+
+ @Override
+ public Optional> getConnectorDeleteTableHandleCodec()
+ {
+ return connectorDeleteTableHandleCodec;
+ }
+
+ 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();
+
+ public Builder setThriftCodecManager(ThriftCodecManager thriftCodecManager)
+ {
+ this.thriftCodecManager = thriftCodecManager;
+ return this;
+ }
+
+ public Builder setConnectorSplitType(Class extends ConnectorSplit> type)
+ {
+ this.connectorSplitType = Optional.ofNullable(type);
+ return this;
+ }
+
+ public Builder setConnectorTransactionHandle(Class extends ConnectorTransactionHandle> type)
+ {
+ this.connectorTransactionHandle = Optional.ofNullable(type);
+ return this;
+ }
+
+ public Builder setConnectorTableLayoutHandle(Class extends ConnectorTableLayoutHandle> type)
+ {
+ this.connectorTableLayoutHandle = Optional.ofNullable(type);
+ return this;
+ }
+
+ public Builder setConnectorTableHandle(Class extends ConnectorTableHandle> type)
+ {
+ this.connectorTableHandle = Optional.ofNullable(type);
+ return this;
+ }
+
+ public Builder setConnectorOutputTableHandle(Class extends ConnectorOutputTableHandle> type)
+ {
+ this.connectorOutputTableHandle = Optional.ofNullable(type);
+ return this;
+ }
+
+ public Builder setConnectorInsertTableHandle(Class extends ConnectorInsertTableHandle> type)
+ {
+ this.connectorInsertTableHandle = Optional.ofNullable(type);
+ return this;
+ }
+
+ public Builder setConnectorDeleteTableHandle(Class extends ConnectorDeleteTableHandle> type)
+ {
+ this.connectorDeleteTableHandle = Optional.ofNullable(type);
+ return this;
+ }
+
+ public ThriftCodecProvider build()
+ {
+ requireNonNull(thriftCodecManager, "ThriftCodecManager not set");
+ 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-connector-toolkit/src/test/java/com/facebook/presto/thrift/codec/TestGenericThriftCodec.java b/presto-thrift-connector-toolkit/src/test/java/com/facebook/presto/thrift/codec/TestGenericThriftCodec.java
new file mode 100644
index 0000000000000..21a147f047290
--- /dev/null
+++ b/presto-thrift-connector-toolkit/src/test/java/com/facebook/presto/thrift/codec/TestGenericThriftCodec.java
@@ -0,0 +1,114 @@
+/*
+ * 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.annotations.ThriftConstructor;
+import com.facebook.drift.annotations.ThriftField;
+import com.facebook.drift.annotations.ThriftStruct;
+import com.facebook.drift.codec.ThriftCodecManager;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+import java.util.Objects;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
+
+@Test(singleThreaded = true)
+public class TestGenericThriftCodec
+{
+ private ThriftCodecManager codecManager;
+
+ @BeforeMethod
+ public void setUp()
+ {
+ codecManager = new ThriftCodecManager();
+ }
+
+ @Test
+ public void testSerializeAndDeserialize()
+ {
+ GenericThriftCodec codec = new GenericThriftCodec<>(codecManager, TestThriftObject.class);
+
+ TestThriftObject original = new TestThriftObject(42, "test-value");
+
+ byte[] serialized = codec.serialize(original);
+
+ assertNotNull(serialized);
+ TestThriftObject deserialized = codec.deserialize(serialized);
+
+ assertEquals(deserialized.getId(), original.getId());
+ assertEquals(deserialized.getName(), original.getName());
+ }
+
+ @Test
+ public void testSerializeNull()
+ {
+ GenericThriftCodec codec = new GenericThriftCodec<>(codecManager, TestThriftObject.class);
+
+ TestThriftObject original = new TestThriftObject(0, null);
+ byte[] serialized = codec.serialize(original);
+
+ assertNotNull(serialized);
+ TestThriftObject deserialized = codec.deserialize(serialized);
+
+ assertEquals(deserialized.getId(), original.getId());
+ assertEquals(deserialized.getName(), original.getName());
+ }
+
+ @ThriftStruct
+ public static class TestThriftObject
+ {
+ private final int id;
+ private final String name;
+
+ @ThriftConstructor
+ public TestThriftObject(int id, String name)
+ {
+ this.id = id;
+ this.name = name;
+ }
+
+ @ThriftField(1)
+ public int getId()
+ {
+ return id;
+ }
+
+ @ThriftField(2)
+ public String getName()
+ {
+ return name;
+ }
+
+ @Override
+ public boolean equals(Object o)
+ {
+ if (this == o) {
+ return true;
+ }
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+ TestThriftObject that = (TestThriftObject) o;
+ return id == that.id && Objects.equals(name, that.name);
+ }
+
+ @Override
+ public int hashCode()
+ {
+ return Objects.hash(id, name);
+ }
+ }
+}
diff --git a/presto-thrift-connector-toolkit/src/test/java/com/facebook/presto/thrift/codec/TestThriftCodecProvider.java b/presto-thrift-connector-toolkit/src/test/java/com/facebook/presto/thrift/codec/TestThriftCodecProvider.java
new file mode 100644
index 0000000000000..955e942203217
--- /dev/null
+++ b/presto-thrift-connector-toolkit/src/test/java/com/facebook/presto/thrift/codec/TestThriftCodecProvider.java
@@ -0,0 +1,341 @@
+/*
+ * 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.annotations.ThriftConstructor;
+import com.facebook.drift.annotations.ThriftField;
+import com.facebook.drift.annotations.ThriftStruct;
+import com.facebook.drift.codec.ThriftCodecManager;
+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.HostAddress;
+import com.facebook.presto.spi.NodeProvider;
+import com.facebook.presto.spi.connector.ConnectorTransactionHandle;
+import com.facebook.presto.spi.schedule.NodeSelectionStrategy;
+import com.google.common.collect.ImmutableList;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+import java.util.List;
+import java.util.Optional;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertTrue;
+
+@Test(singleThreaded = true)
+public class TestThriftCodecProvider
+{
+ private ThriftCodecProvider provider;
+
+ @BeforeMethod
+ public void setUp()
+ {
+ provider = new ThriftCodecProvider.Builder()
+ .setThriftCodecManager(new ThriftCodecManager())
+ .setConnectorSplitType(TestSplit.class)
+ .setConnectorTransactionHandle(TestTransactionHandle.class)
+ .setConnectorTableLayoutHandle(TestTableLayoutHandle.class)
+ .setConnectorTableHandle(TestTableHandle.class)
+ .setConnectorOutputTableHandle(TestOutputTableHandle.class)
+ .setConnectorInsertTableHandle(TestInsertTableHandle.class)
+ .setConnectorDeleteTableHandle(TestDeleteTableHandle.class)
+ .build();
+ }
+
+ @Test
+ public void testSplitCodecRoundTrip()
+ {
+ Optional> codecOptional = provider.getConnectorSplitCodec();
+ assertTrue(codecOptional.isPresent());
+
+ ConnectorCodec codec = codecOptional.get();
+ TestSplit original = new TestSplit("test-id", 100);
+ byte[] serialized = codec.serialize(original);
+ ConnectorSplit deserialized = codec.deserialize(serialized);
+
+ assertTrue(deserialized instanceof TestSplit);
+ TestSplit deserializedSplit = (TestSplit) deserialized;
+ assertEquals(deserializedSplit.getSplitId(), original.getSplitId());
+ assertEquals(deserializedSplit.getSplitSize(), original.getSplitSize());
+ }
+
+ @Test
+ public void testTableHandleCodecRoundTrip()
+ {
+ Optional> codecOptional = provider.getConnectorTableHandleCodec();
+ assertTrue(codecOptional.isPresent());
+
+ ConnectorCodec codec = codecOptional.get();
+ TestTableHandle original = new TestTableHandle("test-schema", "test-table");
+ byte[] serialized = codec.serialize(original);
+ ConnectorTableHandle deserialized = codec.deserialize(serialized);
+
+ assertTrue(deserialized instanceof TestTableHandle);
+ TestTableHandle deserializedHandle = (TestTableHandle) deserialized;
+ assertEquals(deserializedHandle.getSchemaName(), original.getSchemaName());
+ assertEquals(deserializedHandle.getTableName(), original.getTableName());
+ }
+
+ @Test
+ public void testTransactionHandleCodecRoundTrip()
+ {
+ Optional> codecOptional = provider.getConnectorTransactionHandleCodec();
+ assertTrue(codecOptional.isPresent());
+
+ ConnectorCodec codec = codecOptional.get();
+ TestTransactionHandle original = new TestTransactionHandle("test");
+ byte[] serialized = codec.serialize(original);
+ ConnectorTransactionHandle deserialized = codec.deserialize(serialized);
+
+ assertTrue(deserialized instanceof TestTransactionHandle);
+ TestTransactionHandle deserializedHandle = (TestTransactionHandle) deserialized;
+ assertEquals(deserializedHandle.getTransactionId(), original.getTransactionId());
+ }
+
+ @Test
+ public void testTableLayoutHandleCodecRoundTrip()
+ {
+ Optional> codecOptional = provider.getConnectorTableLayoutHandleCodec();
+ assertTrue(codecOptional.isPresent());
+
+ ConnectorCodec codec = codecOptional.get();
+ TestTableLayoutHandle original = new TestTableLayoutHandle("test");
+ byte[] serialized = codec.serialize(original);
+ ConnectorTableLayoutHandle deserialized = codec.deserialize(serialized);
+
+ assertTrue(deserialized instanceof TestTableLayoutHandle);
+ TestTableLayoutHandle deserializedHandle = (TestTableLayoutHandle) deserialized;
+ assertEquals(deserializedHandle.getLayoutId(), original.getLayoutId());
+ }
+
+ @Test
+ public void testOutputTableHandleCodecRoundTrip()
+ {
+ Optional> codecOptional = provider.getConnectorOutputTableHandleCodec();
+ assertTrue(codecOptional.isPresent());
+
+ ConnectorCodec codec = codecOptional.get();
+ TestOutputTableHandle original = new TestOutputTableHandle("test");
+ byte[] serialized = codec.serialize(original);
+ ConnectorOutputTableHandle deserialized = codec.deserialize(serialized);
+
+ assertTrue(deserialized instanceof TestOutputTableHandle);
+ TestOutputTableHandle deserializedHandle = (TestOutputTableHandle) deserialized;
+ assertEquals(deserializedHandle.getOutputId(), original.getOutputId());
+ }
+
+ @Test
+ public void testInsertTableHandleCodecRoundTrip()
+ {
+ Optional> codecOptional = provider.getConnectorInsertTableHandleCodec();
+ assertTrue(codecOptional.isPresent());
+
+ ConnectorCodec codec = codecOptional.get();
+ TestInsertTableHandle original = new TestInsertTableHandle("test");
+ byte[] serialized = codec.serialize(original);
+ ConnectorInsertTableHandle deserialized = codec.deserialize(serialized);
+
+ assertTrue(deserialized instanceof TestInsertTableHandle);
+ TestInsertTableHandle deserializedHandle = (TestInsertTableHandle) deserialized;
+ assertEquals(deserializedHandle.getInsertId(), original.getInsertId());
+ }
+
+ @Test
+ public void testDeleteTableHandleCodecRoundTrip()
+ {
+ Optional> codecOptional = provider.getConnectorDeleteTableHandleCodec();
+ assertTrue(codecOptional.isPresent());
+
+ ConnectorCodec codec = codecOptional.get();
+ TestDeleteTableHandle original = new TestDeleteTableHandle("test");
+ byte[] serialized = codec.serialize(original);
+ ConnectorDeleteTableHandle deserialized = codec.deserialize(serialized);
+
+ assertTrue(deserialized instanceof TestDeleteTableHandle);
+ TestDeleteTableHandle deserializedHandle = (TestDeleteTableHandle) deserialized;
+ assertEquals(deserializedHandle.getDeleteId(), original.getDeleteId());
+ }
+
+ @ThriftStruct
+ public static class TestSplit
+ implements ConnectorSplit
+ {
+ private final String splitId;
+ private final long splitSize;
+
+ @ThriftConstructor
+ public TestSplit(String splitId, long splitSize)
+ {
+ this.splitId = splitId;
+ this.splitSize = splitSize;
+ }
+
+ @ThriftField(1)
+ public String getSplitId()
+ {
+ return splitId;
+ }
+
+ @ThriftField(2)
+ public long getSplitSize()
+ {
+ return splitSize;
+ }
+
+ @Override
+ public NodeSelectionStrategy getNodeSelectionStrategy()
+ {
+ return null;
+ }
+
+ @Override
+ public List getPreferredNodes(NodeProvider nodeProvider)
+ {
+ return ImmutableList.of();
+ }
+
+ @Override
+ public Object getInfo()
+ {
+ return this;
+ }
+ }
+
+ @ThriftStruct
+ public static class TestTransactionHandle
+ implements ConnectorTransactionHandle
+ {
+ private final String transactionId;
+
+ @ThriftConstructor
+ public TestTransactionHandle(String transactionId)
+ {
+ this.transactionId = transactionId;
+ }
+
+ @ThriftField(1)
+ public String getTransactionId()
+ {
+ return transactionId;
+ }
+ }
+
+ @ThriftStruct
+ public static class TestTableLayoutHandle
+ implements ConnectorTableLayoutHandle
+ {
+ private final String layoutId;
+
+ @ThriftConstructor
+ public TestTableLayoutHandle(String layoutId)
+ {
+ this.layoutId = layoutId;
+ }
+
+ @ThriftField(1)
+ public String getLayoutId()
+ {
+ return layoutId;
+ }
+ }
+
+ @ThriftStruct
+ public static class TestTableHandle
+ implements ConnectorTableHandle
+ {
+ private final String schemaName;
+ private final String tableName;
+
+ @ThriftConstructor
+ public TestTableHandle(String schemaName, String tableName)
+ {
+ this.schemaName = schemaName;
+ this.tableName = tableName;
+ }
+
+ @ThriftField(1)
+ public String getSchemaName()
+ {
+ return schemaName;
+ }
+
+ @ThriftField(2)
+ public String getTableName()
+ {
+ return tableName;
+ }
+ }
+
+ @ThriftStruct
+ public static class TestOutputTableHandle
+ implements ConnectorOutputTableHandle
+ {
+ private final String outputId;
+
+ @ThriftConstructor
+ public TestOutputTableHandle(String outputId)
+ {
+ this.outputId = outputId;
+ }
+
+ @ThriftField(1)
+ public String getOutputId()
+ {
+ return outputId;
+ }
+ }
+
+ @ThriftStruct
+ public static class TestInsertTableHandle
+ implements ConnectorInsertTableHandle
+ {
+ private final String insertId;
+
+ @ThriftConstructor
+ public TestInsertTableHandle(String insertId)
+ {
+ this.insertId = insertId;
+ }
+
+ @ThriftField(1)
+ public String getInsertId()
+ {
+ return insertId;
+ }
+ }
+
+ @ThriftStruct
+ public static class TestDeleteTableHandle
+ implements ConnectorDeleteTableHandle
+ {
+ private final String deleteId;
+
+ @ThriftConstructor
+ public TestDeleteTableHandle(String deleteId)
+ {
+ this.deleteId = deleteId;
+ }
+
+ @ThriftField(1)
+ public String getDeleteId()
+ {
+ return deleteId;
+ }
+ }
+}
diff --git a/presto-tpcds/pom.xml b/presto-tpcds/pom.xml
index ef412d28496e0..498ada559f041 100644
--- a/presto-tpcds/pom.xml
+++ b/presto-tpcds/pom.xml
@@ -49,8 +49,8 @@
- com.facebook.airlift.drift
- drift-protocol
+ com.facebook.presto
+ presto-thrift-connector-toolkit
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 428fa9ab70d26..796f7ee5faf9f 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 fa3a678b550a8..0000000000000
--- a/presto-tpcds/src/main/java/com/facebook/presto/tpcds/thrift/TpcdsCodecProvider.java
+++ /dev/null
@@ -1,68 +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 ConnectorCodec splitCodec;
- private final ConnectorCodec transactionHandleCodec;
- private final ConnectorCodec tableLayoutHandleCodec;
- private final ConnectorCodec tableHandleCodec;
-
- public TpcdsCodecProvider(ThriftCodecManager thriftCodecManager)
- {
- requireNonNull(thriftCodecManager, "thriftCodecManager is null");
- this.splitCodec = new TpcdsSplitCodec(thriftCodecManager);
- this.transactionHandleCodec = new TpcdsTransactionHandleCodec(thriftCodecManager);
- this.tableLayoutHandleCodec = new TpcdsTableLayoutHandleCodec(thriftCodecManager);
- this.tableHandleCodec = new TpcdsTableHandleCodec(thriftCodecManager);
- }
-
- @Override
- public Optional> getConnectorSplitCodec()
- {
- return Optional.of(splitCodec);
- }
-
- @Override
- public Optional> getConnectorTransactionHandleCodec()
- {
- return Optional.of(transactionHandleCodec);
- }
-
- @Override
- public Optional> getConnectorTableLayoutHandleCodec()
- {
- return Optional.of(tableLayoutHandleCodec);
- }
-
- @Override
- public Optional> getConnectorTableHandleCodec()
- {
- return Optional.of(tableHandleCodec);
- }
-}
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/test/java/com/facebook/presto/tpcds/TestTpcdsWithThrift.java b/presto-tpcds/src/test/java/com/facebook/presto/tpcds/TestTpcdsWithThrift.java
index 5ce27d9be0815..8011499acd6ac 100644
--- a/presto-tpcds/src/test/java/com/facebook/presto/tpcds/TestTpcdsWithThrift.java
+++ b/presto-tpcds/src/test/java/com/facebook/presto/tpcds/TestTpcdsWithThrift.java
@@ -15,7 +15,9 @@
import com.facebook.presto.testing.QueryRunner;
import com.google.common.collect.ImmutableMap;
+import org.testng.annotations.Test;
+@Test(singleThreaded = true)
public class TestTpcdsWithThrift
extends AbstractTestTpcds
{