Skip to content

Commit 4e82ca3

Browse files
youngsofunyoungsofun
committed
support arrow body.
1 parent 197843a commit 4e82ca3

File tree

21 files changed

+498
-159
lines changed

21 files changed

+498
-159
lines changed

‎Cargo.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,11 +30,13 @@ databend-driver-macros = { path = "macros", version = "0.30.3" }
3030
jsonb = { version = "0.5.3" }
3131
tokio-stream = "0.1"
3232
chrono = { version = "0.4.40", default-features = false, features = ["clock"] }
33+
chrono-tz = { version = "0.10.4" }
3334
arrow = { version = "55.0" }
3435
arrow-array = { version = "55.0" }
3536
arrow-buffer = { version = "55.0" }
3637
arrow-schema = { version = "55.0" }
3738
arrow-flight = { version = "55.0", features = ["flight-sql-experimental"] }
39+
arrow-ipc = { version = "55.0", features = ["lz4", "zstd"]}
3840
tonic = { version = "0.12", default-features = false, features = [
3941
"transport",
4042
"codegen",

‎bindings/nodejs/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ doc = false
1414

1515
[dependencies]
1616
chrono = { workspace = true }
17+
chrono-tz = { workspace = true }
1718
databend-driver = { workspace = true, features = ["rustls", "flight-sql"] }
1819
env_logger = "0.11.8"
1920
tokio-stream = { workspace = true }

‎bindings/nodejs/src/lib.rs‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,13 +15,14 @@
1515
#[macro_use]
1616
extern crate napi_derive;
1717

18-
use chrono::{NaiveDate, NaiveDateTime, NaiveTime};
18+
use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime};
1919
use databend_driver::LoadMethod;
2020
use napi::{bindgen_prelude::*, Env};
2121
use once_cell::sync::Lazy;
2222
use std::str::FromStr;
2323
use std::sync::Arc;
2424
use std::{collections::HashMap, path::Path};
25+
use chrono_tz::Tz;
2526
use tokio_stream::StreamExt;
2627

2728
static VERSION: Lazy<String> = Lazy::new(|| {
@@ -307,10 +308,10 @@ impl ToNapiValue for Value<'_> {
307308
databend_driver::Value::Number(n) => {
308309
NumberValue::to_napi_value(env, NumberValue(n.clone()))
309310
}
310-
databend_driver::Value::Timestamp(_) => {
311+
databend_driver::Value::Timestamp(_, _tz) => {
311312
let inner = val.inner.clone();
312-
let v = NaiveDateTime::try_from(inner).map_err(format_napi_error)?;
313-
NaiveDateTime::to_napi_value(env, v)
313+
let v = DateTime::<Tz>::try_from(inner).map_err(format_napi_error)?;
314+
DateTime::to_napi_value(env, v)
314315
}
315316
databend_driver::Value::Date(_) => {
316317
let inner = val.inner.clone();

‎bindings/python/Cargo.toml‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ doc = false
1414

1515
[dependencies]
1616
chrono = { workspace = true }
17+
chrono-tz = { workspace = true }
1718
databend-client = { workspace = true }
1819
databend-driver = { workspace = true, features = ["rustls", "flight-sql"] }
1920
databend-driver-core = { workspace = true }
@@ -23,6 +24,6 @@ ctor = "0.2"
2324
env_logger = "0.11.8"
2425
http = "1.0"
2526
once_cell = "1.21"
26-
pyo3 = { version = "0.24.2", features = ["abi3-py37", "chrono"] }
27+
pyo3 = { version = "0.24.2", features = ["abi3-py37", "chrono", "chrono-tz"] }
2728
pyo3-async-runtimes = { version = "0.24", features = ["tokio-runtime"] }
2829
tokio = "1.44"

‎bindings/python/src/types.rs‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,9 @@
1414

1515
use std::sync::Arc;
1616

17-
use chrono::{Duration, NaiveDate, NaiveDateTime};
17+
use chrono::offset::Offset;
18+
use chrono::{DateTime, Duration, FixedOffset, NaiveDate};
19+
use chrono_tz::Tz;
1820
use once_cell::sync::Lazy;
1921
use pyo3::exceptions::{PyAttributeError, PyException, PyStopAsyncIteration, PyStopIteration};
2022
use pyo3::sync::GILOnceCell;
@@ -73,10 +75,16 @@ impl<'py> IntoPyObject<'py> for Value {
7375
let v = NumberValue(n);
7476
v.into_bound_py_any(py)?
7577
}
76-
databend_driver::Value::Timestamp(_) => {
77-
let t = NaiveDateTime::try_from(self.0).map_err(|e| {
78+
databend_driver::Value::Timestamp(_, _) => {
79+
let t = DateTime::<Tz>::try_from(self.0).map_err(|e| {
7880
PyException::new_err(format!("failed to convert timestamp: {e}"))
7981
})?;
82+
// impl of IntoPyObject for chrono_tz::Tz is gated by #![cfg(all(Py_3_9, feature = "chrono-tz"))]
83+
// Convert to a fixed-offset datetime so it works on Python < 3.9 (no zoneinfo).
84+
// other options:
85+
// 1. abi3-py37 -> abi3-py39: abandon py3.8
86+
// 2. rm abi3-py37: a wheel for each version of python
87+
let t: DateTime<FixedOffset> = t.with_timezone(&t.offset().fix());
8088
t.into_bound_py_any(py)?
8189
}
8290
databend_driver::Value::Date(_) => {

‎bindings/python/tests/asyncio/steps/binding.py‎

Lines changed: 90 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
import os
1616
import gc
1717
import time
18-
from datetime import datetime, date, timedelta
18+
from datetime import datetime, date, timedelta, timezone
1919
from decimal import Decimal
2020

2121
from behave import given, when, then
@@ -37,6 +37,11 @@
3737
else:
3838
DRIVER_VERSION = (100, 0, 0)
3939

40+
if DRIVER_VERSION > (0, 30, 3) and DB_VERSION >= (1, 2, 836):
41+
default_tzinfo = timezone.utc
42+
else:
43+
default_tzinfo = None
44+
4045

4146
@given("A new Databend Driver Client")
4247
async def _(context):
@@ -130,9 +135,9 @@ async def _(context):
130135
row = await context.conn.query_row(
131136
"select (10, '20', to_datetime('2024-04-16 12:34:56.789'))"
132137
)
133-
assert row.values() == ((10, "20", datetime(2024, 4, 16, 12, 34, 56, 789000)),), (
134-
f"Tuple: {row.values()}"
135-
)
138+
assert row.values() == (
139+
(10, "20", datetime(2024, 4, 16, 12, 34, 56, 789000, tzinfo=default_tzinfo)),
140+
), f"Tuple: {row.values()}"
136141

137142

138143
@then("Select numbers should iterate all rows")
@@ -156,9 +161,33 @@ async def _(context):
156161
rows = await context.conn.query_iter("SELECT * FROM test")
157162
ret = [row.values() for row in rows]
158163
expected = [
159-
(-1, 1, 1.0, "'", None, date(2011, 3, 6), datetime(2011, 3, 6, 6, 20)),
160-
(-2, 2, 2.0, '"', "", date(2012, 5, 31), datetime(2012, 5, 31, 11, 20)),
161-
(-3, 3, 3.0, "\\", "NULL", date(2016, 4, 4), datetime(2016, 4, 4, 11, 30)),
164+
(
165+
-1,
166+
1,
167+
1.0,
168+
"'",
169+
None,
170+
date(2011, 3, 6),
171+
datetime(2011, 3, 6, 6, 20, tzinfo=default_tzinfo),
172+
),
173+
(
174+
-2,
175+
2,
176+
2.0,
177+
'"',
178+
"",
179+
date(2012, 5, 31),
180+
datetime(2012, 5, 31, 11, 20, tzinfo=default_tzinfo),
181+
),
182+
(
183+
-3,
184+
3,
185+
3.0,
186+
"\\",
187+
"NULL",
188+
date(2016, 4, 4),
189+
datetime(2016, 4, 4, 11, 30, tzinfo=default_tzinfo),
190+
),
162191
]
163192
assert ret == expected, f"ret: {ret}"
164193

@@ -176,9 +205,33 @@ async def _(context):
176205
rows = await context.conn.query_iter("SELECT * FROM test")
177206
ret = [row.values() for row in rows]
178207
expected = [
179-
(-1, 1, 1.0, "'", None, date(2011, 3, 6), datetime(2011, 3, 6, 6, 20)),
180-
(-2, 2, 2.0, '"', None, date(2012, 5, 31), datetime(2012, 5, 31, 11, 20)),
181-
(-3, 3, 3.0, "\\", "NULL", date(2016, 4, 4), datetime(2016, 4, 4, 11, 30)),
208+
(
209+
-1,
210+
1,
211+
1.0,
212+
"'",
213+
None,
214+
date(2011, 3, 6),
215+
datetime(2011, 3, 6, 6, 20, tzinfo=default_tzinfo),
216+
),
217+
(
218+
-2,
219+
2,
220+
2.0,
221+
'"',
222+
None,
223+
date(2012, 5, 31),
224+
datetime(2012, 5, 31, 11, 20, tzinfo=default_tzinfo),
225+
),
226+
(
227+
-3,
228+
3,
229+
3.0,
230+
"\\",
231+
"NULL",
232+
date(2016, 4, 4),
233+
datetime(2016, 4, 4, 11, 30, tzinfo=default_tzinfo),
234+
),
182235
]
183236
assert ret == expected, f"ret: {ret}"
184237

@@ -213,9 +266,33 @@ async def test_load_file(context, load_method):
213266
rows = await context.conn.query_iter("SELECT * FROM test1")
214267
ret = [row.values() for row in rows]
215268
expected = [
216-
(-1, 1, 1.0, "'", None, date(2011, 3, 6), datetime(2011, 3, 6, 6, 20)),
217-
(-2, 2, 2.0, '"', None, date(2012, 5, 31), datetime(2012, 5, 31, 11, 20)),
218-
(-3, 3, 3.0, "\\", "NULL", date(2016, 4, 4), datetime(2016, 4, 4, 11, 30)),
269+
(
270+
-1,
271+
1,
272+
1.0,
273+
"'",
274+
None,
275+
date(2011, 3, 6),
276+
datetime(2011, 3, 6, 6, 20, tzinfo=default_tzinfo),
277+
),
278+
(
279+
-2,
280+
2,
281+
2.0,
282+
'"',
283+
None,
284+
date(2012, 5, 31),
285+
datetime(2012, 5, 31, 11, 20, tzinfo=default_tzinfo),
286+
),
287+
(
288+
-3,
289+
3,
290+
3.0,
291+
"\\",
292+
"NULL",
293+
date(2016, 4, 4),
294+
datetime(2016, 4, 4, 11, 30, tzinfo=default_tzinfo),
295+
),
219296
]
220297
assert ret == expected, f"{load_method} ret: {ret}"
221298

‎bindings/python/tests/blocking/steps/binding.py‎

Lines changed: 91 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414

1515
import os
1616
import gc
17-
from datetime import datetime, date, timedelta
17+
from datetime import datetime, date, timedelta, timezone
1818
from decimal import Decimal
1919
import time
2020
from time import sleep
@@ -38,12 +38,17 @@
3838
else:
3939
DRIVER_VERSION = (100, 0, 0)
4040

41+
if DRIVER_VERSION > (0, 30, 3) and DB_VERSION >= (1, 2, 836):
42+
default_tzinfo = timezone.utc
43+
else:
44+
default_tzinfo = None
45+
4146

4247
@given("A new Databend Driver Client")
4348
def _(context):
4449
dsn = os.getenv(
4550
"TEST_DATABEND_DSN",
46-
"databend://root:root@localhost:8000/?sslmode=disable",
51+
"databend://root:root@localhost:8000/?sslmode=disable&arrow_data=disable",
4752
)
4853
client = databend_driver.BlockingDatabendClient(dsn)
4954
context.conn = client.get_conn()
@@ -126,9 +131,9 @@ def _(context):
126131
row = context.conn.query_row(
127132
"select (10, '20', to_datetime('2024-04-16 12:34:56.789'))"
128133
)
129-
assert row.values() == ((10, "20", datetime(2024, 4, 16, 12, 34, 56, 789000)),), (
130-
f"Tuple: {row.values()}"
131-
)
134+
assert row.values() == (
135+
(10, "20", datetime(2024, 4, 16, 12, 34, 56, 789000, tzinfo=default_tzinfo)),
136+
), f"Tuple: {row.values()}"
132137

133138

134139
@then("Select numbers should iterate all rows")
@@ -152,9 +157,33 @@ def _(context):
152157
rows = context.conn.query_iter("SELECT * FROM test")
153158
ret = [row.values() for row in rows]
154159
expected = [
155-
(-1, 1, 1.0, "'", None, date(2011, 3, 6), datetime(2011, 3, 6, 6, 20)),
156-
(-2, 2, 2.0, '"', "", date(2012, 5, 31), datetime(2012, 5, 31, 11, 20)),
157-
(-3, 3, 3.0, "\\", "NULL", date(2016, 4, 4), datetime(2016, 4, 4, 11, 30)),
160+
(
161+
-1,
162+
1,
163+
1.0,
164+
"'",
165+
None,
166+
date(2011, 3, 6),
167+
datetime(2011, 3, 6, 6, 20, tzinfo=default_tzinfo),
168+
),
169+
(
170+
-2,
171+
2,
172+
2.0,
173+
'"',
174+
"",
175+
date(2012, 5, 31),
176+
datetime(2012, 5, 31, 11, 20, tzinfo=default_tzinfo),
177+
),
178+
(
179+
-3,
180+
3,
181+
3.0,
182+
"\\",
183+
"NULL",
184+
date(2016, 4, 4),
185+
datetime(2016, 4, 4, 11, 30, tzinfo=default_tzinfo),
186+
),
158187
]
159188
assert ret == expected, f"ret: {ret}"
160189

@@ -172,9 +201,33 @@ def _(context):
172201
rows = context.conn.query_iter("SELECT * FROM test")
173202
ret = [row.values() for row in rows]
174203
expected = [
175-
(-1, 1, 1.0, "'", None, date(2011, 3, 6), datetime(2011, 3, 6, 6, 20)),
176-
(-2, 2, 2.0, '"', None, date(2012, 5, 31), datetime(2012, 5, 31, 11, 20)),
177-
(-3, 3, 3.0, "\\", "NULL", date(2016, 4, 4), datetime(2016, 4, 4, 11, 30)),
204+
(
205+
-1,
206+
1,
207+
1.0,
208+
"'",
209+
None,
210+
date(2011, 3, 6),
211+
datetime(2011, 3, 6, 6, 20, tzinfo=default_tzinfo),
212+
),
213+
(
214+
-2,
215+
2,
216+
2.0,
217+
'"',
218+
None,
219+
date(2012, 5, 31),
220+
datetime(2012, 5, 31, 11, 20, tzinfo=default_tzinfo),
221+
),
222+
(
223+
-3,
224+
3,
225+
3.0,
226+
"\\",
227+
"NULL",
228+
date(2016, 4, 4),
229+
datetime(2016, 4, 4, 11, 30, tzinfo=default_tzinfo),
230+
),
178231
]
179232
assert ret == expected, f"ret: {ret}"
180233

@@ -207,9 +260,33 @@ def test_load_file(context, load_method):
207260
rows = context.conn.query_iter("SELECT * FROM test1")
208261
ret = [row.values() for row in rows]
209262
expected = [
210-
(-1, 1, 1.0, "'", None, date(2011, 3, 6), datetime(2011, 3, 6, 6, 20)),
211-
(-2, 2, 2.0, '"', None, date(2012, 5, 31), datetime(2012, 5, 31, 11, 20)),
212-
(-3, 3, 3.0, "\\", "NULL", date(2016, 4, 4), datetime(2016, 4, 4, 11, 30)),
263+
(
264+
-1,
265+
1,
266+
1.0,
267+
"'",
268+
None,
269+
date(2011, 3, 6),
270+
datetime(2011, 3, 6, 6, 20, tzinfo=default_tzinfo),
271+
),
272+
(
273+
-2,
274+
2,
275+
2.0,
276+
'"',
277+
None,
278+
date(2012, 5, 31),
279+
datetime(2012, 5, 31, 11, 20, tzinfo=default_tzinfo),
280+
),
281+
(
282+
-3,
283+
3,
284+
3.0,
285+
"\\",
286+
"NULL",
287+
date(2016, 4, 4),
288+
datetime(2016, 4, 4, 11, 30, tzinfo=default_tzinfo),
289+
),
213290
]
214291
assert ret == expected, f"{load_method} ret: {ret}"
215292

0 commit comments

Comments
 (0)