diff --git a/c/pom.xml b/c/pom.xml
index c90b6dc0ef..27b6619c4c 100644
--- a/c/pom.xml
+++ b/c/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-java-root
- 19.0.0-SNAPSHOT
+ 20.0.0-SNAPSHOT
arrow-c-data
diff --git a/c/src/main/java/org/apache/arrow/c/ReferenceCountedArrowArray.java b/c/src/main/java/org/apache/arrow/c/ReferenceCountedArrowArray.java
index cf50f9417b..f51fb25105 100644
--- a/c/src/main/java/org/apache/arrow/c/ReferenceCountedArrowArray.java
+++ b/c/src/main/java/org/apache/arrow/c/ReferenceCountedArrowArray.java
@@ -64,13 +64,18 @@ void release() {
*/
ArrowBuf unsafeAssociateAllocation(
BufferAllocator trackingAllocator, long capacity, long memoryAddress) {
+ // Retain only after wrapForeignAllocation succeeds. On the allocator-limit OOM path,
+ // wrapForeignAllocation throws before the ForeignAllocation is associated, so release0()
+ // is not called; retaining first would leave the count elevated with no matching release0().
+ ArrowBuf buf =
+ trackingAllocator.wrapForeignAllocation(
+ new ForeignAllocation(capacity, memoryAddress) {
+ @Override
+ protected void release0() {
+ ReferenceCountedArrowArray.this.release();
+ }
+ });
retain();
- return trackingAllocator.wrapForeignAllocation(
- new ForeignAllocation(capacity, memoryAddress) {
- @Override
- protected void release0() {
- ReferenceCountedArrowArray.this.release();
- }
- });
+ return buf;
}
}
diff --git a/c/src/test/java/org/apache/arrow/c/ImportOutOfMemoryTest.java b/c/src/test/java/org/apache/arrow/c/ImportOutOfMemoryTest.java
new file mode 100644
index 0000000000..7c099f2ef0
--- /dev/null
+++ b/c/src/test/java/org/apache/arrow/c/ImportOutOfMemoryTest.java
@@ -0,0 +1,139 @@
+/*
+ * 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.c;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.OutOfMemoryException;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.VarCharVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.types.pojo.Schema;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+/**
+ * Regression test: a mid-import {@link OutOfMemoryException} must not leak the imported array.
+ *
+ * A "producer" allocator owns the exported batch; if the C Data release callback fires, the
+ * producer drains to zero. A too-small consumer allocator forces an OOM part-way through the
+ * import. The test asserts the producer drains, confirming the release callback fired despite the
+ * failure.
+ */
+final class ImportOutOfMemoryTest {
+ private static final int ROWS = 1024;
+ private static final int VALUE_BYTES = 256;
+ private static final int COLUMNS = 4;
+ // Far smaller than the exported batch, so the import OOMs part-way through the buffers.
+ private static final long TINY_LIMIT = 16 * 1024;
+
+ private RootAllocator root;
+
+ @BeforeEach
+ public void setUp() {
+ root = new RootAllocator(Long.MAX_VALUE);
+ }
+
+ @AfterEach
+ public void tearDown() {
+ root.close();
+ }
+
+ @Test
+ public void importOomDoesNotLeakExportedArray() {
+ // "producer" owns only the exported batch buffers; the C Data struct containers live on a
+ // separate allocator (they are consumed/closed by import, which would otherwise muddy the
+ // producer's balance). So producer draining to zero is an exact signal that the array's release
+ // callback fired.
+ try (BufferAllocator producer = root.newChildAllocator("producer", 0, Long.MAX_VALUE);
+ BufferAllocator structs = root.newChildAllocator("structs", 0, Long.MAX_VALUE)) {
+ try (ArrowArray array = ArrowArray.allocateNew(structs);
+ ArrowSchema schema = ArrowSchema.allocateNew(structs)) {
+ exportBatch(producer, array, schema);
+ assertTrue(
+ producer.getAllocatedMemory() > 0, "producer holds the exported batch before import");
+
+ // A consumer allocator far too small to hold the batch: the import throws part-way through.
+ try (BufferAllocator consumer = root.newChildAllocator("consumer", 0, TINY_LIMIT);
+ CDataDictionaryProvider provider = new CDataDictionaryProvider()) {
+ Schema importSchema = Data.importSchema(consumer, schema, provider);
+ try (VectorSchemaRoot importRoot = VectorSchemaRoot.create(importSchema, consumer)) {
+ Exception thrown =
+ assertThrows(
+ Exception.class,
+ () -> Data.importIntoVectorSchemaRoot(consumer, array, importRoot, provider));
+ assertTrue(
+ hasOutOfMemoryCause(thrown),
+ "mid-import failure must be an allocator OOM: " + thrown);
+ }
+ }
+
+ // The array's release callback must have fired despite the mid-import OOM, freeing the
+ // whole exported batch. On the unfixed retain-before-wrap code the batch is stranded.
+ assertEquals(
+ 0L,
+ producer.getAllocatedMemory(),
+ "import OOM leaked the exported batch (producer not drained)");
+ }
+ }
+ }
+
+ /** True if {@code t} is, or is caused by, an Arrow {@link OutOfMemoryException}. */
+ private static boolean hasOutOfMemoryCause(Throwable t) {
+ for (Throwable cause = t; cause != null; cause = cause.getCause()) {
+ if (cause instanceof OutOfMemoryException) {
+ return true;
+ }
+ }
+ return false;
+ }
+
+ /**
+ * Builds a wide multi-column VarChar batch on {@code alloc} and exports it into the C structs.
+ */
+ private void exportBatch(BufferAllocator alloc, ArrowArray array, ArrowSchema schema) {
+ byte[] value = new byte[VALUE_BYTES];
+ for (int i = 0; i < value.length; i++) {
+ value[i] = (byte) 'x';
+ }
+ List vectors = new ArrayList<>(COLUMNS);
+ for (int c = 0; c < COLUMNS; c++) {
+ VarCharVector vector = new VarCharVector("col" + c, alloc);
+ vector.allocateNew((long) ROWS * VALUE_BYTES, ROWS);
+ for (int r = 0; r < ROWS; r++) {
+ vector.setSafe(r, value);
+ }
+ vector.setValueCount(ROWS);
+ vectors.add(vector);
+ }
+ try (VectorSchemaRoot source = new VectorSchemaRoot(vectors)) {
+ long total = 0;
+ for (FieldVector vector : source.getFieldVectors()) {
+ total += vector.getBufferSize();
+ }
+ assertTrue(total > TINY_LIMIT, "test setup: batch must exceed the consumer limit");
+ Data.exportVectorSchemaRoot(alloc, source, null, array, schema);
+ }
+ }
+}
diff --git a/c/src/test/java/org/apache/arrow/c/RoundtripTest.java b/c/src/test/java/org/apache/arrow/c/RoundtripTest.java
index 010a305495..f6ff88571e 100644
--- a/c/src/test/java/org/apache/arrow/c/RoundtripTest.java
+++ b/c/src/test/java/org/apache/arrow/c/RoundtripTest.java
@@ -35,7 +35,6 @@
import org.apache.arrow.memory.ArrowBuf;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
-import org.apache.arrow.memory.util.hash.ArrowBufHasher;
import org.apache.arrow.vector.BaseLargeVariableWidthVector;
import org.apache.arrow.vector.BaseVariableWidthVector;
import org.apache.arrow.vector.BigIntVector;
@@ -44,7 +43,6 @@
import org.apache.arrow.vector.DateMilliVector;
import org.apache.arrow.vector.DecimalVector;
import org.apache.arrow.vector.DurationVector;
-import org.apache.arrow.vector.ExtensionTypeVector;
import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.FixedSizeBinaryVector;
import org.apache.arrow.vector.Float2Vector;
@@ -74,6 +72,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;
@@ -92,6 +91,7 @@
import org.apache.arrow.vector.complex.StructVector;
import org.apache.arrow.vector.complex.UnionVector;
import org.apache.arrow.vector.complex.impl.UnionMapWriter;
+import org.apache.arrow.vector.extension.UuidType;
import org.apache.arrow.vector.holders.IntervalDayHolder;
import org.apache.arrow.vector.holders.NullableLargeVarBinaryHolder;
import org.apache.arrow.vector.holders.NullableUInt4Holder;
@@ -100,7 +100,6 @@
import org.apache.arrow.vector.types.Types.MinorType;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType;
-import org.apache.arrow.vector.types.pojo.ExtensionTypeRegistry;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.FieldType;
import org.apache.arrow.vector.types.pojo.Schema;
@@ -810,9 +809,8 @@ public void testEmptyRunEndEncodedVector() {
@Test
public void testExtensionTypeVector() {
- ExtensionTypeRegistry.register(new UuidType());
final Schema schema =
- new Schema(Collections.singletonList(Field.nullable("a", new UuidType())));
+ new Schema(Collections.singletonList(Field.nullable("a", UuidType.INSTANCE)));
try (final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
// Fill with data
UUID u1 = UUID.randomUUID();
@@ -830,13 +828,12 @@ public void testExtensionTypeVector() {
assertEquals(root.getSchema(), importedRoot.getSchema());
final Field field = importedRoot.getSchema().getFields().get(0);
- final UuidType expectedType = new UuidType();
assertEquals(
field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_NAME),
- expectedType.extensionName());
+ UuidType.INSTANCE.extensionName());
assertEquals(
field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_METADATA),
- expectedType.serialize());
+ UuidType.INSTANCE.serialize());
final UuidVector deserialized = (UuidVector) importedRoot.getFieldVectors().get(0);
assertEquals(vector.getValueCount(), deserialized.getValueCount());
@@ -1115,72 +1112,4 @@ private VectorSchemaRoot createTestVSR() {
return new VectorSchemaRoot(fields, vectors);
}
-
- static class UuidType extends ExtensionType {
-
- @Override
- public ArrowType storageType() {
- return new ArrowType.FixedSizeBinary(16);
- }
-
- @Override
- public String extensionName() {
- return "uuid";
- }
-
- @Override
- public boolean extensionEquals(ExtensionType other) {
- return other instanceof UuidType;
- }
-
- @Override
- public ArrowType deserialize(ArrowType storageType, String serializedData) {
- if (!storageType.equals(storageType())) {
- throw new UnsupportedOperationException(
- "Cannot construct UuidType from underlying type " + storageType);
- }
- return new UuidType();
- }
-
- @Override
- public String serialize() {
- return "";
- }
-
- @Override
- public FieldVector getNewVector(String name, FieldType fieldType, BufferAllocator allocator) {
- return new UuidVector(name, allocator, new FixedSizeBinaryVector(name, allocator, 16));
- }
- }
-
- static class UuidVector extends ExtensionTypeVector {
-
- public UuidVector(
- String name, BufferAllocator allocator, FixedSizeBinaryVector underlyingVector) {
- super(name, allocator, underlyingVector);
- }
-
- @Override
- public UUID getObject(int index) {
- final ByteBuffer bb = ByteBuffer.wrap(getUnderlyingVector().getObject(index));
- return new UUID(bb.getLong(), bb.getLong());
- }
-
- @Override
- public int hashCode(int index) {
- return hashCode(index, null);
- }
-
- @Override
- public int hashCode(int index, ArrowBufHasher hasher) {
- return getUnderlyingVector().hashCode(index, hasher);
- }
-
- public void set(int index, UUID uuid) {
- ByteBuffer bb = ByteBuffer.allocate(16);
- bb.putLong(uuid.getMostSignificantBits());
- bb.putLong(uuid.getLeastSignificantBits());
- getUnderlyingVector().set(index, bb.array());
- }
- }
}
diff --git a/ci/docker/conda-jni.dockerfile b/ci/docker/conda-jni.dockerfile
index e14db73688..3f31b74052 100644
--- a/ci/docker/conda-jni.dockerfile
+++ b/ci/docker/conda-jni.dockerfile
@@ -17,7 +17,7 @@
FROM ghcr.io/mamba-org/micromamba:ubuntu24.04
-ARG jdk=11
+ARG jdk=17
ARG maven=3.9.9
RUN micromamba install -y \
diff --git a/ci/docker/vcpkg-jni.dockerfile b/ci/docker/vcpkg-jni.dockerfile
index d6bd322a39..f2f5d0d45a 100644
--- a/ci/docker/vcpkg-jni.dockerfile
+++ b/ci/docker/vcpkg-jni.dockerfile
@@ -20,7 +20,7 @@ FROM ${base}
# Install Java
# We need Java for JNI headers, but we don't invoke Maven in this build.
-ARG java=11
+ARG java=17
RUN dnf install -y java-$java-openjdk-devel && dnf clean all
# For ci/scripts/{cpp,java}_*.sh
diff --git a/ci/scripts/jni_macos_build.sh b/ci/scripts/jni_macos_build.sh
index 13c0675d38..65ab450666 100755
--- a/ci/scripts/jni_macos_build.sh
+++ b/ci/scripts/jni_macos_build.sh
@@ -77,7 +77,9 @@ cmake \
cmake --build "${build_dir}/cpp" --target install
github_actions_group_end
-export JAVA_JNI_CMAKE_ARGS="-DProtobuf_ROOT=${build_dir}/cpp/protobuf_ep-install"
+JAVA_JNI_CMAKE_ARGS="-DProtobuf_ROOT=${build_dir}/cpp/_deps/protobuf-build"
+JAVA_JNI_CMAKE_ARGS+=" -DProtobuf_SRC_ROOT_FOLDER=${build_dir}/cpp/_deps/protobuf-src"
+export JAVA_JNI_CMAKE_ARGS
"${source_dir}/ci/scripts/jni_build.sh" \
"${source_dir}" \
"${install_dir}" \
diff --git a/ci/scripts/test.sh b/ci/scripts/test.sh
index cacc20034e..8061ee455d 100755
--- a/ci/scripts/test.sh
+++ b/ci/scripts/test.sh
@@ -34,10 +34,11 @@ fi
mvn="mvn -B -Dorg.slf4j.simpleLogger.log.org.apache.maven.cli.transfer.Slf4jMavenTransferListener=warn"
# Use `2 * ncores` threads
mvn="${mvn} -T 2C"
+mvn="${mvn} -Denforcer.skip=true"
pushd "${build_dir}"
-${mvn} -Darrow.test.dataRoot="${source_dir}/testing/data" clean test
+${mvn} -Darrow.test.dataRoot="${source_dir}/testing/data" test
projects=()
if [ "${ARROW_JAVA_JNI}" = "ON" ]; then
@@ -46,7 +47,7 @@ if [ "${ARROW_JAVA_JNI}" = "ON" ]; then
projects+=(gandiva)
fi
if [ "${#projects[@]}" -gt 0 ]; then
- ${mvn} clean test \
+ ${mvn} test \
-Parrow-jni \
-pl "$(
IFS=,
@@ -56,7 +57,7 @@ if [ "${#projects[@]}" -gt 0 ]; then
fi
if [ "${ARROW_JAVA_CDATA}" = "ON" ]; then
- ${mvn} clean test -Parrow-c-data -pl c -Darrow.c.jni.dist.dir="${java_jni_dist_dir}"
+ ${mvn} test -Parrow-c-data -pl c -Darrow.c.jni.dist.dir="${java_jni_dist_dir}"
fi
popd
diff --git a/compose.yaml b/compose.yaml
index f5082a22aa..4fd825e5a5 100644
--- a/compose.yaml
+++ b/compose.yaml
@@ -40,7 +40,7 @@ services:
# docker compose run ubuntu
# Parameters:
# MAVEN: 3.9.9
- # JDK: 11, 17, 21
+ # JDK: 17, 21
image: ${ARCH}/maven:${MAVEN}-eclipse-temurin-${JDK}
volumes:
- .:/arrow-java:delegated
@@ -60,7 +60,7 @@ services:
# docker compose run conda-jni-cdata
# Parameters:
# MAVEN: 3.9.9
- # JDK: 11, 17, 21
+ # JDK: 17, 21
image: ${REPO}:${ARCH}-conda-java-${JDK}-maven-${MAVEN}-jni-integration
build:
context: .
@@ -109,5 +109,6 @@ services:
ARROW_JAVA_CDATA: "ON"
CCACHE_DIR: "/ccache"
command:
- ["git config --global --add safe.directory /arrow-java && \
+ ["/bin/bash", "-c",
+ "git config --global --add safe.directory /arrow-java &&
/arrow-java/ci/scripts/jni_manylinux_build.sh /arrow-java /arrow /build/java /arrow-java/jni"]
diff --git a/compression/pom.xml b/compression/pom.xml
index ba13156243..aa7dee6f89 100644
--- a/compression/pom.xml
+++ b/compression/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-java-root
- 19.0.0-SNAPSHOT
+ 20.0.0-SNAPSHOT
arrow-compression
Arrow Compression
@@ -50,12 +50,12 @@ under the License.
org.apache.commons
commons-compress
- 1.27.1
+ 1.28.0
com.github.luben
zstd-jni
- 1.5.7-6
+ 1.5.7-11
diff --git a/compression/src/main/java/org/apache/arrow/compression/Lz4CompressionCodec.java b/compression/src/main/java/org/apache/arrow/compression/Lz4CompressionCodec.java
index 91cefc2a9e..f268e815fe 100644
--- a/compression/src/main/java/org/apache/arrow/compression/Lz4CompressionCodec.java
+++ b/compression/src/main/java/org/apache/arrow/compression/Lz4CompressionCodec.java
@@ -80,6 +80,13 @@ protected ArrowBuf doDecompress(BufferAllocator allocator, ArrowBuf compressedBu
}
byte[] outBytes = out.toByteArray();
+ if (outBytes.length != decompressedLength) {
+ throw new RuntimeException(
+ "Expected != actual decompressed length: "
+ + decompressedLength
+ + " != "
+ + outBytes.length);
+ }
ArrowBuf decompressedBuffer = allocator.buffer(outBytes.length);
decompressedBuffer.setBytes(/* index= */ 0, outBytes);
decompressedBuffer.writerIndex(decompressedLength);
diff --git a/compression/src/main/java/org/apache/arrow/compression/ZstdCompressionCodec.java b/compression/src/main/java/org/apache/arrow/compression/ZstdCompressionCodec.java
index 290723608d..ed46fe81b4 100644
--- a/compression/src/main/java/org/apache/arrow/compression/ZstdCompressionCodec.java
+++ b/compression/src/main/java/org/apache/arrow/compression/ZstdCompressionCodec.java
@@ -44,7 +44,7 @@ protected ArrowBuf doCompress(BufferAllocator allocator, ArrowBuf uncompressedBu
long bytesWritten =
Zstd.compressUnsafe(
compressedBuffer.memoryAddress() + CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH,
- dstSize,
+ maxSize,
/*src*/ uncompressedBuffer.memoryAddress(),
/* srcSize= */ uncompressedBuffer.writerIndex(),
/* level= */ this.compressionLevel);
diff --git a/compression/src/test/java/org/apache/arrow/compression/TestCompressionCodec.java b/compression/src/test/java/org/apache/arrow/compression/TestCompressionCodec.java
index b8fb4e28b9..d2d2921649 100644
--- a/compression/src/test/java/org/apache/arrow/compression/TestCompressionCodec.java
+++ b/compression/src/test/java/org/apache/arrow/compression/TestCompressionCodec.java
@@ -20,6 +20,7 @@
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import java.io.ByteArrayOutputStream;
@@ -59,6 +60,7 @@
import org.apache.arrow.vector.util.ByteArrayReadableSeekableByteChannel;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.MethodSource;
@@ -231,6 +233,26 @@ void testEmptyBuffer(int vectorLength, CompressionCodec codec) throws Exception
AutoCloseables.close(decompressedBuffers);
}
+ @Test
+ void testLz4DecompressRejectsWrongLength() {
+ byte[] data = new byte[512]; // all zeros, highly compressible
+ ArrowBuf orig = allocator.buffer(data.length);
+ orig.setBytes(0, data);
+ orig.writerIndex(data.length);
+
+ CompressionCodec codec = new Lz4CompressionCodec();
+ ArrowBuf compressed = codec.compress(allocator, orig);
+
+ // tamper with the 8-byte uncompressed-length prefix so it no longer matches
+ // the real decompressed size
+ compressed.setLong(0, 1_000_000L);
+
+ RuntimeException e =
+ assertThrows(RuntimeException.class, () -> codec.decompress(allocator, compressed));
+ assertTrue(e.getMessage().contains("decompressed length"));
+ compressed.close();
+ }
+
private static Stream codecTypes() {
return Arrays.stream(CompressionUtil.CodecType.values());
}
diff --git a/dataset/pom.xml b/dataset/pom.xml
index 66233c3970..5acc837860 100644
--- a/dataset/pom.xml
+++ b/dataset/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-java-root
- 19.0.0-SNAPSHOT
+ 20.0.0-SNAPSHOT
arrow-dataset
@@ -32,8 +32,8 @@ under the License.
../../../cpp/release-build/
- 1.15.2
- 1.12.0
+ 1.17.1
+ 1.12.1
@@ -130,7 +130,7 @@ under the License.
org.apache.orc
orc-core
- 2.2.1
+ 2.3.0
test
@@ -156,7 +156,7 @@ under the License.
commons-io
commons-io
- 2.19.0
+ 2.22.0
test
diff --git a/dev/release/bump_version.sh b/dev/release/bump_version.sh
index 458e930f98..68cafb99bd 100755
--- a/dev/release/bump_version.sh
+++ b/dev/release/bump_version.sh
@@ -37,6 +37,7 @@ if [ ! -f "${SOURCE_DIR}/.env" ]; then
exit 1
fi
. "${SOURCE_DIR}/.env"
+export GH_TOKEN
cd "${SOURCE_TOP_DIR}"
diff --git a/dev/release/release.sh b/dev/release/release.sh
index f08a618c4f..d1db7ad05a 100755
--- a/dev/release/release.sh
+++ b/dev/release/release.sh
@@ -36,6 +36,7 @@ if [ ! -f "${SOURCE_DIR}/.env" ]; then
exit 1
fi
. "${SOURCE_DIR}/.env"
+export GH_TOKEN
git_origin_url="$(git remote get-url origin)"
repository="${git_origin_url#*github.com?}"
diff --git a/dev/release/release_rc.sh b/dev/release/release_rc.sh
index ff77718b8d..0920edbe35 100755
--- a/dev/release/release_rc.sh
+++ b/dev/release/release_rc.sh
@@ -42,6 +42,7 @@ if [ ! -f "${SOURCE_DIR}/.env" ]; then
exit 1
fi
. "${SOURCE_DIR}/.env"
+export GH_TOKEN
cd "${SOURCE_TOP_DIR}"
diff --git a/docs/source/cdata.rst b/docs/source/cdata.rst
index 9643d88df3..7b2924d259 100644
--- a/docs/source/cdata.rst
+++ b/docs/source/cdata.rst
@@ -101,8 +101,8 @@ without writing JNI bindings ourselves.
1.0-SNAPSHOT
- 8
- 8
+ 17
+ 17
9.0.0
@@ -237,8 +237,8 @@ For this example, we will build a JAR with all dependencies bundled.
cpptojava
1.0-SNAPSHOT
- 8
- 8
+ 17
+ 17
9.0.0
diff --git a/docs/source/developers/building.rst b/docs/source/developers/building.rst
index f9ef7daea8..b682957714 100644
--- a/docs/source/developers/building.rst
+++ b/docs/source/developers/building.rst
@@ -32,7 +32,7 @@ Arrow Java uses the `Maven `_ build system.
Building requires:
-* JDK 11+
+* JDK 17+
* Maven 3+
.. note::
@@ -345,7 +345,7 @@ configuration file usually located under ``${HOME}/.m2`` with the following snip
jdk
- 21
+ 21
temurin
@@ -383,11 +383,11 @@ Arrow repository, and update the following settings:
right click the directory, and select Mark Directory as > Generated Sources
Root. There is no need to mark other generated sources directories, as only
the ``vector`` module generates sources.
-* For JDK 11, due to an `IntelliJ bug
- `__, you must go into
+* Due to an `IntelliJ bug
+ `__, you may need to go into
Settings > Build, Execution, Deployment > Compiler > Java Compiler and disable
"Use '--release' option for cross-compilation (Java 9 and later)". Otherwise
- you will get an error like "package sun.misc does not exist".
+ you may get an error like "package sun.misc does not exist".
* You may want to disable error-prone entirely if it gives spurious
warnings (disable both error-prone profiles in the Maven tool window
and "Reload All Maven Projects").
@@ -397,7 +397,7 @@ Arrow repository, and update the following settings:
* To enable debugging JNI-based modules like ``dataset``,
activate specific profiles in the Maven tab under "Profiles".
Ensure the profiles ``arrow-c-data``, ``arrow-jni``, ``generate-libs-cdata-all-os``,
- ``generate-libs-jni-macos-linux``, and ``jdk11+`` are enabled, so that the
+ ``generate-libs-jni-macos-linux``, and ``jdk17+`` are enabled, so that the
IDE can build them and enable debugging.
You may not need to update all of these settings if you build/test with the
@@ -478,8 +478,8 @@ Installing Manually
.. code-block:: xml
- 8
- 8
+ 17
+ 17
9.0.0.dev501
diff --git a/docs/source/flight.rst b/docs/source/flight.rst
index fabced8094..fd0fdf07bc 100644
--- a/docs/source/flight.rst
+++ b/docs/source/flight.rst
@@ -232,8 +232,8 @@ Servers can add other gRPC services. For example, to add the `Health Check servi
See the :external+arrow:ref:`best practices for C++ `.
-.. _`FlightClient`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/flight/FlightClient.html
-.. _`FlightProducer`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/flight/FlightProducer.html
-.. _`FlightServer`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/flight/FlightServer.html
-.. _`NoOpFlightProducer`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/flight/NoOpFlightProducer.html
-.. _`Location`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/flight/Location.html
+.. _`FlightClient`: https://arrow.apache.org/java/current/reference/org.apache.arrow.flight.core/org/apache/arrow/flight/FlightClient.html
+.. _`FlightProducer`: https://arrow.apache.org/java/current/reference/org.apache.arrow.flight.core/org/apache/arrow/flight/FlightProducer.html
+.. _`FlightServer`: https://arrow.apache.org/java/current/reference/org.apache.arrow.flight.core/org/apache/arrow/flight/FlightServer.html
+.. _`NoOpFlightProducer`: https://arrow.apache.org/java/current/reference/org.apache.arrow.flight.core/org/apache/arrow/flight/NoOpFlightProducer.html
+.. _`Location`: https://arrow.apache.org/java/current/reference/org.apache.arrow.flight.core/org/apache/arrow/flight/Location.html
diff --git a/docs/source/flight_sql.rst b/docs/source/flight_sql.rst
index 169a0e24bf..09ce1dda0d 100644
--- a/docs/source/flight_sql.rst
+++ b/docs/source/flight_sql.rst
@@ -29,4 +29,4 @@ over the network.
For usage information, see the `API documentation`_.
-.. _API documentation: https://arrow.apache.org/docs/java/reference/org/apache/arrow/flight/sql/package-summary.html
+.. _API documentation: https://arrow.apache.org/java/current/reference/org.apache.arrow.flight.sql/org/apache/arrow/flight/sql/package-summary.html
diff --git a/docs/source/flight_sql_jdbc_driver.rst b/docs/source/flight_sql_jdbc_driver.rst
index 1806930943..6d40434a22 100644
--- a/docs/source/flight_sql_jdbc_driver.rst
+++ b/docs/source/flight_sql_jdbc_driver.rst
@@ -27,7 +27,7 @@ Flight SQL.
Installation and Requirements
=============================
-The driver is compatible with JDK 11+. Note that the following JVM
+The driver is compatible with JDK 17+. Note that the following JVM
parameter is required:
.. code-block:: shell
@@ -173,3 +173,126 @@ DriverManager#getConnection()
`_,
the username and password supplied on the URI supercede the username and
password arguments to the function call.
+
+OAuth 2.0 Authentication
+========================
+
+The driver supports OAuth 2.0 authentication for obtaining access tokens
+from an authorization server. Two OAuth flows are currently supported:
+
+* **Client Credentials** - For service-to-service authentication where no
+ user interaction is required. The application authenticates using its own
+ credentials (client ID and client secret).
+
+* **Token Exchange** (RFC 8693) - For exchanging one token for another,
+ commonly used for federated authentication, delegation, or impersonation
+ scenarios.
+
+OAuth Connection Properties
+---------------------------
+
+The following properties configure OAuth authentication. These properties
+should be provided via the ``Properties`` object when connecting, as they
+may contain special characters that are difficult to encode in a URI.
+
+**Common OAuth Properties**
+
+.. list-table::
+ :header-rows: 1
+
+ * - Parameter
+ - Type
+ - Required
+ - Default
+ - Description
+
+ * - oauth.flow
+ - String
+ - Yes (to enable OAuth)
+ - null
+ - The OAuth grant type. Supported values: ``client_credentials``,
+ ``token_exchange``
+
+ * - oauth.tokenUri
+ - String
+ - Yes
+ - null
+ - The OAuth 2.0 token endpoint URL (e.g.,
+ ``https://auth.example.com/oauth/token``)
+
+ * - oauth.clientId
+ - String
+ - Conditional
+ - null
+ - The OAuth 2.0 client ID. Required for ``client_credentials`` flow,
+ optional for ``token_exchange``
+
+ * - oauth.clientSecret
+ - String
+ - Conditional
+ - null
+ - The OAuth 2.0 client secret. Required for ``client_credentials`` flow,
+ optional for ``token_exchange``
+
+ * - oauth.scope
+ - String
+ - No
+ - null
+ - Space-separated list of OAuth scopes to request
+
+ * - oauth.resource
+ - String
+ - No
+ - null
+ - The resource indicator for the token request (RFC 8707)
+
+**Token Exchange Properties**
+
+These properties are specific to the ``token_exchange`` flow:
+
+.. list-table::
+ :header-rows: 1
+
+ * - Parameter
+ - Type
+ - Required
+ - Default
+ - Description
+
+ * - oauth.exchange.subjectToken
+ - String
+ - Yes
+ - null
+ - The subject token to exchange (e.g., a JWT from an identity provider)
+
+ * - oauth.exchange.subjectTokenType
+ - String
+ - Yes
+ - null
+ - The token type URI of the subject token. Common values:
+ ``urn:ietf:params:oauth:token-type:access_token``,
+ ``urn:ietf:params:oauth:token-type:jwt``
+
+ * - oauth.exchange.actorToken
+ - String
+ - No
+ - null
+ - The actor token for delegation/impersonation scenarios
+
+ * - oauth.exchange.actorTokenType
+ - String
+ - No
+ - null
+ - The token type URI of the actor token
+
+ * - oauth.exchange.aud
+ - String
+ - No
+ - null
+ - The target audience for the exchanged token
+
+ * - oauth.exchange.requestedTokenType
+ - String
+ - No
+ - null
+ - The desired token type for the exchanged token
diff --git a/docs/source/install.rst b/docs/source/install.rst
index b2b1c7163f..e0b34515ef 100644
--- a/docs/source/install.rst
+++ b/docs/source/install.rst
@@ -27,8 +27,8 @@ Java modules are regularly built and tested on macOS and Linux distributions.
Java Compatibility
==================
-Java modules are compatible with JDK 11 and above. Currently, JDK versions
-11, 17, 21, and latest are tested in CI.
+Java modules are compatible with JDK 17 and above. Currently, JDK versions
+17, 21, and latest are tested in CI.
Note that some JDK internals must be exposed by
adding ``--add-opens=java.base/java.nio=org.apache.arrow.memory.core,ALL-UNNAMED`` to the ``java`` command:
diff --git a/docs/source/jdbc.rst b/docs/source/jdbc.rst
index c0477cb06d..e054127faa 100644
--- a/docs/source/jdbc.rst
+++ b/docs/source/jdbc.rst
@@ -95,7 +95,7 @@ Type Mapping
The JDBC to Arrow type mapping can be obtained at runtime from
`JdbcToArrowUtils.getArrowTypeFromJdbcType`_.
-.. _JdbcToArrowUtils.getArrowTypeFromJdbcType: https://arrow.apache.org/docs/java/reference/org/apache/arrow/adapter/jdbc/JdbcToArrowUtils.html#getArrowTypeFromJdbcType-org.apache.arrow.adapter.jdbc.JdbcFieldInfo-java.util.Calendar-
+.. _JdbcToArrowUtils.getArrowTypeFromJdbcType: https://arrow.apache.org/java/current/reference/org.apache.arrow.adapter.jdbc/org/apache/arrow/adapter/jdbc/JdbcToArrowUtils.html#getArrowTypeFromJdbcType-org.apache.arrow.adapter.jdbc.JdbcFieldInfo-java.util.Calendar-
+--------------------+--------------------+-------+
| JDBC Type | Arrow Type | Notes |
@@ -171,8 +171,8 @@ The JDBC to Arrow type mapping can be obtained at runtime from
timezone of the calendar, else it will be a timestamp without
timezone.
-.. _setArraySubTypeByColumnIndexMap: https://arrow.apache.org/docs/java/reference/org/apache/arrow/adapter/jdbc/JdbcToArrowConfigBuilder.html#setArraySubTypeByColumnIndexMap-java.util.Map-
-.. _setArraySubTypeByColumnNameMap: https://arrow.apache.org/docs/java/reference/org/apache/arrow/adapter/jdbc/JdbcToArrowConfigBuilder.html#setArraySubTypeByColumnNameMap-java.util.Map-
+.. _setArraySubTypeByColumnIndexMap: https://arrow.apache.org/java/current/reference/org.apache.arrow.adapter.jdbc/org/apache/arrow/adapter/jdbc/JdbcToArrowConfigBuilder.html#setArraySubTypeByColumnIndexMap-java.util.Map-
+.. _setArraySubTypeByColumnNameMap: https://arrow.apache.org/java/current/reference/org.apache.arrow.adapter.jdbc/org/apache/arrow/adapter/jdbc/JdbcToArrowConfigBuilder.html#setArraySubTypeByColumnNameMap-java.util.Map-
.. _ARROW-17006: https://issues.apache.org/jira/browse/ARROW-17006
VectorSchemaRoot to PreparedStatement Parameter Conversion
@@ -213,7 +213,8 @@ Type Mapping
------------
The Arrow to JDBC type mapping can be obtained at runtime via
-a method on ColumnBinder.
+a method on ColumnBinder. The Flight SQL JDBC driver follows the same
+mapping, with additional support for the UUID extension type noted below.
+----------------------------+----------------------------+-------+
| Arrow Type | JDBC Type | Notes |
@@ -232,6 +233,8 @@ a method on ColumnBinder.
+----------------------------+----------------------------+-------+
| FixedSizeBinary | BINARY (setBytes) | |
+----------------------------+----------------------------+-------+
+| Uuid (extension) | OTHER (setObject) | \(3) |
++----------------------------+----------------------------+-------+
| Float32 | REAL (setFloat) | |
+----------------------------+----------------------------+-------+
| Int8 | TINYINT (setByte) | |
@@ -273,6 +276,9 @@ a method on ColumnBinder.
JDBC value, because a JDBC Timestamp is in UTC, and we have no
timezone information. In this case, the default binder will call
`setTimestamp(int, Timestamp)
- `_,
+ `_,
which will lead to the driver using the "default timezone" (that of
the Java VM).
+* \(3) For the Flight SQL JDBC driver, the Arrow UUID extension type
+ (``arrow.uuid``) maps to JDBC ``OTHER`` and is surfaced as
+ ``java.util.UUID`` values.
diff --git a/docs/source/memory.rst b/docs/source/memory.rst
index 58ef382dc9..5b9148223a 100644
--- a/docs/source/memory.rst
+++ b/docs/source/memory.rst
@@ -333,18 +333,18 @@ How this works:
}
}
-.. _`ArrowBuf`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/memory/ArrowBuf.html
-.. _`ArrowBuf.print()`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/memory/ArrowBuf.html#print-java.lang.StringBuilder-int-org.apache.arrow.memory.BaseAllocator.Verbosity-
-.. _`BufferAllocator`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/memory/BufferAllocator.html
-.. _`BufferLedger`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/memory/BufferLedger.html
-.. _`RootAllocator`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/memory/RootAllocator.html
-.. _`newChildAllocator`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/memory/RootAllocator.html#newChildAllocator-java.lang.String-org.apache.arrow.memory.AllocationListener-long-long-
+.. _`ArrowBuf`: https://arrow.apache.org/java/current/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/ArrowBuf.html
+.. _`ArrowBuf.print()`: https://arrow.apache.org/java/current/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/ArrowBuf.html#print-java.lang.StringBuilder-int-org.apache.arrow.memory.BaseAllocator.Verbosity-
+.. _`BufferAllocator`: https://arrow.apache.org/java/current/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/BufferAllocator.html
+.. _`BufferLedger`: https://arrow.apache.org/java/current/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/BufferLedger.html
+.. _`RootAllocator`: https://arrow.apache.org/java/current/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/RootAllocator.html
+.. _`newChildAllocator`: https://arrow.apache.org/java/current/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/RootAllocator.html#newChildAllocator-java.lang.String-org.apache.arrow.memory.AllocationListener-long-long-
.. _`Netty`: https://netty.io/wiki/
.. _`sun.misc.unsafe`: https://web.archive.org/web/20210929024401/http://www.docjar.com/html/api/sun/misc/Unsafe.java.html
-.. _`Direct Memory`: https://docs.oracle.com/en/java/javase/11/docs/api/java.base/java/nio/ByteBuffer.html
-.. _`ReferenceManager`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/memory/ReferenceManager.html
-.. _`ReferenceManager.release`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/memory/ReferenceManager.html#release--
-.. _`ReferenceManager.retain`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/memory/ReferenceManager.html#retain--
+.. _`Direct Memory`: https://docs.oracle.com/en/java/javase/17/docs/api/java.base/java/nio/ByteBuffer.html
+.. _`ReferenceManager`: https://arrow.apache.org/java/current/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/ReferenceManager.html
+.. _`ReferenceManager.release`: https://arrow.apache.org/java/current/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/ReferenceManager.html#release--
+.. _`ReferenceManager.retain`: https://arrow.apache.org/java/current/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/ReferenceManager.html#retain--
Arrow Memory In-Depth
=====================
diff --git a/docs/source/overview.rst b/docs/source/overview.rst
index be579c1495..1188054114 100644
--- a/docs/source/overview.rst
+++ b/docs/source/overview.rst
@@ -45,6 +45,9 @@ but some modules are JNI bindings to the C++ library.
* - arrow-vector
- An off-heap reference implementation for Arrow columnar data format.
- Native
+ * - arrow-vector-codegen
+ - Template files for Arrow datatypes suitable for code generation.
+ - Native
* - arrow-tools
- Java applications for working with Arrow ValueVectors.
- Native
diff --git a/docs/source/table.rst b/docs/source/table.rst
index 5aa95e153c..880ef84d29 100644
--- a/docs/source/table.rst
+++ b/docs/source/table.rst
@@ -364,15 +364,15 @@ If the table contains dictionary-encoded vectors and was constructed with a ``Di
Data.exportTable(bufferAllocator, table, outArrowArray);
-.. _`ArrowBuf`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/memory/ArrowBuf.html
-.. _`Data`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/c/Data.html
-.. _`DictionaryProvider`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/dictionary/DictionaryProvider.html
-.. _`Field`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/types/pojo/Field.html
-.. _`FieldReader`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/complex/reader/FieldReader.html
-.. _`FieldVector`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/FieldVector.html
-.. _`Row`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/table/Row.html
-.. _`Schema`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/types/pojo/Schema.html
-.. _`Table`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/table/Table.html
-.. _`ValueHolder`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/holders/ValueHolder.html
-.. _`ValueVector`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/ValueVector.html
-.. _`VectorSchemaRoot`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/VectorSchemaRoot.html
+.. _`ArrowBuf`: https://arrow.apache.org/java/current/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/ArrowBuf.html
+.. _`Data`: https://arrow.apache.org/java/current/reference/org.apache.arrow.c/org/apache/arrow/c/Data.html
+.. _`DictionaryProvider`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/dictionary/DictionaryProvider.html
+.. _`Field`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/types/pojo/Field.html
+.. _`FieldReader`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/complex/reader/FieldReader.html
+.. _`FieldVector`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/FieldVector.html
+.. _`Row`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/table/Row.html
+.. _`Schema`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/types/pojo/Schema.html
+.. _`Table`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/table/Table.html
+.. _`ValueHolder`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/holders/ValueHolder.html
+.. _`ValueVector`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/ValueVector.html
+.. _`VectorSchemaRoot`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/VectorSchemaRoot.html
diff --git a/docs/source/vector_schema_root.rst b/docs/source/vector_schema_root.rst
index 3119122d9a..f4a497c4e5 100644
--- a/docs/source/vector_schema_root.rst
+++ b/docs/source/vector_schema_root.rst
@@ -153,11 +153,11 @@ A `Table`_ is an immutable tabular data structure, very similar to VectorSchemaR
See the :doc:`table` documentation for more information.
-.. _`ArrowRecordBatch`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/ipc/message/ArrowRecordBatch.html
-.. _`Field`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/types/pojo/Field.html
-.. _`Flight`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/flight/package-summary.html
-.. _`Schema`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/types/pojo/Schema.html
-.. _`Table`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/table/Table.html
-.. _`VectorLoader`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/VectorLoader.html
-.. _`VectorSchemaRoot`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/VectorSchemaRoot.html
-.. _`VectorUnloader`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/VectorUnloader.html
+.. _`ArrowRecordBatch`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/ipc/message/ArrowRecordBatch.html
+.. _`Field`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/types/pojo/Field.html
+.. _`Flight`: https://arrow.apache.org/java/current/reference/org.apache.arrow.flight.core/org/apache/arrow/flight/package-summary.html
+.. _`Schema`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/types/pojo/Schema.html
+.. _`Table`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/table/Table.html
+.. _`VectorLoader`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/VectorLoader.html
+.. _`VectorSchemaRoot`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/VectorSchemaRoot.html
+.. _`VectorUnloader`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/VectorUnloader.html
diff --git a/flight/flight-core/pom.xml b/flight/flight-core/pom.xml
index 24beac391e..9ae402cdc4 100644
--- a/flight/flight-core/pom.xml
+++ b/flight/flight-core/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-flight
- 19.0.0-SNAPSHOT
+ 20.0.0-SNAPSHOT
flight-core
@@ -134,7 +134,7 @@ under the License.
com.google.api.grpc
proto-google-common-protos
- 2.56.0
+ 2.72.0
test
diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/CallHeaders.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/CallHeaders.java
index f4f6486a3c..0939d232cf 100644
--- a/flight/flight-core/src/main/java/org/apache/arrow/flight/CallHeaders.java
+++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/CallHeaders.java
@@ -26,10 +26,20 @@ public interface CallHeaders {
/** Get the value of a metadata key. If multiple values are present, then get the last one. */
byte[] getByte(String key);
- /** Get all values present for the given metadata key. */
+ /**
+ * Get all values present for the given metadata key.
+ *
+ * @param key the metadata key
+ * @return an iterable of all values for the key. Returns an empty iterable if no value to return.
+ */
Iterable getAll(String key);
- /** Get all values present for the given metadata key. */
+ /**
+ * Get all values present for the given metadata key.
+ *
+ * @param key the metadata key
+ * @return an iterable of all values for the key. Returns an empty iterable if no value to return.
+ */
Iterable getAllByte(String key);
/**
diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/ServerSessionMiddleware.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/ServerSessionMiddleware.java
index 47fd6f1366..5ec01b9c83 100644
--- a/flight/flight-core/src/main/java/org/apache/arrow/flight/ServerSessionMiddleware.java
+++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/ServerSessionMiddleware.java
@@ -80,20 +80,18 @@ public ServerSessionMiddleware onCallStarted(
String sessionId = null;
final Iterable it = incomingHeaders.getAll("cookie");
- if (it != null) {
- findIdCookie:
- for (final String headerValue : it) {
- for (final String cookie : headerValue.split(" ;")) {
- final String[] cookiePair = cookie.split("=");
- if (cookiePair.length != 2) {
- // Soft failure: Ignore invalid cookie list field
- break;
- }
-
- if (sessionCookieName.equals(cookiePair[0]) && cookiePair[1].length() > 0) {
- sessionId = cookiePair[1];
- break findIdCookie;
- }
+ findIdCookie:
+ for (final String headerValue : it) {
+ for (final String cookie : headerValue.split(" ;")) {
+ final String[] cookiePair = cookie.split("=");
+ if (cookiePair.length != 2) {
+ // Soft failure: Ignore invalid cookie list field
+ break;
+ }
+
+ if (sessionCookieName.equals(cookiePair[0]) && cookiePair[1].length() > 0) {
+ sessionId = cookiePair[1];
+ break findIdCookie;
}
}
}
diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/client/ClientCookieMiddleware.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/client/ClientCookieMiddleware.java
index e5eb934001..b33e6b7ecc 100644
--- a/flight/flight-core/src/main/java/org/apache/arrow/flight/client/ClientCookieMiddleware.java
+++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/client/ClientCookieMiddleware.java
@@ -100,10 +100,7 @@ public void onBeforeSendingHeaders(CallHeaders outgoingHeaders) {
@Override
public void onHeadersReceived(CallHeaders incomingHeaders) {
- final Iterable setCookieHeaders = incomingHeaders.getAll(SET_COOKIE_HEADER);
- if (setCookieHeaders != null) {
- factory.updateCookies(setCookieHeaders);
- }
+ factory.updateCookies(incomingHeaders.getAll(SET_COOKIE_HEADER));
}
@Override
diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/CallCredentialAdapter.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/CallCredentialAdapter.java
index f33e9b2f94..fe81f3fb23 100644
--- a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/CallCredentialAdapter.java
+++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/CallCredentialAdapter.java
@@ -18,6 +18,7 @@
import io.grpc.CallCredentials;
import io.grpc.Metadata;
+import io.grpc.Status;
import java.util.concurrent.Executor;
import java.util.function.Consumer;
import org.apache.arrow.flight.CallHeaders;
@@ -36,9 +37,14 @@ public void applyRequestMetadata(
RequestInfo requestInfo, Executor executor, MetadataApplier metadataApplier) {
executor.execute(
() -> {
- final Metadata headers = new Metadata();
- credentialWriter.accept(new MetadataAdapter(headers));
- metadataApplier.apply(headers);
+ try {
+ final Metadata headers = new Metadata();
+ credentialWriter.accept(new MetadataAdapter(headers));
+ metadataApplier.apply(headers);
+ } catch (Throwable t) {
+ metadataApplier.fail(
+ Status.UNAUTHENTICATED.withCause(t).withDescription(t.getMessage()));
+ }
});
}
diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java
index 45c32a86c6..fcba88d212 100644
--- a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java
+++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java
@@ -87,13 +87,13 @@ public static void readIntoBuffer(
final InputStream stream, final ArrowBuf buf, final int size, final boolean fastPath)
throws IOException {
ReadableBuffer readableBuffer = fastPath ? getReadableBuffer(stream) : null;
+ byte[] heapBytes = new byte[size];
if (readableBuffer != null) {
- readableBuffer.readBytes(buf.nioBuffer(0, size));
+ readableBuffer.readBytes(heapBytes, 0, size);
} else {
- byte[] heapBytes = new byte[size];
ByteStreams.readFully(stream, heapBytes);
- buf.writeBytes(heapBytes);
}
+ buf.writeBytes(heapBytes);
buf.writerIndex(size);
}
}
diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/MetadataAdapter.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/MetadataAdapter.java
index a1de16ede6..64a0769d63 100644
--- a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/MetadataAdapter.java
+++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/MetadataAdapter.java
@@ -18,6 +18,7 @@
import io.grpc.Metadata;
import java.nio.charset.StandardCharsets;
+import java.util.Collections;
import java.util.HashSet;
import java.util.Set;
import java.util.stream.Collectors;
@@ -53,13 +54,17 @@ public byte[] getByte(String key) {
@Override
public Iterable getAll(String key) {
- return this.metadata.getAll(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER));
+ final Iterable all =
+ this.metadata.getAll(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER));
+ return all != null ? all : Collections.emptyList();
}
@Override
public Iterable getAllByte(String key) {
if (key.endsWith(Metadata.BINARY_HEADER_SUFFIX)) {
- return this.metadata.getAll(Metadata.Key.of(key, Metadata.BINARY_BYTE_MARSHALLER));
+ final Iterable all =
+ this.metadata.getAll(Metadata.Key.of(key, Metadata.BINARY_BYTE_MARSHALLER));
+ return all != null ? all : Collections.emptyList();
}
return StreamSupport.stream(getAll(key).spliterator(), false)
.map(String::getBytes)
diff --git a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestCallOptions.java b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestCallOptions.java
index a54ce69812..8aef9c69a1 100644
--- a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestCallOptions.java
+++ b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestCallOptions.java
@@ -21,6 +21,7 @@
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assertions.fail;
@@ -110,6 +111,16 @@ public void mixedProperties() {
testHeaders(headers);
}
+ @Test
+ public void getAllReturnsEmptyIterableForMissingKey() {
+ FlightCallHeaders headers = new FlightCallHeaders();
+
+ assertNotNull(headers.getAll("missing"));
+ assertFalse(headers.getAll("missing").iterator().hasNext());
+ assertNotNull(headers.getAllByte("missing-bin"));
+ assertFalse(headers.getAllByte("missing-bin").iterator().hasNext());
+ }
+
private void testHeaders(CallHeaders headers) {
try (BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
HeaderProducer producer = new HeaderProducer();
diff --git a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestErrorMetadata.java b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestErrorMetadata.java
index a9a3e355bc..214614defd 100644
--- a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestErrorMetadata.java
+++ b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestErrorMetadata.java
@@ -20,6 +20,7 @@
import static org.apache.arrow.flight.Location.forGrpcInsecure;
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assertions.fail;
@@ -119,6 +120,16 @@ public void testFlightMetadata() throws Exception {
}
}
+ @Test
+ public void getAllReturnsEmptyIterableForMissingKey() {
+ ErrorFlightMetadata metadata = new ErrorFlightMetadata();
+
+ assertNotNull(metadata.getAll("missing"));
+ assertFalse(metadata.getAll("missing").iterator().hasNext());
+ assertNotNull(metadata.getAllByte("missing-bin"));
+ assertFalse(metadata.getAllByte("missing-bin").iterator().hasNext());
+ }
+
private static class StatusRuntimeExceptionProducer extends NoOpFlightProducer {
private final PerfOuterClass.Perf perf;
diff --git a/flight/flight-core/src/test/java/org/apache/arrow/flight/auth/TestBasicAuth.java b/flight/flight-core/src/test/java/org/apache/arrow/flight/auth/TestBasicAuth.java
index 0c63785c88..0f202ba2d9 100644
--- a/flight/flight-core/src/test/java/org/apache/arrow/flight/auth/TestBasicAuth.java
+++ b/flight/flight-core/src/test/java/org/apache/arrow/flight/auth/TestBasicAuth.java
@@ -178,6 +178,12 @@ public static void shutdown() throws Exception {
AutoCloseables.close(server);
allocator.getChildAllocators().forEach(BufferAllocator::close);
+
+ // gRPC/Netty may still be releasing Arrow buffers asynchronously after server shutdown.
+ // Poll briefly to allow in-flight buffer releases to complete before closing the allocator.
+ for (int i = 0; i < 20 && allocator.getAllocatedMemory() > 0; i++) {
+ Thread.sleep(100);
+ }
AutoCloseables.close(allocator);
}
}
diff --git a/flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/TestMetadataAdapter.java b/flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/TestMetadataAdapter.java
new file mode 100644
index 0000000000..b0f5dcfcfc
--- /dev/null
+++ b/flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/TestMetadataAdapter.java
@@ -0,0 +1,36 @@
+/*
+ * 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.flight.grpc;
+
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+
+import io.grpc.Metadata;
+import org.junit.jupiter.api.Test;
+
+public class TestMetadataAdapter {
+
+ @Test
+ public void getAllReturnsEmptyIterableForMissingKey() {
+ MetadataAdapter headers = new MetadataAdapter(new Metadata());
+
+ assertNotNull(headers.getAll("missing"));
+ assertFalse(headers.getAll("missing").iterator().hasNext());
+ assertNotNull(headers.getAllByte("missing-bin"));
+ assertFalse(headers.getAllByte("missing-bin").iterator().hasNext());
+ }
+}
diff --git a/flight/flight-integration-tests/pom.xml b/flight/flight-integration-tests/pom.xml
index 78a2d08ee1..f6ae8e16a5 100644
--- a/flight/flight-integration-tests/pom.xml
+++ b/flight/flight-integration-tests/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-flight
- 19.0.0-SNAPSHOT
+ 20.0.0-SNAPSHOT
flight-integration-tests
@@ -58,7 +58,7 @@ under the License.
commons-cli
commons-cli
- 1.9.0
+ 1.11.0
org.slf4j
@@ -101,7 +101,7 @@ under the License.
-
+
META-INF/LICENSE.txt
src/shade/LICENSE.txt
diff --git a/flight/flight-sql-jdbc-core/pom.xml b/flight/flight-sql-jdbc-core/pom.xml
index 965e071e72..be2ee32868 100644
--- a/flight/flight-sql-jdbc-core/pom.xml
+++ b/flight/flight-sql-jdbc-core/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-flight
- 19.0.0-SNAPSHOT
+ 20.0.0-SNAPSHOT
flight-sql-jdbc-core
@@ -105,7 +105,7 @@ under the License.
commons-io
commons-io
- 2.19.0
+ 2.22.0
test
@@ -120,6 +120,31 @@ under the License.
test
+
+ com.squareup.okhttp3
+ mockwebserver3
+ 5.4.0
+ test
+
+
+ com.squareup.okhttp3
+ mockwebserver3-junit5
+ 5.4.0
+ test
+
+
+ com.squareup.okhttp3
+ okhttp-jvm
+ 5.4.0
+ test
+
+
+ com.squareup.okio
+ okio-jvm
+ 3.17.0
+ test
+
+
io.netty
netty-common
@@ -134,13 +159,13 @@ under the License.
org.apache.calcite.avatica
avatica
- 1.26.0
+ 1.27.0
org.bouncycastle
bcpkix-jdk18on
- 1.82
+ 1.84
@@ -151,8 +176,15 @@ under the License.
com.github.ben-manes.caffeine
caffeine
- 3.2.0
+ 3.2.4
+
+
+ com.nimbusds
+ oauth2-oidc-sdk
+ 11.37.2
+
+
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..0110525fea 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
@@ -45,6 +45,7 @@
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.Arrays;
+import java.util.Collections;
import java.util.EnumMap;
import java.util.HashMap;
import java.util.List;
@@ -75,18 +76,23 @@
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;
import org.apache.calcite.avatica.AvaticaConnection;
import org.apache.calcite.avatica.AvaticaDatabaseMetaData;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
/** Arrow Flight JDBC's implementation of {@link DatabaseMetaData}. */
public class ArrowDatabaseMetadata extends AvaticaDatabaseMetaData {
+ private static final Logger LOGGER = LoggerFactory.getLogger(ArrowDatabaseMetadata.class);
private static final String JAVA_REGEX_SPECIALS = "[]()|^-+*?{}$\\.";
private static final Charset CHARSET = StandardCharsets.UTF_8;
private static final byte[] EMPTY_BYTE_ARRAY = new byte[0];
@@ -164,6 +170,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) {
@@ -769,7 +778,34 @@ private T getSqlInfoAndCacheIfCacheIsEmpty(
}
}
}
- return desiredType.cast(cachedSqlInfo.get(sqlInfoCommand));
+ T value = desiredType.cast(cachedSqlInfo.get(sqlInfoCommand));
+ if (value != null) {
+ return value;
+ }
+ LOGGER.debug(
+ "SqlInfo {} not provided by server, returning default for type {}",
+ sqlInfoCommand.name(),
+ desiredType.getSimpleName());
+
+ // Return sensible defaults when SqlInfo is unavailable
+ if (desiredType == Long.class) {
+ return desiredType.cast(0L);
+ } else if (desiredType == Integer.class) {
+ return desiredType.cast(0);
+ } else if (desiredType == Boolean.class) {
+ return desiredType.cast(false);
+ } else if (desiredType == String.class) {
+ return desiredType.cast("");
+ } else if (desiredType == Map.class) {
+ return desiredType.cast(Collections.emptyMap());
+ } else if (desiredType == List.class) {
+ return desiredType.cast(Collections.emptyList());
+ }
+
+ throw new SQLException(
+ String.format(
+ "The value of the SqlInfo %s is null and it could not be cast to %s.",
+ sqlInfoCommand.name(), desiredType.getName()));
}
private Optional convertListSqlInfoToString(final List sqlInfoList) {
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java
index f6f17770f1..623c2b81be 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java
@@ -20,6 +20,9 @@
import io.netty.util.concurrent.DefaultThreadFactory;
import java.sql.SQLException;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
import java.util.Properties;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
@@ -41,6 +44,8 @@ public final class ArrowFlightConnection extends AvaticaConnection {
private final ArrowFlightSqlClientHandler clientHandler;
private final ArrowFlightConnectionConfigImpl config;
private ExecutorService executorService;
+ private int metadataResultSetCount;
+ private Map metadataResultSetMap = new HashMap<>();
/**
* Creates a new {@link ArrowFlightConnection}.
@@ -65,6 +70,7 @@ private ArrowFlightConnection(
this.config = Preconditions.checkNotNull(config, "Config cannot be null.");
this.allocator = Preconditions.checkNotNull(allocator, "Allocator cannot be null.");
this.clientHandler = Preconditions.checkNotNull(clientHandler, "Handler cannot be null.");
+ this.metadataResultSetCount = 0;
}
/**
@@ -121,6 +127,7 @@ private static ArrowFlightSqlClientHandler createNewClientHandler(
.withClientCache(config.useClientCache() ? new FlightClientCache() : null)
.withConnectTimeout(config.getConnectTimeout())
.withDriverVersion(driverVersion)
+ .withOAuthConfiguration(config.getOauthConfiguration())
.build();
} catch (final SQLException e) {
try {
@@ -171,6 +178,31 @@ synchronized ExecutorService getExecutorService() {
: executorService;
}
+ /**
+ * Registers a new metadata ResultSet and assigns it a unique ID. Metadata ResultSets are those
+ * created without an associated Statement.
+ *
+ * @param resultSet the ResultSet to register
+ * @return the assigned ID
+ */
+ int getNewMetadataResultSetId(ArrowFlightJdbcFlightStreamResultSet resultSet) {
+ metadataResultSetMap.put(metadataResultSetCount, resultSet);
+ return metadataResultSetCount++;
+ }
+
+ /**
+ * Unregisters a metadata ResultSet when it is closed. This method is called by metadata
+ * ResultSets during their close operation to remove themselves from the tracking map.
+ *
+ * @param id the ID of the ResultSet to unregister, or null if not a metadata ResultSet
+ */
+ void onResultSetClose(Integer id) {
+ if (id == null) {
+ return;
+ }
+ metadataResultSetMap.remove(id);
+ }
+
@Override
public Properties getClientInfo() {
final Properties copy = new Properties();
@@ -180,19 +212,41 @@ public Properties getClientInfo() {
@Override
public void close() throws SQLException {
- clientHandler.close();
- if (executorService != null) {
- executorService.shutdown();
+ Exception topLevelException = null;
+ try {
+ if (executorService != null) {
+ executorService.shutdown();
+ }
+ } catch (final Exception e) {
+ topLevelException = e;
+ }
+ // copies of the collections are used to avoid concurrent modification problems
+ ArrayList closeables = new ArrayList<>(statementMap.values());
+ closeables.addAll(new ArrayList<>(metadataResultSetMap.values()));
+ closeables.add(clientHandler);
+ closeables.addAll(allocator.getChildAllocators());
+ closeables.add(allocator);
+ try {
+ AutoCloseables.close(closeables);
+ } catch (final Exception e) {
+ if (topLevelException == null) {
+ topLevelException = e;
+ } else {
+ topLevelException.addSuppressed(e);
+ }
}
-
try {
- AutoCloseables.close(clientHandler);
- allocator.getChildAllocators().forEach(AutoCloseables::closeNoChecked);
- AutoCloseables.close(allocator);
-
super.close();
} catch (final Exception e) {
- throw AvaticaConnection.HELPER.createException(e.getMessage(), e);
+ if (topLevelException == null) {
+ topLevelException = e;
+ } else {
+ topLevelException.addSuppressed(e);
+ }
+ }
+ if (topLevelException != null) {
+ throw AvaticaConnection.HELPER.createException(
+ topLevelException.getMessage(), topLevelException);
}
}
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcArray.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcArray.java
index 9b9eba51e5..f3d76ace92 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcArray.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcArray.java
@@ -26,6 +26,7 @@
import org.apache.arrow.driver.jdbc.utils.SqlTypes;
import org.apache.arrow.memory.util.LargeMemoryUtil;
import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.ValueVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.types.pojo.ArrowType;
@@ -135,12 +136,22 @@ public ResultSet getResultSet(long index, int count) throws SQLException {
private static ResultSet getResultSetNoBoundariesCheck(
ValueVector dataVector, long start, long count) throws SQLException {
+ int intStart = LargeMemoryUtil.checkedCastToInt(start);
+ int intCount = LargeMemoryUtil.checkedCastToInt(count);
+
+ // Create an index vector with 1-based indices (per JDBC spec) to return with value vector
+ IntVector indexVector = new IntVector("INDEX", dataVector.getAllocator());
+ indexVector.allocateNew(intCount);
+ for (int i = 0; i < intCount; i++) {
+ indexVector.set(i, i + 1);
+ }
+ indexVector.setValueCount(intCount);
+
TransferPair transferPair = dataVector.getTransferPair(dataVector.getAllocator());
- transferPair.splitAndTransfer(
- LargeMemoryUtil.checkedCastToInt(start), LargeMemoryUtil.checkedCastToInt(count));
- FieldVector vectorSlice = (FieldVector) transferPair.getTo();
+ transferPair.splitAndTransfer(intStart, intCount);
+ FieldVector valueVector = (FieldVector) transferPair.getTo();
- VectorSchemaRoot vectorSchemaRoot = VectorSchemaRoot.of(vectorSlice);
+ VectorSchemaRoot vectorSchemaRoot = VectorSchemaRoot.of(indexVector, valueVector);
return ArrowFlightJdbcVectorSchemaRootResultSet.fromVectorSchemaRoot(vectorSchemaRoot);
}
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriver.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriver.java
index 53e6120f62..12ef8030d7 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriver.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriver.java
@@ -75,7 +75,9 @@ public Logger getParentLogger() {
public ArrowFlightConnection connect(final String url, final Properties info)
throws SQLException {
final Properties properties = new Properties(info);
- properties.putAll(info);
+ if (info != null) {
+ properties.putAll(info);
+ }
if (url != null) {
final Optional> maybeProperties = getUrlsArgs(url);
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java
index aabaf01e63..376e5b11e7 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java
@@ -54,6 +54,7 @@ public final class ArrowFlightJdbcFlightStreamResultSet
private VectorSchemaRoot currentVectorSchemaRoot;
private Schema schema;
+ private Integer id = null; // used for metadata result sets only
/** Public constructor used by ArrowFlightJdbcFactory. */
ArrowFlightJdbcFlightStreamResultSet(
@@ -82,6 +83,7 @@ private ArrowFlightJdbcFlightStreamResultSet(
super(null, state, signature, resultSetMetaData, timeZone, firstFrame);
this.connection = connection;
this.flightInfo = flightInfo;
+ this.id = connection.getNewMetadataResultSetId(this);
}
/**
@@ -104,7 +106,7 @@ static ArrowFlightJdbcFlightStreamResultSet fromFlightInfo(
final TimeZone timeZone = TimeZone.getDefault();
final QueryState state = new QueryState();
- final Meta.Signature signature = ArrowFlightMetaImpl.newSignature(null, null, null);
+ final Meta.Signature signature = ArrowFlightMetaImpl.newSignature(null, null, null, null);
final AvaticaResultSetMetaData resultSetMetaData =
new AvaticaResultSetMetaData(null, null, signature);
@@ -234,7 +236,12 @@ protected void cancel() {
@Override
public synchronized void close() {
+
try {
+ if (isClosed()) {
+ return;
+ }
+ this.connection.onResultSetClose(id);
if (flightEndpointDataQueue != null) {
// flightStreamQueue should close currentFlightStream internally
flightEndpointDataQueue.close();
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcVectorSchemaRootResultSet.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcVectorSchemaRootResultSet.java
index 622e5fe7f6..ad6670a001 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcVectorSchemaRootResultSet.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcVectorSchemaRootResultSet.java
@@ -22,7 +22,6 @@
import java.sql.Types;
import java.util.HashSet;
import java.util.List;
-import java.util.Objects;
import java.util.Set;
import java.util.TimeZone;
import org.apache.arrow.driver.jdbc.utils.ConvertUtils;
@@ -74,7 +73,7 @@ public static ArrowFlightJdbcVectorSchemaRootResultSet fromVectorSchemaRoot(
final TimeZone timeZone = TimeZone.getDefault();
final QueryState state = new QueryState();
- final Meta.Signature signature = ArrowFlightMetaImpl.newSignature(null, null, null);
+ final Meta.Signature signature = ArrowFlightMetaImpl.newSignature(null, null, null, null);
final AvaticaResultSetMetaData resultSetMetaData =
new AvaticaResultSetMetaData(null, null, signature);
@@ -129,7 +128,7 @@ public Object getObject(int columnIndex) throws SQLException {
if (metaData.type.id == Types.TIMESTAMP_WITH_TIMEZONE) {
return accessor.getTimestamp(localCalendar);
} else {
- return AvaticaSite.get(accessor, metaData.type.id, localCalendar);
+ return AvaticaSite.get(accessor, metaData.type.id, true, localCalendar);
}
}
@@ -159,12 +158,10 @@ public void close() {
} catch (final Exception e) {
exceptions.add(e);
}
- if (!Objects.isNull(statement)) {
- try {
- super.close();
- } catch (final Exception e) {
- exceptions.add(e);
- }
+ try {
+ super.close();
+ } catch (final Exception e) {
+ exceptions.add(e);
}
exceptions.parallelStream().forEach(e -> LOGGER.error(e.getMessage(), e));
exceptions.stream()
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java
index 21cc3e431f..0d85b5eddb 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java
@@ -53,7 +53,8 @@ public ArrowFlightMetaImpl(final AvaticaConnection connection) {
}
/** Construct a signature. */
- static Signature newSignature(final String sql, Schema resultSetSchema, Schema parameterSchema) {
+ static Signature newSignature(
+ final String sql, Schema resultSetSchema, Schema parameterSchema, Boolean isUpdate) {
List columnMetaData =
resultSetSchema == null
? new ArrayList<>()
@@ -62,10 +63,17 @@ static Signature newSignature(final String sql, Schema resultSetSchema, Schema p
parameterSchema == null
? new ArrayList<>()
: ConvertUtils.convertArrowFieldsToAvaticaParameters(parameterSchema.getFields());
- StatementType statementType =
- resultSetSchema == null || resultSetSchema.getFields().isEmpty()
- ? StatementType.IS_DML
- : StatementType.SELECT;
+ // If the server provided the is_update field, use it to determine the statement type
+ StatementType statementType;
+ if (isUpdate != null) {
+ statementType = isUpdate ? StatementType.IS_DML : StatementType.SELECT;
+ } else {
+ // Fall back to the legacy logic: check if the result set schema is empty
+ statementType =
+ resultSetSchema == null || resultSetSchema.getFields().isEmpty()
+ ? StatementType.IS_DML
+ : StatementType.SELECT;
+ }
return new Signature(
columnMetaData,
sql,
@@ -79,7 +87,8 @@ static Signature newSignature(final String sql, Schema resultSetSchema, Schema p
public void closeStatement(final StatementHandle statementHandle) {
PreparedStatement preparedStatement =
statementHandlePreparedStatementMap.remove(new StatementHandleKey(statementHandle));
- // Testing if the prepared statement was created because the statement can be not created until
+ // Testing if the prepared statement was created because the statement can be
+ // not created until
// this moment
if (preparedStatement != null) {
preparedStatement.close();
@@ -177,7 +186,10 @@ private PreparedStatement prepareForHandle(final String query, StatementHandle h
((ArrowFlightConnection) connection).getClientHandler().prepare(query);
handle.signature =
newSignature(
- query, preparedStatement.getDataSetSchema(), preparedStatement.getParameterSchema());
+ query,
+ preparedStatement.getDataSetSchema(),
+ preparedStatement.getParameterSchema(),
+ preparedStatement.isUpdate());
statementHandlePreparedStatementMap.put(new StatementHandleKey(handle), preparedStatement);
return preparedStatement;
}
@@ -224,7 +236,8 @@ public ExecuteResult prepareAndExecute(
MetaResultSet.create(handle.connectionId, handle.id, false, handle.signature, null);
return new ExecuteResult(Collections.singletonList(metaResultSet));
} catch (SQLTimeoutException e) {
- // So far AvaticaStatement(executeInternal) only handles NoSuchStatement and Runtime
+ // So far AvaticaStatement(executeInternal) only handles NoSuchStatement and
+ // Runtime
// Exceptions.
throw new RuntimeException(e);
} catch (SQLException e) {
@@ -253,6 +266,20 @@ public boolean syncResults(
return false;
}
+ @Override
+ public ConnectionProperties connectionSync(ConnectionHandle ch, ConnectionProperties connProps) {
+ final ConnectionProperties result = super.connectionSync(ch, connProps);
+ final String newCatalog = this.connProps.getCatalog();
+ if (newCatalog != null) {
+ try {
+ ((ArrowFlightConnection) connection).getClientHandler().setCatalog(newCatalog);
+ } catch (SQLException e) {
+ throw new RuntimeException(e);
+ }
+ }
+ return result;
+ }
+
void setDefaultConnectionProperties() {
// TODO Double-check this.
connProps
@@ -268,7 +295,8 @@ PreparedStatement getPreparedStatement(StatementHandle statementHandle) {
return statementHandlePreparedStatementMap.get(new StatementHandleKey(statementHandle));
}
- // Helper used to look up prepared statement instances later. Avatica doesn't give us the
+ // Helper used to look up prepared statement instances later. Avatica doesn't
+ // give us the
// signature in
// an UPDATE code path so we can't directly use StatementHandle as a map key.
private static final class StatementHandleKey {
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessor.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessor.java
index f0fa55fa82..cd762fb1ac 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessor.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessor.java
@@ -36,6 +36,10 @@
import java.util.Calendar;
import java.util.Map;
import java.util.function.IntSupplier;
+import org.joou.UByte;
+import org.joou.UInteger;
+import org.joou.ULong;
+import org.joou.UShort;
/** Base Jdbc Accessor. */
public abstract class ArrowFlightJdbcAccessor implements Accessor {
@@ -99,6 +103,26 @@ public long getLong() throws SQLException {
throw getOperationNotSupported(this.getClass());
}
+ @Override
+ public UByte getUByte() throws SQLException {
+ throw getOperationNotSupported(this.getClass());
+ }
+
+ @Override
+ public UShort getUShort() throws SQLException {
+ throw getOperationNotSupported(this.getClass());
+ }
+
+ @Override
+ public UInteger getUInt() throws SQLException {
+ throw getOperationNotSupported(this.getClass());
+ }
+
+ @Override
+ public ULong getULong() throws SQLException {
+ throw getOperationNotSupported(this.getClass());
+ }
+
@Override
public float getFloat() throws SQLException {
throw getOperationNotSupported(this.getClass());
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:
+ *
+ *
+ * - {@link #getObject()} returns {@link java.util.UUID}
+ *
- {@link #getString()} returns the hyphenated string format (e.g.,
+ * "550e8400-e29b-41d4-a716-446655440000")
+ *
- {@link #getBytes()} returns the 16-byte binary representation
+ *
+ */
+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/client/ArrowFlightSqlClientHandler.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java
index 5dc7e0e2e9..719cc38a2b 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java
@@ -32,6 +32,9 @@
import java.util.Map;
import java.util.Optional;
import java.util.Set;
+import org.apache.arrow.driver.jdbc.client.oauth.OAuthConfiguration;
+import org.apache.arrow.driver.jdbc.client.oauth.OAuthCredentialWriter;
+import org.apache.arrow.driver.jdbc.client.oauth.OAuthTokenProvider;
import org.apache.arrow.driver.jdbc.client.utils.ClientAuthenticationUtils;
import org.apache.arrow.driver.jdbc.client.utils.FlightClientCache;
import org.apache.arrow.driver.jdbc.client.utils.FlightLocationQueue;
@@ -47,7 +50,6 @@
import org.apache.arrow.flight.FlightStatusCode;
import org.apache.arrow.flight.Location;
import org.apache.arrow.flight.LocationSchemes;
-import org.apache.arrow.flight.SessionOptionValue;
import org.apache.arrow.flight.SessionOptionValueFactory;
import org.apache.arrow.flight.SetSessionOptionsRequest;
import org.apache.arrow.flight.SetSessionOptionsResult;
@@ -147,20 +149,26 @@ public List getStreams(final FlightInfo flightInfo)
try {
for (FlightEndpoint endpoint : flightInfo.getEndpoints()) {
if (endpoint.getLocations().isEmpty()) {
- // Create a stream using the current client only and do not close the client at the end.
+ // Create a stream using the current client only and do not close the client at
+ // the end.
endpoints.add(
new CloseableEndpointStreamPair(
sqlClient.getStream(endpoint.getTicket(), getOptions()), null));
} else {
// Clone the builder and then set the new endpoint on it.
- // GH-38574: Currently a new FlightClient will be made for each partition that returns a
- // non-empty Location then disposed of. It may be better to cache clients because a server
- // may report the same Locations. It would also be good to identify when the reported
+ // GH-38574: Currently a new FlightClient will be made for each partition that
+ // returns a
+ // non-empty Location then disposed of. It may be better to cache clients
+ // because a server
+ // may report the same Locations. It would also be good to identify when the
+ // reported
// location
- // is the same as the original connection's Location and skip creating a FlightClient in
+ // is the same as the original connection's Location and skip creating a
+ // FlightClient in
// that scenario.
- // Also copy the cache to the client so we can share a cache. Cache needs to cache
+ // Also copy the cache to the client so we can share a cache. Cache needs to
+ // cache
// negative attempts too.
List exceptions = new ArrayList<>();
CloseableEndpointStreamPair stream = null;
@@ -337,7 +345,8 @@ private boolean isBenignCloseException(FlightRuntimeException fre) {
*/
private void logSuppressedCloseException(
FlightRuntimeException fre, String operationDescription) {
- // ARROW-17785 and GH-863: suppress exceptions caused by flaky gRPC layer during shutdown
+ // ARROW-17785 and GH-863: suppress exceptions caused by flaky gRPC layer during
+ // shutdown
LOGGER.debug("Suppressed error {}", operationDescription, fre);
}
@@ -379,6 +388,14 @@ public interface PreparedStatement extends AutoCloseable {
*/
Schema getParameterSchema();
+ /**
+ * Gets whether this {@link PreparedStatement} is an update statement.
+ *
+ * @return {@code true} if this is an update statement, {@code false} if it's a query, or {@code
+ * null} if the server did not provide this information.
+ */
+ @Nullable Boolean isUpdate();
+
void setParameters(VectorSchemaRoot parameters);
@Override
@@ -388,25 +405,40 @@ public interface PreparedStatement extends AutoCloseable {
/** A connection is created with catalog set as a session option. */
private void setSetCatalogInSessionIfPresent() {
if (catalog.isPresent()) {
- final SetSessionOptionsRequest setSessionOptionRequest =
- new SetSessionOptionsRequest(
- ImmutableMap.builder()
- .put(CATALOG, SessionOptionValueFactory.makeSessionOptionValue(catalog.get()))
- .build());
- final SetSessionOptionsResult result =
- sqlClient.setSessionOptions(setSessionOptionRequest, getOptions());
+ try {
+ setCatalog(catalog.get());
+ } catch (SQLException e) {
+ throw CallStatus.INVALID_ARGUMENT
+ .withDescription(e.getMessage())
+ .withCause(e)
+ .toRuntimeException();
+ }
+ }
+ }
+ /**
+ * Sets the catalog for the current session.
+ *
+ * @param catalog the catalog to set.
+ * @throws SQLException if an error occurs while setting the catalog.
+ */
+ public void setCatalog(final String catalog) throws SQLException {
+ final SetSessionOptionsRequest request =
+ new SetSessionOptionsRequest(
+ ImmutableMap.of(CATALOG, SessionOptionValueFactory.makeSessionOptionValue(catalog)));
+ try {
+ final SetSessionOptionsResult result = sqlClient.setSessionOptions(request, getOptions());
if (result.hasErrors()) {
- Map errors = result.getErrors();
- for (Map.Entry error : errors.entrySet()) {
+ final Map errors = result.getErrors();
+ for (final Map.Entry error : errors.entrySet()) {
LOGGER.warn(error.toString());
}
- throw CallStatus.INVALID_ARGUMENT
- .withDescription(
- String.format(
- "Cannot set session option for catalog = %s. Check log for details.", catalog))
- .toRuntimeException();
+ throw new SQLException(
+ String.format(
+ "Cannot set session option for catalog = %s. Check log for details.", catalog));
}
+ } catch (final FlightRuntimeException e) {
+ throw new SQLException(e);
}
}
@@ -432,6 +464,12 @@ public long executeUpdate() {
@Override
public StatementType getType() {
+ // If the server provided the is_update field, use it to determine the statement type
+ final Boolean isUpdate = preparedStatement.isUpdate();
+ if (isUpdate != null) {
+ return isUpdate ? StatementType.UPDATE : StatementType.SELECT;
+ }
+ // Fall back to the legacy logic: check if the result set schema is empty
final Schema schema = preparedStatement.getResultSetSchema();
return schema.getFields().isEmpty() ? StatementType.UPDATE : StatementType.SELECT;
}
@@ -451,6 +489,11 @@ public void setParameters(VectorSchemaRoot parameters) {
preparedStatement.setParameters(parameters);
}
+ @Override
+ public Boolean isUpdate() {
+ return preparedStatement.isUpdate();
+ }
+
@Override
public void close() {
try {
@@ -654,7 +697,10 @@ public static final class Builder {
@VisibleForTesting @Nullable Duration connectTimeout;
- // These two middleware are for internal use within build() and should not be exposed by builder
+ @VisibleForTesting @Nullable OAuthConfiguration oauthConfig;
+
+ // These two middleware are for internal use within build() and should not be
+ // exposed by builder
// APIs.
// Note that these middleware may not necessarily be registered.
@VisibleForTesting
@@ -692,6 +738,7 @@ public Builder() {}
this.clientKeyPath = original.clientKeyPath;
this.allocator = original.allocator;
this.catalog = original.catalog;
+ this.oauthConfig = original.oauthConfig;
if (original.retainCookies) {
this.cookieFactory = original.cookieFactory;
@@ -961,6 +1008,17 @@ public Builder withDriverVersion(DriverVersion driverVersion) {
return this;
}
+ /**
+ * Sets the OAuth configuration for this handler.
+ *
+ * @param oauthConfig the OAuth configuration
+ * @return this builder instance
+ */
+ public Builder withOAuthConfiguration(final OAuthConfiguration oauthConfig) {
+ this.oauthConfig = oauthConfig;
+ return this;
+ }
+
public String getCacheKey() {
return getLocation().toString();
}
@@ -980,7 +1038,8 @@ public Location getLocation() {
* @throws SQLException on error.
*/
public ArrowFlightSqlClientHandler build() throws SQLException {
- // Copy middleware so that the build method doesn't change the state of the builder fields
+ // Copy middleware so that the build method doesn't change the state of the
+ // builder fields
// itself.
Set buildTimeMiddlewareFactories =
new HashSet<>(this.middlewareFactories);
@@ -988,7 +1047,8 @@ public ArrowFlightSqlClientHandler build() throws SQLException {
boolean isUsingUserPasswordAuth = username != null && token == null;
try {
- // Token should take priority since some apps pass in a username/password even when a token
+ // Token should take priority since some apps pass in a username/password even
+ // when a token
// is provided
if (isUsingUserPasswordAuth) {
buildTimeMiddlewareFactories.add(authFactory);
@@ -1046,9 +1106,15 @@ public ArrowFlightSqlClientHandler build() throws SQLException {
FlightGrpcUtils.createFlightClient(
allocator, channelBuilder.build(), clientBuilder.middleware());
final ArrayList credentialOptions = new ArrayList<>();
- if (isUsingUserPasswordAuth) {
- // If the authFactory has already been used for a handshake, use the existing token.
- // This can occur if the authFactory is being re-used for a new connection spawned for
+ // Authentication priority: OAuth > token > username/password
+ if (oauthConfig != null) {
+ OAuthTokenProvider tokenProvider = oauthConfig.createTokenProvider();
+ credentialOptions.add(new CredentialCallOption(new OAuthCredentialWriter(tokenProvider)));
+ } else if (isUsingUserPasswordAuth) {
+ // If the authFactory has already been used for a handshake, use the existing
+ // token.
+ // This can occur if the authFactory is being re-used for a new connection
+ // spawned for
// getStream().
if (authFactory.getCredentialCallOption() != null) {
credentialOptions.add(authFactory.getCredentialCallOption());
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/AbstractOAuthTokenProvider.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/AbstractOAuthTokenProvider.java
new file mode 100644
index 0000000000..9c377a5850
--- /dev/null
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/AbstractOAuthTokenProvider.java
@@ -0,0 +1,108 @@
+/*
+ * 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.client.oauth;
+
+import com.nimbusds.oauth2.sdk.ParseException;
+import com.nimbusds.oauth2.sdk.Scope;
+import com.nimbusds.oauth2.sdk.TokenErrorResponse;
+import com.nimbusds.oauth2.sdk.TokenRequest;
+import com.nimbusds.oauth2.sdk.TokenResponse;
+import com.nimbusds.oauth2.sdk.auth.ClientAuthentication;
+import com.nimbusds.oauth2.sdk.token.AccessToken;
+import java.io.IOException;
+import java.net.URI;
+import java.sql.SQLException;
+import java.time.Instant;
+import org.apache.arrow.util.VisibleForTesting;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/**
+ * Abstract base class for OAuth token providers that handles token caching, refresh logic, and
+ * common request/response handling.
+ */
+public abstract class AbstractOAuthTokenProvider implements OAuthTokenProvider {
+ protected static final int EXPIRATION_BUFFER_SECONDS = 30;
+ protected static final int DEFAULT_EXPIRATION_SECONDS = 3600;
+
+ private final Object tokenLock = new Object();
+ private volatile @Nullable TokenInfo cachedToken;
+
+ @VisibleForTesting URI tokenUri;
+
+ @VisibleForTesting @Nullable ClientAuthentication clientAuth;
+
+ @VisibleForTesting @Nullable Scope scope;
+
+ @Override
+ public String getValidToken() throws SQLException {
+ TokenInfo token = cachedToken;
+ if (token != null && !token.isExpired(EXPIRATION_BUFFER_SECONDS)) {
+ return token.getAccessToken();
+ }
+
+ synchronized (tokenLock) {
+ token = cachedToken;
+ if (token != null && !token.isExpired(EXPIRATION_BUFFER_SECONDS)) {
+ return token.getAccessToken();
+ }
+ cachedToken = fetchNewToken();
+ return cachedToken.getAccessToken();
+ }
+ }
+
+ /**
+ * Fetches a new token from the authorization server. This method handles the common
+ * request/response logic while delegating flow-specific request building to subclasses.
+ *
+ * @return the new token information
+ * @throws SQLException if token cannot be obtained
+ */
+ protected TokenInfo fetchNewToken() throws SQLException {
+ try {
+ TokenRequest request = buildTokenRequest();
+ TokenResponse response = TokenResponse.parse(request.toHTTPRequest().send());
+
+ if (!response.indicatesSuccess()) {
+ TokenErrorResponse errorResponse = response.toErrorResponse();
+ String errorMsg =
+ String.format(
+ "OAuth request failed: %s - %s",
+ errorResponse.getErrorObject().getCode(),
+ errorResponse.getErrorObject().getDescription());
+ throw new SQLException(errorMsg);
+ }
+
+ AccessToken accessToken = response.toSuccessResponse().getTokens().getAccessToken();
+ long expiresIn =
+ accessToken.getLifetime() > 0 ? accessToken.getLifetime() : DEFAULT_EXPIRATION_SECONDS;
+ Instant expiresAt = Instant.now().plusSeconds(expiresIn);
+
+ return new TokenInfo(accessToken.getValue(), expiresAt);
+ } catch (ParseException e) {
+ throw new SQLException("Failed to parse OAuth token response", e);
+ } catch (IOException e) {
+ throw new SQLException("Failed to send OAuth token request", e);
+ }
+ }
+
+ /**
+ * Builds the flow-specific token request.
+ *
+ * @return the token request to send to the authorization server
+ */
+ protected abstract TokenRequest buildTokenRequest();
+}
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/ClientCredentialsTokenProvider.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/ClientCredentialsTokenProvider.java
new file mode 100644
index 0000000000..7e6289819c
--- /dev/null
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/ClientCredentialsTokenProvider.java
@@ -0,0 +1,58 @@
+/*
+ * 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.client.oauth;
+
+import com.nimbusds.oauth2.sdk.ClientCredentialsGrant;
+import com.nimbusds.oauth2.sdk.Scope;
+import com.nimbusds.oauth2.sdk.TokenRequest;
+import com.nimbusds.oauth2.sdk.auth.ClientSecretBasic;
+import com.nimbusds.oauth2.sdk.auth.Secret;
+import com.nimbusds.oauth2.sdk.id.ClientID;
+import java.net.URI;
+import java.util.Objects;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/**
+ * OAuth 2.0 Client Credentials flow token provider (RFC 6749 Section 4.4).
+ *
+ * This provider handles service-to-service authentication where no user interaction is required.
+ * Tokens are cached and automatically refreshed before expiration.
+ */
+public class ClientCredentialsTokenProvider extends AbstractOAuthTokenProvider {
+
+ /**
+ * Creates a new ClientCredentialsTokenProvider.
+ *
+ * @param tokenUri the OAuth token endpoint URI
+ * @param clientId the OAuth client ID
+ * @param clientSecret the OAuth client secret
+ * @param scope optional OAuth scopes (space-separated)
+ */
+ ClientCredentialsTokenProvider(
+ URI tokenUri, String clientId, String clientSecret, @Nullable String scope) {
+ this.tokenUri = Objects.requireNonNull(tokenUri, "tokenUri cannot be null");
+ Objects.requireNonNull(clientId, "clientId cannot be null");
+ Objects.requireNonNull(clientSecret, "clientSecret cannot be null");
+ this.clientAuth = new ClientSecretBasic(new ClientID(clientId), new Secret(clientSecret));
+ this.scope = (scope != null && !scope.isEmpty()) ? Scope.parse(scope) : null;
+ }
+
+ @Override
+ protected TokenRequest buildTokenRequest() {
+ return new TokenRequest(tokenUri, clientAuth, new ClientCredentialsGrant(), scope);
+ }
+}
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthConfiguration.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthConfiguration.java
new file mode 100644
index 0000000000..cba9d4c2e6
--- /dev/null
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthConfiguration.java
@@ -0,0 +1,240 @@
+/*
+ * 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.client.oauth;
+
+import com.nimbusds.oauth2.sdk.GrantType;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.sql.SQLException;
+import java.util.Locale;
+import java.util.Objects;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Configuration class for OAuth settings parsed from connection properties. */
+public class OAuthConfiguration {
+
+ private final GrantType grantType;
+ private final URI tokenUri;
+ private final @Nullable String clientId;
+ private final @Nullable String clientSecret;
+ private final @Nullable String scope;
+ private final @Nullable String subjectToken;
+ private final @Nullable String subjectTokenType;
+ private final @Nullable String actorToken;
+ private final @Nullable String actorTokenType;
+ private final @Nullable String audience;
+ private final @Nullable String resource;
+ private final @Nullable String requestedTokenType;
+
+ private OAuthConfiguration(Builder builder) throws SQLException {
+ this.grantType = builder.grantType;
+ this.tokenUri = builder.tokenUri;
+ this.clientId = builder.clientId;
+ this.clientSecret = builder.clientSecret;
+ this.scope = builder.scope;
+ this.subjectToken = builder.subjectToken;
+ this.subjectTokenType = builder.subjectTokenType;
+ this.actorToken = builder.actorToken;
+ this.actorTokenType = builder.actorTokenType;
+ this.audience = builder.audience;
+ this.resource = builder.resource;
+ this.requestedTokenType = builder.requestedTokenType;
+
+ validate();
+ }
+
+ private void validate() throws SQLException {
+ Objects.requireNonNull(grantType, "OAuth grant type is required");
+ Objects.requireNonNull(tokenUri, "Token URI is required");
+
+ if (GrantType.CLIENT_CREDENTIALS.equals(grantType)) {
+ if (clientId == null || clientId.isEmpty()) {
+ throw new SQLException("clientId is required for client_credentials flow");
+ }
+ if (clientSecret == null || clientSecret.isEmpty()) {
+ throw new SQLException("clientSecret is required for client_credentials flow");
+ }
+ } else if (GrantType.TOKEN_EXCHANGE.equals(grantType)) {
+ if (subjectToken == null || subjectToken.isEmpty()) {
+ throw new SQLException("subjectToken is required for token_exchange flow");
+ }
+ if (subjectTokenType == null || subjectTokenType.isEmpty()) {
+ throw new SQLException("subjectTokenType is required for token_exchange flow");
+ }
+ } else {
+ throw new SQLException("Unsupported OAuth grant type: " + grantType);
+ }
+ }
+
+ /**
+ * Creates an OAuthTokenProvider based on the configured grant type.
+ *
+ * @return the token provider
+ * @throws SQLException if the grant type is not supported or configuration is invalid
+ */
+ public OAuthTokenProvider createTokenProvider() throws SQLException {
+ if (GrantType.CLIENT_CREDENTIALS.equals(grantType)) {
+ return OAuthTokenProviders.clientCredentials()
+ .tokenUri(tokenUri)
+ .clientId(clientId)
+ .clientSecret(clientSecret)
+ .scope(scope)
+ .build();
+ } else if (GrantType.TOKEN_EXCHANGE.equals(grantType)) {
+ OAuthTokenProviders.TokenExchangeBuilder builder =
+ OAuthTokenProviders.tokenExchange()
+ .tokenUri(tokenUri)
+ .subjectToken(subjectToken)
+ .subjectTokenType(subjectTokenType)
+ .actorToken(actorToken)
+ .actorTokenType(actorTokenType)
+ .audience(audience)
+ .requestedTokenType(requestedTokenType)
+ .scope(scope)
+ .resource(resource);
+
+ if (clientId != null && clientSecret != null) {
+ builder.clientCredentials(clientId, clientSecret);
+ }
+
+ return builder.build();
+ } else {
+ throw new SQLException("Unsupported OAuth grant type: " + grantType);
+ }
+ }
+
+ /** Builder for OAuthConfiguration. */
+ public static class Builder {
+ private GrantType grantType;
+ private URI tokenUri;
+ private @Nullable String clientId;
+ private @Nullable String clientSecret;
+ private @Nullable String scope;
+ private @Nullable String subjectToken;
+ private @Nullable String subjectTokenType;
+ private @Nullable String actorToken;
+ private @Nullable String actorTokenType;
+ private @Nullable String audience;
+ private @Nullable String resource;
+ private @Nullable String requestedTokenType;
+
+ /**
+ * Sets the OAuth grant type from a string value.
+ *
+ *
Accepts either user-friendly names ("client_credentials", "token_exchange") or the full
+ * URN format as defined in RFC 6749 and RFC 8693.
+ *
+ * @param flowStr the flow type string (e.g., "client_credentials", "token_exchange")
+ * @return this builder
+ * @throws SQLException if the flow string is invalid
+ */
+ public Builder flow(String flowStr) throws SQLException {
+ if (flowStr == null || flowStr.isEmpty()) {
+ throw new SQLException("OAuth flow cannot be null or empty");
+ }
+ try {
+ String normalized = flowStr.toLowerCase(Locale.ROOT);
+ // Map user-friendly names to URN format for token_exchange
+ if ("token_exchange".equals(normalized)) {
+ normalized = GrantType.TOKEN_EXCHANGE.getValue();
+ }
+ GrantType parsed = GrantType.parse(normalized);
+ if (!parsed.equals(GrantType.CLIENT_CREDENTIALS)
+ && !parsed.equals(GrantType.TOKEN_EXCHANGE)) {
+ throw new SQLException("Unsupported OAuth flow: " + flowStr);
+ }
+ this.grantType = parsed;
+ } catch (com.nimbusds.oauth2.sdk.ParseException e) {
+ throw new SQLException("Invalid OAuth flow: " + flowStr, e);
+ }
+ return this;
+ }
+
+ /**
+ * Sets the token URI.
+ *
+ * @param tokenUri the OAuth token endpoint URI
+ * @return this builder
+ * @throws SQLException if the URI is invalid
+ */
+ public Builder tokenUri(String tokenUri) throws SQLException {
+ if (tokenUri == null || tokenUri.isEmpty()) {
+ throw new SQLException("Token URI cannot be null or empty");
+ }
+ try {
+ this.tokenUri = new URI(tokenUri);
+ } catch (URISyntaxException e) {
+ throw new SQLException("Invalid token URI: " + tokenUri, e);
+ }
+ return this;
+ }
+
+ public Builder clientId(@Nullable String clientId) {
+ this.clientId = clientId;
+ return this;
+ }
+
+ public Builder clientSecret(@Nullable String clientSecret) {
+ this.clientSecret = clientSecret;
+ return this;
+ }
+
+ public Builder scope(@Nullable String scope) {
+ this.scope = scope;
+ return this;
+ }
+
+ public Builder subjectToken(@Nullable String subjectToken) {
+ this.subjectToken = subjectToken;
+ return this;
+ }
+
+ public Builder subjectTokenType(@Nullable String subjectTokenType) {
+ this.subjectTokenType = subjectTokenType;
+ return this;
+ }
+
+ public Builder actorToken(@Nullable String actorToken) {
+ this.actorToken = actorToken;
+ return this;
+ }
+
+ public Builder actorTokenType(@Nullable String actorTokenType) {
+ this.actorTokenType = actorTokenType;
+ return this;
+ }
+
+ public Builder audience(@Nullable String audience) {
+ this.audience = audience;
+ return this;
+ }
+
+ public Builder resource(@Nullable String resource) {
+ this.resource = resource;
+ return this;
+ }
+
+ public Builder requestedTokenType(@Nullable String requestedTokenType) {
+ this.requestedTokenType = requestedTokenType;
+ return this;
+ }
+
+ public OAuthConfiguration build() throws SQLException {
+ return new OAuthConfiguration(this);
+ }
+ }
+}
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthCredentialWriter.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthCredentialWriter.java
new file mode 100644
index 0000000000..0d4ad4689f
--- /dev/null
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthCredentialWriter.java
@@ -0,0 +1,42 @@
+/*
+ * 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.client.oauth;
+
+import java.sql.SQLException;
+import java.util.Objects;
+import java.util.function.Consumer;
+import org.apache.arrow.flight.CallHeaders;
+import org.apache.arrow.flight.auth2.Auth2Constants;
+
+/** Writes OAuth bearer tokens to Flight call headers. */
+public class OAuthCredentialWriter implements Consumer {
+ private final OAuthTokenProvider tokenProvider;
+
+ public OAuthCredentialWriter(OAuthTokenProvider tokenProvider) {
+ this.tokenProvider = Objects.requireNonNull(tokenProvider, "tokenProvider cannot be null");
+ }
+
+ @Override
+ public void accept(CallHeaders headers) {
+ try {
+ String token = tokenProvider.getValidToken();
+ headers.insert(Auth2Constants.AUTHORIZATION_HEADER, Auth2Constants.BEARER_PREFIX + token);
+ } catch (SQLException e) {
+ throw new OAuthTokenException("Failed to obtain OAuth token", e);
+ }
+ }
+}
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthTokenException.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthTokenException.java
new file mode 100644
index 0000000000..aceadb327b
--- /dev/null
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthTokenException.java
@@ -0,0 +1,31 @@
+/*
+ * 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.client.oauth;
+
+/**
+ * Runtime exception thrown when OAuth token operations fail. Used to wrap checked exceptions in
+ * contexts that don't allow them.
+ */
+public class OAuthTokenException extends RuntimeException {
+ public OAuthTokenException(String message) {
+ super(message);
+ }
+
+ public OAuthTokenException(String message, Throwable cause) {
+ super(message, cause);
+ }
+}
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthTokenProvider.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthTokenProvider.java
new file mode 100644
index 0000000000..241611e432
--- /dev/null
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthTokenProvider.java
@@ -0,0 +1,33 @@
+/*
+ * 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.client.oauth;
+
+import java.sql.SQLException;
+
+/**
+ * Interface for OAuth token providers that handle token acquisition and refresh. Implementations
+ * should cache tokens and automatically refresh them before expiration.
+ */
+public interface OAuthTokenProvider {
+ /**
+ * Gets a valid OAuth access token, refreshing if necessary.
+ *
+ * @return a valid access token string
+ * @throws SQLException if token cannot be obtained
+ */
+ String getValidToken() throws SQLException;
+}
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthTokenProviders.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthTokenProviders.java
new file mode 100644
index 0000000000..bbf7072d39
--- /dev/null
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthTokenProviders.java
@@ -0,0 +1,419 @@
+/*
+ * 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.client.oauth;
+
+import com.nimbusds.oauth2.sdk.ParseException;
+import com.nimbusds.oauth2.sdk.Scope;
+import com.nimbusds.oauth2.sdk.auth.ClientAuthentication;
+import com.nimbusds.oauth2.sdk.auth.ClientSecretBasic;
+import com.nimbusds.oauth2.sdk.auth.Secret;
+import com.nimbusds.oauth2.sdk.id.Audience;
+import com.nimbusds.oauth2.sdk.id.ClientID;
+import com.nimbusds.oauth2.sdk.token.TokenTypeURI;
+import com.nimbusds.oauth2.sdk.token.TypelessAccessToken;
+import com.nimbusds.oauth2.sdk.tokenexchange.TokenExchangeGrant;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.Collections;
+import java.util.List;
+import java.util.Objects;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/**
+ * Unified factory for creating OAuth token providers.
+ *
+ * This class provides a single entry point for creating all OAuth token providers with a
+ * consistent builder API. It supports:
+ *
+ *
+ * - Client Credentials flow (RFC 6749 Section 4.4)
+ *
- Token Exchange flow (RFC 8693)
+ *
+ *
+ * Example usage:
+ *
+ *
{@code
+ * // Client Credentials flow
+ * OAuthTokenProvider provider = OAuthTokenProviders.clientCredentials()
+ * .tokenUri("https://auth.example.com/token")
+ * .clientId("my-client")
+ * .clientSecret("my-secret")
+ * .scope("read write")
+ * .build();
+ *
+ * // Token Exchange flow
+ * OAuthTokenProvider provider = OAuthTokenProviders.tokenExchange()
+ * .tokenUri("https://auth.example.com/token")
+ * .subjectToken("user-token")
+ * .subjectTokenType("urn:ietf:params:oauth:token-type:access_token")
+ * .build();
+ * }
+ */
+public final class OAuthTokenProviders {
+
+ private OAuthTokenProviders() {}
+
+ /**
+ * Creates a new builder for Client Credentials flow.
+ *
+ * @return a new ClientCredentialsBuilder instance
+ */
+ public static ClientCredentialsBuilder clientCredentials() {
+ return new ClientCredentialsBuilder();
+ }
+
+ /**
+ * Creates a new builder for Token Exchange flow.
+ *
+ * @return a new TokenExchangeBuilder instance
+ */
+ public static TokenExchangeBuilder tokenExchange() {
+ return new TokenExchangeBuilder();
+ }
+
+ /** Builder for creating {@link ClientCredentialsTokenProvider} instances. */
+ public static class ClientCredentialsBuilder {
+ private @Nullable URI tokenUri;
+ private @Nullable String clientId;
+ private @Nullable String clientSecret;
+ private @Nullable String scope;
+
+ ClientCredentialsBuilder() {}
+
+ /**
+ * Sets the OAuth token endpoint URI (required).
+ *
+ * @param tokenUri the token endpoint URI
+ * @return this builder
+ */
+ public ClientCredentialsBuilder tokenUri(URI tokenUri) {
+ this.tokenUri = Objects.requireNonNull(tokenUri, "tokenUri cannot be null");
+ return this;
+ }
+
+ /**
+ * Sets the OAuth token endpoint URI from a string (required).
+ *
+ * @param tokenUri the token endpoint URI string
+ * @return this builder
+ * @throws IllegalArgumentException if the URI is invalid
+ */
+ public ClientCredentialsBuilder tokenUri(String tokenUri) {
+ Objects.requireNonNull(tokenUri, "tokenUri cannot be null");
+ try {
+ this.tokenUri = new URI(tokenUri);
+ } catch (URISyntaxException e) {
+ throw new IllegalArgumentException("Invalid token URI: " + tokenUri, e);
+ }
+ return this;
+ }
+
+ /**
+ * Sets the OAuth client ID (required).
+ *
+ * @param clientId the client ID
+ * @return this builder
+ */
+ public ClientCredentialsBuilder clientId(String clientId) {
+ this.clientId = Objects.requireNonNull(clientId, "clientId cannot be null");
+ return this;
+ }
+
+ /**
+ * Sets the OAuth client secret (required).
+ *
+ * @param clientSecret the client secret
+ * @return this builder
+ */
+ public ClientCredentialsBuilder clientSecret(String clientSecret) {
+ this.clientSecret = Objects.requireNonNull(clientSecret, "clientSecret cannot be null");
+ return this;
+ }
+
+ /**
+ * Sets the OAuth scopes (optional).
+ *
+ * @param scope the space-separated scope string
+ * @return this builder
+ */
+ public ClientCredentialsBuilder scope(@Nullable String scope) {
+ this.scope = scope;
+ return this;
+ }
+
+ /**
+ * Builds a new ClientCredentialsTokenProvider instance.
+ *
+ * @return the configured ClientCredentialsTokenProvider
+ * @throws IllegalStateException if required parameters are missing
+ */
+ public ClientCredentialsTokenProvider build() {
+ if (tokenUri == null) {
+ throw new IllegalStateException("tokenUri is required");
+ }
+ if (clientId == null) {
+ throw new IllegalStateException("clientId is required");
+ }
+ if (clientSecret == null) {
+ throw new IllegalStateException("clientSecret is required");
+ }
+ return new ClientCredentialsTokenProvider(tokenUri, clientId, clientSecret, scope);
+ }
+ }
+
+ /** Builder for creating {@link TokenExchangeTokenProvider} instances. */
+ public static class TokenExchangeBuilder {
+ private @Nullable URI tokenUri;
+ private @Nullable String subjectToken;
+ private @Nullable String subjectTokenType;
+ private @Nullable String actorToken;
+ private @Nullable String actorTokenType;
+ private @Nullable String audience;
+ private @Nullable String requestedTokenType;
+ private @Nullable Scope scope;
+ private @Nullable List resources;
+ private @Nullable ClientAuthentication clientAuth;
+
+ TokenExchangeBuilder() {}
+
+ /**
+ * Sets the OAuth token endpoint URI (required).
+ *
+ * @param tokenUri the token endpoint URI
+ * @return this builder
+ */
+ public TokenExchangeBuilder tokenUri(URI tokenUri) {
+ this.tokenUri = Objects.requireNonNull(tokenUri, "tokenUri cannot be null");
+ return this;
+ }
+
+ /**
+ * Sets the OAuth token endpoint URI from a string (required).
+ *
+ * @param tokenUri the token endpoint URI string
+ * @return this builder
+ * @throws IllegalArgumentException if the URI is invalid
+ */
+ public TokenExchangeBuilder tokenUri(String tokenUri) {
+ Objects.requireNonNull(tokenUri, "tokenUri cannot be null");
+ try {
+ this.tokenUri = new URI(tokenUri);
+ } catch (URISyntaxException e) {
+ throw new IllegalArgumentException("Invalid token URI: " + tokenUri, e);
+ }
+ return this;
+ }
+
+ /**
+ * Sets the subject token to exchange (required).
+ *
+ * @param subjectToken the subject token value
+ * @return this builder
+ */
+ public TokenExchangeBuilder subjectToken(String subjectToken) {
+ this.subjectToken = Objects.requireNonNull(subjectToken, "subjectToken cannot be null");
+ return this;
+ }
+
+ /**
+ * Sets the type of the subject token (required).
+ *
+ * @param subjectTokenType the subject token type URI
+ * @return this builder
+ */
+ public TokenExchangeBuilder subjectTokenType(String subjectTokenType) {
+ this.subjectTokenType =
+ Objects.requireNonNull(subjectTokenType, "subjectTokenType cannot be null");
+ return this;
+ }
+
+ /**
+ * Sets the optional actor token for delegation scenarios.
+ *
+ * @param actorToken the actor token value
+ * @return this builder
+ */
+ public TokenExchangeBuilder actorToken(@Nullable String actorToken) {
+ this.actorToken = actorToken;
+ return this;
+ }
+
+ /**
+ * Sets the type of the actor token.
+ *
+ * @param actorTokenType the actor token type URI
+ * @return this builder
+ */
+ public TokenExchangeBuilder actorTokenType(@Nullable String actorTokenType) {
+ this.actorTokenType = actorTokenType;
+ return this;
+ }
+
+ /**
+ * Sets the target audience for the exchanged token.
+ *
+ * @param audience the target audience
+ * @return this builder
+ */
+ public TokenExchangeBuilder audience(@Nullable String audience) {
+ this.audience = audience;
+ return this;
+ }
+
+ /**
+ * Sets the requested token type for the exchanged token.
+ *
+ * @param requestedTokenType the requested token type URI
+ * @return this builder
+ */
+ public TokenExchangeBuilder requestedTokenType(@Nullable String requestedTokenType) {
+ this.requestedTokenType = requestedTokenType;
+ return this;
+ }
+
+ /**
+ * Sets the OAuth scopes for the token request.
+ *
+ * @param scope the OAuth scope object
+ * @return this builder
+ */
+ public TokenExchangeBuilder scope(@Nullable Scope scope) {
+ this.scope = scope;
+ return this;
+ }
+
+ /**
+ * Sets the OAuth scopes from a space-separated string.
+ *
+ * @param scope the space-separated scope string
+ * @return this builder
+ */
+ public TokenExchangeBuilder scope(@Nullable String scope) {
+ this.scope = (scope != null && !scope.isEmpty()) ? Scope.parse(scope) : null;
+ return this;
+ }
+
+ /**
+ * Sets the target resource URIs (RFC 8707).
+ *
+ * @param resources the list of resource URIs
+ * @return this builder
+ */
+ public TokenExchangeBuilder resources(@Nullable List resources) {
+ this.resources = resources;
+ return this;
+ }
+
+ /**
+ * Sets a single target resource URI (RFC 8707).
+ *
+ * @param resource the resource URI
+ * @return this builder
+ */
+ public TokenExchangeBuilder resource(@Nullable URI resource) {
+ this.resources = resource != null ? Collections.singletonList(resource) : null;
+ return this;
+ }
+
+ /**
+ * Sets a single target resource URI from a string (RFC 8707).
+ *
+ * @param resource the resource URI string
+ * @return this builder
+ */
+ public TokenExchangeBuilder resource(@Nullable String resource) {
+ if (resource != null && !resource.isEmpty()) {
+ this.resources = Collections.singletonList(URI.create(resource));
+ } else {
+ this.resources = null;
+ }
+ return this;
+ }
+
+ /**
+ * Sets the client authentication.
+ *
+ * @param clientAuth the client authentication object
+ * @return this builder
+ */
+ public TokenExchangeBuilder clientAuthentication(@Nullable ClientAuthentication clientAuth) {
+ this.clientAuth = clientAuth;
+ return this;
+ }
+
+ /**
+ * Sets client authentication using client ID and secret.
+ *
+ * @param clientId the client ID
+ * @param clientSecret the client secret
+ * @return this builder
+ */
+ public TokenExchangeBuilder clientCredentials(String clientId, String clientSecret) {
+ Objects.requireNonNull(clientId, "clientId cannot be null");
+ Objects.requireNonNull(clientSecret, "clientSecret cannot be null");
+ this.clientAuth = new ClientSecretBasic(new ClientID(clientId), new Secret(clientSecret));
+ return this;
+ }
+
+ /**
+ * Builds a new TokenExchangeTokenProvider instance.
+ *
+ * @return the configured TokenExchangeTokenProvider
+ * @throws IllegalStateException if required parameters are missing
+ */
+ public TokenExchangeTokenProvider build() {
+ if (tokenUri == null) {
+ throw new IllegalStateException("tokenUri is required");
+ }
+ if (subjectToken == null) {
+ throw new IllegalStateException("subjectToken is required");
+ }
+ if (subjectTokenType == null) {
+ throw new IllegalStateException("subjectTokenType is required");
+ }
+
+ TokenExchangeGrant grant = createGrant();
+ return new TokenExchangeTokenProvider(tokenUri, grant, clientAuth, scope, resources);
+ }
+
+ private TokenExchangeGrant createGrant() {
+ try {
+ TypelessAccessToken subjectAccessToken = new TypelessAccessToken(subjectToken);
+ TokenTypeURI subjectTypeUri = TokenTypeURI.parse(subjectTokenType);
+
+ TypelessAccessToken actorAccessToken =
+ actorToken != null ? new TypelessAccessToken(actorToken) : null;
+ TokenTypeURI actorTypeUri =
+ actorTokenType != null ? TokenTypeURI.parse(actorTokenType) : null;
+ TokenTypeURI requestedTypeUri =
+ requestedTokenType != null ? TokenTypeURI.parse(requestedTokenType) : null;
+ List audienceList =
+ audience != null ? Collections.singletonList(new Audience(audience)) : null;
+
+ return new TokenExchangeGrant(
+ subjectAccessToken,
+ subjectTypeUri,
+ actorAccessToken,
+ actorTypeUri,
+ requestedTypeUri,
+ audienceList);
+ } catch (ParseException e) {
+ throw new IllegalStateException("Failed to create TokenExchangeGrant", e);
+ }
+ }
+ }
+}
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/TokenExchangeTokenProvider.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/TokenExchangeTokenProvider.java
new file mode 100644
index 0000000000..af433a2712
--- /dev/null
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/TokenExchangeTokenProvider.java
@@ -0,0 +1,81 @@
+/*
+ * 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.client.oauth;
+
+import com.nimbusds.oauth2.sdk.Scope;
+import com.nimbusds.oauth2.sdk.TokenRequest;
+import com.nimbusds.oauth2.sdk.auth.ClientAuthentication;
+import com.nimbusds.oauth2.sdk.tokenexchange.TokenExchangeGrant;
+import java.net.URI;
+import java.util.List;
+import java.util.Objects;
+import org.apache.arrow.util.VisibleForTesting;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/**
+ * OAuth 2.0 Token Exchange flow token provider (RFC 8693).
+ *
+ * This provider exchanges one token for another, commonly used for federated authentication,
+ * delegation, or impersonation scenarios. Tokens are cached and automatically refreshed.
+ */
+public class TokenExchangeTokenProvider extends AbstractOAuthTokenProvider {
+
+ @VisibleForTesting TokenExchangeGrant grant;
+
+ @VisibleForTesting @Nullable List resources;
+
+ /**
+ * Creates a new TokenExchangeTokenProvider with full configuration.
+ *
+ * @param tokenUri the OAuth token endpoint URI
+ * @param grant the token exchange grant containing subject/actor token information
+ * @param clientAuth optional client authentication
+ * @param scope optional OAuth scopes
+ * @param resource optional target resource URI (RFC 8707)
+ */
+ TokenExchangeTokenProvider(
+ URI tokenUri,
+ TokenExchangeGrant grant,
+ @Nullable ClientAuthentication clientAuth,
+ @Nullable Scope scope,
+ @Nullable List resource) {
+ this.tokenUri = Objects.requireNonNull(tokenUri, "tokenUri cannot be null");
+ this.grant = Objects.requireNonNull(grant, "grant cannot be null");
+ this.scope = scope;
+ this.resources = resource;
+ this.clientAuth = clientAuth;
+ }
+
+ @Override
+ protected TokenRequest buildTokenRequest() {
+ TokenRequest.Builder builder;
+ if (clientAuth != null) {
+ builder = new TokenRequest.Builder(tokenUri, clientAuth, grant);
+ } else {
+ builder = new TokenRequest.Builder(tokenUri, grant);
+ }
+
+ if (scope != null) {
+ builder.scope(scope);
+ }
+ if (resources != null) {
+ builder.resources(resources.toArray(new URI[0]));
+ }
+
+ return builder.build();
+ }
+}
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/TokenInfo.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/TokenInfo.java
new file mode 100644
index 0000000000..f47cc8b053
--- /dev/null
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/TokenInfo.java
@@ -0,0 +1,45 @@
+/*
+ * 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.client.oauth;
+
+import java.time.Instant;
+import java.util.Objects;
+
+/** Holds OAuth token information including the access token and expiration time. */
+public class TokenInfo {
+ private final String accessToken;
+ private final Instant expiresAt;
+
+ public TokenInfo(String accessToken, Instant expiresAt) {
+ this.accessToken = Objects.requireNonNull(accessToken, "accessToken cannot be null");
+ this.expiresAt = Objects.requireNonNull(expiresAt, "expiresAt cannot be null");
+ }
+
+ public String getAccessToken() {
+ return accessToken;
+ }
+
+ /**
+ * Checks if the token is expired or will expire within the buffer period.
+ *
+ * @param bufferSeconds seconds before actual expiration to consider token expired
+ * @return true if token should be refreshed
+ */
+ public boolean isExpired(int bufferSeconds) {
+ return Instant.now().plusSeconds(bufferSeconds).isAfter(expiresAt);
+ }
+}
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/ArrowFlightConnectionConfigImpl.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ArrowFlightConnectionConfigImpl.java
index 76ba964a53..d0ba74dbcc 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ArrowFlightConnectionConfigImpl.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ArrowFlightConnectionConfigImpl.java
@@ -16,6 +16,7 @@
*/
package org.apache.arrow.driver.jdbc.utils;
+import java.sql.SQLException;
import java.time.Duration;
import java.util.Arrays;
import java.util.HashMap;
@@ -23,6 +24,7 @@
import java.util.Objects;
import java.util.Properties;
import org.apache.arrow.driver.jdbc.ArrowFlightConnection;
+import org.apache.arrow.driver.jdbc.client.oauth.OAuthConfiguration;
import org.apache.arrow.flight.CallHeaders;
import org.apache.arrow.flight.CallOption;
import org.apache.arrow.flight.FlightCallHeaders;
@@ -31,6 +33,7 @@
import org.apache.calcite.avatica.ConnectionConfig;
import org.apache.calcite.avatica.ConnectionConfigImpl;
import org.apache.calcite.avatica.ConnectionProperty;
+import org.checkerframework.checker.nullness.qual.Nullable;
/** A {@link ConnectionConfig} for the {@link ArrowFlightConnection}. */
public final class ArrowFlightConnectionConfigImpl extends ConnectionConfigImpl {
@@ -211,6 +214,38 @@ public Map getHeaderAttributes() {
return headers;
}
+ /**
+ * Returns OAuth configuration if oauth.flow is specified, null otherwise.
+ *
+ * @return the OAuth configuration or null
+ * @throws SQLException if the OAuth configuration is invalid
+ */
+ public @Nullable OAuthConfiguration getOauthConfiguration() throws SQLException {
+ String flow = ArrowFlightConnectionProperty.OAUTH_FLOW.getString(properties);
+ if (flow == null) {
+ return null;
+ }
+
+ return new OAuthConfiguration.Builder()
+ .flow(flow)
+ .clientId(ArrowFlightConnectionProperty.OAUTH_CLIENT_ID.getString(properties))
+ .clientSecret(ArrowFlightConnectionProperty.OAUTH_CLIENT_SECRET.getString(properties))
+ .tokenUri(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.getString(properties))
+ .scope(ArrowFlightConnectionProperty.OAUTH_SCOPE.getString(properties))
+ .resource(ArrowFlightConnectionProperty.OAUTH_RESOURCE.getString(properties))
+ .subjectToken(
+ ArrowFlightConnectionProperty.OAUTH_EXCHANGE_SUBJECT_TOKEN.getString(properties))
+ .subjectTokenType(
+ ArrowFlightConnectionProperty.OAUTH_EXCHANGE_SUBJECT_TOKEN_TYPE.getString(properties))
+ .actorToken(ArrowFlightConnectionProperty.OAUTH_EXCHANGE_ACTOR_TOKEN.getString(properties))
+ .actorTokenType(
+ ArrowFlightConnectionProperty.OAUTH_EXCHANGE_ACTOR_TOKEN_TYPE.getString(properties))
+ .audience(ArrowFlightConnectionProperty.OAUTH_EXCHANGE_AUDIENCE.getString(properties))
+ .requestedTokenType(
+ ArrowFlightConnectionProperty.OAUTH_EXCHANGE_REQUESTED_TOKEN_TYPE.getString(properties))
+ .build();
+ }
+
/** Custom {@link ConnectionProperty} for the {@link ArrowFlightConnectionConfigImpl}. */
public enum ArrowFlightConnectionProperty implements ConnectionProperty {
HOST("host", null, Type.STRING, true),
@@ -232,6 +267,23 @@ public enum ArrowFlightConnectionProperty implements ConnectionProperty {
CATALOG("catalog", null, Type.STRING, false),
CONNECT_TIMEOUT_MILLIS("connectTimeoutMs", 10000, Type.NUMBER, false),
USE_CLIENT_CACHE("useClientCache", true, Type.BOOLEAN, false),
+
+ // OAuth configuration properties
+ OAUTH_FLOW("oauth.flow", null, Type.STRING, false),
+ OAUTH_CLIENT_ID("oauth.clientId", null, Type.STRING, false),
+ OAUTH_CLIENT_SECRET("oauth.clientSecret", null, Type.STRING, false),
+ OAUTH_TOKEN_URI("oauth.tokenUri", null, Type.STRING, false),
+ OAUTH_SCOPE("oauth.scope", null, Type.STRING, false),
+ OAUTH_RESOURCE("oauth.resource", null, Type.STRING, false),
+
+ // Token exchange specific properties
+ OAUTH_EXCHANGE_SUBJECT_TOKEN("oauth.exchange.subjectToken", null, Type.STRING, false),
+ OAUTH_EXCHANGE_SUBJECT_TOKEN_TYPE("oauth.exchange.subjectTokenType", null, Type.STRING, false),
+ OAUTH_EXCHANGE_ACTOR_TOKEN("oauth.exchange.actorToken", null, Type.STRING, false),
+ OAUTH_EXCHANGE_ACTOR_TOKEN_TYPE("oauth.exchange.actorTokenType", null, Type.STRING, false),
+ OAUTH_EXCHANGE_AUDIENCE("oauth.exchange.aud", null, Type.STRING, false),
+ OAUTH_EXCHANGE_REQUESTED_TOKEN_TYPE(
+ "oauth.exchange.requestedTokenType", null, Type.STRING, false),
;
private final String camelName;
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..8f40d6698e 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
@@ -41,10 +41,14 @@
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.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/ArrowDatabaseMetadataTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadataTest.java
index 81579cc387..3ab1460b27 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadataTest.java
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadataTest.java
@@ -1543,11 +1543,83 @@ public void testEmptySqlInfo() throws Exception {
try (final Connection testConnection =
FLIGHT_SERVER_EMPTY_SQLINFO_TEST_RULE.getConnection(false)) {
final DatabaseMetaData metaData = testConnection.getMetaData();
+
assertThat(metaData.getSQLKeywords(), is(""));
assertThat(metaData.getNumericFunctions(), is(""));
assertThat(metaData.getStringFunctions(), is(""));
assertThat(metaData.getSystemFunctions(), is(""));
assertThat(metaData.getTimeDateFunctions(), is(""));
+
+ assertThat(metaData.getMaxBinaryLiteralLength(), is(0));
+ assertThat(metaData.getMaxCharLiteralLength(), is(0));
+ assertThat(metaData.getMaxColumnNameLength(), is(0));
+ assertThat(metaData.getMaxColumnsInGroupBy(), is(0));
+ assertThat(metaData.getMaxColumnsInIndex(), is(0));
+ assertThat(metaData.getMaxColumnsInOrderBy(), is(0));
+ assertThat(metaData.getMaxColumnsInSelect(), is(0));
+ assertThat(metaData.getMaxColumnsInTable(), is(0));
+ assertThat(metaData.getMaxConnections(), is(0));
+ assertThat(metaData.getMaxCursorNameLength(), is(0));
+ assertThat(metaData.getMaxIndexLength(), is(0));
+ assertThat(metaData.getMaxSchemaNameLength(), is(0));
+ assertThat(metaData.getMaxProcedureNameLength(), is(0));
+ assertThat(metaData.getMaxCatalogNameLength(), is(0));
+ assertThat(metaData.getMaxRowSize(), is(0));
+ assertThat(metaData.getMaxStatementLength(), is(0));
+ assertThat(metaData.getMaxStatements(), is(0));
+ assertThat(metaData.getMaxTableNameLength(), is(0));
+ assertThat(metaData.getMaxTablesInSelect(), is(0));
+ assertThat(metaData.getMaxUserNameLength(), is(0));
+
+ assertThat(metaData.supportsColumnAliasing(), is(false));
+ assertThat(metaData.nullPlusNonNullIsNull(), is(false));
+ assertThat(metaData.supportsTableCorrelationNames(), is(false));
+ assertThat(metaData.supportsDifferentTableCorrelationNames(), is(false));
+ assertThat(metaData.supportsExpressionsInOrderBy(), is(false));
+ assertThat(metaData.supportsOrderByUnrelated(), is(false));
+ assertThat(metaData.supportsLikeEscapeClause(), is(false));
+ assertThat(metaData.supportsNonNullableColumns(), is(false));
+ assertThat(metaData.supportsIntegrityEnhancementFacility(), is(false));
+ assertThat(metaData.isCatalogAtStart(), is(false));
+ assertThat(metaData.supportsSelectForUpdate(), is(false));
+ assertThat(metaData.supportsStoredProcedures(), is(false));
+ assertThat(metaData.supportsCorrelatedSubqueries(), is(false));
+ assertThat(metaData.doesMaxRowSizeIncludeBlobs(), is(false));
+ assertThat(metaData.supportsTransactions(), is(false));
+ assertThat(metaData.dataDefinitionCausesTransactionCommit(), is(false));
+ assertThat(metaData.dataDefinitionIgnoredInTransactions(), is(false));
+ assertThat(metaData.supportsBatchUpdates(), is(false));
+ assertThat(metaData.supportsSavepoints(), is(false));
+ assertThat(metaData.supportsNamedParameters(), is(false));
+ assertThat(metaData.locatorsUpdateCopy(), is(false));
+ assertThat(metaData.supportsStoredFunctionsUsingCallSyntax(), is(false));
+ assertThat(metaData.supportsGroupBy(), is(false));
+ assertThat(metaData.supportsGroupByUnrelated(), is(false));
+ assertThat(metaData.supportsMinimumSQLGrammar(), is(false));
+ assertThat(metaData.supportsCoreSQLGrammar(), is(false));
+ assertThat(metaData.supportsExtendedSQLGrammar(), is(false));
+ assertThat(metaData.supportsANSI92EntryLevelSQL(), is(false));
+ assertThat(metaData.supportsANSI92IntermediateSQL(), is(false));
+ assertThat(metaData.supportsANSI92FullSQL(), is(false));
+ assertThat(metaData.supportsOuterJoins(), is(false));
+ assertThat(metaData.supportsFullOuterJoins(), is(false));
+ assertThat(metaData.supportsLimitedOuterJoins(), is(false));
+ assertThat(metaData.supportsSchemasInProcedureCalls(), is(false));
+ assertThat(metaData.supportsSchemasInIndexDefinitions(), is(false));
+ assertThat(metaData.supportsSchemasInPrivilegeDefinitions(), is(false));
+ assertThat(metaData.supportsCatalogsInIndexDefinitions(), is(false));
+ assertThat(metaData.supportsCatalogsInPrivilegeDefinitions(), is(false));
+ assertThat(metaData.supportsPositionedDelete(), is(false));
+ assertThat(metaData.supportsPositionedUpdate(), is(false));
+ assertThat(metaData.supportsSubqueriesInComparisons(), is(false));
+ assertThat(metaData.supportsSubqueriesInExists(), is(false));
+ assertThat(metaData.supportsSubqueriesInIns(), is(false));
+ assertThat(metaData.supportsSubqueriesInQuantifieds(), is(false));
+ assertThat(metaData.supportsUnion(), is(false));
+ assertThat(metaData.supportsUnionAll(), is(false));
+ assertThat(metaData.supportsConvert(), is(false));
+
+ assertThat(metaData.getDefaultTransactionIsolation(), is(Connection.TRANSACTION_NONE));
}
}
}
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcArrayTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcArrayTest.java
index 06d101724c..cb6abacb2f 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcArrayTest.java
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcArrayTest.java
@@ -129,7 +129,7 @@ public void testShouldGetResultSetReturnValidResultSet() throws SQLException {
try (ResultSet resultSet = arrowFlightJdbcArray.getResultSet()) {
int count = 0;
while (resultSet.next()) {
- assertEquals((Object) resultSet.getInt(1), dataVector.getObject(count));
+ assertEquals((Object) resultSet.getInt(2), dataVector.getObject(count));
count++;
}
}
@@ -142,7 +142,7 @@ public void testShouldGetResultSetReturnValidResultSetWithOffsets() throws SQLEx
try (ResultSet resultSet = arrowFlightJdbcArray.getResultSet(3, 5)) {
int count = 0;
while (resultSet.next()) {
- assertEquals((Object) resultSet.getInt(1), dataVector.getObject(count + 3));
+ assertEquals((Object) resultSet.getInt(2), dataVector.getObject(count + 3));
count++;
}
assertEquals(5, count);
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriverTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriverTest.java
index ae355829d7..88fb9889b6 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriverTest.java
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriverTest.java
@@ -201,6 +201,30 @@ public void testConnectWithInsensitiveCasePropertyKeys2() throws Exception {
}
}
+ /**
+ * Tests whether the {@link ArrowFlightJdbcDriver} can establish a successful connection to the
+ * Arrow Flight client when provided with null properties.
+ */
+ @Test
+ public void testConnectWithNullProperties() throws Exception {
+ final Driver driver = new ArrowFlightJdbcDriver();
+ try (Connection connection =
+ driver.connect(
+ "jdbc:arrow-flight://"
+ + dataSource.getConfig().getHost()
+ + ":"
+ + dataSource.getConfig().getPort()
+ + "?"
+ + "useEncryption=false"
+ + "&user="
+ + dataSource.getConfig().getUser()
+ + "&password="
+ + dataSource.getConfig().getPassword(),
+ null)) {
+ assertTrue(connection.isValid(300));
+ }
+ }
+
/**
* Tests whether an exception is thrown upon attempting to connect to a malformed URI.
*
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java
index 0369c3a162..078837adf3 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java
@@ -98,6 +98,40 @@ public void testSimpleQueryNoParameterBindingWithExecute() throws SQLException {
}
}
+ @Test
+ public void testSimpleQueryNoParameterBindingWithExecuteV2() throws SQLException {
+ final String query = "SELECT * FROM TEST_V2";
+ final Schema schema =
+ new Schema(Collections.singletonList(Field.nullable("", Types.MinorType.INT.getType())));
+ PRODUCER.addSelectQuery(
+ query,
+ schema,
+ Collections.singletonList(
+ listener -> {
+ try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
+ final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
+ root.allocateNew();
+ ((IntVector) root.getVector(0)).setSafe(0, 123);
+ root.setRowCount(1);
+ listener.start(root);
+ listener.putNext();
+ } finally {
+ listener.completed();
+ }
+ }),
+ false);
+ try (final PreparedStatement preparedStatement = connection.prepareStatement(query)) {
+ boolean isResultSet = preparedStatement.execute();
+ assertTrue(isResultSet);
+ final ResultSet resultSet = preparedStatement.getResultSet();
+ assertTrue(resultSet.next());
+ assertEquals(123, resultSet.getInt(1));
+ assertFalse(resultSet.next());
+ assertFalse(preparedStatement.getMoreResults());
+ assertEquals(-1, preparedStatement.getUpdateCount());
+ }
+ }
+
@Test
public void testQueryWithParameterBinding() throws SQLException {
final String query = "Fake query with parameters";
@@ -203,6 +237,20 @@ public void testUpdateQueryWithExecute() throws SQLException {
}
}
+ @Test
+ public void testUpdateQueryWithExecuteV2() throws SQLException {
+ String query = "Fake update with execute V2";
+ PRODUCER.addUpdateQuery(query, /*updatedRows*/ 99, true);
+ try (final PreparedStatement stmt = connection.prepareStatement(query)) {
+ boolean isResultSet = stmt.execute();
+ assertFalse(isResultSet);
+ int updated = stmt.getUpdateCount();
+ assertEquals(99, updated);
+ assertFalse(stmt.getMoreResults());
+ assertEquals(-1, stmt.getUpdateCount());
+ }
+ }
+
@Test
public void testUpdateQueryWithParameters() throws SQLException {
String query = "Fake update with parameters";
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementExecuteTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementExecuteTest.java
index 632cb0ba56..6acce9c2a6 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementExecuteTest.java
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementExecuteTest.java
@@ -62,6 +62,9 @@ public class ArrowFlightStatementExecuteTest {
private static final String SAMPLE_LARGE_UPDATE_QUERY =
"UPDATE this_large_table SET this_large_field = that_large_field FROM this_large_test WHERE this_large_condition";
private static final long SAMPLE_LARGE_UPDATE_COUNT = Long.MAX_VALUE;
+ private static final String SAMPLE_QUERY_CMD_V2 = "SELECT * FROM this_test_v2";
+ private static final String SAMPLE_LARGE_UPDATE_QUERY_V2 =
+ "UPDATE this_large_table_v2 SET this_large_field = that_large_field FROM this_large_test WHERE this_large_condition";
private static final MockFlightSqlProducer PRODUCER = new MockFlightSqlProducer();
@RegisterExtension
@@ -96,6 +99,31 @@ public static void setUpBeforeClass() {
}));
PRODUCER.addUpdateQuery(SAMPLE_UPDATE_QUERY, SAMPLE_UPDATE_COUNT);
PRODUCER.addUpdateQuery(SAMPLE_LARGE_UPDATE_QUERY, SAMPLE_LARGE_UPDATE_COUNT);
+
+ // V2 queries with is_update field set
+ PRODUCER.addSelectQuery(
+ SAMPLE_QUERY_CMD_V2,
+ SAMPLE_QUERY_SCHEMA,
+ Collections.singletonList(
+ listener -> {
+ try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
+ final VectorSchemaRoot root =
+ VectorSchemaRoot.create(SAMPLE_QUERY_SCHEMA, allocator)) {
+ final UInt1Vector vector = (UInt1Vector) root.getVector(VECTOR_NAME);
+ IntStream.range(0, SAMPLE_QUERY_ROWS)
+ .forEach(index -> vector.setSafe(index, index));
+ vector.setValueCount(SAMPLE_QUERY_ROWS);
+ root.setRowCount(SAMPLE_QUERY_ROWS);
+ listener.start(root);
+ listener.putNext();
+ } catch (final Throwable throwable) {
+ listener.error(throwable);
+ } finally {
+ listener.completed();
+ }
+ }),
+ false);
+ PRODUCER.addUpdateQuery(SAMPLE_LARGE_UPDATE_QUERY_V2, SAMPLE_LARGE_UPDATE_COUNT, true);
}
@BeforeEach
@@ -168,4 +196,42 @@ public void testUpdateCountShouldStartOnZero() throws SQLException {
is(allOf(equalTo(statement.getLargeUpdateCount()), equalTo(0L))));
assertThat(statement.getResultSet(), is(nullValue()));
}
+
+ @Test
+ public void testExecuteShouldRunSelectQueryV2() throws SQLException {
+ assertThat(statement.execute(SAMPLE_QUERY_CMD_V2), is(true));
+ final Set numbers =
+ IntStream.range(0, SAMPLE_QUERY_ROWS)
+ .boxed()
+ .map(Integer::byteValue)
+ .collect(Collectors.toCollection(HashSet::new));
+ try (final ResultSet resultSet = statement.getResultSet()) {
+ final int columnCount = resultSet.getMetaData().getColumnCount();
+ assertThat(columnCount, is(1));
+ int rowCount = 0;
+ for (; resultSet.next(); rowCount++) {
+ assertThat(numbers.remove(resultSet.getByte(1)), is(true));
+ }
+ assertThat(rowCount, is(equalTo(SAMPLE_QUERY_ROWS)));
+ }
+ assertThat(numbers, is(Collections.emptySet()));
+ assertThat(
+ (long) statement.getUpdateCount(),
+ is(allOf(equalTo(statement.getLargeUpdateCount()), equalTo(-1L))));
+ }
+
+ @Test
+ public void testExecuteShouldRunUpdateQueryForLargeUpdateV2() throws SQLException {
+ assertThat(statement.execute(SAMPLE_LARGE_UPDATE_QUERY_V2), is(false)); // UPDATE query.
+ final long updateCountSmall = statement.getUpdateCount();
+ final long updateCountLarge = statement.getLargeUpdateCount();
+ assertThat(updateCountLarge, is(equalTo(SAMPLE_LARGE_UPDATE_COUNT)));
+ assertThat(
+ updateCountSmall,
+ is(
+ allOf(
+ equalTo((long) AvaticaUtils.toSaturatedInt(updateCountLarge)),
+ not(equalTo(updateCountLarge)))));
+ assertThat(statement.getResultSet(), is(nullValue()));
+ }
}
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionTest.java
index 72e4b222a3..55722f60fb 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionTest.java
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionTest.java
@@ -16,25 +16,42 @@
*/
package org.apache.arrow.driver.jdbc;
+import static java.lang.String.format;
+import static java.util.stream.IntStream.range;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assertions.fail;
+import com.google.protobuf.Message;
import java.net.URISyntaxException;
import java.sql.Connection;
import java.sql.Driver;
import java.sql.DriverManager;
+import java.sql.ResultSet;
import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.Map;
import java.util.Properties;
+import java.util.function.Consumer;
import org.apache.arrow.driver.jdbc.authentication.UserPasswordAuthentication;
import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler;
import org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty;
import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer;
import org.apache.arrow.flight.FlightMethod;
+import org.apache.arrow.flight.FlightProducer.ServerStreamListener;
+import org.apache.arrow.flight.NoOpSessionOptionValueVisitor;
+import org.apache.arrow.flight.SessionOptionValue;
+import org.apache.arrow.flight.sql.FlightSqlProducer.Schemas;
+import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetTableTypes;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.util.AutoCloseables;
+import org.apache.arrow.vector.VarCharVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.util.Text;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
@@ -614,7 +631,8 @@ public void testJdbcDriverVersionIntegration() throws Exception {
var expectedUserAgent =
"JDBC Flight SQL Driver " + driverVersion.getDriverVersion().versionString;
- // Driver appends version to grpc user-agent header. Assert the header starts with the
+ // Driver appends version to grpc user-agent header. Assert the header starts
+ // with the
// expected
// value and ignored grpc version.
assertTrue(
@@ -622,4 +640,132 @@ public void testJdbcDriverVersionIntegration() throws Exception {
"Expected: " + expectedUserAgent + " but found: " + actualUserAgent);
}
}
+
+ @Test
+ public void testSetCatalogShouldUpdateSessionOptions() throws Exception {
+ final Properties properties = new Properties();
+ properties.put(ArrowFlightConnectionProperty.USER.camelName(), userTest);
+ properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest);
+ properties.put("useEncryption", false);
+
+ try (Connection connection =
+ DriverManager.getConnection(
+ "jdbc:arrow-flight-sql://"
+ + FLIGHT_SERVER_TEST_EXTENSION.getHost()
+ + ":"
+ + FLIGHT_SERVER_TEST_EXTENSION.getPort(),
+ properties)) {
+ final String catalog = "new_catalog";
+ connection.setCatalog(catalog);
+
+ final Map options = PRODUCER.getSessionOptions();
+ assertTrue(options.containsKey("catalog"));
+ String actualCatalog =
+ options
+ .get("catalog")
+ .acceptVisitor(
+ new NoOpSessionOptionValueVisitor() {
+ @Override
+ public String visit(String value) {
+ return value;
+ }
+ });
+ assertEquals(catalog, actualCatalog);
+ }
+ }
+
+ @Test
+ public void testStatementsClosedOnConnectionClose() throws Exception {
+ // create a connection
+ final Properties properties = new Properties();
+ properties.put(ArrowFlightConnectionProperty.HOST.camelName(), "localhost");
+ properties.put(
+ ArrowFlightConnectionProperty.PORT.camelName(), FLIGHT_SERVER_TEST_EXTENSION.getPort());
+ properties.put(ArrowFlightConnectionProperty.USER.camelName(), userTest);
+ properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest);
+ properties.put("useEncryption", false);
+
+ Connection connection =
+ DriverManager.getConnection(
+ "jdbc:arrow-flight-sql://"
+ + FLIGHT_SERVER_TEST_EXTENSION.getHost()
+ + ":"
+ + FLIGHT_SERVER_TEST_EXTENSION.getPort(),
+ properties);
+
+ // create some statements
+ int numStatements = 3;
+ Statement[] statements = new Statement[numStatements];
+ for (int i = 0; i < numStatements; i++) {
+ statements[i] = connection.createStatement();
+ assertFalse(statements[i].isClosed());
+ }
+
+ // close the connection
+ connection.close();
+
+ // assert the statements are closed
+ for (int i = 0; i < numStatements; i++) {
+ assertTrue(statements[i].isClosed());
+ }
+ }
+
+ @Test
+ public void testResultSetsFromDatabaseMetadataClosedOnConnectionClose() throws Exception {
+ // set up the FlightProducer to respond to metadata queries
+ // getTableTypes() is being used, but any other method would work
+ int rowCount = 3;
+ final Message commandGetTableTypes = CommandGetTableTypes.getDefaultInstance();
+ final Consumer commandGetTableTypesResultProducer =
+ listener -> {
+ try (final BufferAllocator allocator = new RootAllocator();
+ final VectorSchemaRoot root =
+ VectorSchemaRoot.create(Schemas.GET_TABLE_TYPES_SCHEMA, allocator)) {
+ final VarCharVector tableType = (VarCharVector) root.getVector("table_type");
+ range(0, rowCount)
+ .forEach(i -> tableType.setSafe(i, new Text(format("table_type #%d", i))));
+ root.setRowCount(rowCount);
+ listener.start(root);
+ listener.putNext();
+ } catch (final Throwable throwable) {
+ listener.error(throwable);
+ } finally {
+ listener.completed();
+ }
+ };
+ PRODUCER.addCatalogQuery(commandGetTableTypes, commandGetTableTypesResultProducer);
+
+ // create a connection
+ final Properties properties = new Properties();
+ properties.put(ArrowFlightConnectionProperty.HOST.camelName(), "localhost");
+ properties.put(
+ ArrowFlightConnectionProperty.PORT.camelName(), FLIGHT_SERVER_TEST_EXTENSION.getPort());
+ properties.put(ArrowFlightConnectionProperty.USER.camelName(), userTest);
+ properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest);
+ properties.put("useEncryption", false);
+
+ Connection connection =
+ DriverManager.getConnection(
+ "jdbc:arrow-flight-sql://"
+ + FLIGHT_SERVER_TEST_EXTENSION.getHost()
+ + ":"
+ + FLIGHT_SERVER_TEST_EXTENSION.getPort(),
+ properties);
+
+ // create ResultSets from DatabaseMetadata
+ int numResultSets = 3;
+ ResultSet[] resultSets = new ResultSet[numResultSets];
+ for (int i = 0; i < numResultSets; i++) {
+ resultSets[i] = connection.getMetaData().getTableTypes();
+ assertFalse(resultSets[i].isClosed());
+ }
+
+ // close the connection
+ connection.close();
+
+ // assert the ResultSets are closed
+ for (int i = 0; i < numResultSets; i++) {
+ assertTrue(resultSets[i].isClosed());
+ }
+ }
}
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/OAuthIntegrationTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/OAuthIntegrationTest.java
new file mode 100644
index 0000000000..5e782db031
--- /dev/null
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/OAuthIntegrationTest.java
@@ -0,0 +1,474 @@
+/*
+ * 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;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.net.URI;
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.SQLException;
+import java.util.Properties;
+import java.util.concurrent.TimeUnit;
+import mockwebserver3.MockResponse;
+import mockwebserver3.MockWebServer;
+import mockwebserver3.RecordedRequest;
+import mockwebserver3.junit5.StartStop;
+import org.apache.arrow.driver.jdbc.authentication.TokenAuthentication;
+import org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty;
+import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer;
+import org.apache.arrow.flight.sql.FlightSqlProducer.Schemas;
+import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetCatalogs;
+import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetDbSchemas;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.util.AutoCloseables;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+/**
+ * Integration tests for OAuth authentication flows in the JDBC driver.
+ *
+ * These tests verify that OAuth tokens obtained from an OAuth server are correctly used in
+ * Flight SQL requests.
+ */
+public class OAuthIntegrationTest {
+
+ private static final String VALID_ACCESS_TOKEN = "valid-oauth-access-token-12345";
+ private static final String CLIENT_ID = "test-client-id";
+ private static final String CLIENT_SECRET = "test-client-secret";
+ private static final String SUBJECT_TOKEN = "original-subject-token";
+ private static final String SUBJECT_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:jwt";
+ private static final String TEST_SCOPE = "dremio.all";
+
+ private static final MockFlightSqlProducer FLIGHT_SQL_PRODUCER = new MockFlightSqlProducer();
+
+ @RegisterExtension public static FlightServerTestExtension FLIGHT_SERVER_TEST_EXTENSION;
+
+ static {
+ FLIGHT_SERVER_TEST_EXTENSION =
+ new FlightServerTestExtension.Builder()
+ .authentication(new TokenAuthentication.Builder().token(VALID_ACCESS_TOKEN).build())
+ .producer(FLIGHT_SQL_PRODUCER)
+ .build();
+ }
+
+ @StartStop private final MockWebServer oauthServer = new MockWebServer();
+ private URI tokenEndpoint;
+
+ @BeforeAll
+ public static void setUpClass() {
+ // Register a simple catalog query handler
+ FLIGHT_SQL_PRODUCER.addCatalogQuery(
+ CommandGetCatalogs.getDefaultInstance(),
+ listener -> {
+ try (BufferAllocator allocator = new RootAllocator();
+ VectorSchemaRoot root =
+ VectorSchemaRoot.create(Schemas.GET_CATALOGS_SCHEMA, allocator)) {
+ root.setRowCount(0);
+ listener.start(root);
+ listener.putNext();
+ } catch (Throwable t) {
+ listener.error(t);
+ } finally {
+ listener.completed();
+ }
+ });
+
+ // Register a simple schema query handler for getSchemas()
+ FLIGHT_SQL_PRODUCER.addCatalogQuery(
+ CommandGetDbSchemas.getDefaultInstance(),
+ listener -> {
+ try (BufferAllocator allocator = new RootAllocator();
+ VectorSchemaRoot root =
+ VectorSchemaRoot.create(Schemas.GET_SCHEMAS_SCHEMA, allocator)) {
+ root.setRowCount(0);
+ listener.start(root);
+ listener.putNext();
+ } catch (Throwable t) {
+ listener.error(t);
+ } finally {
+ listener.completed();
+ }
+ });
+ }
+
+ @AfterAll
+ public static void tearDownClass() {
+ AutoCloseables.closeNoChecked(FLIGHT_SQL_PRODUCER);
+ }
+
+ @BeforeEach
+ public void setUp() {
+ tokenEndpoint = oauthServer.url("/oauth/token").uri();
+ }
+
+ @AfterEach
+ public void tearDown() {
+ oauthServer.close();
+ }
+
+ // Helper methods for mock OAuth responses
+
+ private void enqueueSuccessfulTokenResponse() {
+ enqueueSuccessfulTokenResponse(VALID_ACCESS_TOKEN, 3600);
+ }
+
+ private void enqueueSuccessfulTokenResponse(String token, int expiresIn) {
+ String body =
+ String.format(
+ "{\"access_token\":\"%s\",\"token_type\":\"Bearer\",\"expires_in\":%d}",
+ token, expiresIn);
+ oauthServer.enqueue(
+ new MockResponse.Builder()
+ .code(200)
+ .setHeader("Content-Type", "application/json")
+ .body(body)
+ .build());
+ }
+
+ private void enqueueErrorResponse(String error, String description) {
+ String body =
+ String.format("{\"error\":\"%s\",\"error_description\":\"%s\"}", error, description);
+ oauthServer.enqueue(
+ new MockResponse.Builder()
+ .code(400)
+ .setHeader("Content-Type", "application/json")
+ .body(body)
+ .build());
+ }
+
+ private Properties createBaseProperties() {
+ Properties props = new Properties();
+ props.put(ArrowFlightConnectionProperty.HOST.camelName(), "localhost");
+ props.put(
+ ArrowFlightConnectionProperty.PORT.camelName(), FLIGHT_SERVER_TEST_EXTENSION.getPort());
+ props.put(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), false);
+ return props;
+ }
+
+ private String getJdbcUrl() {
+ return String.format(
+ "jdbc:arrow-flight-sql://localhost:%d", FLIGHT_SERVER_TEST_EXTENSION.getPort());
+ }
+
+ // ==================== Client Credentials Flow Tests ====================
+
+ @Test
+ public void testClientCredentialsFlowSuccess() throws Exception {
+ enqueueSuccessfulTokenResponse();
+
+ Properties props = createBaseProperties();
+ props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "client_credentials");
+ props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString());
+ props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_ID.camelName(), CLIENT_ID);
+ props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_SECRET.camelName(), CLIENT_SECRET);
+ props.put(ArrowFlightConnectionProperty.OAUTH_SCOPE.camelName(), TEST_SCOPE);
+
+ try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) {
+ assertFalse(conn.isClosed());
+ // Trigger a Flight call to force OAuth token retrieval
+ conn.getMetaData().getCatalogs().close();
+ }
+
+ // Verify OAuth request was made
+ RecordedRequest oauthRequest = oauthServer.takeRequest(5, TimeUnit.SECONDS);
+ assertNotNull(oauthRequest, "OAuth request should have been made");
+ assertEquals("POST", oauthRequest.getMethod());
+ String body = oauthRequest.getBody().utf8();
+ assertTrue(body.contains("grant_type=client_credentials"));
+ assertTrue(body.contains("scope=" + TEST_SCOPE));
+ }
+
+ @Test
+ public void testClientCredentialsFlowWithUrlParameters() throws Exception {
+ enqueueSuccessfulTokenResponse();
+
+ String url =
+ String.format(
+ "jdbc:arrow-flight-sql://localhost:%d?useEncryption=false"
+ + "&oauth.flow=client_credentials"
+ + "&oauth.tokenUri=%s"
+ + "&oauth.clientId=%s"
+ + "&oauth.clientSecret=%s",
+ FLIGHT_SERVER_TEST_EXTENSION.getPort(),
+ tokenEndpoint.toString(),
+ CLIENT_ID,
+ CLIENT_SECRET);
+
+ try (Connection conn = DriverManager.getConnection(url)) {
+ conn.getMetaData().getCatalogs().close();
+ }
+
+ assertEquals(1, oauthServer.getRequestCount());
+ }
+
+ @Test
+ public void testClientCredentialsFlowInvalidCredentials() throws Exception {
+ enqueueErrorResponse("invalid_client", "Client authentication failed");
+
+ Properties props = createBaseProperties();
+ props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "client_credentials");
+ props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString());
+ props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_ID.camelName(), "wrong-client");
+ props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_SECRET.camelName(), "wrong-secret");
+
+ Exception ex =
+ assertThrows(
+ Exception.class,
+ () -> {
+ try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) {
+ conn.getMetaData().getCatalogs().close();
+ }
+ });
+ // Verify the error message contains the OAuth error somewhere in the exception chain
+ assertTrue(
+ containsInExceptionChain(ex, "invalid_client"),
+ "Exception chain should contain 'invalid_client'");
+ }
+
+ private boolean containsInExceptionChain(Throwable t, String message) {
+ while (t != null) {
+ if (t.getMessage() != null && t.getMessage().contains(message)) {
+ return true;
+ }
+ t = t.getCause();
+ }
+ return false;
+ }
+
+ // ==================== Token Exchange Flow Tests ====================
+
+ @Test
+ public void testTokenExchangeFlowMinimalParameters() throws Exception {
+ enqueueSuccessfulTokenResponse();
+
+ Properties props = createBaseProperties();
+ props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "token_exchange");
+ props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString());
+ props.put(
+ ArrowFlightConnectionProperty.OAUTH_EXCHANGE_SUBJECT_TOKEN.camelName(), SUBJECT_TOKEN);
+ props.put(
+ ArrowFlightConnectionProperty.OAUTH_EXCHANGE_SUBJECT_TOKEN_TYPE.camelName(),
+ SUBJECT_TOKEN_TYPE);
+
+ try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) {
+ conn.getMetaData().getCatalogs().close();
+ }
+
+ RecordedRequest oauthRequest = oauthServer.takeRequest(5, TimeUnit.SECONDS);
+ assertNotNull(oauthRequest, "OAuth request should have been made");
+ String body = oauthRequest.getBody().utf8();
+ assertTrue(
+ body.contains("grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Atoken-exchange"),
+ "Should contain token exchange grant type");
+ assertTrue(body.contains("subject_token=" + SUBJECT_TOKEN));
+ }
+
+ @Test
+ public void testTokenExchangeFlowWithAllParameters() throws Exception {
+ enqueueSuccessfulTokenResponse();
+
+ String actorToken = "actor-token-value";
+ String actorTokenType = "urn:ietf:params:oauth:token-type:access_token";
+ String audience = "https://api.example.com";
+ String resource = "https://api.example.com/resource";
+ String requestedTokenType = "urn:ietf:params:oauth:token-type:access_token";
+
+ Properties props = createBaseProperties();
+ props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "token_exchange");
+ props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString());
+ props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_ID.camelName(), CLIENT_ID);
+ props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_SECRET.camelName(), CLIENT_SECRET);
+ props.put(ArrowFlightConnectionProperty.OAUTH_SCOPE.camelName(), TEST_SCOPE);
+ props.put(
+ ArrowFlightConnectionProperty.OAUTH_EXCHANGE_SUBJECT_TOKEN.camelName(), SUBJECT_TOKEN);
+ props.put(
+ ArrowFlightConnectionProperty.OAUTH_EXCHANGE_SUBJECT_TOKEN_TYPE.camelName(),
+ SUBJECT_TOKEN_TYPE);
+ props.put(ArrowFlightConnectionProperty.OAUTH_EXCHANGE_ACTOR_TOKEN.camelName(), actorToken);
+ props.put(
+ ArrowFlightConnectionProperty.OAUTH_EXCHANGE_ACTOR_TOKEN_TYPE.camelName(), actorTokenType);
+ props.put(ArrowFlightConnectionProperty.OAUTH_EXCHANGE_AUDIENCE.camelName(), audience);
+ props.put(ArrowFlightConnectionProperty.OAUTH_RESOURCE.camelName(), resource);
+ props.put(
+ ArrowFlightConnectionProperty.OAUTH_EXCHANGE_REQUESTED_TOKEN_TYPE.camelName(),
+ requestedTokenType);
+
+ try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) {
+ conn.getMetaData().getCatalogs().close();
+ }
+
+ RecordedRequest oauthRequest = oauthServer.takeRequest(5, TimeUnit.SECONDS);
+ assertNotNull(oauthRequest, "OAuth request should have been made");
+ String body = oauthRequest.getBody().utf8();
+ assertTrue(body.contains("subject_token=" + SUBJECT_TOKEN));
+ assertTrue(body.contains("actor_token=" + actorToken));
+ }
+
+ @Test
+ public void testTokenExchangeFlowWithClientAuthentication() throws Exception {
+ enqueueSuccessfulTokenResponse();
+
+ Properties props = createBaseProperties();
+ props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "token_exchange");
+ props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString());
+ props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_ID.camelName(), CLIENT_ID);
+ props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_SECRET.camelName(), CLIENT_SECRET);
+ props.put(
+ ArrowFlightConnectionProperty.OAUTH_EXCHANGE_SUBJECT_TOKEN.camelName(), SUBJECT_TOKEN);
+ props.put(
+ ArrowFlightConnectionProperty.OAUTH_EXCHANGE_SUBJECT_TOKEN_TYPE.camelName(),
+ SUBJECT_TOKEN_TYPE);
+
+ try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) {
+ conn.getMetaData().getCatalogs().close();
+ }
+
+ RecordedRequest oauthRequest = oauthServer.takeRequest(5, TimeUnit.SECONDS);
+ assertNotNull(oauthRequest, "OAuth request should have been made");
+ String authHeader = oauthRequest.getHeaders().get("Authorization");
+ assertNotNull(authHeader, "Should have Basic auth header for client authentication");
+ assertTrue(authHeader.startsWith("Basic "));
+ }
+
+ // ==================== Token Caching Tests ====================
+
+ @Test
+ public void testTokenCachingAcrossMultipleOperations() throws Exception {
+ enqueueSuccessfulTokenResponse();
+
+ Properties props = createBaseProperties();
+ props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "client_credentials");
+ props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString());
+ props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_ID.camelName(), CLIENT_ID);
+ props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_SECRET.camelName(), CLIENT_SECRET);
+
+ try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) {
+ // Execute multiple operations
+ conn.isValid(5);
+ conn.getMetaData().getCatalogs().close();
+ conn.getMetaData().getSchemas().close();
+ }
+
+ // Should only have made one OAuth request due to caching
+ assertEquals(1, oauthServer.getRequestCount());
+ }
+
+ @Test
+ public void testTokenRefreshAfterExpiration() throws Exception {
+ enqueueSuccessfulTokenResponse(VALID_ACCESS_TOKEN, 1);
+ enqueueSuccessfulTokenResponse(VALID_ACCESS_TOKEN, 3600);
+
+ Properties props = createBaseProperties();
+ props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "token_exchange");
+ props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString());
+ props.put(
+ ArrowFlightConnectionProperty.OAUTH_EXCHANGE_SUBJECT_TOKEN.camelName(), SUBJECT_TOKEN);
+ props.put(
+ ArrowFlightConnectionProperty.OAUTH_EXCHANGE_SUBJECT_TOKEN_TYPE.camelName(),
+ SUBJECT_TOKEN_TYPE);
+
+ try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) {
+ // First operation triggers initial token fetch
+ conn.getMetaData().getCatalogs().close();
+
+ // Token with 1s expiry is immediately considered expired (due to 30s buffer)
+ // so the next operation should trigger a refresh
+ conn.getMetaData().getCatalogs().close();
+ }
+
+ // Should have made exactly 2 OAuth requests: initial + refresh
+ assertEquals(2, oauthServer.getRequestCount());
+ }
+
+ // ==================== Error Handling Tests ====================
+
+ @Test
+ public void testMissingRequiredParametersClientCredentials() {
+ Properties props = createBaseProperties();
+ props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "client_credentials");
+ props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString());
+ // Missing client_id and client_secret
+
+ assertThrows(SQLException.class, () -> DriverManager.getConnection(getJdbcUrl(), props));
+ }
+
+ @Test
+ public void testMissingRequiredParametersTokenExchange() {
+ Properties props = createBaseProperties();
+ props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "token_exchange");
+ props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString());
+ // Missing subject_token and subject_token_type
+
+ assertThrows(SQLException.class, () -> DriverManager.getConnection(getJdbcUrl(), props));
+ }
+
+ @Test
+ public void testInvalidOAuthFlow() {
+ Properties props = createBaseProperties();
+ props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "invalid_flow");
+ props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString());
+
+ assertThrows(SQLException.class, () -> DriverManager.getConnection(getJdbcUrl(), props));
+ }
+
+ @Test
+ public void testMalformedTokenEndpoint() {
+ Properties props = createBaseProperties();
+ props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "client_credentials");
+ props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), "not-a-valid-uri://");
+ props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_ID.camelName(), CLIENT_ID);
+ props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_SECRET.camelName(), CLIENT_SECRET);
+
+ assertThrows(SQLException.class, () -> DriverManager.getConnection(getJdbcUrl(), props));
+ }
+
+ // ==================== Authorization Header Verification ====================
+
+ @Test
+ public void testOAuthTokenSentAsBearer() throws Exception {
+ enqueueSuccessfulTokenResponse();
+
+ Properties props = createBaseProperties();
+ props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "client_credentials");
+ props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString());
+ props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_ID.camelName(), CLIENT_ID);
+ props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_SECRET.camelName(), CLIENT_SECRET);
+
+ try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) {
+ conn.getMetaData().getCatalogs().close();
+ }
+
+ // Verify the Flight server received the bearer token
+ String authHeader =
+ FLIGHT_SERVER_TEST_EXTENSION
+ .getInterceptorFactory()
+ .getHeader(org.apache.arrow.flight.FlightMethod.GET_FLIGHT_INFO, "authorization");
+ assertNotNull(authHeader, "Authorization header should be present in Flight requests");
+ assertEquals("Bearer " + VALID_ACCESS_TOKEN, authHeader);
+ }
+}
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..3a5a39be3d 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,174 @@ 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/accessor/impl/complex/AbstractArrowFlightJdbcListAccessorTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/AbstractArrowFlightJdbcListAccessorTest.java
index ad689837e2..c5eb6e34ef 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/AbstractArrowFlightJdbcListAccessorTest.java
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/AbstractArrowFlightJdbcListAccessorTest.java
@@ -191,7 +191,12 @@ public void testShouldGetArrayGetResultSetReturnValidResultSet(
try (ResultSet rs = array.getResultSet()) {
int count = 0;
while (rs.next()) {
- final int value = rs.getInt(1);
+ // Column 1: 1-based index (per JDBC spec)
+ final int index = rs.getInt(1);
+ assertThat(index, equalTo(count + 1));
+
+ // Column 2: actual value (per JDBC spec)
+ final int value = rs.getInt(2);
assertThat(value, equalTo(currentRow * count));
count++;
}
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/ArrowFlightJdbcMapVectorAccessorTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/ArrowFlightJdbcMapVectorAccessorTest.java
index 696e5afb71..f2d1725fd8 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/ArrowFlightJdbcMapVectorAccessorTest.java
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/ArrowFlightJdbcMapVectorAccessorTest.java
@@ -153,15 +153,15 @@ public void testShouldGetArrayReturnValidArray() throws SQLException {
try (ResultSet resultSet = array.getResultSet()) {
assertTrue(resultSet.next());
- Map entry = resultSet.getObject(1, Map.class);
+ Map entry = resultSet.getObject(2, Map.class);
assertEquals(1, entry.get("key"));
assertEquals(11, entry.get("value"));
assertTrue(resultSet.next());
- entry = resultSet.getObject(1, Map.class);
+ entry = resultSet.getObject(2, Map.class);
assertEquals(2, entry.get("key"));
assertEquals(22, entry.get("value"));
assertTrue(resultSet.next());
- entry = resultSet.getObject(1, Map.class);
+ entry = resultSet.getObject(2, Map.class);
assertEquals(3, entry.get("key"));
assertEquals(33, entry.get("value"));
assertFalse(resultSet.next());
@@ -173,7 +173,7 @@ public void testShouldGetArrayReturnValidArray() throws SQLException {
assertFalse(accessor.wasNull());
try (ResultSet resultSet = array.getResultSet()) {
assertTrue(resultSet.next());
- Map entry = resultSet.getObject(1, Map.class);
+ Map entry = resultSet.getObject(2, Map.class);
assertEquals(2, entry.get("key"));
assertNull(entry.get("value"));
assertFalse(resultSet.next());
@@ -185,19 +185,19 @@ public void testShouldGetArrayReturnValidArray() throws SQLException {
assertFalse(accessor.wasNull());
try (ResultSet resultSet = array.getResultSet()) {
assertTrue(resultSet.next());
- Map entry = resultSet.getObject(1, Map.class);
+ Map entry = resultSet.getObject(2, Map.class);
assertEquals(0, entry.get("key"));
assertEquals(2000, entry.get("value"));
assertTrue(resultSet.next());
- entry = resultSet.getObject(1, Map.class);
+ entry = resultSet.getObject(2, Map.class);
assertEquals(1, entry.get("key"));
assertEquals(2001, entry.get("value"));
assertTrue(resultSet.next());
- entry = resultSet.getObject(1, Map.class);
+ entry = resultSet.getObject(2, Map.class);
assertEquals(2, entry.get("key"));
assertEquals(2002, entry.get("value"));
assertTrue(resultSet.next());
- entry = resultSet.getObject(1, Map.class);
+ entry = resultSet.getObject(2, Map.class);
assertEquals(3, entry.get("key"));
assertEquals(2003, entry.get("value"));
assertFalse(resultSet.next());
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthConfigurationTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthConfigurationTest.java
new file mode 100644
index 0000000000..c258a7c652
--- /dev/null
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthConfigurationTest.java
@@ -0,0 +1,296 @@
+/*
+ * 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.client.oauth;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import com.nimbusds.oauth2.sdk.Scope;
+import java.net.URI;
+import java.sql.SQLException;
+import java.util.Collections;
+import java.util.stream.Stream;
+import org.junit.jupiter.api.Named;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+
+/** Tests for {@link OAuthConfiguration}. */
+public class OAuthConfigurationTest {
+
+ private static final String TOKEN_URI = "https://auth.example.com/token";
+ private static final String CLIENT_ID = "test-client-id";
+ private static final String CLIENT_SECRET = "test-client-secret";
+ private static final String SCOPE = "read write";
+ private static final String SUBJECT_TOKEN = "subject-token-value";
+ public static final String RESOURCE = "https://api.example.com/resource";
+
+ @FunctionalInterface
+ interface BuilderConfigurer {
+ void configure(OAuthConfiguration.Builder builder) throws SQLException;
+ }
+
+ static Stream createFlowCases() {
+ return Stream.of(
+ Arguments.of(
+ Named.of(
+ "string flow", (BuilderConfigurer) builder -> builder.flow("client_credentials"))),
+ Arguments.of(
+ Named.of(
+ "uppercase string flow",
+ (BuilderConfigurer) builder -> builder.flow("CLIENT_CREDENTIALS"))));
+ }
+
+ @ParameterizedTest
+ @MethodSource("createFlowCases")
+ public void testCreateFlowConfiguration(BuilderConfigurer flowConfigurer) throws SQLException {
+ OAuthConfiguration.Builder builder = new OAuthConfiguration.Builder();
+ flowConfigurer.configure(builder);
+ OAuthConfiguration config =
+ builder.tokenUri(TOKEN_URI).clientId(CLIENT_ID).clientSecret(CLIENT_SECRET).build();
+
+ // Verify configuration creates correct provider type
+ OAuthTokenProvider provider = config.createTokenProvider();
+ assertInstanceOf(ClientCredentialsTokenProvider.class, provider);
+ }
+
+ @Test
+ public void testCreateClientCredentialsTokenProvider() throws SQLException {
+ OAuthConfiguration config =
+ new OAuthConfiguration.Builder()
+ .flow("client_credentials")
+ .tokenUri(TOKEN_URI)
+ .clientId(CLIENT_ID)
+ .clientSecret(CLIENT_SECRET)
+ .scope(SCOPE)
+ .build();
+
+ OAuthTokenProvider provider = config.createTokenProvider();
+
+ assertNotNull(provider);
+ assertInstanceOf(ClientCredentialsTokenProvider.class, provider);
+
+ ClientCredentialsTokenProvider ccProvider = (ClientCredentialsTokenProvider) provider;
+ assertEquals(URI.create(TOKEN_URI), ccProvider.tokenUri);
+ assertEquals(CLIENT_ID, ccProvider.clientAuth.getClientID().getValue());
+ assertEquals(Scope.parse(SCOPE), ccProvider.scope);
+ }
+
+ @Test
+ public void testCreateTokenExchangeTokenProviderWithAllOptions() throws SQLException {
+ String subjectTokenType = "urn:ietf:params:oauth:token-type:access_token";
+ String actorToken = "actor-token-value";
+ String actorTokenType = "urn:ietf:params:oauth:token-type:jwt";
+ String audience = "https://api.example.com";
+ String requestedTokenType = "urn:ietf:params:oauth:token-type:access_token";
+
+ OAuthConfiguration config =
+ new OAuthConfiguration.Builder()
+ .flow("token_exchange")
+ .tokenUri(TOKEN_URI)
+ .scope(SCOPE)
+ .clientId(CLIENT_ID)
+ .clientSecret(CLIENT_SECRET)
+ .resource(RESOURCE)
+ .subjectToken(SUBJECT_TOKEN)
+ .subjectTokenType(subjectTokenType)
+ .actorToken(actorToken)
+ .actorTokenType(actorTokenType)
+ .audience(audience)
+ .requestedTokenType(requestedTokenType)
+ .build();
+
+ OAuthTokenProvider provider = config.createTokenProvider();
+
+ assertNotNull(provider);
+ assertInstanceOf(TokenExchangeTokenProvider.class, provider);
+
+ TokenExchangeTokenProvider teProvider = (TokenExchangeTokenProvider) provider;
+ assertEquals(URI.create(TOKEN_URI), teProvider.tokenUri);
+ assertNotNull(teProvider.grant);
+ assertEquals(SUBJECT_TOKEN, teProvider.grant.getSubjectToken().getValue());
+ assertEquals(subjectTokenType, teProvider.grant.getSubjectTokenType().getURI().toString());
+ assertEquals(actorToken, teProvider.grant.getActorToken().getValue());
+ assertEquals(actorTokenType, teProvider.grant.getActorTokenType().getURI().toString());
+ assertNotNull(teProvider.grant.getAudience());
+ assertEquals(1, teProvider.grant.getAudience().size());
+ assertEquals(audience, teProvider.grant.getAudience().get(0).getValue());
+ assertEquals(requestedTokenType, teProvider.grant.getRequestedTokenType().getURI().toString());
+ assertEquals(Scope.parse(SCOPE), teProvider.scope);
+ assertEquals(Collections.singletonList(URI.create(RESOURCE)), teProvider.resources);
+
+ assertEquals(CLIENT_ID, teProvider.clientAuth.getClientID().getValue());
+ }
+
+ static Stream generalValidationErrorCases() {
+ return Stream.of(
+ Arguments.of(
+ Named.of(
+ "null flow",
+ (BuilderConfigurer) builder -> builder.flow((String) null).tokenUri(TOKEN_URI)),
+ "OAuth flow cannot be null or empty"),
+ Arguments.of(
+ Named.of(
+ "empty flow", (BuilderConfigurer) builder -> builder.flow("").tokenUri(TOKEN_URI)),
+ "OAuth flow cannot be null or empty"),
+ Arguments.of(
+ Named.of(
+ "invalid flow",
+ (BuilderConfigurer) builder -> builder.flow("invalid_flow").tokenUri(TOKEN_URI)),
+ "Unsupported OAuth flow: invalid_flow"),
+ Arguments.of(
+ Named.of(
+ "null tokenUri",
+ (BuilderConfigurer)
+ builder ->
+ builder
+ .flow("client_credentials")
+ .tokenUri((String) null)
+ .clientId(CLIENT_ID)
+ .clientSecret(CLIENT_SECRET)),
+ "Token URI cannot be null or empty"),
+ Arguments.of(
+ Named.of(
+ "empty tokenUri",
+ (BuilderConfigurer)
+ builder ->
+ builder
+ .flow("client_credentials")
+ .tokenUri("")
+ .clientId(CLIENT_ID)
+ .clientSecret(CLIENT_SECRET)),
+ "Token URI cannot be null or empty"),
+ Arguments.of(
+ Named.of(
+ "invalid tokenUri",
+ (BuilderConfigurer)
+ builder ->
+ builder
+ .flow("client_credentials")
+ .tokenUri("not a valid uri ://")
+ .clientId(CLIENT_ID)
+ .clientSecret(CLIENT_SECRET)),
+ null),
+ Arguments.of(
+ Named.of(
+ "invalid tokenUri",
+ (BuilderConfigurer)
+ builder ->
+ builder.flow("client_credentials").tokenUri(TOKEN_URI).clientId(CLIENT_ID)),
+ // null means verify exception has message and cause
+ "clientSecret is required for client_credentials flow"));
+ }
+
+ @ParameterizedTest
+ @MethodSource("generalValidationErrorCases")
+ public void testGeneralValidationErrors(BuilderConfigurer configurer, String expectedMessage) {
+ SQLException exception =
+ assertThrows(
+ SQLException.class,
+ () -> {
+ OAuthConfiguration.Builder builder = new OAuthConfiguration.Builder();
+ configurer.configure(builder);
+ builder.build();
+ });
+
+ if (expectedMessage != null) {
+ assertEquals(expectedMessage, exception.getMessage());
+ } else {
+ assertNotNull(exception.getMessage());
+ assertNotNull(exception.getCause());
+ }
+ }
+
+ static Stream flowSpecificValidationErrorCases() {
+ return Stream.of(
+ // client_credentials flow validation
+ Arguments.of(
+ Named.of(
+ "client_credentials: missing clientId",
+ (BuilderConfigurer)
+ builder ->
+ builder
+ .flow("client_credentials")
+ .tokenUri(TOKEN_URI)
+ .clientSecret(CLIENT_SECRET)),
+ "clientId is required for client_credentials flow"),
+ Arguments.of(
+ Named.of(
+ "client_credentials: missing clientSecret",
+ (BuilderConfigurer)
+ builder ->
+ builder.flow("client_credentials").tokenUri(TOKEN_URI).clientId(CLIENT_ID)),
+ "clientSecret is required for client_credentials flow"),
+ // token_exchange flow validation
+ Arguments.of(
+ Named.of(
+ "token_exchange: missing subjectToken",
+ (BuilderConfigurer) builder -> builder.flow("token_exchange").tokenUri(TOKEN_URI)),
+ "subjectToken is required for token_exchange flow"),
+ Arguments.of(
+ Named.of(
+ "token_exchange: empty subjectToken",
+ (BuilderConfigurer)
+ builder ->
+ builder
+ .flow("token_exchange")
+ .tokenUri(TOKEN_URI)
+ .subjectToken("")
+ .subjectTokenType("urn:ietf:params:oauth:token-type:access_token")),
+ "subjectToken is required for token_exchange flow"),
+ Arguments.of(
+ Named.of(
+ "token_exchange: missing subjectTokenType",
+ (BuilderConfigurer)
+ builder ->
+ builder
+ .flow("token_exchange")
+ .tokenUri(TOKEN_URI)
+ .subjectToken(SUBJECT_TOKEN)),
+ "subjectTokenType is required for token_exchange flow"),
+ Arguments.of(
+ Named.of(
+ "token_exchange: empty subjectTokenType",
+ (BuilderConfigurer)
+ builder ->
+ builder
+ .flow("token_exchange")
+ .tokenUri(TOKEN_URI)
+ .subjectToken(SUBJECT_TOKEN)
+ .subjectTokenType("")),
+ "subjectTokenType is required for token_exchange flow"));
+ }
+
+ @ParameterizedTest
+ @MethodSource("flowSpecificValidationErrorCases")
+ public void testFlowSpecificValidationErrors(
+ BuilderConfigurer configurer, String expectedMessage) {
+ SQLException exception =
+ assertThrows(
+ SQLException.class,
+ () -> {
+ OAuthConfiguration.Builder builder = new OAuthConfiguration.Builder();
+ configurer.configure(builder);
+ builder.build();
+ });
+
+ assertEquals(expectedMessage, exception.getMessage());
+ }
+}
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthCredentialWriterTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthCredentialWriterTest.java
new file mode 100644
index 0000000000..1a33f7f0ae
--- /dev/null
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthCredentialWriterTest.java
@@ -0,0 +1,95 @@
+/*
+ * 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.client.oauth;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import java.sql.SQLException;
+import org.apache.arrow.flight.CallHeaders;
+import org.apache.arrow.flight.FlightCallHeaders;
+import org.apache.arrow.flight.auth2.Auth2Constants;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
+
+/** Tests for {@link OAuthCredentialWriter}. */
+@ExtendWith(MockitoExtension.class)
+public class OAuthCredentialWriterTest {
+
+ @Mock private OAuthTokenProvider mockTokenProvider;
+
+ @Test
+ public void testConstructorRejectsNullTokenProvider() {
+ assertThrows(NullPointerException.class, () -> new OAuthCredentialWriter(null));
+ }
+
+ @Test
+ public void testAcceptWritesBearerTokenToHeaders() throws SQLException {
+ String testToken = "test-access-token-12345";
+ when(mockTokenProvider.getValidToken()).thenReturn(testToken);
+
+ OAuthCredentialWriter writer = new OAuthCredentialWriter(mockTokenProvider);
+ CallHeaders headers = new FlightCallHeaders();
+
+ writer.accept(headers);
+
+ verify(mockTokenProvider).getValidToken();
+ assertEquals(
+ Auth2Constants.BEARER_PREFIX + testToken, headers.get(Auth2Constants.AUTHORIZATION_HEADER));
+ }
+
+ @Test
+ public void testAcceptThrowsOAuthTokenExceptionOnSQLException() throws SQLException {
+ SQLException sqlException = new SQLException("Token fetch failed");
+ when(mockTokenProvider.getValidToken()).thenThrow(sqlException);
+
+ OAuthCredentialWriter writer = new OAuthCredentialWriter(mockTokenProvider);
+ CallHeaders headers = new FlightCallHeaders();
+
+ OAuthTokenException exception =
+ assertThrows(OAuthTokenException.class, () -> writer.accept(headers));
+
+ assertEquals("Failed to obtain OAuth token", exception.getMessage());
+ assertEquals(sqlException, exception.getCause());
+ }
+
+ @Test
+ public void testAcceptCallsTokenProviderEachTime() throws SQLException {
+ when(mockTokenProvider.getValidToken())
+ .thenReturn("token1")
+ .thenReturn("token2")
+ .thenReturn("token3");
+
+ OAuthCredentialWriter writer = new OAuthCredentialWriter(mockTokenProvider);
+
+ CallHeaders headers1 = new FlightCallHeaders();
+ writer.accept(headers1);
+ assertEquals("Bearer token1", headers1.get(Auth2Constants.AUTHORIZATION_HEADER));
+
+ CallHeaders headers2 = new FlightCallHeaders();
+ writer.accept(headers2);
+ assertEquals("Bearer token2", headers2.get(Auth2Constants.AUTHORIZATION_HEADER));
+
+ CallHeaders headers3 = new FlightCallHeaders();
+ writer.accept(headers3);
+ assertEquals("Bearer token3", headers3.get(Auth2Constants.AUTHORIZATION_HEADER));
+ }
+}
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..7c17755693 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
@@ -28,8 +28,10 @@
import java.sql.SQLException;
import java.sql.Timestamp;
import java.util.ArrayList;
+import java.util.Arrays;
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,22 @@ 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 +100,109 @@ 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/MockFlightSqlProducer.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java
index a8874c4869..6627d91ab6 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java
@@ -52,6 +52,9 @@
import org.apache.arrow.flight.PutResult;
import org.apache.arrow.flight.Result;
import org.apache.arrow.flight.SchemaResult;
+import org.apache.arrow.flight.SessionOptionValue;
+import org.apache.arrow.flight.SetSessionOptionsRequest;
+import org.apache.arrow.flight.SetSessionOptionsResult;
import org.apache.arrow.flight.Ticket;
import org.apache.arrow.flight.sql.FlightSqlProducer;
import org.apache.arrow.flight.sql.SqlInfoBuilder;
@@ -84,6 +87,7 @@
import org.apache.arrow.vector.types.pojo.Schema;
import org.apache.arrow.vector.util.JsonStringArrayList;
import org.apache.calcite.avatica.Meta.StatementType;
+import org.checkerframework.checker.nullness.qual.Nullable;
/** An ad-hoc {@link FlightSqlProducer} for tests. */
public final class MockFlightSqlProducer implements FlightSqlProducer {
@@ -98,6 +102,7 @@ public final class MockFlightSqlProducer implements FlightSqlProducer {
private final SqlInfoBuilder sqlInfoBuilder = new SqlInfoBuilder();
private final Map parameterSchemas = new HashMap<>();
private final Map>> expectedParameterValues = new HashMap<>();
+ private final Map isUpdateMap = new HashMap<>();
private final Map actionTypeCounter = new HashMap<>();
@@ -173,6 +178,40 @@ public void addUpdateQuery(final String sqlCommand, final long updatedRows) {
});
}
+ /**
+ * Registers a new {@link StatementType#SELECT} SQL query, optionally setting the is_update field.
+ *
+ * @param sqlCommand the SQL command under which to register the new query.
+ * @param schema the schema to use for the query result.
+ * @param resultProviders the result provider for this query.
+ * @param isUpdate value to report for the is_update field, or {@code null} to leave it unset.
+ */
+ public void addSelectQuery(
+ final String sqlCommand,
+ final Schema schema,
+ final List> resultProviders,
+ final @Nullable Boolean isUpdate) {
+ addSelectQuery(sqlCommand, schema, resultProviders);
+ if (isUpdate != null) {
+ isUpdateMap.put(sqlCommand, isUpdate);
+ }
+ }
+
+ /**
+ * Registers a new {@link StatementType#UPDATE} SQL query, optionally setting the is_update field.
+ *
+ * @param sqlCommand the SQL command.
+ * @param updatedRows the number of rows affected.
+ * @param isUpdate value to report for the is_update field, or {@code null} to leave it unset.
+ */
+ public void addUpdateQuery(
+ final String sqlCommand, final long updatedRows, final @Nullable Boolean isUpdate) {
+ addUpdateQuery(sqlCommand, updatedRows);
+ if (isUpdate != null) {
+ isUpdateMap.put(sqlCommand, isUpdate);
+ }
+ }
+
/**
* Adds a catalog query to the results.
*
@@ -244,6 +283,12 @@ public void createPreparedStatement(
resultBuilder.setParameterSchema(ByteString.copyFrom(outputStream.toByteArray()));
}
+ // Set is_update field if present
+ final Boolean isUpdate = isUpdateMap.get(query);
+ if (isUpdate != null) {
+ resultBuilder.setIsUpdate(isUpdate);
+ }
+
listener.onNext(new Result(pack(resultBuilder.build()).toByteArray()));
} catch (final Throwable t) {
listener.onError(t);
@@ -664,6 +709,22 @@ public SqlInfoBuilder getSqlInfoBuilder() {
return sqlInfoBuilder;
}
+ private final Map sessionOptions = new HashMap<>();
+
+ @Override
+ public void setSessionOptions(
+ final SetSessionOptionsRequest request,
+ final CallContext context,
+ final StreamListener listener) {
+ sessionOptions.putAll(request.getSessionOptions());
+ listener.onNext(new SetSessionOptionsResult(Collections.emptyMap()));
+ listener.onCompleted();
+ }
+
+ public Map getSessionOptions() {
+ return sessionOptions;
+ }
+
private static final class TicketConversionUtils {
private TicketConversionUtils() {
// Prevent instantiation.
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..c4858d787d 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(UuidType.INSTANCE));
}
@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(UuidType.INSTANCE));
}
}
diff --git a/flight/flight-sql-jdbc-driver/pom.xml b/flight/flight-sql-jdbc-driver/pom.xml
index 559c42597d..801089d090 100644
--- a/flight/flight-sql-jdbc-driver/pom.xml
+++ b/flight/flight-sql-jdbc-driver/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-flight
- 19.0.0-SNAPSHOT
+ 20.0.0-SNAPSHOT
flight-sql-jdbc-driver
@@ -138,7 +138,7 @@ under the License.
-
+
META-INF/LICENSE.txt
src/shade/LICENSE.txt
@@ -159,6 +159,7 @@ under the License.
org.apache.calcite.avatica:*
META-INF/services/java.sql.Driver
+ META-INF/README.txt
diff --git a/flight/flight-sql-jdbc-driver/src/shade/LICENSE.txt b/flight/flight-sql-jdbc-driver/src/shade/LICENSE.txt
index 8bc43cbe0f..8476bd9995 100644
--- a/flight/flight-sql-jdbc-driver/src/shade/LICENSE.txt
+++ b/flight/flight-sql-jdbc-driver/src/shade/LICENSE.txt
@@ -345,6 +345,14 @@ License: https://www.apache.org/licenses/LICENSE-2.0
--------------------------------------------------------------------------------
+This binary artifact contains Nimbus OAuth 2.0 SDK with OpenID Connect extensions 11.20.1.
+
+Copyright: Copyright 2012-2024 Connect2id Ltd.
+Home page: https://connect2id.com/products/nimbus-oauth-openid-connect-sdk
+License: https://www.apache.org/licenses/LICENSE-2.0
+
+--------------------------------------------------------------------------------
+
This binary artifact contains Bouncycastle 1.80.
Copyright: Copyright (c) 2000-2024 The Legion of the Bouncy Castle Inc. (https://www.bouncycastle.org).
diff --git a/flight/flight-sql-jdbc-driver/src/test/java/org/apache/arrow/driver/jdbc/ITDriverJarValidation.java b/flight/flight-sql-jdbc-driver/src/test/java/org/apache/arrow/driver/jdbc/ITDriverJarValidation.java
index c1bd111fb9..145744ad38 100644
--- a/flight/flight-sql-jdbc-driver/src/test/java/org/apache/arrow/driver/jdbc/ITDriverJarValidation.java
+++ b/flight/flight-sql-jdbc-driver/src/test/java/org/apache/arrow/driver/jdbc/ITDriverJarValidation.java
@@ -70,6 +70,10 @@ public class ITDriverJarValidation {
"LICENSE.txt",
"NOTICE.txt",
"arrow-git.properties",
+ "iso3166_1alpha2-codes.properties",
+ "iso3166_1alpha3-codes.properties",
+ "iso3166_1alpha-2-3-map.properties",
+ "iso3166_3-codes.properties",
"properties/flight.properties",
"META-INF/io.netty.versions.properties",
"META-INF/MANIFEST.MF",
diff --git a/flight/flight-sql/pom.xml b/flight/flight-sql/pom.xml
index 15d00e3e18..b7c8931391 100644
--- a/flight/flight-sql/pom.xml
+++ b/flight/flight-sql/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-flight
- 19.0.0-SNAPSHOT
+ 20.0.0-SNAPSHOT
flight-sql
@@ -95,7 +95,7 @@ under the License.
org.apache.commons
commons-dbcp2
- 2.13.0
+ 2.14.0
test
@@ -107,19 +107,19 @@ under the License.
org.apache.commons
commons-pool2
- 2.12.1
+ 2.13.1
test
org.apache.commons
commons-text
- 1.13.1
+ 1.15.0
test
commons-cli
commons-cli
- 1.9.0
+ 1.11.0
true
diff --git a/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java b/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java
index 623f9311e8..0af09faee1 100644
--- a/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java
+++ b/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java
@@ -1284,6 +1284,19 @@ public Schema getParameterSchema() {
return parameterSchema;
}
+ /**
+ * Returns whether the server indicated this prepared statement is an update query.
+ *
+ * @return true if the server indicated this is an update query, false if the server indicated
+ * this is a select query, or null if the server did not provide this information.
+ */
+ public Boolean isUpdate() {
+ if (preparedStatementResult.hasIsUpdate()) {
+ return preparedStatementResult.getIsUpdate();
+ }
+ return null;
+ }
+
/** Get the schema of the result set (should be identical to {@link #getResultSetSchema()}). */
public SchemaResult fetchSchema(CallOption... options) {
checkOpen();
diff --git a/flight/pom.xml b/flight/pom.xml
index 2fc3e89ef8..30f75fa27e 100644
--- a/flight/pom.xml
+++ b/flight/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-java-root
- 19.0.0-SNAPSHOT
+ 20.0.0-SNAPSHOT
arrow-flight
diff --git a/format/pom.xml b/format/pom.xml
index d3578b63d2..8c2f2d3387 100644
--- a/format/pom.xml
+++ b/format/pom.xml
@@ -23,7 +23,7 @@ under the License.
org.apache.arrow
arrow-java-root
- 19.0.0-SNAPSHOT
+ 20.0.0-SNAPSHOT
arrow-format
diff --git a/gandiva/pom.xml b/gandiva/pom.xml
index 5367bfdedf..190bf016ce 100644
--- a/gandiva/pom.xml
+++ b/gandiva/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-java-root
- 19.0.0-SNAPSHOT
+ 20.0.0-SNAPSHOT
org.apache.arrow.gandiva
diff --git a/gandiva/src/main/cpp/jni_common.cc b/gandiva/src/main/cpp/jni_common.cc
index 2851250072..ec4888a512 100644
--- a/gandiva/src/main/cpp/jni_common.cc
+++ b/gandiva/src/main/cpp/jni_common.cc
@@ -221,7 +221,7 @@ DataTypePtr ProtoTypeToDataType(const gandiva::types::ExtGandivaType& ext_type)
return arrow::date64();
case gandiva::types::DECIMAL:
// TODO: error handling
- return arrow::decimal(ext_type.precision(), ext_type.scale());
+ return arrow::decimal128(ext_type.precision(), ext_type.scale());
case gandiva::types::TIME32:
return ProtoTypeToTime32(ext_type);
case gandiva::types::TIME64:
diff --git a/memory/memory-core/pom.xml b/memory/memory-core/pom.xml
index 72ee69d60a..825b3dae4b 100644
--- a/memory/memory-core/pom.xml
+++ b/memory/memory-core/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-memory
- 19.0.0-SNAPSHOT
+ 20.0.0-SNAPSHOT
arrow-memory-core
@@ -100,8 +100,8 @@ under the License.
test
-
-
+
+
**/TestOpens.java
diff --git a/memory/memory-core/src/main/java/org/apache/arrow/memory/Accountant.java b/memory/memory-core/src/main/java/org/apache/arrow/memory/Accountant.java
index 5d052c2cde..d4d76f57f4 100644
--- a/memory/memory-core/src/main/java/org/apache/arrow/memory/Accountant.java
+++ b/memory/memory-core/src/main/java/org/apache/arrow/memory/Accountant.java
@@ -16,7 +16,7 @@
*/
package org.apache.arrow.memory;
-import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicLongFieldUpdater;
import org.apache.arrow.util.Preconditions;
import org.checkerframework.checker.nullness.qual.Nullable;
@@ -37,16 +37,24 @@ class Accountant implements AutoCloseable {
*/
protected final long reservation;
- private final AtomicLong peakAllocation = new AtomicLong();
+ // AtomicLongFieldUpdaters for memory accounting fields to reduce memory overhead
+ private static final AtomicLongFieldUpdater PEAK_ALLOCATION_UPDATER =
+ AtomicLongFieldUpdater.newUpdater(Accountant.class, "peakAllocation");
+ private static final AtomicLongFieldUpdater ALLOCATION_LIMIT_UPDATER =
+ AtomicLongFieldUpdater.newUpdater(Accountant.class, "allocationLimit");
+ private static final AtomicLongFieldUpdater LOCALLY_HELD_MEMORY_UPDATER =
+ AtomicLongFieldUpdater.newUpdater(Accountant.class, "locallyHeldMemory");
+
+ private volatile long peakAllocation = 0;
/**
* Maximum local memory that can be held. This can be externally updated. Changing it won't cause
* past memory to change but will change responses to future allocation efforts
*/
- private final AtomicLong allocationLimit = new AtomicLong();
+ private volatile long allocationLimit = 0;
/** Currently allocated amount of memory. */
- private final AtomicLong locallyHeldMemory = new AtomicLong();
+ private volatile long locallyHeldMemory = 0;
public Accountant(
@Nullable Accountant parent, String name, long reservation, long maxAllocation) {
@@ -64,7 +72,7 @@ public Accountant(
this.parent = parent;
this.name = name;
this.reservation = reservation;
- this.allocationLimit.set(maxAllocation);
+ ALLOCATION_LIMIT_UPDATER.set(this, maxAllocation);
if (reservation != 0) {
Preconditions.checkArgument(parent != null, "parent must not be null");
@@ -117,12 +125,12 @@ private AllocationOutcome.Status allocateBytesInternal(long size) {
}
private void updatePeak() {
- final long currentMemory = locallyHeldMemory.get();
+ final long currentMemory = locallyHeldMemory;
while (true) {
- final long previousPeak = peakAllocation.get();
+ final long previousPeak = peakAllocation;
if (currentMemory > previousPeak) {
- if (!peakAllocation.compareAndSet(previousPeak, currentMemory)) {
+ if (!PEAK_ALLOCATION_UPDATER.compareAndSet(this, previousPeak, currentMemory)) {
// peak allocation changed underneath us. try again.
continue;
}
@@ -166,7 +174,7 @@ private AllocationOutcome.Status allocate(
final boolean incomingUpdatePeak,
final boolean forceAllocation,
@Nullable AllocationOutcomeDetails details) {
- final long oldLocal = locallyHeldMemory.getAndAdd(size);
+ final long oldLocal = LOCALLY_HELD_MEMORY_UPDATER.getAndAdd(this, size);
final long newLocal = oldLocal + size;
// Borrowed from Math.addExact (but avoid exception here)
// Overflow if result has opposite sign of both arguments
@@ -174,7 +182,7 @@ private AllocationOutcome.Status allocate(
// failure
final boolean overflow = ((oldLocal ^ newLocal) & (size ^ newLocal)) < 0;
final long beyondReservation = newLocal - reservation;
- final boolean beyondLimit = overflow || newLocal > allocationLimit.get();
+ final boolean beyondLimit = overflow || newLocal > allocationLimit;
final boolean updatePeak = forceAllocation || (incomingUpdatePeak && !beyondLimit);
if (details != null) {
@@ -214,7 +222,7 @@ private AllocationOutcome.Status allocate(
public void releaseBytes(long size) {
// reduce local memory. all memory released above reservation should be released up the tree.
- final long newSize = locallyHeldMemory.addAndGet(-size);
+ final long newSize = LOCALLY_HELD_MEMORY_UPDATER.addAndGet(this, -size);
Preconditions.checkArgument(newSize >= 0, "Accounted size went negative.");
@@ -255,7 +263,7 @@ public String getName() {
* @return Limit in bytes.
*/
public long getLimit() {
- return allocationLimit.get();
+ return allocationLimit;
}
/**
@@ -274,7 +282,7 @@ public long getInitReservation() {
* @param newLimit The limit in bytes.
*/
public void setLimit(long newLimit) {
- allocationLimit.set(newLimit);
+ ALLOCATION_LIMIT_UPDATER.set(this, newLimit);
}
/**
@@ -284,7 +292,7 @@ public void setLimit(long newLimit) {
* @return Currently allocate memory in bytes.
*/
public long getAllocatedMemory() {
- return locallyHeldMemory.get();
+ return locallyHeldMemory;
}
/**
@@ -293,17 +301,17 @@ public long getAllocatedMemory() {
* @return The peak allocated memory in bytes.
*/
public long getPeakMemoryAllocation() {
- return peakAllocation.get();
+ return peakAllocation;
}
public long getHeadroom() {
- long localHeadroom = allocationLimit.get() - locallyHeldMemory.get();
+ long localHeadroom = allocationLimit - locallyHeldMemory;
if (parent == null) {
return localHeadroom;
}
// Amount of reserved memory left on top of what parent has
- long reservedHeadroom = Math.max(0, reservation - locallyHeldMemory.get());
+ long reservedHeadroom = Math.max(0, reservation - locallyHeldMemory);
return Math.min(localHeadroom, parent.getHeadroom() + reservedHeadroom);
}
}
diff --git a/memory/memory-core/src/main/java/org/apache/arrow/memory/ArrowBuf.java b/memory/memory-core/src/main/java/org/apache/arrow/memory/ArrowBuf.java
index 775a8925ad..9712be34d7 100644
--- a/memory/memory-core/src/main/java/org/apache/arrow/memory/ArrowBuf.java
+++ b/memory/memory-core/src/main/java/org/apache/arrow/memory/ArrowBuf.java
@@ -24,7 +24,6 @@
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.ReadOnlyBufferException;
-import java.util.concurrent.atomic.AtomicLong;
import org.apache.arrow.memory.BaseAllocator.Verbosity;
import org.apache.arrow.memory.util.CommonUtil;
import org.apache.arrow.memory.util.HistoricalLog;
@@ -57,9 +56,8 @@ public final class ArrowBuf implements AutoCloseable {
private static final int DOUBLE_SIZE = Double.BYTES;
private static final int LONG_SIZE = Long.BYTES;
- private static final AtomicLong idGenerator = new AtomicLong(0);
private static final int LOG_BYTES_PER_ROW = 10;
- private final long id = idGenerator.incrementAndGet();
+
private final ReferenceManager referenceManager;
private final @Nullable BufferManager bufferManager;
private final long addr;
@@ -67,7 +65,8 @@ public final class ArrowBuf implements AutoCloseable {
private long writerIndex;
private final @Nullable HistoricalLog historicalLog =
BaseAllocator.DEBUG
- ? new HistoricalLog(BaseAllocator.DEBUG_LOG_LENGTH, "ArrowBuf[%d]", id)
+ ? new HistoricalLog(
+ BaseAllocator.DEBUG_LOG_LENGTH, "ArrowBuf[%d]", System.identityHashCode(this))
: null;
private volatile long capacity;
@@ -136,7 +135,7 @@ public long capacity() {
/**
* Adjusts the capacity of this buffer. Size increases are NOT supported.
*
- * @param newCapacity Must be in in the range [0, length).
+ * @param newCapacity Must be in the range [0, length).
*/
public synchronized ArrowBuf capacity(long newCapacity) {
@@ -218,7 +217,8 @@ public long memoryAddress() {
@Override
public String toString() {
- return String.format("ArrowBuf[%d], address:%d, capacity:%d", id, memoryAddress(), capacity);
+ return String.format(
+ "ArrowBuf[%d], address:%d, capacity:%d", getId(), memoryAddress(), capacity);
}
@Override
@@ -1080,12 +1080,15 @@ public String toHexString(final long start, final int length) {
}
/**
- * Get the integer id assigned to this ArrowBuf for debugging purposes.
+ * Get the id assigned to this ArrowBuf for debugging purposes.
+ *
+ * Returns {@link System#identityHashCode(Object)} which provides a unique identifier for this
+ * buffer without any per-instance memory overhead.
*
- * @return integer id
+ * @return the identity hash code for this buffer
*/
public long getId() {
- return id;
+ return System.identityHashCode(this);
}
/**
diff --git a/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferAllocator.java b/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferAllocator.java
index 4f9d3c61c6..dbd6da3291 100644
--- a/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferAllocator.java
+++ b/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferAllocator.java
@@ -25,9 +25,10 @@
public interface BufferAllocator extends AutoCloseable {
/**
- * Allocate a new or reused buffer of the provided size. Note that the buffer may technically be
- * larger than the requested size for rounding purposes. However, the buffer's capacity will be
- * set to the configured size.
+ * Allocate a new or reused buffer of the provided size. The buffer may be larger than the
+ * requested size for rounding purposes (e.g. to a power of two), and the buffer's capacity will
+ * reflect the actual allocated size. Use {@link ArrowBuf#capacity(long)} to set the capacity to
+ * the requested size if needed.
*
* @param size The size in bytes.
* @return a new ArrowBuf, or null if the request can't be satisfied
@@ -36,9 +37,10 @@ public interface BufferAllocator extends AutoCloseable {
ArrowBuf buffer(long size);
/**
- * Allocate a new or reused buffer of the provided size. Note that the buffer may technically be
- * larger than the requested size for rounding purposes. However, the buffer's capacity will be
- * set to the configured size.
+ * Allocate a new or reused buffer of the provided size. The buffer may be larger than the
+ * requested size for rounding purposes (e.g. to a power of two), and the buffer's capacity will
+ * reflect the actual allocated size. Use {@link ArrowBuf#capacity(long)} to set the capacity to
+ * the requested size if needed.
*
* @param size The size in bytes.
* @param manager A buffer manager to manage reallocation.
diff --git a/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferLedger.java b/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferLedger.java
index b562a421e7..eb90efcbb5 100644
--- a/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferLedger.java
+++ b/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferLedger.java
@@ -17,8 +17,7 @@
package org.apache.arrow.memory;
import java.util.IdentityHashMap;
-import java.util.concurrent.atomic.AtomicInteger;
-import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
import org.apache.arrow.memory.util.CommonUtil;
import org.apache.arrow.memory.util.HistoricalLog;
import org.apache.arrow.util.Preconditions;
@@ -32,12 +31,13 @@
public class BufferLedger implements ValueWithKeyIncluded, ReferenceManager {
private final @Nullable IdentityHashMap buffers =
BaseAllocator.DEBUG ? new IdentityHashMap<>() : null;
- private static final AtomicLong LEDGER_ID_GENERATOR = new AtomicLong(0);
- // unique ID assigned to each ledger
- private final long ledgerId = LEDGER_ID_GENERATOR.incrementAndGet();
- private final AtomicInteger bufRefCnt = new AtomicInteger(0); // start at zero so we can
- // manage request for retain
- // correctly
+
+ // AtomicIntegerFieldUpdater for bufRefCnt to reduce memory overhead
+ private static final AtomicIntegerFieldUpdater BUF_REF_CNT_UPDATER =
+ AtomicIntegerFieldUpdater.newUpdater(BufferLedger.class, "bufRefCnt");
+ // start at zero so we can manage request for retain correctly
+ private volatile int bufRefCnt = 0;
+
private final long lCreationTime = System.nanoTime();
private final BufferAllocator allocator;
private final AllocationManager allocationManager;
@@ -78,7 +78,7 @@ public BufferAllocator getAllocator() {
*/
@Override
public int getRefCount() {
- return bufRefCnt.get();
+ return bufRefCnt;
}
/**
@@ -86,7 +86,7 @@ public int getRefCount() {
* ArrowBufs managed by this ledger will share the ref count.
*/
void increment() {
- bufRefCnt.incrementAndGet();
+ BUF_REF_CNT_UPDATER.incrementAndGet(this);
}
/**
@@ -144,7 +144,7 @@ private int decrement(int decrement) {
allocator.assertOpen();
final int outcome;
synchronized (allocationManager) {
- outcome = bufRefCnt.addAndGet(-decrement);
+ outcome = BUF_REF_CNT_UPDATER.addAndGet(this, -decrement);
if (outcome == 0) {
lDestructionTime = System.nanoTime();
// refcount of this reference manager has dropped to 0
@@ -174,7 +174,7 @@ public void retain(int increment) {
if (historicalLog != null) {
historicalLog.recordEvent("retain(%d)", increment);
}
- final int originalReferenceCount = bufRefCnt.getAndAdd(increment);
+ final int originalReferenceCount = BUF_REF_CNT_UPDATER.getAndAdd(this, increment);
Preconditions.checkArgument(originalReferenceCount > 0);
}
@@ -472,13 +472,13 @@ public long getAccountedSize() {
void print(StringBuilder sb, int indent, BaseAllocator.Verbosity verbosity) {
CommonUtil.indent(sb, indent)
.append("ledger[")
- .append(ledgerId)
+ .append(System.identityHashCode(this))
.append("] allocator: ")
.append(allocator.getName())
.append("), isOwning: ")
.append(", size: ")
.append(", references: ")
- .append(bufRefCnt.get())
+ .append(bufRefCnt)
.append(", life: ")
.append(lCreationTime)
.append("..")
diff --git a/memory/memory-core/src/main/java/org/apache/arrow/memory/util/MemoryUtil.java b/memory/memory-core/src/main/java/org/apache/arrow/memory/util/MemoryUtil.java
index 91bd7cd905..be0749a215 100644
--- a/memory/memory-core/src/main/java/org/apache/arrow/memory/util/MemoryUtil.java
+++ b/memory/memory-core/src/main/java/org/apache/arrow/memory/util/MemoryUtil.java
@@ -18,6 +18,7 @@
import java.lang.reflect.Constructor;
import java.lang.reflect.Field;
+import java.lang.reflect.InaccessibleObjectException;
import java.lang.reflect.InvocationTargetException;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
@@ -81,9 +82,18 @@ public Object run() {
BYTE_ARRAY_BASE_OFFSET = UNSAFE.arrayBaseOffset(byte[].class);
// get the offset of the address field in a java.nio.Buffer object
+ long maybeOffset;
Field addressField = java.nio.Buffer.class.getDeclaredField("address");
- addressField.setAccessible(true);
- BYTE_BUFFER_ADDRESS_OFFSET = UNSAFE.objectFieldOffset(addressField);
+ try {
+ addressField.setAccessible(true);
+ maybeOffset = UNSAFE.objectFieldOffset(addressField);
+ } catch (InaccessibleObjectException e) {
+ maybeOffset = -1;
+ logger.debug(
+ "Cannot access the address field of java.nio.Buffer. DirectBuffer operations wont be available",
+ e);
+ }
+ BYTE_BUFFER_ADDRESS_OFFSET = maybeOffset;
Constructor directBufferConstructor;
long address = -1;
@@ -109,6 +119,9 @@ public Object run() {
} catch (SecurityException e) {
logger.debug("Cannot get constructor for direct buffer allocation", e);
return e;
+ } catch (InaccessibleObjectException e) {
+ logger.debug("Cannot get constructor for direct buffer allocation", e);
+ return e;
}
}
});
@@ -156,7 +169,11 @@ public Object run() {
* @return address of the underlying memory.
*/
public static long getByteBufferAddress(ByteBuffer buf) {
- return UNSAFE.getLong(buf, BYTE_BUFFER_ADDRESS_OFFSET);
+ if (BYTE_BUFFER_ADDRESS_OFFSET != -1) {
+ return UNSAFE.getLong(buf, BYTE_BUFFER_ADDRESS_OFFSET);
+ }
+ throw new UnsupportedOperationException(
+ "Byte buffer address cannot be obtained because sun.misc.Unsafe or java.nio.DirectByteBuffer.(long, int) is not available");
}
private MemoryUtil() {}
diff --git a/memory/memory-core/src/main/java/org/apache/arrow/util/AutoCloseables.java b/memory/memory-core/src/main/java/org/apache/arrow/util/AutoCloseables.java
index a39004a9d0..ba5a539a87 100644
--- a/memory/memory-core/src/main/java/org/apache/arrow/util/AutoCloseables.java
+++ b/memory/memory-core/src/main/java/org/apache/arrow/util/AutoCloseables.java
@@ -22,7 +22,9 @@
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
+import java.util.stream.Stream;
import java.util.stream.StreamSupport;
+import org.checkerframework.checker.nullness.qual.Nullable;
/** Utilities for AutoCloseable classes. */
public final class AutoCloseables {
@@ -33,7 +35,8 @@ private AutoCloseables() {}
* Returns a new {@link AutoCloseable} that calls {@link #close(Iterable)} on autoCloseables
* when close is called.
*/
- public static AutoCloseable all(final Collection autoCloseables) {
+ public static AutoCloseable all(
+ final @Nullable Collection autoCloseables) {
return new AutoCloseable() {
@Override
public void close() throws Exception {
@@ -48,7 +51,10 @@ public void close() throws Exception {
* @param t the throwable to add suppressed exception to
* @param autoCloseables the closeables to close
*/
- public static void close(Throwable t, AutoCloseable... autoCloseables) {
+ public static void close(Throwable t, @Nullable AutoCloseable... autoCloseables) {
+ if (autoCloseables == null) {
+ return;
+ }
close(t, Arrays.asList(autoCloseables));
}
@@ -58,7 +64,8 @@ public static void close(Throwable t, AutoCloseable... autoCloseables) {
* @param t the throwable to add suppressed exception to
* @param autoCloseables the closeables to close
*/
- public static void close(Throwable t, Iterable autoCloseables) {
+ public static void close(
+ Throwable t, @Nullable Iterable autoCloseables) {
try {
close(autoCloseables);
} catch (Exception e) {
@@ -71,7 +78,10 @@ public static void close(Throwable t, Iterable autoClos
*
* @param autoCloseables the closeables to close
*/
- public static void close(AutoCloseable... autoCloseables) throws Exception {
+ public static void close(@Nullable AutoCloseable... autoCloseables) throws Exception {
+ if (autoCloseables == null) {
+ return;
+ }
close(Arrays.asList(autoCloseables));
}
@@ -80,7 +90,8 @@ public static void close(AutoCloseable... autoCloseables) throws Exception {
*
* @param ac the closeables to close
*/
- public static void close(Iterable ac) throws Exception {
+ public static void close(@Nullable Iterable ac)
+ throws Exception {
// this method can be called on a single object if it implements Iterable
// like for example VectorContainer make sure we handle that properly
if (ac == null) {
@@ -111,12 +122,17 @@ public static void close(Iterable ac) throws Exception
/** Calls {@link #close(Iterable)} on the flattened list of closeables. */
@SafeVarargs
- public static void close(Iterable... closeables) throws Exception {
+ public static void close(@Nullable Iterable... closeables)
+ throws Exception {
+ if (closeables == null) {
+ return;
+ }
close(flatten(closeables));
}
@SafeVarargs
- private static Iterable flatten(Iterable... closeables) {
+ private static Iterable flatten(
+ Iterable... closeables) {
return new Iterable() {
// Cast from Iterable to Iterable is safe in this
// context
@@ -127,16 +143,18 @@ public Iterator iterator() {
return Arrays.stream(closeables)
.flatMap(
(Iterable i) ->
- StreamSupport.stream(
- ((Iterable) i).spliterator(), /* parallel= */ false))
+ i == null
+ ? Stream.empty()
+ : StreamSupport.stream(
+ ((Iterable) i).spliterator(), /* parallel= */ false))
.iterator();
}
};
}
/** Converts ac to a {@link Iterable} filtering out any null values. */
- public static Iterable iter(AutoCloseable... ac) {
- if (ac.length == 0) {
+ public static Iterable iter(@Nullable AutoCloseable... ac) {
+ if (ac == null || ac.length == 0) {
return Collections.emptyList();
} else {
final List nonNullAc = new ArrayList<>();
@@ -153,10 +171,11 @@ public static Iterable iter(AutoCloseable... ac) {
public static class RollbackCloseable implements AutoCloseable {
private boolean commit = false;
- private List closeables;
+ private final List closeables;
- public RollbackCloseable(AutoCloseable... closeables) {
- this.closeables = new ArrayList<>(Arrays.asList(closeables));
+ public RollbackCloseable(@Nullable AutoCloseable... closeables) {
+ this.closeables =
+ closeables == null ? new ArrayList<>() : new ArrayList<>(Arrays.asList(closeables));
}
public T add(T t) {
@@ -165,12 +184,18 @@ public T add(T t) {
}
/** Add all of list to the rollback list. */
- public void addAll(AutoCloseable... list) {
+ public void addAll(@Nullable AutoCloseable... list) {
+ if (list == null) {
+ return;
+ }
closeables.addAll(Arrays.asList(list));
}
/** Add all of list to the rollback list. */
- public void addAll(Iterable list) {
+ public void addAll(@Nullable Iterable list) {
+ if (list == null) {
+ return;
+ }
for (AutoCloseable ac : list) {
closeables.add(ac);
}
@@ -189,7 +214,7 @@ public void close() throws Exception {
}
/** Creates an {@link RollbackCloseable} from the given closeables. */
- public static RollbackCloseable rollbackable(AutoCloseable... closeables) {
+ public static RollbackCloseable rollbackable(@Nullable AutoCloseable... closeables) {
return new RollbackCloseable(closeables);
}
@@ -203,7 +228,7 @@ public static RollbackCloseable rollbackable(AutoCloseable... closeables) {
* @throws RuntimeException if an Exception occurs; the Exception is wrapped by the
* RuntimeException
*/
- public static void closeNoChecked(final AutoCloseable autoCloseable) {
+ public static void closeNoChecked(final @Nullable AutoCloseable autoCloseable) {
if (autoCloseable != null) {
try {
autoCloseable.close();
diff --git a/memory/memory-core/src/test/java/org/apache/arrow/memory/TestOpens.java b/memory/memory-core/src/test/java/org/apache/arrow/memory/TestOpens.java
index b5e0a71e7e..f74bf63f82 100644
--- a/memory/memory-core/src/test/java/org/apache/arrow/memory/TestOpens.java
+++ b/memory/memory-core/src/test/java/org/apache/arrow/memory/TestOpens.java
@@ -20,32 +20,27 @@
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.condition.JRE.JAVA_16;
+import org.apache.arrow.memory.util.MemoryUtil;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledForJreRange;
public class TestOpens {
- /** Instantiating the RootAllocator should poke MemoryUtil and fail. */
+ /** Accessing MemoryUtil.directBuffer should fail as add-opens is not configured. */
@Test
@EnabledForJreRange(min = JAVA_16)
public void testMemoryUtilFailsLoudly() {
// This test is configured by Maven to run WITHOUT add-opens. So this should fail on JDK16+
// (where JEP396 means that add-opens is required to access JDK internals).
// The test will likely fail in your IDE if it doesn't correctly pick this up.
- Throwable e =
- assertThrows(
- Throwable.class,
- () -> {
- BufferAllocator allocator = new RootAllocator();
- allocator.close();
- });
+ Throwable e = assertThrows(Throwable.class, () -> MemoryUtil.directBuffer(0, 10));
boolean found = false;
while (e != null) {
- e = e.getCause();
- if (e instanceof RuntimeException
- && e.getMessage().contains("Failed to initialize MemoryUtil")) {
+ if (e instanceof UnsupportedOperationException
+ && e.getMessage().contains("java.nio.DirectByteBuffer.(long, int) not available")) {
found = true;
break;
}
+ e = e.getCause();
}
assertTrue(found, "Expected exception was not thrown");
}
diff --git a/memory/memory-core/src/test/java/org/apache/arrow/util/TestAutoCloseables.java b/memory/memory-core/src/test/java/org/apache/arrow/util/TestAutoCloseables.java
new file mode 100644
index 0000000000..ba5b78178a
--- /dev/null
+++ b/memory/memory-core/src/test/java/org/apache/arrow/util/TestAutoCloseables.java
@@ -0,0 +1,268 @@
+/*
+ * 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.util;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.List;
+import org.junit.jupiter.api.Test;
+
+public class TestAutoCloseables {
+
+ /** Closeable that records that it was closed and can optionally throw. */
+ private static final class TrackCloseable implements AutoCloseable {
+ private boolean closed;
+ private final Exception toThrow;
+
+ TrackCloseable() {
+ this.toThrow = null;
+ }
+
+ TrackCloseable(Exception toThrow) {
+ this.toThrow = toThrow;
+ }
+
+ @Override
+ public void close() throws Exception {
+ closed = true;
+ if (toThrow != null) {
+ throw toThrow;
+ }
+ }
+
+ boolean isClosed() {
+ return closed;
+ }
+ }
+
+ @Test
+ public void testCloseVarargsIgnoresNulls() throws Exception {
+ TrackCloseable a = new TrackCloseable();
+ TrackCloseable b = new TrackCloseable();
+ AutoCloseables.close(a, null, b);
+ assertTrue(a.isClosed());
+ assertTrue(b.isClosed());
+ }
+
+ @Test
+ public void testCloseVarargsThrowsFirstExceptionAndSuppressesRest() throws Exception {
+ Exception e1 = new Exception("first");
+ Exception e2 = new Exception("second");
+ TrackCloseable c1 = new TrackCloseable(e1);
+ TrackCloseable c2 = new TrackCloseable(e2);
+ Exception thrown = assertThrows(Exception.class, () -> AutoCloseables.close(c1, c2));
+ assertEquals("first", thrown.getMessage());
+ assertTrue(Arrays.asList(thrown.getSuppressed()).contains(e2));
+ }
+
+ @Test
+ public void testCloseIterableNullIterableReturns() throws Exception {
+ AutoCloseables.close((List) null); // no exception
+ }
+
+ @Test
+ public void testCloseIterableIgnoresNullElements() throws Exception {
+ TrackCloseable a = new TrackCloseable();
+ TrackCloseable b = new TrackCloseable();
+ List list = Arrays.asList(a, null, b);
+ AutoCloseables.close(list);
+ assertTrue(a.isClosed());
+ assertTrue(b.isClosed());
+ }
+
+ @Test
+ public void testCloseIterableWhenIterableIsAlsoAutoCloseable() throws Exception {
+ TrackCloseable iter = new TrackCloseable();
+ TrackCloseable inner = new TrackCloseable();
+ // When the Iterable itself implements AutoCloseable (e.g. VectorContainer),
+ // close(Iterable) calls close() on it and does not iterate over elements
+ class IterableCloseable implements Iterable, AutoCloseable {
+ @Override
+ @SuppressWarnings("unchecked")
+ public Iterator iterator() {
+ return (Iterator) Collections.singletonList(inner);
+ }
+
+ @Override
+ public void close() throws Exception {
+ iter.close();
+ }
+ }
+ AutoCloseables.close(new IterableCloseable());
+ assertTrue(iter.isClosed());
+ assertFalse(inner.isClosed());
+ }
+
+ @Test
+ public void testCloseIterableVarargsWithNullIterables() throws Exception {
+ TrackCloseable a = new TrackCloseable();
+ TrackCloseable b = new TrackCloseable();
+ TrackCloseable c = new TrackCloseable();
+ List list1 = Arrays.asList(null, a, b);
+ List list2 = Collections.singletonList(c);
+ AutoCloseables.close(list1, null, list2);
+ assertTrue(a.isClosed());
+ assertTrue(b.isClosed());
+ assertTrue(c.isClosed());
+ }
+
+ @Test
+ public void testCloseThrowableSuppressesException() {
+ Exception e = new Exception("from close");
+ TrackCloseable c = new TrackCloseable(e);
+ Exception main = new Exception("main");
+ AutoCloseables.close(main, c);
+ assertTrue(c.isClosed());
+ assertEquals(1, main.getSuppressed().length);
+ assertEquals(e, main.getSuppressed()[0]);
+ }
+
+ @Test
+ public void testCloseThrowableWithNullCloseables() {
+ Exception main = new Exception("main");
+ AutoCloseables.close(main, (AutoCloseable) null);
+ assertEquals(0, main.getSuppressed().length);
+
+ AutoCloseables.close(main, (AutoCloseable[]) null); // no exception
+ }
+
+ @Test
+ public void testIterFiltersNulls() {
+ TrackCloseable a = new TrackCloseable();
+ TrackCloseable b = new TrackCloseable();
+ Iterable it = AutoCloseables.iter(a, null, b);
+ List list = new ArrayList<>();
+ it.forEach(list::add);
+ assertEquals(2, list.size());
+ assertTrue(list.contains(a));
+ assertTrue(list.contains(b));
+ }
+
+ @Test
+ public void testIterEmptyVarargs() {
+ Iterable it = AutoCloseables.iter();
+ List list = new ArrayList<>();
+ it.forEach(list::add);
+ assertTrue(list.isEmpty());
+ }
+
+ @Test
+ public void testIterWithNull() {
+ AutoCloseables.iter((AutoCloseable) null); // no exception
+ }
+
+ @Test
+ public void testCloseNoCheckedWithNull() {
+ AutoCloseables.closeNoChecked(null); // no exception
+ }
+
+ @Test
+ public void testCloseNoCheckedWrapsException() {
+ Exception e = new Exception("close failed");
+ TrackCloseable c = new TrackCloseable(e);
+ RuntimeException re =
+ assertThrows(RuntimeException.class, () -> AutoCloseables.closeNoChecked(c));
+ assertSame(re.getCause(), e);
+ assertTrue(re.getMessage().contains("close failed"));
+ }
+
+ @Test
+ public void testNoop() throws Exception {
+ AutoCloseable noop = AutoCloseables.noop();
+ assertSame(noop, AutoCloseables.noop());
+ noop.close(); // no exception
+ }
+
+ @Test
+ public void testAllClosesCollectionOnClose() throws Exception {
+ TrackCloseable a = new TrackCloseable();
+ TrackCloseable b = new TrackCloseable();
+ List list = Arrays.asList(a, b);
+ AutoCloseable all = AutoCloseables.all(list);
+ assertFalse(a.isClosed());
+ assertFalse(b.isClosed());
+ all.close();
+ assertTrue(a.isClosed());
+ assertTrue(b.isClosed());
+ }
+
+ @Test
+ public void testAllWithNullCollection() throws Exception {
+ AutoCloseable all = AutoCloseables.all(null);
+ all.close(); // no exception
+ }
+
+ @Test
+ public void testRollbackCloseableClosesWhenNotCommitted() throws Exception {
+ TrackCloseable a = new TrackCloseable();
+ TrackCloseable b = new TrackCloseable();
+ AutoCloseables.RollbackCloseable rb = AutoCloseables.rollbackable(a, b);
+ rb.close();
+ assertTrue(a.isClosed());
+ assertTrue(b.isClosed());
+ }
+
+ @Test
+ public void testRollbackCloseableDoesNotCloseWhenCommitted() throws Exception {
+ TrackCloseable a = new TrackCloseable();
+ TrackCloseable b = new TrackCloseable();
+ AutoCloseables.RollbackCloseable rb = AutoCloseables.rollbackable(a, b);
+ rb.commit();
+ rb.close();
+ assertFalse(a.isClosed());
+ assertFalse(b.isClosed());
+ }
+
+ @Test
+ public void testRollbackCloseableAddAndAddAll() throws Exception {
+ TrackCloseable a = new TrackCloseable();
+ TrackCloseable b = new TrackCloseable();
+ TrackCloseable c = new TrackCloseable();
+ TrackCloseable d = new TrackCloseable();
+ AutoCloseables.RollbackCloseable rb = AutoCloseables.rollbackable(a);
+ rb.add(b);
+ rb.addAll(c, d);
+ rb.addAll((AutoCloseable[]) null); // null varargs shouldn't fail
+ rb.addAll((List) null); // null Iterable shouldn't fail
+ rb.close();
+ assertTrue(a.isClosed());
+ assertTrue(b.isClosed());
+ assertTrue(c.isClosed());
+ assertTrue(d.isClosed());
+ }
+
+ @Test
+ public void testRollbackCloseableWithNull() throws Exception {
+ AutoCloseables.rollbackable((AutoCloseable) null); // no exception
+ }
+
+ @Test
+ public void testRollbackCloseableWithNulls() throws Exception {
+ TrackCloseable a = new TrackCloseable();
+ AutoCloseables.RollbackCloseable rb = AutoCloseables.rollbackable(a, null);
+ rb.close();
+ assertTrue(a.isClosed());
+ }
+}
diff --git a/memory/memory-netty-buffer-patch/pom.xml b/memory/memory-netty-buffer-patch/pom.xml
index 07dc7d2403..039b2aa04a 100644
--- a/memory/memory-netty-buffer-patch/pom.xml
+++ b/memory/memory-netty-buffer-patch/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-memory
- 19.0.0-SNAPSHOT
+ 20.0.0-SNAPSHOT
arrow-memory-netty-buffer-patch
diff --git a/memory/memory-netty/pom.xml b/memory/memory-netty/pom.xml
index 6d660da117..4218910980 100644
--- a/memory/memory-netty/pom.xml
+++ b/memory/memory-netty/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-memory
- 19.0.0-SNAPSHOT
+ 20.0.0-SNAPSHOT
arrow-memory-netty
diff --git a/memory/memory-unsafe/pom.xml b/memory/memory-unsafe/pom.xml
index 92dc0c9fe5..3fafb42802 100644
--- a/memory/memory-unsafe/pom.xml
+++ b/memory/memory-unsafe/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-memory
- 19.0.0-SNAPSHOT
+ 20.0.0-SNAPSHOT
arrow-memory-unsafe
diff --git a/memory/pom.xml b/memory/pom.xml
index bc34c26050..4ea3d1f9ca 100644
--- a/memory/pom.xml
+++ b/memory/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-java-root
- 19.0.0-SNAPSHOT
+ 20.0.0-SNAPSHOT
arrow-memory
pom
diff --git a/performance/pom.xml b/performance/pom.xml
index 3f18188e3a..d994bff45b 100644
--- a/performance/pom.xml
+++ b/performance/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-java-root
- 19.0.0-SNAPSHOT
+ 20.0.0-SNAPSHOT
arrow-performance
jar
@@ -35,10 +35,10 @@ under the License.
true
.*
1
-
+
5
5
-
+
jmh-result.json
json
@@ -75,7 +75,7 @@ under the License.
com.h2database
h2
- 2.3.232
+ 2.4.240
runtime
@@ -143,7 +143,7 @@ under the License.
java
-classpath
-
+
org.openjdk.jmh.Main
${benchmark.filter}
-f
diff --git a/performance/src/main/java/org/apache/arrow/memory/MemoryFootprintBenchmarks.java b/performance/src/main/java/org/apache/arrow/memory/MemoryFootprintBenchmarks.java
new file mode 100644
index 0000000000..395ba13b9d
--- /dev/null
+++ b/performance/src/main/java/org/apache/arrow/memory/MemoryFootprintBenchmarks.java
@@ -0,0 +1,213 @@
+/*
+ * 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.memory;
+
+import java.lang.management.ManagementFactory;
+import java.lang.management.MemoryMXBean;
+import java.lang.management.MemoryUsage;
+import java.util.concurrent.TimeUnit;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.TearDown;
+import org.openjdk.jmh.annotations.Warmup;
+import org.openjdk.jmh.runner.Runner;
+import org.openjdk.jmh.runner.RunnerException;
+import org.openjdk.jmh.runner.options.Options;
+import org.openjdk.jmh.runner.options.OptionsBuilder;
+
+/**
+ * Benchmarks for memory footprint of Arrow memory objects.
+ *
+ * This benchmark measures the heap memory overhead of creating many ArrowBuf instances. The
+ * optimizations using AtomicFieldUpdater instead of AtomicLong/AtomicInteger objects should reduce
+ * memory overhead significantly.
+ *
+ *
Expected savings per instance: - ArrowBuf: 8 bytes (id field removed) - BufferLedger: 28 bytes
+ * (20 from AtomicInteger + 8 from ledgerId) - Accountant: 48 bytes (3 × 16 bytes from AtomicLong
+ * objects)
+ *
+ *
For 1M ArrowBuf instances, this should save approximately 8 MB of heap memory.
+ */
+@State(Scope.Benchmark)
+@Fork(
+ value = 1,
+ jvmArgs = {"-Xms2g", "-Xmx2g"})
+@Warmup(iterations = 3, time = 1, timeUnit = TimeUnit.SECONDS)
+@Measurement(iterations = 5, time = 1, timeUnit = TimeUnit.SECONDS)
+public class MemoryFootprintBenchmarks {
+
+ /** Number of ArrowBuf instances to create for memory footprint measurement. */
+ private static final int NUM_BUFFERS = 100_000;
+
+ /** Size in bytes of each buffer allocation. */
+ private static final int BUFFER_SIZE = 1024;
+
+ /** Root allocator used for all buffer allocations in the benchmark. */
+ private RootAllocator allocator;
+
+ /** Array to hold references to allocated buffers, preventing garbage collection. */
+ private ArrowBuf[] buffers;
+
+ /** JMX bean for querying heap memory usage statistics. */
+ private MemoryMXBean memoryBean;
+
+ /**
+ * Sets up the benchmark state before each trial.
+ *
+ *
Initializes the memory monitoring bean, creates a root allocator with sufficient capacity,
+ * and allocates the buffer reference array.
+ */
+ @Setup(Level.Trial)
+ public void setup() {
+ memoryBean = ManagementFactory.getMemoryMXBean();
+ allocator = new RootAllocator((long) NUM_BUFFERS * BUFFER_SIZE);
+ buffers = new ArrowBuf[NUM_BUFFERS];
+ }
+
+ /**
+ * Cleans up buffers after each benchmark invocation.
+ *
+ *
Closes all allocated buffers to prevent memory leaks and ensure each iteration starts with a
+ * clean slate. This is critical for the memory footprint benchmark which allocates many buffers
+ * that would otherwise accumulate across warmup and measurement iterations.
+ */
+ @TearDown(Level.Invocation)
+ public void tearDown() {
+ for (int i = 0; i < NUM_BUFFERS; i++) {
+ if (buffers[i] != null) {
+ buffers[i].close();
+ buffers[i] = null;
+ }
+ }
+ }
+
+ /**
+ * Cleans up the allocator after the trial completes.
+ *
+ *
Closes the root allocator to release all resources after all warmup and measurement
+ * iterations are complete.
+ */
+ @TearDown(Level.Trial)
+ public void tearDownTrial() {
+ allocator.close();
+ }
+
+ /**
+ * Benchmark that measures heap memory usage when creating many ArrowBuf instances.
+ *
+ *
This benchmark creates {@value #NUM_BUFFERS} ArrowBuf instances and measures the heap memory
+ * used. With the AtomicFieldUpdater optimizations, we expect to save approximately 800 KB of heap
+ * memory (8 bytes × 100,000 instances) just from removing the id field in ArrowBuf.
+ *
+ *
The benchmark performs garbage collection before and after allocation to ensure accurate
+ * measurement of heap memory delta. Results are printed to stdout for analysis.
+ *
+ * @return the total heap memory used by the allocated buffers in bytes
+ */
+ @Benchmark
+ @BenchmarkMode(Mode.SingleShotTime)
+ @OutputTimeUnit(TimeUnit.MILLISECONDS)
+ public long measureArrowBufMemoryFootprint() {
+ // Force GC before measurement
+ System.gc();
+ System.gc();
+ System.gc();
+
+ MemoryUsage heapBefore = memoryBean.getHeapMemoryUsage();
+ long usedBefore = heapBefore.getUsed();
+
+ // Allocate buffers
+ for (int i = 0; i < NUM_BUFFERS; i++) {
+ buffers[i] = allocator.buffer(BUFFER_SIZE);
+ }
+
+ // Force GC to get accurate measurement
+ System.gc();
+ System.gc();
+ System.gc();
+
+ MemoryUsage heapAfter = memoryBean.getHeapMemoryUsage();
+ long usedAfter = heapAfter.getUsed();
+
+ long memoryUsed = usedAfter - usedBefore;
+
+ // Print memory usage for analysis
+ System.out.printf(
+ "Created %d ArrowBuf instances. Heap memory used: %d bytes (%.2f MB)%n",
+ NUM_BUFFERS, memoryUsed, memoryUsed / (1024.0 * 1024.0));
+ System.out.printf(
+ "Average memory per ArrowBuf: %.2f bytes%n", (double) memoryUsed / NUM_BUFFERS);
+
+ return memoryUsed;
+ }
+
+ /**
+ * Benchmark that measures allocation and deallocation performance.
+ *
+ *
This complements the memory footprint benchmark by measuring the time it takes to allocate
+ * and deallocate 1,000 buffers in a tight loop. This helps identify any performance regressions
+ * introduced by memory optimizations.
+ *
+ *
Uses a local buffer array to avoid interference with the shared {@link #buffers} array used
+ * by other benchmarks.
+ */
+ @Benchmark
+ @BenchmarkMode(Mode.AverageTime)
+ @OutputTimeUnit(TimeUnit.MICROSECONDS)
+ public void measureAllocationPerformance() {
+ ArrowBuf[] localBuffers = new ArrowBuf[1000];
+
+ for (int i = 0; i < 1000; i++) {
+ localBuffers[i] = allocator.buffer(BUFFER_SIZE);
+ }
+
+ for (int i = 0; i < 1000; i++) {
+ localBuffers[i].close();
+ }
+ }
+
+ /**
+ * Main entry point for running the benchmarks standalone.
+ *
+ *
This allows running the benchmarks directly from the command line or IDE without using the
+ * Maven JMH plugin. Example usage:
+ *
+ *
{@code
+ * java -cp target/benchmarks.jar org.apache.arrow.memory.MemoryFootprintBenchmarks
+ * }
+ *
+ * @param args command line arguments (not used)
+ * @throws RunnerException if the benchmark runner encounters an error
+ */
+ public static void main(String[] args) throws RunnerException {
+ Options opt =
+ new OptionsBuilder()
+ .include(MemoryFootprintBenchmarks.class.getSimpleName())
+ .forks(1)
+ .build();
+
+ new Runner(opt).run();
+ }
+}
diff --git a/performance/src/main/java/org/apache/arrow/vector/UuidVectorBenchmarks.java b/performance/src/main/java/org/apache/arrow/vector/UuidVectorBenchmarks.java
new file mode 100644
index 0000000000..b5f87e7a75
--- /dev/null
+++ b/performance/src/main/java/org/apache/arrow/vector/UuidVectorBenchmarks.java
@@ -0,0 +1,134 @@
+/*
+ * 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.vector;
+
+import java.util.UUID;
+import java.util.concurrent.TimeUnit;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.complex.impl.UuidWriterImpl;
+import org.apache.arrow.vector.holders.NullableUuidHolder;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.TearDown;
+import org.openjdk.jmh.profile.GCProfiler;
+import org.openjdk.jmh.runner.Runner;
+import org.openjdk.jmh.runner.RunnerException;
+import org.openjdk.jmh.runner.options.Options;
+import org.openjdk.jmh.runner.options.OptionsBuilder;
+
+/** Benchmarks for {@link UuidVector}. */
+@State(Scope.Benchmark)
+public class UuidVectorBenchmarks {
+ // checkstyle:off: MissingJavadocMethod
+
+ private static final int VECTOR_LENGTH = 10_000;
+
+ private static final int ALLOCATOR_CAPACITY = 1024 * 1024;
+
+ private BufferAllocator allocator;
+
+ private UuidVector vector;
+
+ private UUID[] testUuids;
+
+ @Setup
+ public void prepare() {
+ allocator = new RootAllocator(ALLOCATOR_CAPACITY);
+ vector = new UuidVector("vector", allocator);
+ vector.allocateNew(VECTOR_LENGTH);
+ vector.setValueCount(VECTOR_LENGTH);
+
+ // Pre-generate UUIDs for consistent benchmarking
+ testUuids = new UUID[VECTOR_LENGTH];
+ for (int i = 0; i < VECTOR_LENGTH; i++) {
+ testUuids[i] = new UUID(i, i * 2L);
+ }
+ }
+
+ @TearDown
+ public void tearDown() {
+ vector.close();
+ allocator.close();
+ }
+
+ @Benchmark
+ @BenchmarkMode(Mode.AverageTime)
+ @OutputTimeUnit(TimeUnit.MICROSECONDS)
+ public void setWithHolder() {
+ NullableUuidHolder holder = new NullableUuidHolder();
+ for (int i = 0; i < VECTOR_LENGTH; i++) {
+ vector.get(i, holder);
+ vector.setSafe(i, holder);
+ }
+ }
+
+ @Benchmark
+ @BenchmarkMode(Mode.AverageTime)
+ @OutputTimeUnit(TimeUnit.MICROSECONDS)
+ public void setUuidDirectly() {
+ for (int i = 0; i < VECTOR_LENGTH; i++) {
+ vector.setSafe(i, testUuids[i]);
+ }
+ }
+
+ @Benchmark
+ @BenchmarkMode(Mode.AverageTime)
+ @OutputTimeUnit(TimeUnit.MICROSECONDS)
+ public void setWithWriter() {
+ UuidWriterImpl writer = new UuidWriterImpl(vector);
+ for (int i = 0; i < VECTOR_LENGTH; i++) {
+ writer.writeExtension(testUuids[i]);
+ }
+ }
+
+ @Benchmark
+ @BenchmarkMode(Mode.AverageTime)
+ @OutputTimeUnit(TimeUnit.MICROSECONDS)
+ public void getWithUuidHolder() {
+ NullableUuidHolder holder = new NullableUuidHolder();
+ for (int i = 0; i < VECTOR_LENGTH; i++) {
+ vector.get(i, holder);
+ }
+ }
+
+ @Benchmark
+ @BenchmarkMode(Mode.AverageTime)
+ @OutputTimeUnit(TimeUnit.MICROSECONDS)
+ public void getUuidDirectly() {
+ for (int i = 0; i < VECTOR_LENGTH; i++) {
+ UUID uuid = vector.getObject(i);
+ }
+ }
+
+ public static void main(String[] args) throws RunnerException {
+ Options opt =
+ new OptionsBuilder()
+ .include(UuidVectorBenchmarks.class.getSimpleName())
+ .forks(1)
+ .addProfiler(GCProfiler.class)
+ .build();
+
+ new Runner(opt).run();
+ }
+ // checkstyle:on: MissingJavadocMethod
+}
diff --git a/pom.xml b/pom.xml
index 777f7d5cef..78137eb4e9 100644
--- a/pom.xml
+++ b/pom.xml
@@ -23,12 +23,12 @@ under the License.
org.apache
apache
- 35
+ 38
org.apache.arrow
arrow-java-root
- 19.0.0-SNAPSHOT
+ 20.0.0-SNAPSHOT
pom
Apache Arrow Java Root POM
@@ -68,6 +68,7 @@ under the License.
bom
format
memory
+ arrow-variant
vector
tools
adapter/jdbc
@@ -91,36 +92,37 @@ under the License.
- 1695310533
+ 1773644827
${project.build.directory}/generated-sources
1.9.0
- 5.12.2
- 2.0.17
- 33.4.8-jre
- 4.1.127.Final
- 1.73.0
- 4.33.1
- 2.18.3
- 3.4.2
+ 6.1.1
+ 2.0.18
+ 33.6.0-jre
+ 4.2.15.Final
+ 1.82.1
+ 4.35.1
+ 2.22.0
+ 3.5.0
25.2.10
- 1.12.0
- 5.17.0
-
+ 1.12.1
+ 1.17.1
+ 5.23.0
+
2
10.23.0
true
2.42.0
- 3.52.0
- 1.5.21
+ 4.2.1
+ 1.5.37
none
-Xdoclint:none
--add-opens=java.base/java.nio=org.apache.arrow.memory.core,ALL-UNNAMED
- 11
- 11
- 11
- 11
+ 17
+ 17
+ 17
+ 17
+ codegen
+ ${basedir}/src/main/codegen
+
+ **/*.tdd
+ **/*.fmpp
+ **/*.ftl
+
+
+
+
+
diff --git a/vector/src/main/codegen/templates/AbstractFieldReader.java b/vector/src/main/codegen/templates/AbstractFieldReader.java
index c7c5b4d78d..789295e959 100644
--- a/vector/src/main/codegen/templates/AbstractFieldReader.java
+++ b/vector/src/main/codegen/templates/AbstractFieldReader.java
@@ -29,9 +29,9 @@
* Source code generated using FreeMarker template ${.template_name}
*/
@SuppressWarnings("unused")
-abstract class AbstractFieldReader extends AbstractBaseReader implements FieldReader{
+public abstract class AbstractFieldReader extends AbstractBaseReader implements FieldReader{
- AbstractFieldReader(){
+ protected AbstractFieldReader(){
super();
}
@@ -109,10 +109,6 @@ public void copyAsField(String name, ${name}Writer writer) {
- public void copyAsValue(StructWriter writer, ExtensionTypeWriterFactory writerFactory) {
- fail("CopyAsValue StructWriter");
- }
-
public void read(ExtensionHolder holder) {
fail("Extension");
}
@@ -147,4 +143,5 @@ public int size() {
private void fail(String name) {
throw new IllegalArgumentException(String.format("You tried to read a [%s] type when you are using a field reader of type [%s].", name, this.getClass().getSimpleName()));
}
+
}
diff --git a/vector/src/main/codegen/templates/AbstractFieldWriter.java b/vector/src/main/codegen/templates/AbstractFieldWriter.java
index ae5b97faef..4b4a17d932 100644
--- a/vector/src/main/codegen/templates/AbstractFieldWriter.java
+++ b/vector/src/main/codegen/templates/AbstractFieldWriter.java
@@ -107,14 +107,17 @@ public void endEntry() {
throw new IllegalStateException(String.format("You tried to end a map entry when you are using a ValueWriter of type %s.", this.getClass().getSimpleName()));
}
+ @Override
public void write(ExtensionHolder var1) {
- this.fail("ExtensionType");
+ this.fail("Cannot write ExtensionHolder");
}
+ @Override
public void writeExtension(Object var1) {
- this.fail("ExtensionType");
+ this.fail("Cannot write extension object");
}
- public void addExtensionTypeWriterFactory(ExtensionTypeWriterFactory var1) {
- this.fail("ExtensionType");
+ @Override
+ public void writeExtension(Object var1, ArrowType type) {
+ this.fail("Cannot write extension with type " + type);
}
<#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first />
diff --git a/vector/src/main/codegen/templates/ArrowType.java b/vector/src/main/codegen/templates/ArrowType.java
index fd35c1cd2b..b428f09155 100644
--- a/vector/src/main/codegen/templates/ArrowType.java
+++ b/vector/src/main/codegen/templates/ArrowType.java
@@ -27,8 +27,10 @@
import org.apache.arrow.flatbuf.Type;
import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.vector.complex.writer.FieldWriter;
import org.apache.arrow.vector.types.*;
import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.ValueVector;
import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonIgnore;
@@ -331,6 +333,10 @@ public boolean equals(Object obj) {
public