This is an automated email from the ASF dual-hosted git repository.
lidavidm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-adbc.git
The following commit(s) were added to refs/heads/main by this push:
new 6f5b5fc56 feat(java/driver/jni): wire up cancel, getParameterSchema
(#4249)
6f5b5fc56 is described below
commit 6f5b5fc5691f57a8e0e546b68194c99dc82ffb50
Author: David Li <[email protected]>
AuthorDate: Mon Apr 27 09:50:54 2026 +0900
feat(java/driver/jni): wire up cancel, getParameterSchema (#4249)
Closes #4239.
---
.../org/apache/arrow/adbc/core/AdbcConnection.java | 2 +-
.../adbc/driver/jni/PostgresIntegrationTest.java | 37 +++++++++++++++++++-
.../adbc/driver/jni/SqlServerIntegrationTest.java | 37 +++++++++++++++++++-
java/driver/jni/src/main/cpp/jni_wrapper.cc | 39 ++++++++++++++++++++++
.../arrow/adbc/driver/jni/JniConnection.java | 5 +++
.../apache/arrow/adbc/driver/jni/JniStatement.java | 12 ++++++-
.../arrow/adbc/driver/jni/impl/JniLoader.java | 13 ++++++++
.../arrow/adbc/driver/jni/impl/NativeAdbc.java | 15 ++++++---
8 files changed, 152 insertions(+), 8 deletions(-)
diff --git
a/java/core/src/main/java/org/apache/arrow/adbc/core/AdbcConnection.java
b/java/core/src/main/java/org/apache/arrow/adbc/core/AdbcConnection.java
index 060c65e81..0b39a3d6a 100644
--- a/java/core/src/main/java/org/apache/arrow/adbc/core/AdbcConnection.java
+++ b/java/core/src/main/java/org/apache/arrow/adbc/core/AdbcConnection.java
@@ -40,7 +40,7 @@ public interface AdbcConnection extends AutoCloseable,
AdbcOptions {
* @since ADBC API revision 1.1.0
*/
default void cancel() throws AdbcException {
- throw AdbcException.notImplemented("Statement does not support cancel");
+ throw AdbcException.notImplemented("Connection does not support cancel");
}
/** Commit the pending transaction. */
diff --git
a/java/driver/jni-validation/src/test/java/org/apache/arrow/adbc/driver/jni/PostgresIntegrationTest.java
b/java/driver/jni-validation/src/test/java/org/apache/arrow/adbc/driver/jni/PostgresIntegrationTest.java
index b805a420d..125644041 100644
---
a/java/driver/jni-validation/src/test/java/org/apache/arrow/adbc/driver/jni/PostgresIntegrationTest.java
+++
b/java/driver/jni-validation/src/test/java/org/apache/arrow/adbc/driver/jni/PostgresIntegrationTest.java
@@ -316,8 +316,43 @@ class PostgresIntegrationTest {
assertThat(result.getReader().loadNextBatch()).isTrue();
assertThat(result.getReader().loadNextBatch()).isFalse();
}
+
+ assertThat(stmt.getParameterSchema().getFields()).isEmpty();
}
- // TODO(https://github.com/apache/arrow-adbc/issues/4239): test get
parameter schema
+
+ try (var stmt = conn.createStatement()) {
+ stmt.setSqlQuery("SELECT $1 || 'foo'");
+ assertSchema(stmt.getParameterSchema())
+ .isEqualTo(new Schema(List.of(Field.nullable("$1",
Types.MinorType.VARCHAR.getType()))));
+ }
+ }
+
+ @Test
+ void cancelQuery() throws Exception {
+ // There's nothing really we can test reliably; it is wired up but we'd
need a long-running
+ // query and a reliable way to start the cancel at the right time
+ Schema schema = new Schema(List.of(Field.nullable("$1",
Types.MinorType.VARCHAR.getType())));
+ try (var stmt = conn.createStatement();
+ var vsr = VectorSchemaRoot.create(schema, allocator)) {
+ var vcv = (VarCharVector) vsr.getVector(0);
+ vcv.setSafe(0, "test".getBytes(StandardCharsets.UTF_8));
+ vcv.setSafe(1, "bar".getBytes(StandardCharsets.UTF_8));
+
+ stmt.setSqlQuery("SELECT CAST($1 AS VARCHAR) || 'foo'");
+ stmt.bind(vsr);
+
+ try (AdbcStatement.QueryResult result = stmt.executeQuery()) {
+ stmt.cancel();
+ //noinspection StatementWithEmptyBody
+ while (result.getReader().loadNextBatch()) {}
+ }
+ }
+ }
+
+ @Test
+ void cancelConnection() throws Exception {
+ // There's nothing really we can test reliably
+ conn.cancel();
}
@Test
diff --git
a/java/driver/jni-validation/src/test/java/org/apache/arrow/adbc/driver/jni/SqlServerIntegrationTest.java
b/java/driver/jni-validation/src/test/java/org/apache/arrow/adbc/driver/jni/SqlServerIntegrationTest.java
index 8a6737816..0b4bfd452 100644
---
a/java/driver/jni-validation/src/test/java/org/apache/arrow/adbc/driver/jni/SqlServerIntegrationTest.java
+++
b/java/driver/jni-validation/src/test/java/org/apache/arrow/adbc/driver/jni/SqlServerIntegrationTest.java
@@ -322,8 +322,43 @@ class SqlServerIntegrationTest {
assertThat(result.getReader().loadNextBatch()).isTrue();
assertThat(result.getReader().loadNextBatch()).isFalse();
}
+
+ assertThat(stmt.getParameterSchema().getFields()).isEmpty();
}
- // TODO(https://github.com/apache/arrow-adbc/issues/4239): test get
parameter schema
+
+ try (var stmt = conn.createStatement()) {
+ stmt.setSqlQuery("SELECT CONCAT(CAST(@p1 AS NVARCHAR), 'foo')");
+ assertSchema(stmt.getParameterSchema())
+ .isEqualTo(new Schema(List.of(Field.nullable("@p1",
Types.MinorType.VARCHAR.getType()))));
+ }
+ }
+
+ @Test
+ void cancelQuery() throws Exception {
+ // There's nothing really we can test reliably (MSSQL driver doesn't react
to cancel)
+ Schema schema = new Schema(List.of(Field.nullable("$1",
Types.MinorType.VARCHAR.getType())));
+ try (var stmt = conn.createStatement();
+ var vsr = VectorSchemaRoot.create(schema, allocator)) {
+ var vcv = (VarCharVector) vsr.getVector(0);
+ vcv.setSafe(0, "test".getBytes(StandardCharsets.UTF_8));
+ vcv.setSafe(1, "bar".getBytes(StandardCharsets.UTF_8));
+
+ stmt.setSqlQuery("SELECT CAST(@p1 AS NVARCHAR) || 'foo'");
+ stmt.bind(vsr);
+
+ // N.B. the MSSQL driver doesn't appear to react
+ try (AdbcStatement.QueryResult result = stmt.executeQuery()) {
+ stmt.cancel();
+ //noinspection StatementWithEmptyBody
+ while (result.getReader().loadNextBatch()) {}
+ }
+ }
+ }
+
+ @Test
+ void cancelConnection() throws Exception {
+ // There's nothing really we can test reliably
+ conn.cancel();
}
@Test
diff --git a/java/driver/jni/src/main/cpp/jni_wrapper.cc
b/java/driver/jni/src/main/cpp/jni_wrapper.cc
index a830e7304..0cb128f84 100644
--- a/java/driver/jni/src/main/cpp/jni_wrapper.cc
+++ b/java/driver/jni/src/main/cpp/jni_wrapper.cc
@@ -353,6 +353,33 @@ jobject MakeNativeSchemaResult(JNIEnv* env, struct
ArrowSchema* schema) {
static_cast<jlong>(reinterpret_cast<uintptr_t>(schema)));
}
+JNIEXPORT void JNICALL
+Java_org_apache_arrow_adbc_driver_jni_impl_NativeAdbc_statementCancel(
+ JNIEnv* env, [[maybe_unused]] jclass self, jlong handle) {
+ struct AdbcError error = ADBC_ERROR_INIT;
+ auto* ptr = reinterpret_cast<struct
AdbcStatement*>(static_cast<uintptr_t>(handle));
+ try {
+ CHECK_ADBC_ERROR(AdbcStatementCancel(ptr, &error), error);
+ } catch (const AdbcException& e) {
+ e.ThrowJavaException(env);
+ }
+}
+
+JNIEXPORT jobject JNICALL
+Java_org_apache_arrow_adbc_driver_jni_impl_NativeAdbc_statementGetParameterSchema(
+ JNIEnv* env, [[maybe_unused]] jclass self, jlong handle) {
+ struct AdbcError error = ADBC_ERROR_INIT;
+ auto* ptr = reinterpret_cast<struct
AdbcStatement*>(static_cast<uintptr_t>(handle));
+ struct ArrowSchema schema = {};
+ try {
+ CHECK_ADBC_ERROR(AdbcStatementGetParameterSchema(ptr, &schema, &error),
error);
+ return MakeNativeSchemaResult(env, &schema);
+ } catch (const AdbcException& e) {
+ e.ThrowJavaException(env);
+ }
+ return nullptr;
+}
+
JNIEXPORT jobject JNICALL
Java_org_apache_arrow_adbc_driver_jni_impl_NativeAdbc_statementExecuteQuery(
JNIEnv* env, [[maybe_unused]] jclass self, jlong handle) {
@@ -621,6 +648,18 @@
Java_org_apache_arrow_adbc_driver_jni_impl_NativeAdbc_statementSetOptionString(
}
}
+JNIEXPORT void JNICALL
+Java_org_apache_arrow_adbc_driver_jni_impl_NativeAdbc_connectionCancel(
+ JNIEnv* env, [[maybe_unused]] jclass self, jlong handle) {
+ struct AdbcError error = ADBC_ERROR_INIT;
+ auto* ptr = reinterpret_cast<struct
AdbcConnection*>(static_cast<uintptr_t>(handle));
+ try {
+ CHECK_ADBC_ERROR(AdbcConnectionCancel(ptr, &error), error);
+ } catch (const AdbcException& e) {
+ e.ThrowJavaException(env);
+ }
+}
+
JNIEXPORT jobject JNICALL
Java_org_apache_arrow_adbc_driver_jni_impl_NativeAdbc_connectionGetObjects(
JNIEnv* env, [[maybe_unused]] jclass self, jlong handle, jint depth,
jstring catalog,
diff --git
a/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniConnection.java
b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniConnection.java
index 563420fa3..ebfd75957 100644
---
a/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniConnection.java
+++
b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniConnection.java
@@ -45,6 +45,11 @@ public class JniConnection implements AdbcConnection {
return new JniStatement(allocator,
JniLoader.INSTANCE.openStatement(handle));
}
+ @Override
+ public void cancel() throws AdbcException {
+ JniLoader.INSTANCE.connectionCancel(handle);
+ }
+
@Override
public AdbcStatement bulkIngest(String targetTableName, BulkIngestMode mode)
throws AdbcException {
diff --git
a/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniStatement.java
b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniStatement.java
index 40a1eeec1..ec8fd46b2 100644
---
a/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniStatement.java
+++
b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniStatement.java
@@ -40,6 +40,11 @@ public class JniStatement implements AdbcStatement {
this.handle = handle;
}
+ @Override
+ public void cancel() throws AdbcException {
+ JniLoader.INSTANCE.statementCancel(handle);
+ }
+
@Override
public void setSqlQuery(String query) throws AdbcException {
JniLoader.INSTANCE.statementSetSqlQuery(handle, query);
@@ -87,6 +92,11 @@ public class JniStatement implements AdbcStatement {
return
JniLoader.INSTANCE.statementExecuteSchema(handle).importSchema(allocator);
}
+ @Override
+ public Schema getParameterSchema() throws AdbcException {
+ return
JniLoader.INSTANCE.statementGetParameterSchema(handle).importSchema(allocator);
+ }
+
@Override
public void prepare() throws AdbcException {
JniLoader.INSTANCE.statementPrepare(handle);
@@ -143,7 +153,7 @@ public class JniStatement implements AdbcStatement {
JniLoader.INSTANCE.statementSetOptionDouble(handle, key.getKey(),
(Double) value);
} else if (value instanceof Boolean) {
JniLoader.INSTANCE.statementSetOptionString(
- handle, key.getKey(), ((Boolean) value) ? "true" : "false");
+ handle, key.getKey(), ((Boolean) value).toString());
} else if (value instanceof byte[]) {
JniLoader.INSTANCE.statementSetOptionBytes(handle, key.getKey(),
(byte[]) value);
} else {
diff --git
a/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/JniLoader.java
b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/JniLoader.java
index 5b1adaa0d..ebfc3c51c 100644
---
a/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/JniLoader.java
+++
b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/JniLoader.java
@@ -81,6 +81,10 @@ public enum JniLoader {
return NativeAdbc.openStatement(connection.getConnectionHandle());
}
+ public void statementCancel(NativeStatementHandle statement) throws
AdbcException {
+ NativeAdbc.statementCancel(statement.getStatementHandle());
+ }
+
public NativeQueryResult statementExecuteQuery(NativeStatementHandle
statement)
throws AdbcException {
return NativeAdbc.statementExecuteQuery(statement.getStatementHandle());
@@ -110,6 +114,11 @@ public enum JniLoader {
return NativeAdbc.statementExecuteSchema(statement.getStatementHandle());
}
+ public NativeSchemaResult statementGetParameterSchema(NativeStatementHandle
statement)
+ throws AdbcException {
+ return
NativeAdbc.statementGetParameterSchema(statement.getStatementHandle());
+ }
+
public byte[] statementGetOptionBytes(NativeStatementHandle handle, String
key)
throws AdbcException {
return NativeAdbc.statementGetOptionBytes(handle.getStatementHandle(),
key);
@@ -150,6 +159,10 @@ public enum JniLoader {
NativeAdbc.statementSetOptionString(statement.getStatementHandle(), key,
value);
}
+ public void connectionCancel(NativeConnectionHandle connection) throws
AdbcException {
+ NativeAdbc.connectionCancel(connection.getConnectionHandle());
+ }
+
public NativeQueryResult connectionGetObjects(
NativeConnectionHandle connection,
int depth,
diff --git
a/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/NativeAdbc.java
b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/NativeAdbc.java
index 601a4716b..fb93dacc8 100644
---
a/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/NativeAdbc.java
+++
b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/NativeAdbc.java
@@ -34,10 +34,6 @@ class NativeAdbc {
static native void closeStatement(long handle) throws AdbcException;
- static native NativeQueryResult statementExecuteQuery(long handle) throws
AdbcException;
-
- static native void statementSetSqlQuery(long handle, String query) throws
AdbcException;
-
static native void statementBind(long handle, long values, long schema)
throws AdbcException;
// TODO(lidavidm): we need a way to bind an ArrowReader (or some other
suitable interface that
@@ -46,12 +42,21 @@ class NativeAdbc {
@SuppressWarnings("unused")
static native void statementBindStream(long handle, long stream) throws
AdbcException;
+ static native void statementCancel(long handle) throws AdbcException;
+
static native long statementExecuteUpdate(long handle) throws AdbcException;
static native void statementPrepare(long handle) throws AdbcException;
+ static native NativeQueryResult statementExecuteQuery(long handle) throws
AdbcException;
+
static native NativeSchemaResult statementExecuteSchema(long handle) throws
AdbcException;
+ static native NativeSchemaResult statementGetParameterSchema(long
statementHandle)
+ throws AdbcException;
+
+ static native void statementSetSqlQuery(long handle, String query) throws
AdbcException;
+
static native byte[] statementGetOptionBytes(long handle, String key) throws
AdbcException;
static native double statementGetOptionDouble(long handle, String key)
throws AdbcException;
@@ -72,6 +77,8 @@ class NativeAdbc {
static native void statementSetOptionString(long handle, String key, String
value)
throws AdbcException;
+ static native void connectionCancel(long handle) throws AdbcException;
+
static native NativeQueryResult connectionGetObjects(
long handle,
int depth,
