test(java/driver/jni): add more integration/unit tests

java/driver/jni-validation/src/test/java/org/apache/arrow/adbc/driver/jni/MultiDriverIntegrationTest.java

@@ -0,0 +1,145 @@
+/*
+ * 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.adbc.driver.jni;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.arrow.adbc.core.AdbcConnection;
+import org.apache.arrow.adbc.core.AdbcDatabase;
+import org.apache.arrow.adbc.core.AdbcDriver;
+import org.apache.arrow.adbc.core.AdbcException;
+import org.apache.arrow.adbc.driver.testsuite.ArrowToJava;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.util.AutoCloseables;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Assumptions;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+/** Integration test that uses multiple drivers simultaneously. */
+public class MultiDriverIntegrationTest {
+  BufferAllocator allocator;
+  JniDriver driver;
+
+  AdbcDatabase pgDb, mssqlDb, fsqlDb;
+  AdbcConnection pgConn, mssqlConn, fsqlConn;
+
+  @BeforeAll
+  static void beforeAll() {
+    Assumptions.assumeFalse(
+        FlightSqlSqliteIntegrationTest.URI == null || FlightSqlSqliteIntegrationTest.URI.isEmpty(),
+        String.format("Must set %s", FlightSqlSqliteIntegrationTest.URI_ENV));
+    Assumptions.assumeFalse(
+        PostgresIntegrationTest.URI == null || PostgresIntegrationTest.URI.isEmpty(),
+        String.format("Must set %s", PostgresIntegrationTest.URI_ENV));
+    Assumptions.assumeFalse(
+        SqlServerIntegrationTest.URI == null || SqlServerIntegrationTest.URI.isEmpty(),
+        String.format("Must set %s", SqlServerIntegrationTest.URI_ENV));
+  }
+
+  @BeforeEach
+  void beforeEach() throws Exception {
+    allocator = new RootAllocator();
+    driver = new JniDriver(allocator);
+
+    {
+      System.err.println(
+          "Connecting to Flight SQL with URI: " + FlightSqlSqliteIntegrationTest.URI);
+      Map<String, Object> parameters = new HashMap<>();
+      JniDriver.PARAM_DRIVER.set(parameters, "adbc_driver_flightsql");
+      AdbcDriver.PARAM_URI.set(parameters, FlightSqlSqliteIntegrationTest.URI);
+      fsqlDb = driver.open(parameters);
+      fsqlConn = fsqlDb.connect();
+    }
+    {
+      System.err.println("Connecting to MSSQL with URI: " + SqlServerIntegrationTest.URI);
+      Map<String, Object> parameters = new HashMap<>();
+      JniDriver.PARAM_DRIVER.set(parameters, "mssql");
+      AdbcDriver.PARAM_URI.set(parameters, SqlServerIntegrationTest.URI);
+      mssqlDb = driver.open(parameters);
+      mssqlConn = mssqlDb.connect();
+    }
+    {
+      System.err.println("Connecting to PostgreSQL with URI: " + PostgresIntegrationTest.URI);
+      Map<String, Object> parameters = new HashMap<>();
+      JniDriver.PARAM_DRIVER.set(parameters, "adbc_driver_postgresql");
+      AdbcDriver.PARAM_URI.set(parameters, PostgresIntegrationTest.URI);
+      pgDb = driver.open(parameters);
+      pgConn = pgDb.connect();
+    }
+  }
+
+  @AfterEach
+  void afterEach() throws Exception {
+    AutoCloseables.close(pgConn, mssqlConn, fsqlConn, pgDb, mssqlDb, fsqlDb, allocator);
+  }
+
+  @Test
+  void queryAll() throws Exception {
+    try (var pgStmt = pgConn.createStatement();
+        var mssqlStmt = mssqlConn.createStatement();
+        var fsqlStmt = fsqlConn.createStatement()) {
+      pgStmt.setSqlQuery("SELECT 1 as foo");
+      mssqlStmt.setSqlQuery("SELECT 1 as bar");
+      fsqlStmt.setSqlQuery("SELECT 1 as baz");
+
+      try (var pgRes = pgStmt.executeQuery();
+          var mssqlRes = mssqlStmt.executeQuery();
+          var fsqlRes = fsqlStmt.executeQuery()) {
+        assertThat(ArrowToJava.toIntegers(pgRes.getReader(), "foo")).containsExactly(1);
+        assertThat(ArrowToJava.toIntegers(mssqlRes.getReader(), "bar")).containsExactly(1);
+        assertThat(ArrowToJava.toLongs(fsqlRes.getReader(), "baz")).containsExactly(1L);
+      }
+    }
+  }
+
+  @Test
+  void errorDriverDoesNotExist() {
+    Map<String, Object> parameters = new HashMap<>();
+    JniDriver.PARAM_DRIVER.set(parameters, "thisdriverdoesnotexist");
+    assertThatThrownBy(
+            () -> {
+              try (var db = driver.open(parameters)) {}
+            })
+        .isInstanceOf(AdbcException.class);
+  }
+
+  @Test
+  void errorFailedConnection() throws Exception {
+    Map<String, Object> parameters = new HashMap<>();
+    JniDriver.PARAM_DRIVER.set(parameters, "mssql");
+    AdbcDriver.PARAM_URI.set(parameters, "mssql://localhost:9999");
+    try (var db = driver.open(parameters)) {
+      assertThatThrownBy(db::connect).hasMessageContaining("Could not get connection");
+    }
+  }
+
+  @Test
+  void errorBadConnectionParameter() throws Exception {
+    Map<String, Object> parameters = new HashMap<>();
+    JniDriver.PARAM_DRIVER.set(parameters, "mssql");
+    parameters.put("this parameter does not exist", "");
+    AdbcDriver.PARAM_URI.set(parameters, "mssql://localhost:9999");
+    assertThatThrownBy(() -> driver.open(parameters))
+        .hasMessageContaining("Unknown database option");
+  }
+}

