diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadata.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadata.java index 7185ddfe01..502270e1cd 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadata.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadata.java @@ -75,10 +75,12 @@ import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.extension.UuidType; import org.apache.arrow.vector.ipc.ReadChannel; import org.apache.arrow.vector.ipc.message.MessageSerializer; import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.ExtensionTypeRegistry; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.Text; @@ -164,6 +166,9 @@ public class ArrowDatabaseMetadata extends AvaticaDatabaseMetaData { LONGNVARCHAR, SqlSupportsConvert.SQL_CONVERT_LONGVARCHAR_VALUE); sqlTypesToFlightEnumConvertTypes.put(DATE, SqlSupportsConvert.SQL_CONVERT_DATE_VALUE); sqlTypesToFlightEnumConvertTypes.put(TIMESTAMP, SqlSupportsConvert.SQL_CONVERT_TIMESTAMP_VALUE); + + // Register the UUID extension type so it is always available for the driver + ExtensionTypeRegistry.register(UuidType.INSTANCE); } ArrowDatabaseMetadata(final AvaticaConnection connection) { diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessorFactory.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessorFactory.java index bbfe88a78a..8362eb7627 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessorFactory.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessorFactory.java @@ -19,6 +19,7 @@ import java.util.function.IntSupplier; import org.apache.arrow.driver.jdbc.accessor.impl.ArrowFlightJdbcNullVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.binary.ArrowFlightJdbcBinaryVectorAccessor; +import org.apache.arrow.driver.jdbc.accessor.impl.binary.ArrowFlightJdbcUuidVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcDateVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcDurationVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcIntervalVectorAccessor; @@ -65,6 +66,7 @@ import org.apache.arrow.vector.UInt2Vector; import org.apache.arrow.vector.UInt4Vector; import org.apache.arrow.vector.UInt8Vector; +import org.apache.arrow.vector.UuidVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; @@ -138,6 +140,9 @@ public static ArrowFlightJdbcAccessor createAccessor( } else if (vector instanceof LargeVarBinaryVector) { return new ArrowFlightJdbcBinaryVectorAccessor( (LargeVarBinaryVector) vector, getCurrentRow, setCursorWasNull); + } else if (vector instanceof UuidVector) { + return new ArrowFlightJdbcUuidVectorAccessor( + (UuidVector) vector, getCurrentRow, setCursorWasNull); } else if (vector instanceof FixedSizeBinaryVector) { return new ArrowFlightJdbcBinaryVectorAccessor( (FixedSizeBinaryVector) vector, getCurrentRow, setCursorWasNull); diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/binary/ArrowFlightJdbcUuidVectorAccessor.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/binary/ArrowFlightJdbcUuidVectorAccessor.java new file mode 100644 index 0000000000..4bdbcbb63d --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/binary/ArrowFlightJdbcUuidVectorAccessor.java @@ -0,0 +1,88 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You 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 org.apache.arrow.driver.jdbc.accessor.impl.binary; + +import java.util.UUID; +import java.util.function.IntSupplier; +import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessor; +import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory; +import org.apache.arrow.vector.UuidVector; +import org.apache.arrow.vector.util.UuidUtility; + +/** + * Accessor for the Arrow UUID extension type ({@link UuidVector}). + * + *

This accessor provides JDBC-compatible access to UUID values stored in Arrow's canonical UUID + * extension type ('arrow.uuid'). It follows PostgreSQL JDBC driver conventions: + * + *

