14
14
15
15
#include < gtest/gtest.h>
16
16
#include " presto_cpp/main/thrift/ProtocolToThrift.h"
17
+ #include " presto_cpp/main/thrift/ThriftIO.h"
17
18
#include " presto_cpp/main/common/tests/test_json.h"
18
19
#include " presto_cpp/main/connectors/PrestoToVeloxConnector.h"
19
20
@@ -100,7 +101,7 @@ TEST_F(TaskUpdateRequestTest, mapOutputBuffers) {
100
101
ASSERT_EQ (outputBuffers.buffers [" 2" ], 20 );
101
102
}
102
103
103
- TEST_F (TaskUpdateRequestTest, binarySplitFromThrift ) {
104
+ TEST_F (TaskUpdateRequestTest, binaryHiveSplitFromThrift ) {
104
105
thrift::Split thriftSplit;
105
106
thriftSplit.connectorId ()->catalogName_ref () = " hive" ;
106
107
thriftSplit.transactionHandle ()->jsonValue_ref () = R"( {
@@ -127,14 +128,89 @@ TEST_F(TaskUpdateRequestTest, binarySplitFromThrift) {
127
128
protocol::NodeSelectionStrategy::NO_PREFERENCE);
128
129
}
129
130
130
- TEST_F (TaskUpdateRequestTest, binaryTableWriteInfo) {
131
- std::string str = slurp (getDataPath (BASE_DATA_PATH, " TableWriteInfo.json" ));
132
- protocol::TableWriteInfo tableWriteInfo;
131
+ TEST_F (TaskUpdateRequestTest, binaryRemoteSplitFromThrift) {
132
+ thrift::Split thriftSplit;
133
+ thrift::RemoteTransactionHandle thriftTransactionHandle;
134
+ thrift::RemoteSplit thriftRemoteSplit;
135
+
136
+ thriftSplit.connectorId ()->catalogName_ref () = " $remote" ;
137
+ thriftSplit.transactionHandle ()->customSerializedValue_ref () =
138
+ thriftWrite (thriftTransactionHandle);
139
+
140
+ thriftRemoteSplit.location ()->location_ref () = " /test_location" ;
141
+ thriftRemoteSplit.remoteSourceTaskId ()->id_ref () = 100 ;
142
+ thriftRemoteSplit.remoteSourceTaskId ()->attemptNumber_ref () = 200 ;
143
+ thriftRemoteSplit.remoteSourceTaskId ()->stageExecutionId ()->id_ref () = 300 ;
144
+ thriftRemoteSplit.remoteSourceTaskId ()->stageExecutionId ()->stageId ()->id_ref () = 400 ;
145
+ thriftRemoteSplit.remoteSourceTaskId ()->stageExecutionId ()->stageId ()->queryId_ref () = " test_query_id" ;
146
+
147
+ thriftSplit.connectorSplit ()->connectorId_ref () = " $remote" ;
148
+ thriftSplit.connectorSplit ()->customSerializedValue_ref () =
149
+ thriftWrite (thriftRemoteSplit);
150
+
151
+ protocol::Split split;
152
+ thrift::fromThrift (thriftSplit, split);
153
+
154
+ // Verify that connector specific fields are set correctly with thrift codec
155
+ auto remoteSplit = std::dynamic_pointer_cast<protocol::RemoteSplit>(
156
+ split.connectorSplit );
157
+ ASSERT_EQ ((remoteSplit->location ).location , " /test_location" );
158
+ ASSERT_EQ (remoteSplit->remoteSourceTaskId , " test_query_id.400.300.100.200" );
159
+ }
160
+
161
+ TEST_F (TaskUpdateRequestTest, unionExecutionWriterTargetFromThrift) {
162
+ // Construct ExecutionWriterTarget with CreateHandle
163
+ thrift::CreateHandle thriftCreateHandle;
164
+ thrift::ExecutionWriterTargetUnion thriftWriterTarget;
165
+ thriftCreateHandle.schemaTableName ()->schema_ref () = " test_schema" ;
166
+ thriftCreateHandle.schemaTableName ()->table_ref () = " test_table" ;
167
+ thriftCreateHandle.handle ()->connectorId ()->catalogName_ref () = " hive" ;
168
+ thriftCreateHandle.handle ()->transactionHandle ()->jsonValue_ref () = R"( {
169
+ "@type": "hive",
170
+ "uuid": "8a4d6c83-60ee-46de-9715-bc91755619fa"
171
+ })" ;
172
+ thriftCreateHandle.handle ()->connectorHandle ()->jsonValue_ref () = slurp (getDataPath (BASE_DATA_PATH, " HiveOutputTableHandle.json" ));;
173
+ thriftWriterTarget.set_createHandle (std::move (thriftCreateHandle));
174
+
175
+ // Convert from thrift to protocol and verify fields
176
+ auto writerTarget = std::make_shared<protocol::ExecutionWriterTarget>();
177
+ thrift::fromThrift (thriftWriterTarget, writerTarget);
178
+
179
+ ASSERT_EQ (writerTarget->_type , " CreateHandle" );
180
+ auto createHandle = std::dynamic_pointer_cast<protocol::CreateHandle>(writerTarget);
181
+ ASSERT_NE (createHandle, nullptr );
182
+ ASSERT_EQ (createHandle->schemaTableName .schema , " test_schema" );
183
+ ASSERT_EQ (createHandle->schemaTableName .table , " test_table" );
184
+
185
+ auto * hiveTxnHandle = dynamic_cast <protocol::hive::HiveTransactionHandle*>(createHandle->handle .transactionHandle .get ());
186
+ ASSERT_NE (hiveTxnHandle, nullptr );
187
+ ASSERT_EQ (hiveTxnHandle->uuid , " 8a4d6c83-60ee-46de-9715-bc91755619fa" );
188
+
189
+ auto * hiveOutputTableHandle = dynamic_cast <protocol::hive::HiveOutputTableHandle*>(createHandle->handle .connectorHandle .get ());
190
+ ASSERT_NE (hiveOutputTableHandle, nullptr );
191
+ ASSERT_EQ (hiveOutputTableHandle->schemaName , " test_schema" );
192
+ ASSERT_EQ (hiveOutputTableHandle->tableName , " test_table" );
193
+ ASSERT_EQ (hiveOutputTableHandle->tableStorageFormat , protocol::hive::HiveStorageFormat::ORC);
194
+ ASSERT_EQ (hiveOutputTableHandle->locationHandle .targetPath , " /path/to/target" );
195
+ }
196
+
197
+ TEST_F (TaskUpdateRequestTest, unionExecutionWriterTargetToThrift) {
198
+ // Construct thrift ExecutionWriterTarget with CreateHandle
199
+ auto createHandle = std::make_shared<protocol::CreateHandle>();
200
+ createHandle->schemaTableName .schema = " test_schema" ;
201
+ createHandle->schemaTableName .table = " test_table" ;
202
+
203
+ auto writerTarget = std::make_shared<protocol::ExecutionWriterTarget>();
204
+ writerTarget->_type = " CreateHandle" ;
205
+ writerTarget = createHandle;
133
206
134
- thrift::fromThrift (str, tableWriteInfo);
135
- auto hiveTableHandle = std::dynamic_pointer_cast<protocol::hive::HiveTableHandle>((*tableWriteInfo.analyzeTableHandle ).connectorHandle );
136
- ASSERT_EQ (hiveTableHandle->tableName , " test_table" );
137
- ASSERT_EQ (hiveTableHandle->analyzePartitionValues ->size (), 2 );
207
+ // Convert to thrift and verify fields. Note that toThrift functions for connector fields are not implemented.
208
+ thrift::ExecutionWriterTargetUnion thriftWriterTarget;
209
+ thrift::toThrift (writerTarget, thriftWriterTarget);
210
+ ASSERT_TRUE (thriftWriterTarget.createHandle_ref ().has_value ());
211
+ const auto & thriftCreateHandle = thriftWriterTarget.createHandle_ref ().value ();
212
+ ASSERT_EQ (thriftCreateHandle.schemaTableName ()->schema_ref ().value (), " test_schema" );
213
+ ASSERT_EQ (thriftCreateHandle.schemaTableName ()->table_ref ().value (), " test_table" );
138
214
}
139
215
140
216
TEST_F (TaskUpdateRequestTest, fragment) {
0 commit comments