java/driver/jni/CMakeLists.txt

@@ -23,7 +23,7 @@ add_custom_command(OUTPUT ${CMAKE_CURRENT_SOURCE_DIR}/target/headers/org_apache_
                    COMMAND rm -rf ${CMAKE_CURRENT_SOURCE_DIR}/target/headers
                            ${CMAKE_CURRENT_SOURCE_DIR}/target/maven-status
                    COMMAND mvn --file ${CMAKE_CURRENT_SOURCE_DIR}/../.. -Pjni,javah
-                           compile --also-make --projects :adbc-driver-jni
+                           compile --also-make --projects :adbc-driver-jni -Drat.skip=true
                    DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/src/main/java/org/apache/arrow/adbc/driver/jni/impl/NativeAdbc.java
 )
 

java/driver/jni/src/main/cpp/jni_wrapper.cc

@@ -1562,4 +1562,29 @@ Java_org_apache_arrow_adbc_driver_jni_impl_NativeAdbc_databaseSetOptionString(
     e.ThrowJavaException(env);
   }
 }
+
+// wrapper around GetJniByteBuffer for direct unit testing
+JNIEXPORT jbyteArray JNICALL
+Java_org_apache_arrow_adbc_driver_jni_impl_NativeAdbc_internalGetByteBuffer(
+    JNIEnv* env, [[maybe_unused]] jclass self, jobject input) {
+  std::vector<uint8_t> scratch;
+  size_t length = 0;
+  try {
+    const uint8_t* raw = GetJniByteBuffer(env, input, scratch, length);
+    if (env->ExceptionCheck()) return nullptr;
+    // valid for raw to be nullptr if length == 0
+    jbyteArray result = env->NewByteArray(static_cast<jsize>(length));
+    if (result == nullptr || env->ExceptionCheck()) return nullptr;
+    if (length > 0) {
+      env->SetByteArrayRegion(result, 0, static_cast<jsize>(length),
+                              reinterpret_cast<const jbyte*>(raw));
+      if (env->ExceptionCheck()) return nullptr;
+    }
+    return result;
+  } catch (const AdbcException& e) {
+    e.ThrowJavaException(env);
+    return nullptr;
+  }
+  return nullptr;
+}
 }

java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/JniLoader.java

@@ -338,6 +338,11 @@ public void databaseSetOptionString(NativeDatabaseHandle handle, String key, Str
         handle.getDatabaseHandle(), stringToUtf8(key), stringToUtf8(value));
   }
 
+  /** For unit testing only. */
+  byte[] internalGetByteBuffer(ByteBuffer buf) throws AdbcException {
+    return NativeAdbc.internalGetByteBuffer(buf);
+  }
+
   private static byte[] stringToUtf8(String value) {
     return value == null ? null : value.getBytes(StandardCharsets.UTF_8);
   }

java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/NativeAdbc.java

@@ -157,4 +157,7 @@ static native void databaseSetOptionLong(long handle, byte[] key, long value)
 
   static native void databaseSetOptionString(long handle, byte[] key, byte[] value)
       throws AdbcException;
+
+  // Purely for unit testing.
+  static native byte[] internalGetByteBuffer(ByteBuffer input) throws AdbcException;
 }

java/driver/jni/src/test/java/org/apache/arrow/adbc/driver/jni/impl/ImplTest.java