+ */ +public class ArrowFlightJdbcUuidVectorAccessor extends ArrowFlightJdbcAccessor { + + private final UuidVector vector; + + /** + * Creates a new accessor for a UUID vector. + * + * @param vector the UUID vector to access + * @param currentRowSupplier supplier for the current row index + * @param setCursorWasNull consumer to set the wasNull flag + */ + public ArrowFlightJdbcUuidVectorAccessor( + UuidVector vector, + IntSupplier currentRowSupplier, + ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { + super(currentRowSupplier, setCursorWasNull); + this.vector = vector; + } + + @Override + public Object getObject() { + UUID uuid = vector.getObject(getCurrentRow()); + this.wasNull = uuid == null; + this.wasNullConsumer.setWasNull(this.wasNull); + return uuid; + } + + @Override + public Class getObjectClass() { + return UUID.class; + } + + @Override + public String getString() { + UUID uuid = (UUID) getObject(); + if (uuid == null) { + return null; + } + return uuid.toString(); + } + + @Override + public byte[] getBytes() { + UUID uuid = (UUID) getObject(); + if (uuid == null) { + return null; + } + return UuidUtility.getBytesFromUUID(uuid); + } +} diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/UuidAvaticaParameterConverter.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/UuidAvaticaParameterConverter.java new file mode 100644 index 0000000000..b2157890cf --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/UuidAvaticaParameterConverter.java @@ -0,0 +1,103 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You 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 org.apache.arrow.driver.jdbc.converter.impl; + +import static org.apache.arrow.driver.jdbc.utils.SqlTypes.getSqlTypeIdFromArrowType; +import static org.apache.arrow.driver.jdbc.utils.SqlTypes.getSqlTypeNameFromArrowType; + +import java.nio.ByteBuffer; +import java.util.UUID; +import org.apache.arrow.driver.jdbc.converter.AvaticaParameterConverter; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.UuidVector; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.util.UuidUtility; +import org.apache.calcite.avatica.AvaticaParameter; +import org.apache.calcite.avatica.remote.TypedValue; +import org.apache.calcite.avatica.util.ByteString; + +/** + * AvaticaParameterConverter for UUID Arrow extension type. + * + *

Handles conversion of UUID values from JDBC parameters to Arrow's UUID extension type. Accepts + * both {@link UUID} objects and String representations of UUIDs. + */ +public class UuidAvaticaParameterConverter implements AvaticaParameterConverter { + + public UuidAvaticaParameterConverter() {} + + @Override + public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { + if (!(vector instanceof UuidVector)) { + return false; + } + + UuidVector uuidVector = (UuidVector) vector; + Object value = typedValue.toJdbc(null); + + if (value == null) { + uuidVector.setNull(index); + return true; + } + + UUID uuid; + if (value instanceof UUID) { + uuid = (UUID) value; + } else if (value instanceof String) { + uuid = UUID.fromString((String) value); + } else if (value instanceof byte[]) { + byte[] bytes = (byte[]) value; + if (bytes.length != 16) { + throw new IllegalArgumentException("UUID byte array must be 16 bytes, got " + bytes.length); + } + uuid = uuidFromBytes(bytes); + } else if (value instanceof ByteString) { + byte[] bytes = ((ByteString) value).getBytes(); + if (bytes.length != 16) { + throw new IllegalArgumentException("UUID byte array must be 16 bytes, got " + bytes.length); + } + uuid = uuidFromBytes(bytes); + } else { + throw new IllegalArgumentException( + "Cannot convert " + value.getClass().getName() + " to UUID"); + } + + uuidVector.setSafe(index, UuidUtility.getBytesFromUUID(uuid)); + return true; + } + + @Override + public AvaticaParameter createParameter(Field field) { + final String name = field.getName(); + final int jdbcType = getSqlTypeIdFromArrowType(field.getType()); + final String typeName = getSqlTypeNameFromArrowType(field.getType()); + final String className = UUID.class.getCanonicalName(); + return new AvaticaParameter(false, 0, 0, jdbcType, typeName, className, name); + } + + private static UUID uuidFromBytes(byte[] bytes) { + final long mostSignificantBits; + final long leastSignificantBits; + ByteBuffer bb = ByteBuffer.wrap(bytes); + // Reads the first eight bytes + mostSignificantBits = bb.getLong(); + // Reads the first eight bytes at this buffer's current + leastSignificantBits = bb.getLong(); + + return new UUID(mostSignificantBits, leastSignificantBits); + } +} diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/AvaticaParameterBinder.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/AvaticaParameterBinder.java index 8c98ee4077..3c282a096f 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/AvaticaParameterBinder.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/AvaticaParameterBinder.java @@ -40,11 +40,15 @@ import org.apache.arrow.driver.jdbc.converter.impl.TimestampAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.UnionAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.Utf8AvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.UuidAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.Utf8ViewAvaticaParameterConverter; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.extension.UuidType; import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.ArrowType.ArrowTypeVisitor; +import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType; import org.apache.calcite.avatica.remote.TypedValue; import org.checkerframework.checker.nullness.qual.Nullable; @@ -290,5 +294,15 @@ public Boolean visit(ArrowType.RunEndEncoded type) { throw new UnsupportedOperationException( "No Avatica parameter binder implemented for type " + type); } + + @Override + public Boolean visit(ExtensionType type) { + if (type instanceof UuidType) { + return new UuidAvaticaParameterConverter().bindParameter(vector, typedValue, index); + } + + // fallback to default implementation + return ArrowTypeVisitor.super.visit(type); + } } } diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ConvertUtils.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ConvertUtils.java index 5dd4c69c73..dd51ee5361 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ConvertUtils.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ConvertUtils.java @@ -43,8 +43,12 @@ import org.apache.arrow.driver.jdbc.converter.impl.UnionAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.Utf8AvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.Utf8ViewAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.UuidAvaticaParameterConverter; import org.apache.arrow.flight.sql.FlightSqlColumnMetadata; +import org.apache.arrow.vector.extension.UuidType; import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.ArrowType.ArrowTypeVisitor; +import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.calcite.avatica.AvaticaParameter; import org.apache.calcite.avatica.ColumnMetaData; @@ -294,5 +298,15 @@ public AvaticaParameter visit(ArrowType.RunEndEncoded type) { throw new UnsupportedOperationException( "No Avatica parameter binder implemented for type " + type); } + + @Override + public AvaticaParameter visit(ExtensionType type) { + if (type instanceof UuidType) { + return new UuidAvaticaParameterConverter().createParameter(field); + } + + // fallback to default implementation + return ArrowTypeVisitor.super.visit(type); + } } } diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/SqlTypes.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/SqlTypes.java index 5ba3957f8b..7982d5bc73 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/SqlTypes.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/SqlTypes.java @@ -20,11 +20,13 @@ import java.sql.Types; import java.util.HashMap; import java.util.Map; +import org.apache.arrow.vector.extension.UuidType; import org.apache.arrow.vector.types.FloatingPointPrecision; import org.apache.arrow.vector.types.pojo.ArrowType; /** SQL Types utility functions. */ public class SqlTypes { + private static final Map typeIdToName = new HashMap<>(); static { @@ -110,6 +112,9 @@ public static int getSqlTypeIdFromArrowType(ArrowType arrowType) { case BinaryView: return Types.VARBINARY; case FixedSizeBinary: + if (arrowType instanceof UuidType) { + return Types.OTHER; + } return Types.BINARY; case LargeBinary: return Types.LONGVARBINARY; diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java index 569b5495fe..1e8b996551 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java @@ -22,8 +22,10 @@ import static org.hamcrest.CoreMatchers.allOf; import static org.hamcrest.CoreMatchers.anyOf; import static org.hamcrest.CoreMatchers.containsString; +import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.instanceOf; import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.CoreMatchers.nullValue; import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.jupiter.api.Assertions.*; @@ -31,7 +33,9 @@ import java.nio.charset.StandardCharsets; import java.sql.Connection; import java.sql.DriverManager; +import java.sql.PreparedStatement; import java.sql.ResultSet; +import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.sql.SQLTimeoutException; import java.sql.Statement; @@ -42,6 +46,7 @@ import java.util.List; import java.util.Random; import java.util.Set; +import java.util.UUID; import java.util.concurrent.CountDownLatch; import org.apache.arrow.driver.jdbc.utils.CoreMockedSqlProducers; import org.apache.arrow.driver.jdbc.utils.FallbackFlightSqlProducer; @@ -61,6 +66,7 @@ import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.arrow.vector.util.UuidUtility; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; @@ -795,4 +801,172 @@ public void testResultSetAppMetadata() throws Exception { "foo".getBytes(StandardCharsets.UTF_8)); } } + + @Test + public void testSelectQueryWithUuidColumn() throws SQLException { + // Expectations + final int expectedRowCount = 4; + final UUID[] expectedUuids = new UUID[] { + CoreMockedSqlProducers.UUID_1, + CoreMockedSqlProducers.UUID_2, + CoreMockedSqlProducers.UUID_3, + null + }; + + final Integer[] expectedIds = new Integer[] {1, 2, 3, 4}; + + final List actualUuids = new ArrayList<>(expectedRowCount); + final List actualIds = new ArrayList<>(expectedRowCount); + + + // Query + int actualRowCount = 0; + try (Statement statement = connection.createStatement(); + ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.UUID_SQL_CMD)) { + for (; resultSet.next(); actualRowCount++) { + actualIds.add((Integer) resultSet.getObject("id")); + actualUuids.add((UUID) resultSet.getObject("uuid_col")); + } + } + + // Assertions + int finalActualRowCount = actualRowCount; + assertAll( + "UUID ResultSet values are as expected", + () -> assertThat(finalActualRowCount, is(equalTo(expectedRowCount))), + () -> assertThat(actualIds.toArray(new Integer[0]), is(expectedIds)), + () -> assertThat(actualUuids.toArray(new UUID[0]), is(expectedUuids))); + } + + @Test + public void testGetObjectReturnsUuid() throws SQLException { + try (Statement statement = connection.createStatement(); + ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.UUID_SQL_CMD)) { + resultSet.next(); + Object result = resultSet.getObject("uuid_col"); + assertThat(result, instanceOf(UUID.class)); + assertThat(result, is(CoreMockedSqlProducers.UUID_1)); + } + } + + @Test + public void testGetObjectByIndexReturnsUuid() throws SQLException { + try (Statement statement = connection.createStatement(); + ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.UUID_SQL_CMD)) { + resultSet.next(); + Object result = resultSet.getObject(2); + assertThat(result, instanceOf(UUID.class)); + assertThat(result, is(CoreMockedSqlProducers.UUID_1)); + } + } + + @Test + public void testGetStringReturnsHyphenatedFormat() throws SQLException { + try (Statement statement = connection.createStatement(); + ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.UUID_SQL_CMD)) { + resultSet.next(); + String result = resultSet.getString("uuid_col"); + assertThat(result, is(CoreMockedSqlProducers.UUID_1.toString())); + } + } + + @Test + public void testGetBytesReturns16ByteArray() throws SQLException { + try (Statement statement = connection.createStatement(); + ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.UUID_SQL_CMD)) { + resultSet.next(); + byte[] result = resultSet.getBytes("uuid_col"); + assertThat(result.length, is(16)); + assertThat(result, is(UuidUtility.getBytesFromUUID(CoreMockedSqlProducers.UUID_1))); + } + } + + @Test + public void testNullUuidHandling() throws SQLException { + try (Statement statement = connection.createStatement(); + ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.UUID_SQL_CMD)) { + // Skip to row 4 which has NULL UUID + resultSet.next(); // row 1 + resultSet.next(); // row 2 + resultSet.next(); // row 3 + resultSet.next(); // row 4 (NULL UUID) + + Object objResult = resultSet.getObject("uuid_col"); + assertThat(objResult, nullValue()); + assertThat(resultSet.wasNull(), is(true)); + + String strResult = resultSet.getString("uuid_col"); + assertThat(strResult, nullValue()); + assertThat(resultSet.wasNull(), is(true)); + + byte[] bytesResult = resultSet.getBytes("uuid_col"); + assertThat(bytesResult, nullValue()); + assertThat(resultSet.wasNull(), is(true)); + } + } + + @Test + public void testMultipleUuidRows() throws SQLException { + try (Statement statement = connection.createStatement(); + ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.UUID_SQL_CMD)) { + resultSet.next(); + assertThat(resultSet.getObject("uuid_col"), is(CoreMockedSqlProducers.UUID_1)); + + resultSet.next(); + assertThat(resultSet.getObject("uuid_col"), is(CoreMockedSqlProducers.UUID_2)); + + resultSet.next(); + assertThat(resultSet.getObject("uuid_col"), is(CoreMockedSqlProducers.UUID_3)); + + resultSet.next(); + assertThat(resultSet.getObject("uuid_col"), nullValue()); + } + } + + @Test + public void testUuidExtensionTypeInSchema() throws SQLException { + try (Statement statement = connection.createStatement(); + ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.UUID_SQL_CMD)) { + ResultSetMetaData metaData = resultSet.getMetaData(); + + assertThat(metaData.getColumnCount(), is(2)); + assertThat(metaData.getColumnName(1), is("id")); + assertThat(metaData.getColumnName(2), is("uuid_col")); + + assertThat(metaData.getColumnType(2), is(java.sql.Types.OTHER)); + } + } + + @Test + public void testPreparedStatementWithUuidParameter() throws SQLException { + try (PreparedStatement pstmt = + connection.prepareStatement(CoreMockedSqlProducers.UUID_PREPARED_SELECT_SQL_CMD)) { + pstmt.setObject(1, CoreMockedSqlProducers.UUID_1); + try (ResultSet rs = pstmt.executeQuery()) { + rs.next(); + assertThat(rs.getObject("uuid_col"), is(CoreMockedSqlProducers.UUID_1)); + } + } + } + + @Test + public void testPreparedStatementWithUuidStringParameter() throws SQLException { + try (PreparedStatement pstmt = connection.prepareStatement(CoreMockedSqlProducers.UUID_PREPARED_SELECT_SQL_CMD)) { + pstmt.setString(1, CoreMockedSqlProducers.UUID_1.toString()); + try (ResultSet rs = pstmt.executeQuery()) { + rs.next(); + assertThat(rs.getObject("uuid_col"), is(CoreMockedSqlProducers.UUID_1)); + } + } + } + + @Test + public void testPreparedStatementUpdateWithUuid() throws SQLException { + try (PreparedStatement pstmt = connection.prepareStatement(CoreMockedSqlProducers.UUID_PREPARED_UPDATE_SQL_CMD)) { + pstmt.setObject(1, CoreMockedSqlProducers.UUID_3); + pstmt.setInt(2, 1); + int updated = pstmt.executeUpdate(); + assertThat(updated, is(1)); + } + } } diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessorFactoryTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessorFactoryTest.java index 8b39041f0c..1fbd2f86a9 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessorFactoryTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessorFactoryTest.java @@ -16,10 +16,12 @@ */ package org.apache.arrow.driver.jdbc.accessor; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertTrue; import java.util.function.IntSupplier; import org.apache.arrow.driver.jdbc.accessor.impl.binary.ArrowFlightJdbcBinaryVectorAccessor; +import org.apache.arrow.driver.jdbc.accessor.impl.binary.ArrowFlightJdbcUuidVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcDateVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcDurationVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcIntervalVectorAccessor; @@ -497,4 +499,15 @@ public void createAccessorForMapVector() { assertTrue(accessor instanceof ArrowFlightJdbcMapVectorAccessor); } } + + @Test + public void createAccessorForUuidVector() { + try (ValueVector valueVector = rootAllocatorTestExtension.createUuidVector()) { + ArrowFlightJdbcAccessor accessor = + ArrowFlightJdbcAccessorFactory.createAccessor( + valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); + + assertInstanceOf(ArrowFlightJdbcUuidVectorAccessor.class, accessor); + } + } } diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/binary/ArrowFlightJdbcUuidVectorAccessorTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/binary/ArrowFlightJdbcUuidVectorAccessorTest.java new file mode 100644 index 0000000000..b7f341240c --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/binary/ArrowFlightJdbcUuidVectorAccessorTest.java @@ -0,0 +1,188 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You 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 org.apache.arrow.driver.jdbc.accessor.impl.binary; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.CoreMatchers.nullValue; +import static org.hamcrest.MatcherAssert.assertThat; + +import java.util.UUID; +import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory; +import org.apache.arrow.driver.jdbc.utils.RootAllocatorTestExtension; +import org.apache.arrow.vector.UuidVector; +import org.apache.arrow.vector.util.UuidUtility; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; + +/** + * Tests for {@link ArrowFlightJdbcUuidVectorAccessor}. + * + *

Verifies that the accessor correctly handles UUID values from Arrow's UUID extension type, + * following PostgreSQL JDBC driver conventions. + */ +public class ArrowFlightJdbcUuidVectorAccessorTest { + + @RegisterExtension + public static RootAllocatorTestExtension rootAllocatorTestExtension = + new RootAllocatorTestExtension(); + + private static final UUID UUID_1 = UUID.fromString("550e8400-e29b-41d4-a716-446655440000"); + private static final UUID UUID_2 = UUID.fromString("6ba7b810-9dad-11d1-80b4-00c04fd430c8"); + private static final UUID UUID_3 = UUID.fromString("f47ac10b-58cc-4372-a567-0e02b2c3d479"); + + private UuidVector vector; + private ArrowFlightJdbcUuidVectorAccessor accessor; + private boolean wasNullCalled; + private boolean wasNullValue; + + @BeforeEach + public void setUp() { + vector = rootAllocatorTestExtension.createUuidVector(); + wasNullCalled = false; + wasNullValue = false; + ArrowFlightJdbcAccessorFactory.WasNullConsumer wasNullConsumer = + (wasNull) -> { + wasNullCalled = true; + wasNullValue = wasNull; + }; + accessor = new ArrowFlightJdbcUuidVectorAccessor(vector, () -> 0, wasNullConsumer); + } + + @AfterEach + public void tearDown() { + vector.close(); + } + + @Test + public void testGetObjectReturnsUuid() { + accessor = new ArrowFlightJdbcUuidVectorAccessor(vector, () -> 0, (wasNull) -> {}); + Object result = accessor.getObject(); + assertThat(result, is(UUID_1)); + assertThat(accessor.wasNull(), is(false)); + } + + @Test + public void testGetObjectReturnsCorrectUuidForEachRow() { + accessor = new ArrowFlightJdbcUuidVectorAccessor(vector, () -> 0, (wasNull) -> {}); + assertThat(accessor.getObject(), is(UUID_1)); + + accessor = new ArrowFlightJdbcUuidVectorAccessor(vector, () -> 1, (wasNull) -> {}); + assertThat(accessor.getObject(), is(UUID_2)); + + accessor = new ArrowFlightJdbcUuidVectorAccessor(vector, () -> 2, (wasNull) -> {}); + assertThat(accessor.getObject(), is(UUID_3)); + } + + @Test + public void testGetObjectReturnsNullForNullValue() { + vector.reset(); + vector.allocateNew(1); + vector.setNull(0); + vector.setValueCount(1); + + accessor = new ArrowFlightJdbcUuidVectorAccessor(vector, () -> 0, (wasNull) -> {}); + Object result = accessor.getObject(); + assertThat(result, nullValue()); + assertThat(accessor.wasNull(), is(true)); + } + + @Test + public void testGetObjectClassReturnsUuidClass() { + assertThat(accessor.getObjectClass(), equalTo(UUID.class)); + } + + @Test + public void testGetStringReturnsHyphenatedFormat() { + accessor = new ArrowFlightJdbcUuidVectorAccessor(vector, () -> 0, (wasNull) -> {}); + String result = accessor.getString(); + assertThat(result, is("550e8400-e29b-41d4-a716-446655440000")); + assertThat(accessor.wasNull(), is(false)); + } + + @Test + public void testGetStringReturnsNullForNullValue() { + vector.reset(); + vector.allocateNew(1); + vector.setNull(0); + vector.setValueCount(1); + + accessor = new ArrowFlightJdbcUuidVectorAccessor(vector, () -> 0, (wasNull) -> {}); + String result = accessor.getString(); + assertThat(result, nullValue()); + assertThat(accessor.wasNull(), is(true)); + } + + @Test + public void testGetBytesReturns16ByteArray() { + accessor = new ArrowFlightJdbcUuidVectorAccessor(vector, () -> 0, (wasNull) -> {}); + byte[] result = accessor.getBytes(); + assertThat(result.length, is(16)); + assertThat(result, is(UuidUtility.getBytesFromUUID(UUID_1))); + assertThat(accessor.wasNull(), is(false)); + } + + @Test + public void testGetBytesReturnsNullForNullValue() { + vector.reset(); + vector.allocateNew(1); + vector.setNull(0); + vector.setValueCount(1); + + accessor = new ArrowFlightJdbcUuidVectorAccessor(vector, () -> 0, (wasNull) -> {}); + byte[] result = accessor.getBytes(); + assertThat(result, nullValue()); + assertThat(accessor.wasNull(), is(true)); + } + + @Test + public void testWasNullConsumerIsCalled() { + accessor = + new ArrowFlightJdbcUuidVectorAccessor( + vector, + () -> 0, + (wasNull) -> { + wasNullCalled = true; + wasNullValue = wasNull; + }); + accessor.getObject(); + assertThat(wasNullCalled, is(true)); + assertThat(wasNullValue, is(false)); + } + + @Test + public void testWasNullConsumerIsCalledWithTrueForNull() { + vector.reset(); + vector.allocateNew(1); + vector.setNull(0); + vector.setValueCount(1); + + accessor = + new ArrowFlightJdbcUuidVectorAccessor( + vector, + () -> 0, + (wasNull) -> { + wasNullCalled = true; + wasNullValue = wasNull; + }); + accessor.getObject(); + assertThat(wasNullCalled, is(true)); + assertThat(wasNullValue, is(true)); + } +} diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/converter/impl/UuidAvaticaParameterConverterTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/converter/impl/UuidAvaticaParameterConverterTest.java new file mode 100644 index 0000000000..07751f0abc --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/converter/impl/UuidAvaticaParameterConverterTest.java @@ -0,0 +1,160 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You 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 org.apache.arrow.driver.jdbc.converter.impl; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.CoreMatchers.nullValue; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.sql.Types; +import java.util.UUID; +import org.apache.arrow.driver.jdbc.utils.RootAllocatorTestExtension; +import org.apache.arrow.vector.UuidVector; +import org.apache.arrow.vector.extension.UuidType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.util.UuidUtility; +import org.apache.calcite.avatica.AvaticaParameter; +import org.apache.calcite.avatica.ColumnMetaData; +import org.apache.calcite.avatica.remote.TypedValue; +import org.apache.calcite.avatica.util.ByteString; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; + +/** + * Tests for {@link UuidAvaticaParameterConverter}. + * + *

Verifies that the converter correctly handles UUID parameter binding from JDBC to Arrow's UUID + * extension type. + */ +public class UuidAvaticaParameterConverterTest { + + @RegisterExtension + public static RootAllocatorTestExtension rootAllocatorTestExtension = + new RootAllocatorTestExtension(); + + private static final UUID TEST_UUID = UUID.fromString("550e8400-e29b-41d4-a716-446655440000"); + + private UuidVector vector; + private UuidAvaticaParameterConverter converter; + + @BeforeEach + public void setUp() { + vector = new UuidVector("uuid_param", rootAllocatorTestExtension.getRootAllocator()); + vector.allocateNew(5); + converter = new UuidAvaticaParameterConverter(); + } + + @AfterEach + public void tearDown() { + vector.close(); + } + + @Test + public void testBindParameterWithUuidObject() { + TypedValue typedValue = TypedValue.ofLocal(ColumnMetaData.Rep.OBJECT, TEST_UUID); + + boolean result = converter.bindParameter(vector, typedValue, 0); + + assertTrue(result); + assertThat(vector.getObject(0), is(TEST_UUID)); + } + + @Test + public void testBindParameterWithUuidString() { + String uuidString = "550e8400-e29b-41d4-a716-446655440000"; + TypedValue typedValue = TypedValue.ofLocal(ColumnMetaData.Rep.STRING, uuidString); + + boolean result = converter.bindParameter(vector, typedValue, 0); + + assertTrue(result); + assertThat(vector.getObject(0), is(TEST_UUID)); + } + + @Test + public void testBindParameterWithByteArray() { + byte[] uuidBytes = UuidUtility.getBytesFromUUID(TEST_UUID); + ByteString byteString = new ByteString(uuidBytes); + TypedValue typedValue = TypedValue.ofLocal(ColumnMetaData.Rep.BYTE_STRING, byteString); + + boolean result = converter.bindParameter(vector, typedValue, 0); + + assertTrue(result); + assertThat(vector.getObject(0), is(TEST_UUID)); + } + + @Test + public void testBindParameterWithNullValue() { + TypedValue typedValue = TypedValue.ofLocal(ColumnMetaData.Rep.OBJECT, null); + + boolean result = converter.bindParameter(vector, typedValue, 0); + + assertTrue(result); + assertTrue(vector.isNull(0)); + assertThat(vector.getObject(0), nullValue()); + } + + @Test + public void testBindParameterWithInvalidByteArrayLength() { + byte[] invalidBytes = new byte[8]; // Should be 16 bytes + ByteString byteString = new ByteString(invalidBytes); + TypedValue typedValue = TypedValue.ofLocal(ColumnMetaData.Rep.BYTE_STRING, byteString); + + assertThrows( + IllegalArgumentException.class, () -> converter.bindParameter(vector, typedValue, 0)); + } + + @Test + public void testBindParameterWithInvalidType() { + TypedValue typedValue = TypedValue.ofLocal(ColumnMetaData.Rep.INTEGER, 12345); + + assertThrows( + IllegalArgumentException.class, () -> converter.bindParameter(vector, typedValue, 0)); + } + + @Test + public void testBindParameterMultipleValues() { + UUID uuid1 = UUID.fromString("550e8400-e29b-41d4-a716-446655440000"); + UUID uuid2 = UUID.fromString("6ba7b810-9dad-11d1-80b4-00c04fd430c8"); + UUID uuid3 = UUID.fromString("f47ac10b-58cc-4372-a567-0e02b2c3d479"); + + converter.bindParameter(vector, TypedValue.ofLocal(ColumnMetaData.Rep.OBJECT, uuid1), 0); + converter.bindParameter(vector, TypedValue.ofLocal(ColumnMetaData.Rep.OBJECT, uuid2), 1); + converter.bindParameter(vector, TypedValue.ofLocal(ColumnMetaData.Rep.OBJECT, uuid3), 2); + + assertThat(vector.getObject(0), is(uuid1)); + assertThat(vector.getObject(1), is(uuid2)); + assertThat(vector.getObject(2), is(uuid3)); + } + + @Test + public void testCreateParameter() { + Field uuidField = new Field("uuid_col", new FieldType(true, UuidType.INSTANCE, null), null); + + AvaticaParameter parameter = converter.createParameter(uuidField); + + assertThat(parameter.name, is("uuid_col")); + assertThat(parameter.parameterType, is(Types.OTHER)); + assertThat(parameter.typeName, is("OTHER")); + assertThat(parameter.className, equalTo(UUID.class.getCanonicalName())); + } +} diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/CoreMockedSqlProducers.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/CoreMockedSqlProducers.java index 8197d7d95f..70f3c26cd8 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/CoreMockedSqlProducers.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/CoreMockedSqlProducers.java @@ -17,6 +17,7 @@ package org.apache.arrow.driver.jdbc.utils; import static java.lang.String.format; +import java.util.Arrays; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.MatcherAssert.assertThat; @@ -30,6 +31,7 @@ import java.util.ArrayList; import java.util.Collections; import java.util.List; +import java.util.UUID; import java.util.function.Consumer; import java.util.stream.IntStream; import org.apache.arrow.flight.FlightProducer.ServerStreamListener; @@ -40,10 +42,13 @@ import org.apache.arrow.vector.DateDayVector; import org.apache.arrow.vector.Float4Vector; import org.apache.arrow.vector.Float8Vector; +import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.TimeStampMilliVector; import org.apache.arrow.vector.UInt4Vector; +import org.apache.arrow.vector.UuidVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.extension.UuidType; import org.apache.arrow.vector.types.DateUnit; import org.apache.arrow.vector.types.FloatingPointPrecision; import org.apache.arrow.vector.types.TimeUnit; @@ -52,6 +57,7 @@ import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.Text; +import org.apache.arrow.vector.util.UuidUtility; /** Standard {@link MockFlightSqlProducer} instances for tests. */ // TODO Remove this once all tests are refactor to use only the queries they need. @@ -62,6 +68,20 @@ public final class CoreMockedSqlProducers { public static final String LEGACY_CANCELLATION_SQL_CMD = "SELECT * FROM TAKES_FOREVER"; public static final String LEGACY_REGULAR_WITH_EMPTY_SQL_CMD = "SELECT * FROM TEST_EMPTIES"; + public static final String UUID_SQL_CMD = "SELECT * FROM UUID_TABLE"; + public static final String UUID_PREPARED_SELECT_SQL_CMD = "SELECT * FROM UUID_TABLE WHERE uuid_col = ?"; + public static final String UUID_PREPARED_UPDATE_SQL_CMD = "UPDATE UUID_TABLE SET uuid_col = ? WHERE id = ?"; + + public static final UUID UUID_1 = UUID.fromString("550e8400-e29b-41d4-a716-446655440000"); + public static final UUID UUID_2 = UUID.fromString("6ba7b810-9dad-11d1-80b4-00c04fd430c8"); + public static final UUID UUID_3 = UUID.fromString("f47ac10b-58cc-4372-a567-0e02b2c3d479"); + + public static final Schema UUID_SCHEMA = + new Schema( + ImmutableList.of( + new Field("id", new FieldType(true, new ArrowType.Int(32, true), null), null), + new Field("uuid_col", new FieldType(true, UuidType.INSTANCE, null), null))); + private CoreMockedSqlProducers() { // Prevent instantiation. } @@ -78,9 +98,108 @@ public static MockFlightSqlProducer getLegacyProducer() { addLegacyMetadataSqlCmdSupport(producer); addLegacyCancellationSqlCmdSupport(producer); addQueryWithEmbeddedEmptyRoot(producer); + addUuidSqlCmdSupport(producer); + addUuidPreparedSelectSqlCmdSupport(producer); + addUuidPreparedUpdateSqlCmdSupport(producer); return producer; } + /** + * Gets a {@link MockFlightSqlProducer} configured with UUID test data. + * + * @return a new producer with UUID support. + */ + public static MockFlightSqlProducer getUuidProducer() { + final MockFlightSqlProducer producer = new MockFlightSqlProducer(); + addUuidSqlCmdSupport(producer); + return producer; + } + private static void addUuidPreparedUpdateSqlCmdSupport(final MockFlightSqlProducer producer) { + final String query = "UPDATE UUID_TABLE SET uuid_col = ? WHERE id = ?"; + final Schema parameterSchema = + new Schema( + Arrays.asList( + new Field("", new FieldType(true, UuidType.INSTANCE, null), null), + Field.nullable("", new ArrowType.Int(32, true)))); + + producer.addUpdateQuery(query, 1); + producer.addExpectedParameters( + UUID_PREPARED_UPDATE_SQL_CMD, + parameterSchema, + Collections.singletonList(Arrays.asList(CoreMockedSqlProducers.UUID_3, 1))); + } + + private static void addUuidPreparedSelectSqlCmdSupport(final MockFlightSqlProducer producer) { + final Schema parameterSchema = + new Schema( + Collections.singletonList( + new Field("", new FieldType(true, UuidType.INSTANCE, null), null))); + + final Consumer uuidResultProvider = + listener -> { + try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + final VectorSchemaRoot root = VectorSchemaRoot.create(UUID_SCHEMA, allocator)) { + root.allocateNew(); + IntVector idVector = (IntVector) root.getVector("id"); + UuidVector uuidVector = (UuidVector) root.getVector("uuid_col"); + idVector.setSafe(0, 1); + uuidVector.setSafe(0, UuidUtility.getBytesFromUUID(CoreMockedSqlProducers.UUID_1)); + root.setRowCount(1); + listener.start(root); + listener.putNext(); + } catch (final Throwable throwable) { + listener.error(throwable); + } finally { + listener.completed(); + } + }; + + producer.addSelectQuery( + UUID_PREPARED_SELECT_SQL_CMD, UUID_SCHEMA, Collections.singletonList(uuidResultProvider)); + producer.addExpectedParameters( + UUID_PREPARED_SELECT_SQL_CMD, + parameterSchema, + Collections.singletonList(Collections.singletonList(CoreMockedSqlProducers.UUID_1))); + } + + private static void addUuidSqlCmdSupport(final MockFlightSqlProducer producer) { + final Consumer uuidResultProvider = + listener -> { + try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + final VectorSchemaRoot root = VectorSchemaRoot.create(UUID_SCHEMA, allocator)) { + root.allocateNew(); + + IntVector idVector = (IntVector) root.getVector("id"); + UuidVector uuidVector = (UuidVector) root.getVector("uuid_col"); + + // Row 0: id=1, uuid=UUID_1 + idVector.setSafe(0, 1); + uuidVector.setSafe(0, UuidUtility.getBytesFromUUID(UUID_1)); + + // Row 1: id=2, uuid=UUID_2 + idVector.setSafe(1, 2); + uuidVector.setSafe(1, UuidUtility.getBytesFromUUID(UUID_2)); + + // Row 2: id=3, uuid=UUID_3 + idVector.setSafe(2, 3); + uuidVector.setSafe(2, UuidUtility.getBytesFromUUID(UUID_3)); + + // Row 3: id=4, uuid=NULL + idVector.setSafe(3, 4); + uuidVector.setNull(3); + + root.setRowCount(4); + listener.start(root); + listener.putNext(); + } finally { + listener.completed(); + } + }; + + producer.addSelectQuery( + UUID_SQL_CMD, UUID_SCHEMA, Collections.singletonList(uuidResultProvider)); + } + private static void addQueryWithEmbeddedEmptyRoot(final MockFlightSqlProducer producer) { final Schema querySchema = new Schema( diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/RootAllocatorTestExtension.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/RootAllocatorTestExtension.java index 347e92a16c..4b299d63e0 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/RootAllocatorTestExtension.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/RootAllocatorTestExtension.java @@ -19,6 +19,7 @@ import java.math.BigDecimal; import java.nio.charset.StandardCharsets; import java.util.Random; +import java.util.UUID; import java.util.concurrent.TimeUnit; import java.util.stream.IntStream; import org.apache.arrow.memory.BufferAllocator; @@ -53,6 +54,7 @@ import org.apache.arrow.vector.UInt2Vector; import org.apache.arrow.vector.UInt4Vector; import org.apache.arrow.vector.UInt8Vector; +import org.apache.arrow.vector.UuidVector; import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.complex.FixedSizeListVector; import org.apache.arrow.vector.complex.LargeListVector; @@ -60,6 +62,7 @@ import org.apache.arrow.vector.complex.impl.UnionFixedSizeListWriter; import org.apache.arrow.vector.complex.impl.UnionLargeListWriter; import org.apache.arrow.vector.complex.impl.UnionListWriter; +import org.apache.arrow.vector.util.UuidUtility; import org.junit.jupiter.api.extension.AfterAllCallback; import org.junit.jupiter.api.extension.BeforeAllCallback; import org.junit.jupiter.api.extension.ExtensionContext; @@ -811,4 +814,23 @@ public FixedSizeListVector createFixedSizeListVector() { return valueVector; } + + /** + * Create a UuidVector to be used in the accessor tests. + * + * @return UuidVector + */ + public UuidVector createUuidVector() { + UuidVector valueVector = new UuidVector("", this.getRootAllocator()); + valueVector.allocateNew(3); + valueVector.setSafe( + 0, UuidUtility.getBytesFromUUID(UUID.fromString("550e8400-e29b-41d4-a716-446655440000"))); + valueVector.setSafe( + 1, UuidUtility.getBytesFromUUID(UUID.fromString("6ba7b810-9dad-11d1-80b4-00c04fd430c8"))); + valueVector.setSafe( + 2, UuidUtility.getBytesFromUUID(UUID.fromString("f47ac10b-58cc-4372-a567-0e02b2c3d479"))); + valueVector.setValueCount(3); + + return valueVector; + } } diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/SqlTypesTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/SqlTypesTest.java index d69c549296..98459922a6 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/SqlTypesTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/SqlTypesTest.java @@ -21,6 +21,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import java.sql.Types; +import org.apache.arrow.vector.extension.UuidType; import org.apache.arrow.vector.types.DateUnit; import org.apache.arrow.vector.types.FloatingPointPrecision; import org.apache.arrow.vector.types.IntervalUnit; @@ -85,6 +86,8 @@ public void testGetSqlTypeIdFromArrowType() { assertEquals(Types.JAVA_OBJECT, getSqlTypeIdFromArrowType(new ArrowType.Map(true))); assertEquals(Types.NULL, getSqlTypeIdFromArrowType(new ArrowType.Null())); + + assertEquals(Types.OTHER, getSqlTypeIdFromArrowType(new UuidType())); } @Test @@ -140,5 +143,7 @@ public void testGetSqlTypeNameFromArrowType() { assertEquals("JAVA_OBJECT", getSqlTypeNameFromArrowType(new ArrowType.Map(true))); assertEquals("NULL", getSqlTypeNameFromArrowType(new ArrowType.Null())); + + assertEquals("OTHER", getSqlTypeNameFromArrowType(new UuidType())); } }