diff --git a/presto-base-arrow-flight/pom.xml b/presto-base-arrow-flight/pom.xml index 2548ad2d572ea..88d4a47caf963 100644 --- a/presto-base-arrow-flight/pom.xml +++ b/presto-base-arrow-flight/pom.xml @@ -17,7 +17,7 @@ 1.63.0 4.10.0 17.0.0 - 4.1.100.Final + 4.1.110.Final @@ -114,6 +114,13 @@ configuration + + org.apache.arrow + arrow-jdbc + 17.0.0 + + + io.netty netty-codec-http2 @@ -197,6 +204,12 @@ provided + + com.fasterxml.jackson.core + jackson-core + provided + + org.apache.arrow flight-core @@ -208,6 +221,24 @@ + + com.facebook.presto + presto-main + test + + + + com.facebook.presto + presto-tests + test + + + + com.h2database + h2 + 2.2.220 + test + @@ -292,7 +323,6 @@ 4.1.110.Final - io.netty netty-handler @@ -317,6 +347,12 @@ 3.3 + + org.jetbrains.kotlin + kotlin-stdlib-common + 1.6.20 + + org.slf4j slf4j-api @@ -362,6 +398,7 @@ arrow-git.properties + about.html diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/ArrowServer.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/ArrowServer.java deleted file mode 100644 index 334b744f0e881..0000000000000 --- a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/ArrowServer.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.plugin.arrow; - -import org.apache.arrow.flight.Action; -import org.apache.arrow.flight.ActionType; -import org.apache.arrow.flight.Criteria; -import org.apache.arrow.flight.FlightDescriptor; -import org.apache.arrow.flight.FlightInfo; -import org.apache.arrow.flight.FlightProducer; -import org.apache.arrow.flight.FlightStream; -import org.apache.arrow.flight.PutResult; -import org.apache.arrow.flight.Result; -import org.apache.arrow.flight.Ticket; - -public class ArrowServer - implements FlightProducer -{ - @Override - public void getStream(CallContext callContext, Ticket ticket, ServerStreamListener serverStreamListener) - { - } - - @Override - public void listFlights(CallContext callContext, Criteria criteria, StreamListener streamListener) - { - } - - @Override - public FlightInfo getFlightInfo(CallContext callContext, FlightDescriptor flightDescriptor) - { - return null; - } - - @Override - public Runnable acceptPut(CallContext callContext, FlightStream flightStream, StreamListener streamListener) - { - return null; - } - - @Override - public void doAction(CallContext callContext, Action action, StreamListener streamListener) - { - } - - @Override - public void listActions(CallContext callContext, StreamListener streamListener) - { - } -} diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/ArrowTableHandleTest.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/ArrowTableHandleTest.java deleted file mode 100644 index cd8d3984cc533..0000000000000 --- a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/ArrowTableHandleTest.java +++ /dev/null @@ -1,53 +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.plugin.arrow; - -import com.fasterxml.jackson.databind.ObjectMapper; -import org.testng.annotations.Test; - -import static org.testng.Assert.assertEquals; -import static org.testng.Assert.assertNotNull; - -public class ArrowTableHandleTest -{ - @Test - public void testConstructorAndGetters() - { - ArrowTableHandle handle = new ArrowTableHandle("test_schema", "test_table"); - assertEquals(handle.getSchema(), "test_schema"); - assertEquals(handle.getTable(), "test_table"); - } - - @Test - public void testToString() - { - ArrowTableHandle handle = new ArrowTableHandle("test_schema", "test_table"); - assertEquals(handle.toString(), "test_schema:test_table"); - } - - @Test - public void testJsonSerialization() throws Exception - { - ObjectMapper objectMapper = new ObjectMapper(); - - ArrowTableHandle handle = new ArrowTableHandle("test_schema", "test_table"); - String json = objectMapper.writeValueAsString(handle); - - assertNotNull(json); - - ArrowTableHandle deserialized = objectMapper.readValue(json, ArrowTableHandle.class); - assertEquals(deserialized.getSchema(), handle.getSchema()); - assertEquals(deserialized.getTable(), handle.getTable()); - } -} diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowAbstractMetadata.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowAbstractMetadata.java deleted file mode 100644 index 378dc325da791..0000000000000 --- a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowAbstractMetadata.java +++ /dev/null @@ -1,112 +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.plugin.arrow; - -import com.facebook.presto.common.type.BooleanType; -import com.facebook.presto.common.type.DecimalType; -import com.facebook.presto.common.type.IntegerType; -import com.facebook.presto.common.type.VarcharType; -import com.facebook.presto.spi.ColumnHandle; -import com.facebook.presto.spi.ColumnMetadata; -import com.facebook.presto.spi.ConnectorSession; -import com.facebook.presto.spi.ConnectorTableHandle; -import com.facebook.presto.spi.ConnectorTableLayoutResult; -import com.facebook.presto.spi.ConnectorTableMetadata; -import com.facebook.presto.spi.Constraint; -import com.facebook.presto.spi.SchemaTableName; -import org.apache.arrow.vector.types.pojo.ArrowType; -import org.apache.arrow.vector.types.pojo.Field; -import org.apache.arrow.vector.types.pojo.FieldType; -import org.mockito.Mockito; -import org.testng.annotations.Test; - -import java.io.IOException; -import java.util.Arrays; -import java.util.HashSet; -import java.util.List; -import java.util.Optional; -import java.util.Set; - -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; -import static org.testng.Assert.assertEquals; -import static org.testng.Assert.assertNotNull; - -@Test(singleThreaded = true) -public class TestArrowAbstractMetadata -{ - public TestArrowAbstractMetadata() throws IOException - { - } - - @Test - public void testGetTableMetadata() - { - // Mock dependencies - ArrowAbstractMetadata metadata = mock(ArrowAbstractMetadata.class); - Mockito.doCallRealMethod().when(metadata).getTableMetadata(Mockito.any(ConnectorSession.class), Mockito.any(ConnectorTableHandle.class)); - ConnectorSession session = mock(ConnectorSession.class); - ArrowTableHandle tableHandle = new ArrowTableHandle("testSchema", "testTable"); - - // Mock the behavior of getColumnsList - List columnList = Arrays.asList( - new Field("column1", FieldType.notNullable(new ArrowType.Int(32, true)), null), - new Field("column2", FieldType.notNullable(new ArrowType.Decimal(10, 2)), null), - new Field("column3", FieldType.notNullable(new ArrowType.Bool()), null)); - - when(metadata.getColumnsList("testSchema", "testTable", session)).thenReturn(columnList); - // Call the method under test - ConnectorTableMetadata tableMetadata = metadata.getTableMetadata(session, tableHandle); - - // Verify the result - assertNotNull(tableMetadata); - assertEquals(tableMetadata.getTable(), new SchemaTableName("testSchema", "testTable")); - List columns = tableMetadata.getColumns(); - assertEquals(columns.size(), 3); - assertEquals(columns.get(0), new ColumnMetadata("column1", IntegerType.INTEGER)); - assertEquals(columns.get(1), new ColumnMetadata("column2", DecimalType.createDecimalType(10, 2))); - assertEquals(columns.get(2), new ColumnMetadata("column3", BooleanType.BOOLEAN)); - } - - @Test - public void testGetTableLayouts() - { - // Mock dependencies - ConnectorSession session = mock(ConnectorSession.class); - ConnectorTableHandle tableHandle = new ArrowTableHandle("testSchema", "testTable"); - - // Mock the constraint - Constraint trueConstraint = Constraint.alwaysTrue(); - - Set desiredColumns = new HashSet<>(); - desiredColumns.add(new ArrowColumnHandle("column1", IntegerType.INTEGER)); - desiredColumns.add(new ArrowColumnHandle("column2", VarcharType.VARCHAR)); - - // Call the method under test - ArrowAbstractMetadata metadata = mock(ArrowAbstractMetadata.class); - Mockito.doCallRealMethod().when(metadata).getTableLayouts(Mockito.any(ConnectorSession.class), Mockito.any(ConnectorTableHandle.class), Mockito.any(Constraint.class), Mockito.any(Optional.class)); - - List tableLayouts = metadata.getTableLayouts(session, tableHandle, trueConstraint, Optional.of(desiredColumns)); - - // Verify the result - assertNotNull(tableLayouts); - assertEquals(tableLayouts.size(), 1); - ConnectorTableLayoutResult layoutResult = tableLayouts.get(0); - assertNotNull(layoutResult); - assertNotNull(layoutResult.getTableLayout()); - assertEquals(((ArrowTableLayoutHandle) layoutResult.getTableLayout().getHandle()).getTableHandle(), tableHandle); - assertEquals(((ArrowTableLayoutHandle) layoutResult.getTableLayout().getHandle()).getColumnHandles(), desiredColumns); - assertEquals(layoutResult.getTableLayout().getPredicate(), trueConstraint.getSummary()); - } -} diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowColumnHandle.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowColumnHandle.java deleted file mode 100644 index 6eca6de0a85c5..0000000000000 --- a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowColumnHandle.java +++ /dev/null @@ -1,49 +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.plugin.arrow; - -import com.facebook.presto.common.type.Type; -import com.facebook.presto.common.type.VarcharType; -import com.facebook.presto.spi.ColumnMetadata; -import org.testng.annotations.Test; - -import static org.testng.Assert.assertEquals; - -public class TestArrowColumnHandle -{ - @Test - public void testConstructorAndGetters() - { - String columnName = "TestColumn"; - Type columnType = VarcharType.createUnboundedVarcharType(); - - ArrowColumnHandle columnHandle = new ArrowColumnHandle(columnName, columnType); - - assertEquals(columnHandle.getColumnName(), columnName); - assertEquals(columnHandle.getColumnType(), columnType); - } - - @Test - public void testGetColumnMetadata() - { - String columnName = "testcolumn"; - Type columnType = VarcharType.createUnboundedVarcharType(); - - ArrowColumnHandle columnHandle = new ArrowColumnHandle(columnName, columnType); - ColumnMetadata columnMetadata = columnHandle.getColumnMetadata(); - - assertEquals(columnMetadata.getName(), columnName); - assertEquals(columnMetadata.getType(), columnType); - } -} diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightClient.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightClient.java deleted file mode 100644 index ce5039c3e7186..0000000000000 --- a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightClient.java +++ /dev/null @@ -1,88 +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.plugin.arrow; - -import org.apache.arrow.flight.FlightClient; -import org.apache.arrow.memory.RootAllocator; -import org.testng.annotations.Test; - -import java.io.InputStream; -import java.util.Optional; - -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.times; -import static org.mockito.Mockito.verify; -import static org.testng.Assert.assertEquals; -import static org.testng.Assert.assertFalse; -import static org.testng.Assert.assertTrue; - -public class TestArrowFlightClient -{ - @Test - public void testArrowFlightClient() - { - FlightClient flightClient = mock(FlightClient.class); - InputStream certificateStream = mock(InputStream.class); - Optional trustedCertificate = Optional.of(certificateStream); - RootAllocator allocator = mock(RootAllocator.class); - - ArrowFlightClient arrowFlightClient = new ArrowFlightClient(flightClient, trustedCertificate, allocator); - - assertEquals(arrowFlightClient.getFlightClient(), flightClient); - assertTrue(arrowFlightClient.getTrustedCertificate().isPresent()); - assertEquals(arrowFlightClient.getTrustedCertificate().get(), certificateStream); - } - - @Test - public void testArrowFlightClientWithoutCertificate() - { - FlightClient flightClient = mock(FlightClient.class); - Optional trustedCertificate = Optional.empty(); - RootAllocator allocator = mock(RootAllocator.class); - ArrowFlightClient arrowFlightClient = new ArrowFlightClient(flightClient, trustedCertificate, allocator); - - assertEquals(arrowFlightClient.getFlightClient(), flightClient); - assertFalse(arrowFlightClient.getTrustedCertificate().isPresent()); - } - - @Test - public void testClose() throws Exception - { - FlightClient flightClient = mock(FlightClient.class); - InputStream certificateStream = mock(InputStream.class); - Optional trustedCertificate = Optional.of(certificateStream); - RootAllocator allocator = mock(RootAllocator.class); - - ArrowFlightClient arrowFlightClient = new ArrowFlightClient(flightClient, trustedCertificate, allocator); - - arrowFlightClient.close(); - - verify(flightClient, times(1)).close(); - verify(certificateStream, times(1)).close(); - } - - @Test - public void testCloseWithoutCertificate() throws Exception - { - FlightClient flightClient = mock(FlightClient.class); - Optional trustedCertificate = Optional.empty(); - RootAllocator allocator = mock(RootAllocator.class); - - ArrowFlightClient arrowFlightClient = new ArrowFlightClient(flightClient, trustedCertificate, allocator); - - arrowFlightClient.close(); - - verify(flightClient, times(1)).close(); - } -} diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightClientHandler.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightClientHandler.java new file mode 100644 index 0000000000000..7ff77852e9b0f --- /dev/null +++ b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightClientHandler.java @@ -0,0 +1,36 @@ +/* + * 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.plugin.arrow; + +import com.facebook.presto.spi.ConnectorSession; +import org.apache.arrow.flight.auth2.BearerCredentialWriter; +import org.apache.arrow.flight.grpc.CredentialCallOption; + +import javax.inject.Inject; + +public class TestArrowFlightClientHandler + extends ArrowFlightClientHandler +{ + @Inject + public TestArrowFlightClientHandler(ArrowFlightConfig config) + { + super(config); + } + + @Override + protected CredentialCallOption getCallOptions(ConnectorSession connectorSession) + { + return new CredentialCallOption(new BearerCredentialWriter(null)); + } +} diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightConfig.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightConfig.java index 3e9e643ebdd29..a909f7c9be484 100644 --- a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightConfig.java +++ b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightConfig.java @@ -13,43 +13,101 @@ */ package com.facebook.plugin.arrow; -import com.facebook.airlift.configuration.testing.ConfigAssertions; -import com.google.common.collect.ImmutableMap; -import org.testng.annotations.Test; - -import java.util.Map; +import com.facebook.airlift.configuration.Config; +import com.facebook.airlift.configuration.ConfigSecuritySensitive; public class TestArrowFlightConfig { - @Test - public void testDefaults() - { - ConfigAssertions.assertRecordedDefaults(ConfigAssertions.recordDefaults(ArrowFlightConfig.class) - .setFlightServerName(null) - .setVerifyServer(null) - .setFlightServerSSLCertificate(null) - .setArrowFlightServerSslEnabled(null) - .setArrowFlightPort(null)); - } - - @Test - public void testExplicitPropertyMappings() - { - Map properties = new ImmutableMap.Builder() - .put("arrow-flight.server", "127.0.0.1") - .put("arrow-flight.server.verify", "true") - .put("arrow-flight.server-ssl-certificate", "cert") - .put("arrow-flight.server-ssl-enabled", "true") - .put("arrow-flight.server.port", "443") - .build(); - - ArrowFlightConfig expected = new ArrowFlightConfig() - .setFlightServerName("127.0.0.1") - .setVerifyServer(true) - .setFlightServerSSLCertificate("cert") - .setArrowFlightServerSslEnabled(true) - .setArrowFlightPort(443); - - ConfigAssertions.assertFullMapping(properties, expected); + private String host; // non-static field + private String database; // non-static field + private String username; // non-static field + private String password; // non-static field + private String name; // non-static field + private Integer port; // non-static field + private Boolean ssl; + + public String getDataSourceHost() + { // non-static getter + return host; + } + + public String getDataSourceDatabase() + { // non-static getter + return database; + } + + public String getDataSourceUsername() + { // non-static getter + return username; + } + + public String getDataSourcePassword() + { // non-static getter + return password; + } + + public String getDataSourceName() + { // non-static getter + return name; + } + + public Integer getDataSourcePort() + { // non-static getter + return port; + } + + public Boolean getDataSourceSSL() + { // non-static getter + return ssl; + } + + @Config("data-source.host") + public TestArrowFlightConfig setDataSourceHost(String host) + { // non-static setter + this.host = host; + return this; + } + + @Config("data-source.database") + public TestArrowFlightConfig setDataSourceDatabase(String database) + { // non-static setter + this.database = database; + return this; + } + + @Config("data-source.username") + public TestArrowFlightConfig setDataSourceUsername(String username) + { // non-static setter + this.username = username; + return this; + } + + @Config("data-source.password") + @ConfigSecuritySensitive + public TestArrowFlightConfig setDataSourcePassword(String password) + { // non-static setter + this.password = password; + return this; + } + + @Config("data-source.name") + public TestArrowFlightConfig setDataSourceName(String name) + { // non-static setter + this.name = name; + return this; + } + + @Config("data-source.port") + public TestArrowFlightConfig setDataSourcePort(Integer port) + { // non-static setter + this.port = port; + return this; + } + + @Config("data-source.ssl") + public TestArrowFlightConfig setDataSourceSSL(Boolean ssl) + { // non-static setter + this.ssl = ssl; + return this; } } diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightRequest.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightRequest.java index b546da8735b67..50257fc0ae0f0 100644 --- a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightRequest.java +++ b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightRequest.java @@ -13,12 +13,102 @@ */ package com.facebook.plugin.arrow; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.SerializationFeature; + +import java.nio.charset.StandardCharsets; +import java.util.Optional; + public class TestArrowFlightRequest implements ArrowFlightRequest { + private final String schema; + private final String table; + private final Optional query; + private final ArrowFlightConfig config; + private final int noOfPartitions; + + private final TestArrowFlightConfig testconfig; + + public TestArrowFlightRequest(ArrowFlightConfig config, TestArrowFlightConfig testconfig, String schema, String table, Optional query, int noOfPartitions) + { + this.config = config; + this.schema = schema; + this.table = table; + this.query = query; + this.testconfig = testconfig; + this.noOfPartitions = noOfPartitions; + } + + public TestArrowFlightRequest(ArrowFlightConfig config, String schema, int noOfPartitions, TestArrowFlightConfig testconfig) + { + this.schema = schema; + this.table = null; + this.query = Optional.empty(); + this.config = config; + this.testconfig = testconfig; + this.noOfPartitions = noOfPartitions; + } + + public String getSchema() + { + return schema; + } + + public String getTable() + { + return table; + } + + public Optional getQuery() + { + return query; + } + + public TestRequestData build() + { + TestRequestData requestData = new TestRequestData(); + requestData.setConnectionProperties(getConnectionProperties()); + requestData.setInteractionProperties(createInteractionProperties()); + return requestData; + } + @Override public byte[] getCommand() { - return new byte[0]; + ObjectMapper objectMapper = new ObjectMapper(); + objectMapper.configure(SerializationFeature.FAIL_ON_EMPTY_BEANS, false); + try { + String jsonString = objectMapper.writeValueAsString(build()); + return jsonString.getBytes(StandardCharsets.UTF_8); + } + catch (JsonProcessingException e) { + throw new ArrowException(ArrowErrorCode.ARROW_FLIGHT_ERROR, "JSON request cannot be created.", e); + } + } + + private TestConnectionProperties getConnectionProperties() + { + TestConnectionProperties properties = new TestConnectionProperties(); + properties.database = testconfig.getDataSourceDatabase(); + properties.host = testconfig.getDataSourceHost(); + properties.port = testconfig.getDataSourcePort(); + properties.username = testconfig.getDataSourceUsername(); + properties.password = testconfig.getDataSourcePassword(); + return properties; + } + + private TestInteractionProperties createInteractionProperties() + { + TestInteractionProperties interactionProperties = new TestInteractionProperties(); + if (getQuery().isPresent()) { + interactionProperties.setSelectStatement(getQuery().get()); + } + else { + interactionProperties.setSchema(getSchema()); + interactionProperties.setTable(getTable()); + } + return interactionProperties; } } diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightSmoke.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightSmoke.java new file mode 100644 index 0000000000000..bab2a5865ecc9 --- /dev/null +++ b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightSmoke.java @@ -0,0 +1,131 @@ +/* + * 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.plugin.arrow; + +import com.facebook.airlift.log.Logger; +import com.facebook.presto.testing.QueryRunner; +import org.apache.arrow.flight.FlightServer; +import org.apache.arrow.flight.Location; +import org.apache.arrow.memory.RootAllocator; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import java.io.File; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; + +public class TestArrowFlightSmoke +{ + private static final Logger logger = Logger.get(TestArrowFlightSmoke.class); + + private static RootAllocator allocator; + private static FlightServer server; + private static Location serverLocation; + private QueryRunner queryRunner; + + @BeforeClass + public void setup() throws Exception + { + File certChainFile = new File("src/test/resources/server.crt"); + File privateKeyFile = new File("src/test/resources/server.key"); + + allocator = new RootAllocator(Long.MAX_VALUE); + serverLocation = Location.forGrpcTls("127.0.0.1", 9443); + server = FlightServer.builder(allocator, serverLocation, new TestArrowServer(allocator)) + .useTls(certChainFile, privateKeyFile) + .build(); + server.start(); + logger.info("Server listening on port " + server.getPort()); + queryRunner = TestArrowQueryRunner.createQueryRunner(); + } + + @AfterClass(alwaysRun = true) + public void close() throws InterruptedException + { + if (queryRunner != null) { + queryRunner.close(); + server.close(); + allocator.close(); + } + } + + @Test + public void testShowSchemas() + { + assertEquals(queryRunner.execute("SHOW SCHEMAS FROM arrow").getRowCount(), 3); + } + + @Test + public void testSelectQuery() + { + String query1 = "SELECT * FROM testdb.example_table1"; + String expected1 = "1, John Doe, 1990-05-15, 50000.00, true"; + String expected2 = "2, Jane Smith, 1985-11-20, 60000.00, false"; + + String result1 = queryRunner.execute(query1).toString(); + assertTrue(result1.contains(expected1), "Expected row not found: " + expected1); + assertTrue(result1.contains(expected2), "Expected row not found: " + expected2); + } + + @Test + public void testShowTables() + { + // Ensure the catalog and schema names are correct + String query = "SHOW TABLES FROM testdb"; + String[] expectedTables = {"example_table1", "example_table2", "example_table3"}; + + String result = queryRunner.execute(query).toString(); + for (String expected : expectedTables) { + assertTrue(result.contains(expected), "Expected table not found: " + expected); + } + } + + @Test + public void testSelectQueryWithWhereClause() + { + String query1 = "SELECT * FROM testdb.example_table3 WHERE created_at = 36000000"; + String expected1 = "1, 36000000, A"; + + String result1 = queryRunner.execute(query1).toString(); + assertTrue(result1.contains(expected1), "Expected row not found: " + expected1); + + String query2 = "SELECT * FROM testdb.example_table2 WHERE price > 20.00"; + String expected2 = "2, Product B, 5, 29.99"; + + String result2 = queryRunner.execute(query2).toString(); + assertTrue(result2.contains(expected2), "Expected row not found: " + expected2); + } + + @Test + public void describeQuery() + { + String query1 = "DESCRIBE testdb.example_table1"; + assertEquals(queryRunner.execute(query1).getRowCount(), 5); + + String query2 = "DESCRIBE testdb.example_table2"; + assertEquals(queryRunner.execute(query2).getRowCount(), 4); + + String query3 = "DESCRIBE testdb.example_table3"; + assertEquals(queryRunner.execute(query3).getRowCount(), 3); + } + + @Test + public void testSelectQueryWithjoinClause() + { + String query = "SELECT t1.id, t1.name, t1.birthdate, t1.salary, t1.active, t2.description, t2.quantity, t2.price FROM example_table1 t1 JOIN example_table2 t2 ON t1.id = t2.id"; + assertEquals(queryRunner.execute(query).getRowCount(), 2); + } +} diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowMetadata.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowMetadata.java new file mode 100644 index 0000000000000..fa6b19aedd3c2 --- /dev/null +++ b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowMetadata.java @@ -0,0 +1,128 @@ +/* + * 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.plugin.arrow; + +import com.facebook.presto.spi.ConnectorSession; +import com.facebook.presto.spi.NodeManager; +import com.facebook.presto.spi.SchemaTableName; +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.google.common.collect.ImmutableList; +import org.apache.arrow.flight.Action; +import org.apache.arrow.flight.Result; + +import javax.inject.Inject; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Iterator; +import java.util.List; +import java.util.Optional; + +import static java.util.Locale.ENGLISH; + +public class TestArrowMetadata + extends ArrowAbstractMetadata +{ + private static final ObjectMapper objectMapper = new ObjectMapper(); + private final NodeManager nodeManager; + private final TestArrowFlightConfig testconfig; + private final ArrowFlightClientHandler clientHandler; + private final ArrowFlightConfig config; + + @Inject + public TestArrowMetadata(ArrowFlightConfig config, ArrowFlightClientHandler clientHandler, NodeManager nodeManager, TestArrowFlightConfig testconfig) + { + super(config, clientHandler); + this.nodeManager = nodeManager; + this.testconfig = testconfig; + this.clientHandler = clientHandler; + this.config = config; + } + + @Override + protected ArrowFlightRequest getArrowFlightRequest(ArrowFlightConfig config, String schema) + { + return new TestArrowFlightRequest(config, schema, nodeManager.getWorkerNodes().size(), testconfig); + } + + @Override + public List listSchemaNames(ConnectorSession session) + { + List listSchemas = extractSchemaAndTableData(Optional.empty(), session); + List names = new ArrayList<>(); + for (String value : listSchemas) { + names.add(value.toLowerCase(ENGLISH)); + } + return ImmutableList.copyOf(names); + } + + @Override + public List listTables(ConnectorSession session, Optional schemaName) + { + String schemaValue = schemaName.orElse(""); + String dataSourceSpecificSchemaName = getDataSourceSpecificSchemaName(config, schemaValue); + List listTables = extractSchemaAndTableData(Optional.ofNullable(dataSourceSpecificSchemaName), session); + List tables = new ArrayList<>(); + for (String value : listTables) { + tables.add(new SchemaTableName(dataSourceSpecificSchemaName.toLowerCase(ENGLISH), value.toLowerCase(ENGLISH))); + } + + return tables; + } + + public List extractSchemaAndTableData(Optional schema, ConnectorSession connectorSession) + { + try { + List names = new ArrayList<>(); + ArrowFlightClient client = clientHandler.getClient(Optional.empty()); + ArrowFlightRequest request = getArrowFlightRequest(config, schema.orElse(null)); + ObjectNode rootNode = (ObjectNode) objectMapper.readTree(request.getCommand()); + + String modifiedQueryJson = objectMapper.writeValueAsString(rootNode); + byte[] queryJsonBytes = modifiedQueryJson.getBytes(StandardCharsets.UTF_8); + Iterator iterator = client.getFlightClient().doAction(new Action("discovery", queryJsonBytes), clientHandler.getCallOptions(connectorSession)); + while (iterator.hasNext()) { + Result result = iterator.next(); + String jsonResult = new String(result.getBody(), StandardCharsets.UTF_8); + List tableNames = objectMapper.readValue(jsonResult, new TypeReference>() {}); + names.addAll(tableNames); + } + return names; + } + catch (IOException e) { + throw new RuntimeException(e); + } + } + + @Override + protected String getDataSourceSpecificSchemaName(ArrowFlightConfig config, String schemaName) + { + return schemaName; + } + + @Override + protected String getDataSourceSpecificTableName(ArrowFlightConfig config, String tableName) + { + return tableName; + } + + @Override + protected ArrowFlightRequest getArrowFlightRequest(ArrowFlightConfig config, Optional query, String schema, String table) + { + return new TestArrowFlightRequest(config, testconfig, schema, table, query, nodeManager.getWorkerNodes().size()); + } +} diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowModule.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowModule.java new file mode 100644 index 0000000000000..1f106ccda209b --- /dev/null +++ b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowModule.java @@ -0,0 +1,35 @@ +/* + * 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.plugin.arrow; + +import com.facebook.presto.spi.connector.ConnectorMetadata; +import com.facebook.presto.spi.connector.ConnectorSplitManager; +import com.google.inject.Binder; +import com.google.inject.Module; +import com.google.inject.Scopes; + +import static com.facebook.airlift.configuration.ConfigBinder.configBinder; + +public class TestArrowModule + implements Module +{ + @Override + public void configure(Binder binder) + { + configBinder(binder).bindConfig(TestArrowFlightConfig.class); + binder.bind(ConnectorSplitManager.class).to(TestArrowSplitManager.class).in(Scopes.SINGLETON); + binder.bind(ArrowFlightClientHandler.class).to(TestArrowFlightClientHandler.class).in(Scopes.SINGLETON); + binder.bind(ConnectorMetadata.class).to(TestArrowMetadata.class).in(Scopes.SINGLETON); + } +} diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowPageSource.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowPageSource.java deleted file mode 100644 index 2e245ecb7b46f..0000000000000 --- a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowPageSource.java +++ /dev/null @@ -1,101 +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.plugin.arrow; - -import com.facebook.presto.spi.ConnectorSession; -import org.apache.arrow.flight.FlightClient; -import org.apache.arrow.flight.FlightStream; -import org.apache.arrow.flight.Ticket; -import org.apache.arrow.vector.VectorSchemaRoot; -import org.mockito.Mock; -import org.mockito.MockitoAnnotations; -import org.testng.annotations.BeforeClass; -import org.testng.annotations.Test; - -import java.util.List; -import java.util.Optional; - -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; -import static org.testng.Assert.assertNotNull; -import static org.testng.Assert.assertNull; - -public class TestArrowPageSource -{ - @Mock - private ArrowSplit mockSplit; - - @Mock - private List mockColumnHandles; - - @Mock - private ArrowFlightClientHandler mockClientHandler; - - @Mock - private ConnectorSession mockSession; - - @Mock - private FlightClient mockFlightClient; - - @Mock - private FlightStream mockFlightStream; - - @Mock - private VectorSchemaRoot mockVectorSchemaRoot; - - @BeforeClass - public void setUp() - { - MockitoAnnotations.openMocks(this); - ArrowFlightClient mockArrowFlightClient = mock(ArrowFlightClient.class); - when(mockClientHandler.getClient(any(Optional.class))).thenReturn(mockArrowFlightClient); - when(mockArrowFlightClient.getFlightClient()).thenReturn(mockFlightClient); - when(mockFlightClient.getStream(any(Ticket.class), any())).thenReturn(mockFlightStream); - } - - @Test - public void testInitialization() - { - ArrowPageSource arrowPageSource = new ArrowPageSource(mockSplit, mockColumnHandles, mockClientHandler, mockSession); - assertNotNull(arrowPageSource); - } - - @Test - public void testGetNextPageWithEmptyFlightStream() - { - when(mockFlightStream.next()).thenReturn(false); - ArrowPageSource arrowPageSource = new ArrowPageSource(mockSplit, mockColumnHandles, mockClientHandler, mockSession); - assertNull(arrowPageSource.getNextPage()); - } - - @Test - public void testGetNextPageWithNonEmptyFlightStream() - { - when(mockFlightStream.next()).thenReturn(true); - when(mockFlightStream.getRoot()).thenReturn(mockVectorSchemaRoot); - when(mockVectorSchemaRoot.getRowCount()).thenReturn(1); - ArrowPageSource arrowPageSource = new ArrowPageSource(mockSplit, mockColumnHandles, mockClientHandler, mockSession); - assertNotNull(arrowPageSource.getNextPage()); - } - - @Test - public void testCloseResources() throws Exception - { - ArrowPageSource arrowPageSource = new ArrowPageSource(mockSplit, mockColumnHandles, mockClientHandler, mockSession); - arrowPageSource.close(); - verify(mockFlightStream).close(); - } -} diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowPageSourceProvider.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowPageSourceProvider.java deleted file mode 100644 index 77b89cf5d17cb..0000000000000 --- a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowPageSourceProvider.java +++ /dev/null @@ -1,58 +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.plugin.arrow; - -import com.facebook.presto.spi.ColumnHandle; -import com.facebook.presto.spi.ConnectorPageSource; -import com.facebook.presto.spi.ConnectorSession; -import com.facebook.presto.spi.SplitContext; -import com.facebook.presto.spi.connector.ConnectorTransactionHandle; -import com.google.common.collect.ImmutableList; -import org.apache.arrow.flight.FlightClient; -import org.apache.arrow.flight.FlightStream; -import org.apache.arrow.flight.Ticket; -import org.testng.annotations.Test; - -import java.util.List; - -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; -import static org.testng.Assert.assertNotNull; -import static org.testng.Assert.assertTrue; - -public class TestArrowPageSourceProvider -{ - @Test - public void testCreatePageSourceWithValidParameters() - { - ArrowFlightClientHandler clientHandler = mock(ArrowFlightClientHandler.class); - ArrowFlightClient flightClient = mock(ArrowFlightClient.class); - FlightClient flightClientInstance = mock(FlightClient.class); - FlightStream flightStream = mock(FlightStream.class); - when(clientHandler.getClient(any())).thenReturn(flightClient); - when(flightClient.getFlightClient()).thenReturn(flightClientInstance); - when(flightClientInstance.getStream(any(Ticket.class), any(), any())).thenReturn(flightStream); - ArrowPageSourceProvider arrowPageSourceProvider = new ArrowPageSourceProvider(clientHandler); - ConnectorTransactionHandle transactionHandle = mock(ConnectorTransactionHandle.class); - ConnectorSession session = mock(ConnectorSession.class); - ArrowSplit split = mock(ArrowSplit.class); - List columns = ImmutableList.of(mock(ArrowColumnHandle.class)); - SplitContext splitContext = mock(SplitContext.class); - when(split.getTicket()).thenReturn(new byte[0]); - ConnectorPageSource pageSource = arrowPageSourceProvider.createPageSource(transactionHandle, session, split, columns, splitContext); - assertNotNull(pageSource); - assertTrue(pageSource instanceof ArrowPageSource); - } -} diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowQueryBuilder.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowQueryBuilder.java new file mode 100644 index 0000000000000..75c9a0cfc92a4 --- /dev/null +++ b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowQueryBuilder.java @@ -0,0 +1,295 @@ +/* + * 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.plugin.arrow; + +import com.facebook.presto.common.predicate.Domain; +import com.facebook.presto.common.predicate.Range; +import com.facebook.presto.common.predicate.TupleDomain; +import com.facebook.presto.common.type.BigintType; +import com.facebook.presto.common.type.BooleanType; +import com.facebook.presto.common.type.CharType; +import com.facebook.presto.common.type.DateType; +import com.facebook.presto.common.type.DoubleType; +import com.facebook.presto.common.type.IntegerType; +import com.facebook.presto.common.type.RealType; +import com.facebook.presto.common.type.SmallintType; +import com.facebook.presto.common.type.TimeType; +import com.facebook.presto.common.type.TimestampType; +import com.facebook.presto.common.type.TinyintType; +import com.facebook.presto.common.type.Type; +import com.facebook.presto.common.type.VarcharType; +import com.facebook.presto.spi.ColumnHandle; +import com.google.common.base.Joiner; +import com.google.common.collect.ImmutableList; +import io.airlift.slice.Slice; + +import java.sql.Time; +import java.sql.Timestamp; +import java.text.SimpleDateFormat; +import java.time.ZoneId; +import java.util.ArrayList; +import java.util.Date; +import java.util.List; +import java.util.Map; +import java.util.TimeZone; +import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; + +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Preconditions.checkState; +import static com.google.common.base.Strings.isNullOrEmpty; +import static com.google.common.collect.Iterables.getOnlyElement; +import static java.lang.String.format; +import static java.util.Objects.requireNonNull; +import static java.util.stream.Collectors.joining; + +public class TestArrowQueryBuilder +{ + // not all databases support booleans, so use 1=1 and 1=0 instead + private static final String ALWAYS_TRUE = "1=1"; + private static final String ALWAYS_FALSE = "1=0"; + public static final String DATE_FORMAT = "yyyy-MM-dd"; + public static final String TIMESTAMP_FORMAT = "yyyy-MM-dd HH:mm:ss.SSS"; + public static final String TIME_FORMAT = "HH:mm:ss"; + public static final TimeZone UTC_TIME_ZONE = TimeZone.getTimeZone(ZoneId.of("UTC")); + + public String buildSql( + String schema, + String table, + List columns, + Map columnExpressions, + TupleDomain tupleDomain) + { + StringBuilder sql = new StringBuilder(); + + sql.append("SELECT "); + sql.append(addColumnExpression(columns, columnExpressions)); + + sql.append(" FROM "); + if (!isNullOrEmpty(schema)) { + sql.append(quote(schema)).append('.'); + } + sql.append(quote(table)); + + List accumulator = new ArrayList<>(); + + if (tupleDomain != null) { + List clauses = toConjuncts(columns, tupleDomain, accumulator); + if (!clauses.isEmpty()) { + sql.append(" WHERE ") + .append(Joiner.on(" AND ").join(clauses)); + } + } + + return sql.toString(); + } + + public static String convertEpochToString(long epochValue, Type type) + { + if (type instanceof DateType) { + long millis = TimeUnit.DAYS.toMillis(epochValue); + Date date = new Date(millis); + SimpleDateFormat dateFormat = new SimpleDateFormat(DATE_FORMAT); + dateFormat.setTimeZone(UTC_TIME_ZONE); + return dateFormat.format(date); + } + else if (type instanceof TimestampType) { + Timestamp timestamp = new Timestamp(epochValue); + SimpleDateFormat timestampFormat = new SimpleDateFormat(TIMESTAMP_FORMAT); + timestampFormat.setTimeZone(UTC_TIME_ZONE); + return timestampFormat.format(timestamp); + } + else if (type instanceof TimeType) { + long millis = TimeUnit.SECONDS.toMillis(epochValue / 1000); + Time time = new Time(millis); + SimpleDateFormat timeFormat = new SimpleDateFormat(TIME_FORMAT); + timeFormat.setTimeZone(UTC_TIME_ZONE); + return timeFormat.format(time); + } + else { + throw new UnsupportedOperationException(type + " is not supported."); + } + } + + protected static class TypeAndValue + { + private final Type type; + private final Object value; + + public TypeAndValue(Type type, Object value) + { + this.type = requireNonNull(type, "type is null"); + this.value = requireNonNull(value, "value is null"); + } + + public Type getType() + { + return type; + } + + public Object getValue() + { + return value; + } + } + + private String addColumnExpression(List columns, Map columnExpressions) + { + if (columns.isEmpty()) { + return "null"; + } + + return columns.stream() + .map(ArrowColumnHandle -> { + String columnAlias = quote(ArrowColumnHandle.getColumnName()); + String expression = columnExpressions.get(ArrowColumnHandle.getColumnName()); + if (expression == null) { + return columnAlias; + } + return format("%s AS %s", expression, columnAlias); + }) + .collect(joining(", ")); + } + + private static boolean isAcceptedType(Type type) + { + Type validType = requireNonNull(type, "type is null"); + return validType.equals(BigintType.BIGINT) || + validType.equals(TinyintType.TINYINT) || + validType.equals(SmallintType.SMALLINT) || + validType.equals(IntegerType.INTEGER) || + validType.equals(DoubleType.DOUBLE) || + validType.equals(RealType.REAL) || + validType.equals(BooleanType.BOOLEAN) || + validType.equals(DateType.DATE) || + validType.equals(TimeType.TIME) || + validType.equals(TimestampType.TIMESTAMP) || + validType instanceof VarcharType || + validType instanceof CharType; + } + private List toConjuncts(List columns, TupleDomain tupleDomain, List accumulator) + { + ImmutableList.Builder builder = ImmutableList.builder(); + for (ArrowColumnHandle column : columns) { + Type type = column.getColumnType(); + //Todo handle tupleDomain null + if (isAcceptedType(type)) { + Domain domain = tupleDomain.getDomains().get().get(column); + if (domain != null) { + builder.add(toPredicate(column.getColumnName(), domain, column, accumulator)); + } + } + } + return builder.build(); + } + + private String toPredicate(String columnName, Domain domain, ArrowColumnHandle columnHandle, List accumulator) + { + checkArgument(domain.getType().isOrderable(), "Domain type must be orderable"); + + if (domain.getValues().isNone()) { + return domain.isNullAllowed() ? quote(columnName) + " IS NULL" : ALWAYS_FALSE; + } + + if (domain.getValues().isAll()) { + return domain.isNullAllowed() ? ALWAYS_TRUE : quote(columnName) + " IS NOT NULL"; + } + + List disjuncts = new ArrayList<>(); + List singleValues = new ArrayList<>(); + for (Range range : domain.getValues().getRanges().getOrderedRanges()) { + checkState(!range.isAll()); // Already checked + if (range.isSingleValue()) { + singleValues.add(range.getSingleValue()); + } + else { + List rangeConjuncts = new ArrayList<>(); + if (!range.isLowUnbounded()) { + rangeConjuncts.add(toPredicate(columnName, range.isLowInclusive() ? ">=" : ">", range.getLowBoundedValue(), columnHandle, accumulator)); + } + if (!range.isHighUnbounded()) { + rangeConjuncts.add(toPredicate(columnName, range.isHighInclusive() ? "<=" : "<", range.getHighBoundedValue(), columnHandle, accumulator)); + } + // If rangeConjuncts is null, then the range was ALL, which should already have been checked for + checkState(!rangeConjuncts.isEmpty()); + disjuncts.add("(" + Joiner.on(" AND ").join(rangeConjuncts) + ")"); + } + } + + // Add back all of the possible single values either as an equality or an IN predicate + if (singleValues.size() == 1) { + disjuncts.add(toPredicate(columnName, "=", getOnlyElement(singleValues), columnHandle, accumulator)); + } + else if (singleValues.size() > 1) { + for (Object value : singleValues) { + bindValue(value, columnHandle, accumulator); + } + String values = Joiner.on(",").join(singleValues.stream().map(v -> + parameterValueToString(columnHandle.getColumnType(), v)).collect(Collectors.toList())); + disjuncts.add(quote(columnName) + " IN (" + values + ")"); + } + + // Add nullability disjuncts + checkState(!disjuncts.isEmpty()); + if (domain.isNullAllowed()) { + disjuncts.add(quote(columnName) + " IS NULL"); + } + + return "(" + Joiner.on(" OR ").join(disjuncts) + ")"; + } + + private String toPredicate(String columnName, String operator, Object value, ArrowColumnHandle columnHandle, List accumulator) + { + bindValue(value, columnHandle, accumulator); + return quote(columnName) + " " + operator + " " + parameterValueToString(columnHandle.getColumnType(), value); + } + private String quote(String name) + { + return name; + } + + private String quoteValue(String name) + { + return "'" + name + "'"; + } + + private void bindValue(Object value, ArrowColumnHandle columnHandle, List accumulator) + { + Type type = columnHandle.getColumnType(); + accumulator.add(new TypeAndValue(type, value)); + } + + private String parameterValueToString(Type type, Object value) + { + Class javaType = type.getJavaType(); + if (type instanceof DateType && javaType == long.class) { + return quoteValue(convertEpochToString((Long) value, type)); + } + else if (type instanceof TimeType && javaType == long.class) { + return quoteValue(convertEpochToString((Long) value, type)); + } + else if (type instanceof TimestampType && javaType == long.class) { + return quoteValue(convertEpochToString((Long) value, type)); + } + else if (javaType == boolean.class || javaType == double.class || javaType == long.class) { + return value.toString(); + } + else if (javaType == Slice.class) { + return quoteValue(((Slice) value).toStringUtf8()); + } + else { + return quoteValue(value.toString()); + } + } +} diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowQueryRunner.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowQueryRunner.java new file mode 100644 index 0000000000000..58ea5f9230957 --- /dev/null +++ b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowQueryRunner.java @@ -0,0 +1,66 @@ +/* + * 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.plugin.arrow; + +import com.facebook.presto.Session; +import com.facebook.presto.testing.LocalQueryRunner; +import com.google.common.collect.ImmutableMap; + +import java.util.Map; + +import static com.facebook.presto.testing.TestingSession.testSessionBuilder; + +public class TestArrowQueryRunner +{ + private TestArrowQueryRunner() {} + + public static LocalQueryRunner createQueryRunner() + { + return createQueryRunner(ImmutableMap.of(), TestingArrowFactory.class); + } + + public static LocalQueryRunner createQueryRunner(Map catalogProperties, Class factoryClass) + { + Session session = testSessionBuilder() + .setCatalog("arrow") + .setSchema("testdb") + .build(); + + LocalQueryRunner queryRunner = new LocalQueryRunner(session); + + try { + ArrowConnectorFactory connectorFactory = new ArrowConnectorFactory(factoryClass.getSimpleName(), new TestArrowModule(), TestArrowQueryRunner.class.getClassLoader()); + ImmutableMap.Builder properties = ImmutableMap.builder() + .putAll(catalogProperties) + .put("arrow-flight.server", "127.0.0.1") + .put("arrow-flight.server-ssl-enabled", "true") + .put("arrow-flight.server.port", "9443") + .put("arrow-flight.server.verify", "false") + .put("data-source.port", "8443") + .put("data-source.host", "9.30.120.2") + .put("data-source.database", "db") + .put("data-source.username", "user") + .put("data-source.password", "password") + .put("data-source.ssl", "true"); + queryRunner.createCatalog("arrow", connectorFactory, properties.build()); + + return queryRunner; + } + catch (Exception e) { + e.printStackTrace(); + throw new RuntimeException("Failed to create ArrowQueryRunner", e); + } + } +} diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowServer.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowServer.java new file mode 100644 index 0000000000000..19b52409e2547 --- /dev/null +++ b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowServer.java @@ -0,0 +1,261 @@ +/* + * 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.plugin.arrow; + +import com.facebook.airlift.log.Logger; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.arrow.adapter.jdbc.ArrowVectorIterator; +import org.apache.arrow.adapter.jdbc.JdbcToArrow; +import org.apache.arrow.flight.Action; +import org.apache.arrow.flight.ActionType; +import org.apache.arrow.flight.Criteria; +import org.apache.arrow.flight.FlightDescriptor; +import org.apache.arrow.flight.FlightEndpoint; +import org.apache.arrow.flight.FlightInfo; +import org.apache.arrow.flight.FlightProducer; +import org.apache.arrow.flight.FlightStream; +import org.apache.arrow.flight.PutResult; +import org.apache.arrow.flight.Result; +import org.apache.arrow.flight.Ticket; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.DateUnit; +import org.apache.arrow.vector.types.FloatingPointPrecision; +import org.apache.arrow.vector.types.TimeUnit; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + +public class TestArrowServer + implements FlightProducer +{ + private final BufferAllocator allocator; + private final ObjectMapper objectMapper = new ObjectMapper(); + + private static Connection connection; + + private static final Logger logger = Logger.get(TestArrowServer.class); + + public TestArrowServer(BufferAllocator allocator) throws Exception + { + this.allocator = allocator; + TestH2DatabaseSetup.setup(); + this.connection = DriverManager.getConnection("jdbc:h2:mem:testdb;DB_CLOSE_DELAY=-1", "sa", ""); + } + + @Override + public void getStream(CallContext callContext, Ticket ticket, ServerStreamListener serverStreamListener) + { + try { + String ticketString = new String(ticket.getBytes(), StandardCharsets.UTF_8); + + JsonNode ticketJson = objectMapper.readTree(ticketString); + String query = ticketJson.get("interactionProperties").get("select_statement").asText(); + + if (query == null || query.trim().isEmpty()) { + throw new IllegalArgumentException("Query cannot be null or empty."); + } + + logger.info("Executing query: " + query); + query = query.toUpperCase(); + try (Statement stmt = connection.createStatement(); + ResultSet rs = stmt.executeQuery(query)) { + ArrowVectorIterator iterator = JdbcToArrow.sqlToArrowVectorIterator(rs, allocator); + + while (iterator.hasNext()) { + try (VectorSchemaRoot root = iterator.next()) { + serverStreamListener.start(root); + serverStreamListener.putNext(); + } + } + serverStreamListener.completed(); + } + catch (SQLException e) { + serverStreamListener.error(e); + throw new RuntimeException("Failed to execute query", e); + } + catch (IOException e) { + serverStreamListener.error(e); + throw new RuntimeException("Failed to process Arrow data", e); + } + } + catch (Exception e) { + serverStreamListener.error(e); + throw new RuntimeException("Failed to process the ticket", e); + } + } + + @Override + public void listFlights(CallContext callContext, Criteria criteria, StreamListener streamListener) + { + } + + @Override + public FlightInfo getFlightInfo(CallContext callContext, FlightDescriptor flightDescriptor) + { + try { + String jsonRequest = new String(flightDescriptor.getCommand(), StandardCharsets.UTF_8); + JsonNode rootNode = objectMapper.readTree(jsonRequest); + + String schemaName = rootNode.get("interactionProperties").get("schema_name").asText(null); + String tableName = rootNode.get("interactionProperties").get("table_name").asText(null); + String selectStatement = rootNode.get("interactionProperties").get("select_statement").asText(null); + + List fields = new ArrayList<>(); + if (schemaName != null && tableName != null) { + String query = "SELECT COLUMN_NAME, DATA_TYPE FROM INFORMATION_SCHEMA.COLUMNS " + + "WHERE TABLE_SCHEMA='" + schemaName.toUpperCase() + "' " + + "AND TABLE_NAME='" + tableName.toUpperCase() + "'"; + + try (ResultSet rs = connection.createStatement().executeQuery(query)) { + while (rs.next()) { + String columnName = rs.getString("COLUMN_NAME"); + String dataType = rs.getString("DATA_TYPE"); + + ArrowType arrowType = convertSqlTypeToArrowType(dataType); + Field field = new Field(columnName, FieldType.nullable(arrowType), null); + fields.add(field); + } + } + } + else if (selectStatement != null) { + selectStatement = selectStatement.toUpperCase(); + logger.info("Executing SELECT query: " + selectStatement); + try (ResultSet rs = connection.createStatement().executeQuery(selectStatement)) { + ResultSetMetaData metaData = rs.getMetaData(); + int columnCount = metaData.getColumnCount(); + + for (int i = 1; i <= columnCount; i++) { + String columnName = metaData.getColumnName(i); + String columnType = metaData.getColumnTypeName(i); + + ArrowType arrowType = convertSqlTypeToArrowType(columnType); + Field field = new Field(columnName, FieldType.nullable(arrowType), null); + fields.add(field); + } + } + } + else { + throw new IllegalArgumentException("Either schema_name/table_name or select_statement must be provided."); + } + + Schema schema = new Schema(fields); + FlightEndpoint endpoint = new FlightEndpoint(new Ticket(flightDescriptor.getCommand())); + return new FlightInfo(schema, flightDescriptor, Collections.singletonList(endpoint), -1, -1); + } + catch (Exception e) { + e.printStackTrace(); + throw new RuntimeException("Failed to retrieve FlightInfo", e); + } + } + + @Override + public Runnable acceptPut(CallContext callContext, FlightStream flightStream, StreamListener streamListener) + { + return null; + } + + @Override + public void doAction(CallContext callContext, Action action, StreamListener streamListener) + { + try { + String jsonRequest = new String(action.getBody(), StandardCharsets.UTF_8); + JsonNode rootNode = objectMapper.readTree(jsonRequest); + String schemaName = rootNode.get("interactionProperties").get("schema_name").asText(null); + + String query; + if (schemaName == null) { + query = "SELECT SCHEMA_NAME FROM INFORMATION_SCHEMA.SCHEMATA"; + } + else { + schemaName = schemaName.toUpperCase(); + query = "SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA='" + schemaName + "'"; + } + ResultSet rs = connection.createStatement().executeQuery(query); + List names = new ArrayList<>(); + while (rs.next()) { + names.add(rs.getString(1)); + } + + String jsonResponse = objectMapper.writeValueAsString(names); + streamListener.onNext(new Result(jsonResponse.getBytes(StandardCharsets.UTF_8))); + streamListener.onCompleted(); + } + catch (Exception e) { + streamListener.onError(e); + } + } + + @Override + public void listActions(CallContext callContext, StreamListener streamListener) + { + } + + private ArrowType convertSqlTypeToArrowType(String sqlType) + { + switch (sqlType.toUpperCase()) { + case "VARCHAR": + case "CHAR": + case "CHARACTER VARYING": + case "CHARACTER": + return new ArrowType.Utf8(); + case "INTEGER": + case "INT": + return new ArrowType.Int(32, true); + case "BIGINT": + return new ArrowType.Int(64, true); + case "SMALLINT": + return new ArrowType.Int(16, true); + case "TINYINT": + return new ArrowType.Int(8, true); + case "DOUBLE": + case "DOUBLE PRECISION": + case "FLOAT": + return new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE); + case "REAL": + return new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE); + case "BOOLEAN": + return new ArrowType.Bool(); + case "DATE": + return new ArrowType.Date(DateUnit.DAY); + case "TIMESTAMP": + return new ArrowType.Timestamp(TimeUnit.MILLISECOND, null); + case "TIME": + return new ArrowType.Time(TimeUnit.MILLISECOND, 32); + case "DECIMAL": + case "NUMERIC": + return new ArrowType.Decimal(5, 2); + case "BINARY": + case "VARBINARY": + return new ArrowType.Binary(); + default: + throw new IllegalArgumentException("Unsupported SQL type: " + sqlType); + } + } +} diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowSplitManager.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowSplitManager.java new file mode 100644 index 0000000000000..16bb2e4bafc86 --- /dev/null +++ b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowSplitManager.java @@ -0,0 +1,48 @@ +/* + * 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.plugin.arrow; + +import com.facebook.presto.spi.NodeManager; +import com.google.common.collect.ImmutableMap; + +import javax.inject.Inject; + +import java.util.Optional; + +public class TestArrowSplitManager + extends ArrowAbstractSplitManager +{ + private TestArrowFlightConfig testconfig; + + private final NodeManager nodeManager; + + @Inject + public TestArrowSplitManager(ArrowFlightClientHandler client, TestArrowFlightConfig testconfig, NodeManager nodeManager) + { + super(client); + this.testconfig = testconfig; + this.nodeManager = nodeManager; + } + + @Override + protected ArrowFlightRequest getArrowFlightRequest(ArrowFlightConfig config, ArrowTableLayoutHandle tableLayoutHandle) + { + ArrowTableHandle tableHandle = tableLayoutHandle.getTableHandle(); + Optional query = Optional.of(new TestArrowQueryBuilder().buildSql(tableHandle.getSchema(), + tableHandle.getTable(), + tableLayoutHandle.getColumnHandles(), ImmutableMap.of(), + tableLayoutHandle.getTupleDomain())); + return new TestArrowFlightRequest(config, testconfig, tableHandle.getSchema(), tableHandle.getTable(), query, nodeManager.getWorkerNodes().size()); + } +} diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestConnectionProperties.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestConnectionProperties.java new file mode 100644 index 0000000000000..d9e4523b3367b --- /dev/null +++ b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestConnectionProperties.java @@ -0,0 +1,24 @@ +/* + * 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.plugin.arrow; + +public class TestConnectionProperties +{ + public String database; + public String password; + public Integer port; + public String host; + public Boolean ssl; + public String username; +} diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestFlightClient.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestFlightClient.java deleted file mode 100644 index 746db7395d6ca..0000000000000 --- a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestFlightClient.java +++ /dev/null @@ -1,90 +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.plugin.arrow; - -import com.facebook.presto.spi.ConnectorSession; -import org.apache.arrow.flight.FlightInfo; -import org.apache.arrow.flight.FlightServer; -import org.apache.arrow.flight.Location; -import org.apache.arrow.flight.grpc.CredentialCallOption; -import org.apache.arrow.memory.RootAllocator; -import org.mockito.Mockito; -import org.testng.annotations.AfterClass; -import org.testng.annotations.BeforeClass; -import org.testng.annotations.Test; - -import java.util.Optional; - -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; -import static org.testng.Assert.assertEquals; -import static org.testng.Assert.assertNotNull; - -public class TestFlightClient -{ - private static RootAllocator allocator; - private static FlightServer server; - - private static Location serverLocation; - - @BeforeClass - public void setup() throws Exception - { - allocator = new RootAllocator(Long.MAX_VALUE); - serverLocation = Location.forGrpcInsecure("127.0.0.1", 9443); - server = FlightServer.builder(allocator, serverLocation, new ArrowServer()) - .build(); - server.start(); - System.out.println("Server listening on port " + server.getPort()); - } - - @AfterClass - public void tearDown() throws Exception - { - server.close(); - allocator.close(); - } - - @Test - public void testInitializeClient() - { - ArrowFlightConfig config = mock(ArrowFlightConfig.class); - when(config.getFlightServerSSLCertificate()).thenReturn(null); - when(config.getVerifyServer()).thenReturn(true); - ConnectorSession connectorSession = mock(ConnectorSession.class); - ArrowFlightClientHandler handler = new ArrowFlightClientHandler(config) { - @Override - protected CredentialCallOption getCallOptions(ConnectorSession connectorSession) - { - return null; - } - }; - Optional uri = Optional.of("grpc://127.0.0.1:9443"); - handler.getClient(uri); - assertNotNull(handler.getClient(uri)); - } - - @Test - public void testGetFlightInfo() - { - ArrowFlightConfig config = Mockito.mock(ArrowFlightConfig.class); - FlightInfo mockFlightInfo = mock(FlightInfo.class); - ArrowFlightClientHandler clientHandler = Mockito.mock(ArrowFlightClientHandler.class); - ArrowFlightRequest request = new TestArrowFlightRequest(); - ConnectorSession session = Mockito.mock(ConnectorSession.class); - when(clientHandler.getFlightInfo(request, session)).thenReturn(mockFlightInfo); - FlightInfo result = clientHandler.getFlightInfo(request, session); - assertEquals(mockFlightInfo, result); - } -} diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestH2DatabaseSetup.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestH2DatabaseSetup.java new file mode 100644 index 0000000000000..5a62ae3540cec --- /dev/null +++ b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestH2DatabaseSetup.java @@ -0,0 +1,46 @@ +/* + * 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.plugin.arrow; + +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.Statement; + +public class TestH2DatabaseSetup +{ + private TestH2DatabaseSetup() + { + throw new UnsupportedOperationException("This is a utility class and cannot be instantiated"); + } + public static void setup() throws Exception + { + Class.forName("org.h2.Driver"); + + Connection conn = DriverManager.getConnection("jdbc:h2:mem:testdb", "sa", ""); + + Statement stmt = conn.createStatement(); + + stmt.execute("CREATE SCHEMA IF NOT EXISTS testdb"); + + stmt.execute("CREATE TABLE IF NOT EXISTS testdb.example_table1 (id INT PRIMARY KEY, name VARCHAR(255), birthdate DATE, salary DECIMAL(10, 2), active BOOLEAN)"); + stmt.execute("INSERT INTO testdb.example_table1 (id, name, birthdate, salary, active) VALUES (1, 'John Doe', '1990-05-15', 50000.00, TRUE), (2, 'Jane Smith', '1985-11-20', 60000.00, FALSE)"); + + stmt.execute("CREATE TABLE IF NOT EXISTS testdb.example_table2 (id INT PRIMARY KEY, description TEXT, quantity INT, price DOUBLE)"); + stmt.execute("INSERT INTO testdb.example_table2 (id, description, quantity, price) VALUES (1, 'Product A', 10, 19.99), (2, 'Product B', 5, 29.99)"); + + stmt.execute("DROP TABLE IF EXISTS testdb.example_table3"); + stmt.execute("CREATE TABLE IF NOT EXISTS testdb.example_table3 (id INT PRIMARY KEY, created_at INT, status VARCHAR(255))"); + stmt.execute("INSERT INTO testdb.example_table3 (id, created_at, status) VALUES (1, 36000000, 'A'), (2, 38000000, 'B')"); + } +} diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestInteractionProperties.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestInteractionProperties.java new file mode 100644 index 0000000000000..aa6f0d5da75ff --- /dev/null +++ b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestInteractionProperties.java @@ -0,0 +1,56 @@ +/* + * 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.plugin.arrow; + +import com.fasterxml.jackson.annotation.JsonProperty; + +public class TestInteractionProperties +{ + @JsonProperty("select_statement") + private String selectStatement; + + @JsonProperty("schema_name") + private String schema; + @JsonProperty("table_name") + private String table; + public String getSelectStatement() + { + return selectStatement; + } + + public String getSchema() + { + return schema; + } + + public String getTable() + { + return table; + } + + public void setSchema(String schema) + { + this.schema = schema; + } + + public void setSelectStatement(String selectStatement) + { + this.selectStatement = selectStatement; + } + + public void setTable(String table) + { + this.table = table; + } +} diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestRequestData.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestRequestData.java new file mode 100644 index 0000000000000..a98841d106ca6 --- /dev/null +++ b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestRequestData.java @@ -0,0 +1,40 @@ +/* + * 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.plugin.arrow; + +public class TestRequestData +{ + private TestConnectionProperties connectionProperties; + private TestInteractionProperties interactionProperties; + + public TestConnectionProperties getConnectionProperties() + { + return connectionProperties; + } + + public TestInteractionProperties getInteractionProperties() + { + return interactionProperties; + } + + public void setConnectionProperties(TestConnectionProperties connectionProperties) + { + this.connectionProperties = connectionProperties; + } + + public void setInteractionProperties(TestInteractionProperties interactionProperties) + { + this.interactionProperties = interactionProperties; + } +} diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestingArrowFactory.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestingArrowFactory.java new file mode 100644 index 0000000000000..88120bf8ebf22 --- /dev/null +++ b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestingArrowFactory.java @@ -0,0 +1,23 @@ +/* + * 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.plugin.arrow; + +public class TestingArrowFactory + extends ArrowConnectorFactory +{ + public TestingArrowFactory() + { + super("arrow", new TestArrowModule(), TestingArrowFactory.class.getClassLoader()); + } +} diff --git a/presto-base-arrow-flight/src/test/resources/server.crt b/presto-base-arrow-flight/src/test/resources/server.crt new file mode 100644 index 0000000000000..9af01752e2959 --- /dev/null +++ b/presto-base-arrow-flight/src/test/resources/server.crt @@ -0,0 +1,21 @@ +-----BEGIN CERTIFICATE----- +MIIDazCCAlOgAwIBAgIUBYcCyz9qphcpgV9wIx2caf4ikoEwDQYJKoZIhvcNAQEL +BQAwRTELMAkGA1UEBhMCQVUxEzARBgNVBAgMClNvbWUtU3RhdGUxITAfBgNVBAoM +GEludGVybmV0IFdpZGdpdHMgUHR5IEx0ZDAeFw0yNDA4MzAxMzM0MjJaFw0yNTA4 +MzAxMzM0MjJaMEUxCzAJBgNVBAYTAkFVMRMwEQYDVQQIDApTb21lLVN0YXRlMSEw +HwYDVQQKDBhJbnRlcm5ldCBXaWRnaXRzIFB0eSBMdGQwggEiMA0GCSqGSIb3DQEB +AQUAA4IBDwAwggEKAoIBAQCTu4vGRFav9/aNuGorx2v+xAevMRFfSi6x2onRCpMp +u+9L3bK3em+vSJcC/+QwvNkvTmf0ac8/G79BXQiUtPSV3tDSFnpO7LIICWpO2gJX +UKjTN+SZ3LfLEXSYAijVUCkWf8RswJLvJ12qNQRiv4IDrpBh/X+ICr+ALMa7y6cn +hBfpFtVyxlLkCp6Q71JgMGDjocv2195pc4uiMEDsn6tZfEKxmw1phX+CLke42srX ++AHoBQvlfcOBVIuMkxU7pZvvDDOuw0tSm0qomjRV+Azjn4oWbPDfkIEZIbBvgLwN +nxtKzSAgW+7xWQ0eVMRq1f5JPk6vLAV/yit8lFk8TCtdAgMBAAGjUzBRMB0GA1Ud +DgQWBBS4K4nHKbp4qw77KxR+tqUWXKCADjAfBgNVHSMEGDAWgBS4K4nHKbp4qw77 +KxR+tqUWXKCADjAPBgNVHRMBAf8EBTADAQH/MA0GCSqGSIb3DQEBCwUAA4IBAQBI +yraxUUgEb53q0aF2ID1ml3/9uMZALK4xQHKYP0dtkqdaS0CvuMsqSe0AT7PMDJOD +Q+s1lxe0gJ/lb7dLSOQpFGAU7nnUsBaL5alCGmYgeF6ayLzNWobsn4amnjFpSlOH +P/nOQvOrRgVbphozhaYBSjMtwy2Uzj8G7ZHn+Vg/fdEUf/mGb7/4M9sL7iQWJmIg +slzNnA/LssAqGvyJ65cSkjGrb82VvDJv6JZ7fam6nrMEpq5D4H4GzOwupNvdItLH +Zf9IBsedkDEu8P+Rzp9kGIMbGTLfL3u/pP1st+pD4PfHYRsbl8aAUqYA7P+d29mf +aVcp98Z1h1s5e005BT1T +-----END CERTIFICATE----- diff --git a/presto-base-arrow-flight/src/test/resources/server.key b/presto-base-arrow-flight/src/test/resources/server.key new file mode 100644 index 0000000000000..894dc1fad8e61 --- /dev/null +++ b/presto-base-arrow-flight/src/test/resources/server.key @@ -0,0 +1,28 @@ +-----BEGIN PRIVATE KEY----- +MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQCTu4vGRFav9/aN +uGorx2v+xAevMRFfSi6x2onRCpMpu+9L3bK3em+vSJcC/+QwvNkvTmf0ac8/G79B +XQiUtPSV3tDSFnpO7LIICWpO2gJXUKjTN+SZ3LfLEXSYAijVUCkWf8RswJLvJ12q +NQRiv4IDrpBh/X+ICr+ALMa7y6cnhBfpFtVyxlLkCp6Q71JgMGDjocv2195pc4ui +MEDsn6tZfEKxmw1phX+CLke42srX+AHoBQvlfcOBVIuMkxU7pZvvDDOuw0tSm0qo +mjRV+Azjn4oWbPDfkIEZIbBvgLwNnxtKzSAgW+7xWQ0eVMRq1f5JPk6vLAV/yit8 +lFk8TCtdAgMBAAECggEAAKgSdrLajMUmFhql9CRafUMbQqLN8DW47+bn+mMY5NRW +O6jUUL7tTKLesu92sOXB9FUdnqcyudXSe4Shk2GbfagEFw7tA7lHEESUcZ3D6WXt +HiUvMaTatz8QXNWTn2EQEa7HLXGMpZ3v61/5cUPnHMOTli/ld3IOyE/KoU6GI2WP +5+8eKJUdJiZDLdywLTF/OefTuHhnbkNxsoPtR1DwO1UZKsX1OFo0b3x6ibqPZSu3 +x01QkYfTlUoVjUBNFQTcyfXKWX5rrAK9C/2h7TVJxbktj7V4DOGLb8ee6oN+yr4e +lIg6Vfz/KAakXKaRVappoxSxvFTZXqGhL8kBnLPFoQKBgQDIfr3ir65pi2Gzc+6e +0DWCkznfDqVvfgzoAxCrIqKp/6MqVUIJmProPjBztSAhiaNhK8sDbLFBssw0AOeZ +YEqMcY69ClyVrtDoGYBNeuCTTz3S8yXf0mzZe+CgWA6vNnbzEgKe7NqB+EbGjLnH +BcNoSE1GuBjt7Juvuor8P9XW/QKBgQC8oXtd0gQmpIyApgTXtE4AdbckPZE6sX+5 +aWWk9RY7yL9ivx86ej+gG2oj0YqmBtCbcabrYdF2Dv1idSpr+UoW5cxMa+jfAwSQ +8WNeUOeWk1OfEMrnbp8XLZMTDMGMTTyHdPMNpXq0zw7bvIReiRzeFb7fEgdE2+q+ +Me+GHU5D4QKBgDK+eTrFciQ+ZbTwk6VYVyK8NnpxD4f/ZC7Yj8BwnLDgBaDyQSuC +r4ZWLxcp8X7rghFW7yPnv5k8Mpi63eMgzt1q5FCOLc6olzEXOzTg87P061XXum9C +p9AHnVuXzeekpkhw937XvZoFh4w7E83+dG2RVxWeBJk7OFAqq4Cae3nVAoGAYXOD +yqqvnk8wj1417kKmcbJfFYgBObNt6xo6ewhrniNOTPO0bH+v00WWhj7BRJkMuOH0 +fHKixj1kRrOFYRb/YekCrRCq1Fw4xbEPxzBBFRe0Ad+pE/ugkVboPtU+QP++H7UZ +xJkTVcoLQRaZxEVN9qaBX7luq/J5yhz+Q+lr/8ECgYEAsDUXgtuRyMvKr1PFzrfQ +n/aVb4wmY39Zyr1sERHwHYLUvX2gsiNMaUq7lFWUNBGy0x0gcaMCWOlxXT3SMUP9 +aE3LCjrkUWP69/RuFjKdcustLGOtJmnr9ioSpZkkrR49tmiu1SWFsK2IunTd4GUg +/p/ITUEsA2oiyDZcpTkWx0Q= +-----END PRIVATE KEY-----