libraryLoaded = new AtomicReference<>(LibraryState.NOT_LOADED);
+
+ static {
+ NativeLibrary.loadLibrary();
+ }
+
+ public static void loadLibrary() {
+ if (libraryLoaded.get() == LibraryState.LOADED) {
+ return;
+ }
+
+ if (libraryLoaded.compareAndSet(LibraryState.NOT_LOADED, LibraryState.LOADING)) {
+ try {
+ doLoadLibrary();
+ } catch (IOException e) {
+ libraryLoaded.set(LibraryState.NOT_LOADED);
+ throw new UncheckedIOException("Unable to load the Databend shared library", e);
+ }
+ libraryLoaded.set(LibraryState.LOADED);
+ return;
+ }
+
+ while (libraryLoaded.get() == LibraryState.LOADING) {
+ try {
+ Thread.sleep(10);
+ } catch (InterruptedException ignore) {
+ }
+ }
+ }
+
+ private static void doLoadLibrary() throws IOException {
+ try {
+ // try dynamic library - the search path can be configured via "-Djava.library.path"
+ System.loadLibrary("bendsql_java");
+ return;
+ } catch (UnsatisfiedLinkError ignore) {
+ // ignore - try from classpath
+ }
+
+ doLoadBundledLibrary();
+ }
+
+ private static void doLoadBundledLibrary() throws IOException {
+ final String libraryPath = bundledLibraryPath();
+ try (final InputStream is = NativeObject.class.getResourceAsStream(libraryPath)) {
+ if (is == null) {
+ throw new IOException("cannot find " + libraryPath);
+ }
+ final int dot = libraryPath.indexOf('.');
+ final File tmpFile = File.createTempFile(libraryPath.substring(0, dot), libraryPath.substring(dot));
+ tmpFile.deleteOnExit();
+ Files.copy(is, tmpFile.toPath(), StandardCopyOption.REPLACE_EXISTING);
+ System.load(tmpFile.getAbsolutePath());
+ }
+ }
+
+ private static String bundledLibraryPath() {
+ final String classifier = Environment.getClassifier();
+ final String libraryName = System.mapLibraryName("bendsql_java");
+ return "/native/" + classifier + "/" + libraryName;
+ }
+}
diff --git a/bindings/java/src/main/java/com/databend/bendsql/jni_utils/NativeObject.java b/bindings/java/src/main/java/com/databend/bendsql/jni_utils/NativeObject.java
new file mode 100644
index 000000000..c1a5bd9d0
--- /dev/null
+++ b/bindings/java/src/main/java/com/databend/bendsql/jni_utils/NativeObject.java
@@ -0,0 +1,99 @@
+/*
+ * 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 com.databend.bendsql.jni_utils;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+
+/**
+ * NativeObject is the base-class of all BendSQL classes that have
+ * a pointer to a native object.
+ *
+ *
+ * NativeObject has the {@link NativeObject#close()} method, which frees its associated
+ * native object.
+ *
+ *
+ * This function should be called manually, or even better, called implicitly using a
+ * try-with-resources
+ * statement, when you are finished with the object. It is no longer called automatically
+ * during the regular Java GC process via {@link NativeObject#finalize()}.
+ *
+ *
+ * Explanatory note
+ *
+ *
+ * When or if the Garbage Collector calls {@link Object#finalize()}
+ * depends on the JVM implementation and system conditions, which the programmer
+ * cannot control. In addition, the GC cannot see through the native reference
+ * long member variable (which is the pointer value to the native object),
+ * and cannot know what other resources depend on it.
+ *
+ *
+ * Finalization is deprecated and subject to removal in a future release.
+ * The use of finalization can lead to problems with security, performance,
+ * and reliability. See JEP 421
+ * for discussion and alternatives.
+ */
+public abstract class NativeObject implements AutoCloseable {
+ static {
+ NativeLibrary.loadLibrary();
+ }
+
+ private final AtomicBoolean disposed = new AtomicBoolean(false);
+
+/**
+ * An immutable reference to the value of the underneath pointer pointing
+ * to some underlying native Databend object.
+ */
+ protected final long nativeHandle;
+ protected final long executorHandle;
+
+ protected NativeObject(long nativeHandle, long executorHandle) {
+ this.nativeHandle = nativeHandle;
+ this.executorHandle = executorHandle;
+ }
+
+ @Override
+ public void close() {
+ if (disposed.compareAndSet(false, true)) {
+ disposeInternal(nativeHandle, executorHandle);
+ }
+ }
+
+ /**
+ * Check if the object has been disposed. Useful for defensive programming.
+ *
+ * @return if the object has been disposed.
+ */
+ public boolean isDisposed() {
+ return disposed.get();
+ }
+
+ /**
+ * Deletes underlying native object pointer.
+ *
+ * @param handle to the native object pointer
+ */
+ protected abstract void disposeInternal(long handle, long executorHandle);
+
+ public long getNativeHandle() {
+ return nativeHandle;
+ }
+}
diff --git a/bindings/java/src/main/resources/META-INF/services/java.sql.Driver b/bindings/java/src/main/resources/META-INF/services/java.sql.Driver
new file mode 100644
index 000000000..f2b3502f9
--- /dev/null
+++ b/bindings/java/src/main/resources/META-INF/services/java.sql.Driver
@@ -0,0 +1 @@
+com.databend.bendsql.DatabendDriver
\ No newline at end of file
diff --git a/bindings/java/src/main/resources/bindings.properties b/bindings/java/src/main/resources/bindings.properties
new file mode 100644
index 000000000..ad06a8d04
--- /dev/null
+++ b/bindings/java/src/main/resources/bindings.properties
@@ -0,0 +1,15 @@
+# Copyright 2021 Datafuse Labs
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+project.version=${project.version}
diff --git a/bindings/java/src/row_batch_iterator.rs b/bindings/java/src/row_batch_iterator.rs
new file mode 100644
index 000000000..eb6bdfa38
--- /dev/null
+++ b/bindings/java/src/row_batch_iterator.rs
@@ -0,0 +1,88 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::jni_utils::executor::executor_or_default;
+use crate::jni_utils::executor::Executor;
+use crate::Result;
+use jni::objects::JClass;
+
+use databend_driver::rest_api::RowBatch;
+use jni::sys::jstring;
+use jni::sys::{jlong, jobject};
+
+use jni::JNIEnv;
+
+#[no_mangle]
+pub extern "system" fn Java_com_databend_bendsql_NativeRowBatchIterator_fetchNextRowBatch(
+ mut env: JNIEnv,
+ _class: JClass,
+ it: *mut RowBatch,
+ executor: *const Executor,
+) -> jstring {
+ fetch_next_row_batch(&mut env, it, executor).unwrap_or_else(|e| {
+ e.throw(&mut env);
+ std::ptr::null_mut()
+ })
+}
+
+fn fetch_next_row_batch(
+ env: &mut JNIEnv,
+ it: *mut RowBatch,
+ executor: *const Executor,
+) -> Result {
+ let batch = unsafe { &mut *it };
+
+ let data = executor_or_default(env, executor)?
+ .block_on(async move { batch.fetch_next_page().await })?;
+
+ if !data.is_empty() {
+ let json = serde_json::to_string(&data).unwrap();
+ let jstring = env.new_string(json)?;
+ Ok(jstring.into_raw())
+ } else {
+ Ok(std::ptr::null_mut())
+ }
+}
+
+#[no_mangle]
+pub extern "system" fn Java_com_databend_bendsql_NativeRowBatchIterator_getSchema(
+ mut env: JNIEnv,
+ _class: JClass,
+ it: *mut RowBatch,
+) -> jstring {
+ get_schema(&mut env, it).unwrap_or_else(|e| {
+ e.throw(&mut env);
+ std::ptr::null_mut()
+ })
+}
+
+fn get_schema(env: &mut JNIEnv, it: *mut RowBatch) -> Result {
+ let batch = unsafe { &mut *it };
+ let schema = batch.schema();
+ let json = serde_json::to_string(&schema).unwrap();
+ let jstring = env.new_string(json)?;
+ Ok(jstring.into_raw())
+}
+
+#[no_mangle]
+pub extern "system" fn Java_com_databend_bendsql_NativeRowBatchIterator_disposeInternal(
+ _env: JNIEnv,
+ _class: JClass,
+ handle: jlong,
+ _executor: *const Executor,
+) {
+ if handle != 0 {
+ let _ = unsafe { Box::from_raw(handle as *mut RowBatch) };
+ }
+}
diff --git a/bindings/java/src/test/java/com/databend/bendsql/DatabendConnectionTest.java b/bindings/java/src/test/java/com/databend/bendsql/DatabendConnectionTest.java
new file mode 100644
index 000000000..b24a6b291
--- /dev/null
+++ b/bindings/java/src/test/java/com/databend/bendsql/DatabendConnectionTest.java
@@ -0,0 +1,55 @@
+package com.databend.bendsql;
+
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import java.sql.ResultSet;
+import java.sql.Statement;
+import java.util.Properties;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+public class DatabendConnectionTest {
+
+ private DatabendConnection connection;
+ private static final String TEST_DSN = "databend://root:@localhost:8000/default?sslmode=disable";
+
+ @BeforeEach
+ void setUp() throws Exception {
+ Properties props = new Properties();
+ connection = new DatabendConnection(TEST_DSN, props);
+ }
+
+ @Test
+ void testCreateStatement() throws Exception {
+ Statement stmt = connection.createStatement();
+ assertNotNull(stmt);
+ assertTrue(stmt instanceof DatabendStatement);
+ }
+
+ @Test
+ void testSimpleQuery() throws Exception {
+ Statement stmt = connection.createStatement();
+
+ ResultSet rs = stmt.executeQuery("SELECT 1");
+
+ assertTrue(rs.next());
+ assertEquals(1, rs.getInt(1));
+ assertFalse(rs.next());
+ }
+
+ @Test
+ void testTableQuery() throws Exception {
+ Statement stmt = connection.createStatement();
+
+ stmt.execute("CREATE OR REPLACE TABLE test_table (id INT, name VARCHAR)");
+ stmt.execute("INSERT INTO test_table VALUES (1, 'test')");
+ ResultSet rs = stmt.executeQuery("SELECT * FROM test_table");
+
+ assertTrue(rs.next());
+ assertEquals(1, rs.getInt("id"));
+ assertEquals("test", rs.getString("name"));
+ assertFalse(rs.next());
+
+ stmt.execute("DROP TABLE IF EXISTS test_table");
+ }
+}
\ No newline at end of file
diff --git a/bindings/java/src/test/java/com/databend/bendsql/DatabendDriverTest.java b/bindings/java/src/test/java/com/databend/bendsql/DatabendDriverTest.java
new file mode 100644
index 000000000..df7f1496f
--- /dev/null
+++ b/bindings/java/src/test/java/com/databend/bendsql/DatabendDriverTest.java
@@ -0,0 +1,64 @@
+package com.databend.bendsql;
+
+import org.junit.jupiter.api.Test;
+import static org.junit.jupiter.api.Assertions.*;
+
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.Statement;
+import java.util.Arrays;
+import java.util.List;
+import static com.databend.bendsql.utils.ResultSetTestUtils.assertResultSet;
+
+public class DatabendDriverTest {
+ private static final String TEST_DSN = "jdbc:databend://root:@localhost:8000/default?sslmode=disable";
+
+ @Test
+ public void testSimpleSelect() throws Exception {
+ try (Connection conn = DriverManager.getConnection(TEST_DSN, null, null);
+ Statement stmt = conn.createStatement();
+ ResultSet rs = stmt.executeQuery("SELECT 1, 'hello'")) {
+
+ assertTrue(rs.next());
+ assertEquals(1, rs.getInt(1));
+ assertEquals("hello", rs.getString(2));
+
+ assertFalse(rs.next());
+ }
+ }
+
+ @Test
+ public void testBatchInsert() throws Exception {
+ try (Connection conn = DriverManager.getConnection(TEST_DSN, null, null);) {
+ try(Statement stmt = conn.createStatement();) {
+ stmt.execute("create or replace table test_prepare_statement (a int, b string)");
+ }
+
+ try(PreparedStatement ps = conn.prepareStatement("insert into test_prepare_statement values");) {
+ ps.setInt(1, 1);
+ ps.setString(2, "a");
+ ps.addBatch();
+ ps.setInt(1, 2);
+ ps.setString(2, "b");
+ ps.addBatch();
+ int[] ans = ps.executeBatch();
+ assertEquals(ans.length, 2);
+ //assertEquals(ans[0], 1);
+ //assertEquals(ans[1], 1);
+ Statement statement = conn.createStatement();
+
+ boolean hasResultSet = statement.execute("SELECT * from test_prepare_statement");
+ assertTrue(hasResultSet);
+ try(ResultSet rs = statement.getResultSet();) {
+ List expected = Arrays.asList(
+ new Object[]{1, "a"},
+ new Object[]{2, "b"}
+ );
+ assertResultSet(rs, expected);
+ }
+ }
+ }
+ }
+}
diff --git a/bindings/java/src/test/java/com/databend/bendsql/NativeConnectionTest.java b/bindings/java/src/test/java/com/databend/bendsql/NativeConnectionTest.java
new file mode 100644
index 000000000..77be84733
--- /dev/null
+++ b/bindings/java/src/test/java/com/databend/bendsql/NativeConnectionTest.java
@@ -0,0 +1,91 @@
+/*
+ * Copyright 2021 Datafuse Labs
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+
+
+package com.databend.bendsql;
+
+import org.junit.jupiter.api.Test;
+
+import com.databend.bendsql.jni_utils.NativeLibrary;
+
+import org.junit.jupiter.api.BeforeEach;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+import java.sql.SQLException;
+import java.util.List;
+
+class NativeConnectionTest {
+ private NativeConnection connection;
+ private static final String TEST_DSN = "databend://root:@localhost:8000/default?sslmode=disable";
+
+ static {
+ NativeLibrary.loadLibrary();
+ }
+
+
+ @BeforeEach
+ void setUp() {
+ connection = NativeConnection.of(TEST_DSN);
+ }
+
+ @Test
+ void testSimpleQuery() {
+ String sql = "SELECT 1, 2";
+
+ NativeRowBatchIterator result = connection.execute(sql);
+
+ assertNotNull(result);
+ assertTrue(result.hasNext());
+ List> batch = result.next();
+ assertEquals(1, batch.size());
+ assertEquals(2, batch.get(0).size());
+ assertEquals("1", batch.get(0).get(0));
+ assertEquals("2", batch.get(0).get(1));
+ assertFalse(result.hasNext());
+ }
+
+ @Test
+ void testQueryInvalidQuery() {
+ String sql = "INVALID SQL QUERY";
+
+ assertThrows(SQLException.class, () -> connection.execute(sql));
+ }
+
+ @Test
+ void testQueryNullQuery() {
+ assertThrows(NullPointerException.class, () -> connection.execute(null));
+ }
+
+ @Test
+ void testQueryEmptyQuery() {
+ String sql = "";
+ assertThrows(SQLException.class, () -> connection.execute(sql));
+ }
+
+ @Test
+ void testMultipleQueriesSequentially() {
+ String sql1 = "SELECT 1";
+ String sql2 = "SELECT 2";
+
+ NativeRowBatchIterator result1 = connection.execute(sql1);
+ NativeRowBatchIterator result2 = connection.execute(sql2);
+
+ assertNotNull(result1);
+ assertNotNull(result2);
+ }
+}
\ No newline at end of file
diff --git a/bindings/java/src/test/java/com/databend/bendsql/utils/ResultSetTestUtils.java b/bindings/java/src/test/java/com/databend/bendsql/utils/ResultSetTestUtils.java
new file mode 100644
index 000000000..0a615e3b7
--- /dev/null
+++ b/bindings/java/src/test/java/com/databend/bendsql/utils/ResultSetTestUtils.java
@@ -0,0 +1,51 @@
+package com.databend.bendsql.utils;
+
+import org.junit.jupiter.api.Assertions;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.util.List;
+
+public class ResultSetTestUtils {
+
+ /**
+ * 验证 ResultSet 的内容是否与预期匹配
+ *
+ * @param rs ResultSet 实例
+ * @param expectedRows 预期的数据行，每行是一个 Object 数组
+ * @throws SQLException 如果访问 ResultSet 出错
+ */
+ public static void assertResultSet(ResultSet rs, List expectedRows) throws SQLException {
+ int rowNum = 0;
+ while (rs.next()) {
+ Assertions.assertTrue(rowNum < expectedRows.size(),
+ "Got more rows than expected. Expected " + expectedRows.size() + " rows");
+
+ Object[] expectedRow = expectedRows.get(rowNum);
+ for (int i = 0; i < expectedRow.length; i++) {
+ Object expected = expectedRow[i];
+ Object actual = rs.getObject(i + 1);
+ Assertions.assertEquals(expected, actual,
+ String.format("Row %d, Column %d mismatch", rowNum + 1, i + 1));
+ }
+ rowNum++;
+ }
+ Assertions.assertEquals(expectedRows.size(), rowNum,
+ "Got fewer rows than expected. Expected " + expectedRows.size() + " rows");
+ }
+
+ /**
+ * 验证 ResultSet 的行数
+ *
+ * @param rs ResultSet 实例
+ * @param expectedCount 预期的行数
+ * @throws SQLException 如果访问 ResultSet 出错
+ */
+ public static void assertRowCount(ResultSet rs, int expectedCount) throws SQLException {
+ int rowCount = 0;
+ while (rs.next()) {
+ rowCount++;
+ }
+ Assertions.assertEquals(expectedCount, rowCount,
+ String.format("Expected %d rows but got %d", expectedCount, rowCount));
+ }
+}
\ No newline at end of file
diff --git a/bindings/java/tools/build.py b/bindings/java/tools/build.py
new file mode 100755
index 000000000..df159e5b2
--- /dev/null
+++ b/bindings/java/tools/build.py
@@ -0,0 +1,106 @@
+#!/usr/bin/env python3
+# 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.
+
+
+from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
+from pathlib import Path
+import shutil
+import subprocess
+
+
+def classifier_to_target(classifier: str) -> str:
+ if classifier == "osx-aarch_64":
+ return "aarch64-apple-darwin"
+ if classifier == "osx-x86_64":
+ return "x86_64-apple-darwin"
+ if classifier == "linux-aarch_64":
+ return "aarch64-unknown-linux-gnu"
+ if classifier == "linux-x86_64":
+ return "x86_64-unknown-linux-gnu"
+ if classifier == "windows-x86_64":
+ return "x86_64-pc-windows-msvc"
+ raise Exception(f"Unsupported classifier: {classifier}")
+
+
+def get_cargo_artifact_name(classifier: str) -> str:
+ if classifier.startswith("osx"):
+ return "libbendsql_java.dylib"
+ if classifier.startswith("linux"):
+ return "libbendsql_java.so"
+ if classifier.startswith("windows"):
+ return "bendsql_java.dll"
+ raise Exception(f"Unsupported classifier: {classifier}")
+
+
+if __name__ == "__main__":
+ basedir = Path(__file__).parent.parent
+
+ parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
+ parser.add_argument("--classifier", type=str, required=True)
+ parser.add_argument("--target", type=str, default="")
+ parser.add_argument("--profile", type=str, default="dev")
+ parser.add_argument("--enable-zigbuild", type=str, default="false")
+ args = parser.parse_args()
+
+ if args.target:
+ target = args.target
+ else:
+ target = classifier_to_target(args.classifier)
+
+ # Setup target.
+ command = ["rustup", "target", "add", target]
+ print("$ " + subprocess.list2cmdline(command))
+ subprocess.run(command, cwd=basedir, check=True)
+
+ # Enable zigbuild if flag enabled and we are building linux target
+ enable_zigbuild = args.enable_zigbuild == "true" and "linux" in target
+
+ cmd = [
+ "cargo",
+ "zigbuild" if enable_zigbuild else "build",
+ "--color=always",
+ f"--profile={args.profile}",
+ ]
+
+
+ if enable_zigbuild:
+ # Pin glibc to 2.17 if zigbuild has been enabled.
+ cmd += ["--target", f"{target}.2.17"]
+ else:
+ cmd += ["--target", target]
+
+ output = basedir / "target" / "bindings"
+ Path(output).mkdir(exist_ok=True, parents=True)
+ cmd += ["--target-dir", str(output)]
+
+ print("$ " + subprocess.list2cmdline(cmd))
+ try:
+ subprocess.run(cmd, cwd=basedir, check=True, capture_output=True, text=True)
+ except subprocess.CalledProcessError as e:
+ print(f"Command failed with exit code {e.returncode}")
+ print("Error output:")
+ print(e.stderr)
+ raise
+
+ # History reason of cargo profiles.
+ profile = "debug" if args.profile in ["dev", "test", "bench"] else args.profile
+ artifact = get_cargo_artifact_name(args.classifier)
+ src = output / target / profile / artifact
+ dst = basedir / "target" / "classes" / "native" / args.classifier / artifact
+ dst.parent.mkdir(exist_ok=True, parents=True)
+ shutil.copy2(src, dst)
diff --git a/core/src/response.rs b/core/src/response.rs
index 048964f58..e7268ea86 100644
--- a/core/src/response.rs
+++ b/core/src/response.rs
@@ -14,7 +14,7 @@
use crate::error_code::ErrorCode;
use crate::session::SessionState;
-use serde::Deserialize;
+use serde::{Deserialize, Serialize};
#[derive(Deserialize, Debug)]
pub struct QueryStats {
@@ -53,7 +53,7 @@ pub struct ProgressValues {
pub bytes: usize,
}
-#[derive(Deserialize, Debug, Clone)]
+#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct SchemaField {
pub name: String,
#[serde(rename = "type")]
diff --git a/driver/src/rest_api.rs b/driver/src/rest_api.rs
index 39ef627c5..cbfb71989 100644
--- a/driver/src/rest_api.rs
+++ b/driver/src/rest_api.rs
@@ -27,9 +27,9 @@ use tokio::fs::File;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio_stream::Stream;
-use databend_client::APIClient;
use databend_client::PresignedResponse;
use databend_client::QueryResponse;
+use databend_client::{APIClient, SchemaField};
use databend_driver_core::error::{Error, Result};
use databend_driver_core::rows::{Row, RowIterator, RowStatsIterator, RowWithStats, ServerStats};
use databend_driver_core::schema::{Schema, SchemaRef};
@@ -82,7 +82,7 @@ impl Connection for RestAPIConnection {
async fn query_iter_ext(&self, sql: &str) -> Result {
info!("query iter ext: {}", sql);
let resp = self.client.start_query(sql).await?;
- let resp = self.wait_for_schema(resp).await?;
+ let resp = self.wait_for_schema(resp, true).await?;
let (schema, rows) = RestAPIRows::from_response(self.client.clone(), resp)?;
Ok(RowStatsIterator::new(Arc::new(schema), Box::pin(rows)))
}
@@ -209,8 +209,14 @@ impl<'o> RestAPIConnection {
})
}
- async fn wait_for_schema(&self, resp: QueryResponse) -> Result {
- if !resp.data.is_empty() || !resp.schema.is_empty() || resp.stats.progresses.has_progress()
+ async fn wait_for_schema(
+ &self,
+ resp: QueryResponse,
+ return_on_progress: bool,
+ ) -> Result {
+ if !resp.data.is_empty()
+ || !resp.schema.is_empty()
+ || (return_on_progress && resp.stats.progresses.has_progress())
{
return Ok(resp);
}
@@ -228,7 +234,7 @@ impl<'o> RestAPIConnection {
if !result.data.is_empty()
|| !result.schema.is_empty()
- || result.stats.progresses.has_progress()
+ || (return_on_progress && result.stats.progresses.has_progress())
{
break;
}
@@ -250,6 +256,12 @@ impl<'o> RestAPIConnection {
fn default_copy_options() -> BTreeMap<&'o str, &'o str> {
vec![("purge", "true")].into_iter().collect()
}
+
+ pub async fn query_row_batch(&self, sql: &str) -> Result {
+ let resp = self.client.start_query(sql).await?;
+ let resp = self.wait_for_schema(resp, false).await?;
+ Ok(RowBatch::from_response(self.client.clone(), resp)?)
+ }
}
type PageFut = Pin> + Send>>;
@@ -341,3 +353,48 @@ impl Stream for RestAPIRows {
}
}
}
+
+pub struct RowBatch {
+ schema: Vec,
+ client: Arc,
+ query_id: String,
+ node_id: Option,
+
+ next_uri: Option,
+ data: Vec>>,
+}
+
+impl RowBatch {
+ pub fn schema(&self) -> Vec {
+ self.schema.clone()
+ }
+
+ fn from_response(client: Arc, mut resp: QueryResponse) -> Result {
+ Ok(Self {
+ schema: std::mem::take(&mut resp.schema),
+ client,
+ query_id: resp.id,
+ node_id: resp.node_id,
+ next_uri: resp.next_uri,
+ data: resp.data,
+ })
+ }
+
+ pub async fn fetch_next_page(&mut self) -> Result>>> {
+ if !self.data.is_empty() {
+ return Ok(std::mem::take(&mut self.data));
+ }
+ while let Some(next_uri) = &self.next_uri {
+ let resp = self
+ .client
+ .query_page(&self.query_id, &next_uri, &self.node_id)
+ .await?;
+
+ self.next_uri = resp.next_uri;
+ if !resp.data.is_empty() {
+ return Ok(resp.data);
+ }
+ }
+ Ok(vec![])
+ }
+}
diff --git a/sql/src/error.rs b/sql/src/error.rs
index 752a80ce2..88ba25cca 100644
--- a/sql/src/error.rs
+++ b/sql/src/error.rs
@@ -48,6 +48,7 @@ pub enum Error {
#[cfg(feature = "flight-sql")]
Arrow(arrow_schema::ArrowError),
Convert(ConvertError),
+ Unexpected(String),
}
impl std::fmt::Display for Error {
@@ -57,7 +58,7 @@ impl std::fmt::Display for Error {
Error::Protocol(msg) => write!(f, "ProtocolError: {}", msg),
Error::Transport(msg) => write!(f, "TransportError: {}", msg),
Error::IO(msg) => write!(f, "IOError: {}", msg),
-
+ Error::Unexpected(msg) => write!(f, "Unexpected: {}", msg),
Error::BadArgument(msg) => write!(f, "BadArgument: {}", msg),
Error::InvalidResponse(msg) => write!(f, "ResponseError: {}", msg),
#[cfg(feature = "flight-sql")]