@@ -0,0 +1,122 @@
+/*
+ * 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.adbc.driver.jni.impl;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import java.nio.ByteBuffer;
+import org.junit.jupiter.api.Test;
+
+public class ImplTest {
+  @Test
+  void emptyHeap() throws Exception {
+    ByteBuffer buf = ByteBuffer.allocate(0);
+    assertThat(JniLoader.INSTANCE.internalGetByteBuffer(buf)).isEmpty();
+
+    buf = ByteBuffer.allocate(16);
+    buf.position(16);
+    assertThat(buf.remaining()).isEqualTo(0);
+    assertThat(JniLoader.INSTANCE.internalGetByteBuffer(buf)).isEmpty();
+  }
+
+  @Test
+  void emptyDirect() throws Exception {
+    ByteBuffer buf = ByteBuffer.allocateDirect(0);
+    assertThat(buf.isDirect()).isTrue();
+    assertThat(JniLoader.INSTANCE.internalGetByteBuffer(buf)).isEmpty();
+
+    buf = ByteBuffer.allocateDirect(16);
+    assertThat(buf.isDirect()).isTrue();
+    buf.position(16);
+    assertThat(buf.remaining()).isEqualTo(0);
+    assertThat(JniLoader.INSTANCE.internalGetByteBuffer(buf)).isEmpty();
+  }
+
+  @Test
+  void emptyArray() throws Exception {
+    ByteBuffer buf = ByteBuffer.wrap(new byte[0]);
+    assertThat(buf.hasArray()).isTrue();
+    assertThat(JniLoader.INSTANCE.internalGetByteBuffer(buf)).isEmpty();
+
+    buf = ByteBuffer.wrap(new byte[16]);
+    buf.position(16);
+    assertThat(buf.hasArray()).isTrue();
+    assertThat(buf.remaining()).isEqualTo(0);
+    assertThat(JniLoader.INSTANCE.internalGetByteBuffer(buf)).isEmpty();
+  }
+
+  @Test
+  void emptySlow() throws Exception {
+    ByteBuffer buf = ByteBuffer.wrap(new byte[0]).asReadOnlyBuffer();
+    // take the slow path
+    assertThat(buf.hasArray()).isFalse();
+    assertThat(buf.isDirect()).isFalse();
+    assertThat(JniLoader.INSTANCE.internalGetByteBuffer(buf)).isEmpty();
+
+    buf = ByteBuffer.wrap(new byte[16]);
+    buf.position(16);
+    buf = buf.asReadOnlyBuffer();
+    assertThat(buf.hasArray()).isFalse();
+    assertThat(buf.isDirect()).isFalse();
+    assertThat(buf.remaining()).isEqualTo(0);
+    assertThat(JniLoader.INSTANCE.internalGetByteBuffer(buf)).isEmpty();
+  }
+
+  @Test
+  void offsetHeap() throws Exception {
+    ByteBuffer buf = ByteBuffer.allocate(4);
+    buf.put((byte) 0);
+    buf.put((byte) 1);
+    buf.put((byte) 2);
+    buf.put((byte) 3);
+    buf.position(1);
+    assertThat(buf.remaining()).isEqualTo(3);
+    assertThat(JniLoader.INSTANCE.internalGetByteBuffer(buf)).isEqualTo(new byte[] {1, 2, 3});
+  }
+
+  @Test
+  void offsetDirect() throws Exception {
+    ByteBuffer buf = ByteBuffer.allocateDirect(4);
+    buf.put((byte) 0);
+    buf.put((byte) 1);
+    buf.put((byte) 2);
+    buf.put((byte) 3);
+    buf.position(1);
+    assertThat(buf.remaining()).isEqualTo(3);
+    assertThat(JniLoader.INSTANCE.internalGetByteBuffer(buf)).isEqualTo(new byte[] {1, 2, 3});
+  }
+
+  @Test
+  void offsetArray() throws Exception {
+    ByteBuffer buf = ByteBuffer.wrap(new byte[] {(byte) 0, (byte) 1, (byte) 2, (byte) 3});
+    buf.position(1);
+    assertThat(buf.hasArray()).isTrue();
+    assertThat(buf.remaining()).isEqualTo(3);
+    assertThat(JniLoader.INSTANCE.internalGetByteBuffer(buf)).isEqualTo(new byte[] {1, 2, 3});
+  }
+
+  @Test
+  void offsetSlow() throws Exception {
+    ByteBuffer buf = ByteBuffer.wrap(new byte[] {(byte) 0, (byte) 1, (byte) 2, (byte) 3});
+    buf.position(1);
+    buf = buf.asReadOnlyBuffer();
+    assertThat(buf.hasArray()).isFalse();
+    assertThat(buf.isDirect()).isFalse();
+    assertThat(buf.remaining()).isEqualTo(3);
+    assertThat(JniLoader.INSTANCE.internalGetByteBuffer(buf)).isEqualTo(new byte[] {1, 2, 3});
+  }
+}