quack-ruby 0.1.0
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- checksums.yaml +7 -0
- data/LICENSE +21 -0
- data/LICENSE-DuckDB +7 -0
- data/LICENSE-QuackProtocol +21 -0
- data/LICENSE-Rust-dependencies.json +2057 -0
- data/README.md +229 -0
- data/ext/quack/Cargo.lock +1803 -0
- data/ext/quack/Cargo.toml +17 -0
- data/ext/quack/extconf.rb +4 -0
- data/ext/quack/src/lib.rs +647 -0
- data/ext/quack/vendor/quack_protocol/Cargo.lock +1907 -0
- data/ext/quack/vendor/quack_protocol/Cargo.toml +41 -0
- data/ext/quack/vendor/quack_protocol/LICENSE +21 -0
- data/ext/quack/vendor/quack_protocol/PATCHES.md +16 -0
- data/ext/quack/vendor/quack_protocol/README.md +272 -0
- data/ext/quack/vendor/quack_protocol/src/arrow/columns.rs +399 -0
- data/ext/quack/vendor/quack_protocol/src/arrow/mod.rs +134 -0
- data/ext/quack/vendor/quack_protocol/src/arrow/schema.rs +103 -0
- data/ext/quack/vendor/quack_protocol/src/arrow/tests.rs +851 -0
- data/ext/quack/vendor/quack_protocol/src/arrow/values.rs +150 -0
- data/ext/quack/vendor/quack_protocol/src/binary.rs +504 -0
- data/ext/quack/vendor/quack_protocol/src/builders.rs +217 -0
- data/ext/quack/vendor/quack_protocol/src/client.rs +1486 -0
- data/ext/quack/vendor/quack_protocol/src/constants.rs +11 -0
- data/ext/quack/vendor/quack_protocol/src/errors.rs +103 -0
- data/ext/quack/vendor/quack_protocol/src/inspection.rs +250 -0
- data/ext/quack/vendor/quack_protocol/src/json.rs +225 -0
- data/ext/quack/vendor/quack_protocol/src/lib.rs +54 -0
- data/ext/quack/vendor/quack_protocol/src/logical_types.rs +977 -0
- data/ext/quack/vendor/quack_protocol/src/messages.rs +843 -0
- data/ext/quack/vendor/quack_protocol/src/pool.rs +741 -0
- data/ext/quack/vendor/quack_protocol/src/protocol_tests.rs +759 -0
- data/ext/quack/vendor/quack_protocol/src/proxy_error.rs +77 -0
- data/ext/quack/vendor/quack_protocol/src/sql.rs +359 -0
- data/ext/quack/vendor/quack_protocol/src/tls.rs +354 -0
- data/ext/quack/vendor/quack_protocol/src/values.rs +77 -0
- data/ext/quack/vendor/quack_protocol/src/vector.rs +1364 -0
- data/ext/quack/vendor/quack_protocol/testdata/tls/key.der +0 -0
- data/ext/quack/vendor/quack_protocol/testdata/tls/spki.der +0 -0
- data/ext/quack/vendor/quack_protocol/testdata/tls/v1-cert.der +0 -0
- data/ext/quack/vendor/quack_protocol/testdata/tls/v3-cert.der +0 -0
- data/lib/quack/codec.rb +263 -0
- data/lib/quack/connection.rb +131 -0
- data/lib/quack/logical_types.rb +72 -0
- data/lib/quack/result.rb +64 -0
- data/lib/quack/sql.rb +183 -0
- data/lib/quack/types.rb +40 -0
- data/lib/quack/variant.rb +298 -0
- data/lib/quack.rb +34 -0
- metadata +136 -0
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
[package]
|
|
2
|
+
name = "quack-ruby"
|
|
3
|
+
version = "0.1.0"
|
|
4
|
+
edition = "2024"
|
|
5
|
+
rust-version = "1.85"
|
|
6
|
+
license = "MIT"
|
|
7
|
+
|
|
8
|
+
[lib]
|
|
9
|
+
name = "quack_native"
|
|
10
|
+
crate-type = ["cdylib"]
|
|
11
|
+
|
|
12
|
+
[dependencies]
|
|
13
|
+
magnus = { version = "0.9", features = ["rb-sys"] }
|
|
14
|
+
rb-sys = "0.9"
|
|
15
|
+
quack_protocol = { version = "=0.4.1", path = "vendor/quack_protocol" }
|
|
16
|
+
tokio = { version = "1", features = ["rt-multi-thread", "time", "net", "sync"] }
|
|
17
|
+
futures-util = "0.3"
|
|
@@ -0,0 +1,647 @@
|
|
|
1
|
+
use futures_util::{StreamExt, stream::BoxStream};
|
|
2
|
+
use magnus::{Error, IntoValue, RArray, RHash, RModule, Ruby, Value, function, method, prelude::*};
|
|
3
|
+
use quack_protocol::{
|
|
4
|
+
ColumnDefinition, DataChunk, ExtraTypeInfo, LogicalType, QuackClient, QuackClientOptions,
|
|
5
|
+
QuackError, Row, Value as QValue,
|
|
6
|
+
};
|
|
7
|
+
use std::{
|
|
8
|
+
ffi::c_void,
|
|
9
|
+
mem::ManuallyDrop,
|
|
10
|
+
panic::{AssertUnwindSafe, catch_unwind},
|
|
11
|
+
sync::{
|
|
12
|
+
Arc, Mutex,
|
|
13
|
+
atomic::{AtomicBool, Ordering},
|
|
14
|
+
},
|
|
15
|
+
time::Duration,
|
|
16
|
+
};
|
|
17
|
+
use tokio::runtime::Runtime;
|
|
18
|
+
|
|
19
|
+
// Ruby unwinding is contained by protect; Rust unwinding is contained by the
|
|
20
|
+
// callback. No Ruby value or API is touched while the GVL is released.
|
|
21
|
+
fn without_gvl<T: Send, F: FnOnce() -> T + Send>(ruby: &Ruby, operation: F) -> Result<T, Error> {
|
|
22
|
+
struct Call<F, T> {
|
|
23
|
+
operation: Option<F>,
|
|
24
|
+
result: Option<std::thread::Result<T>>,
|
|
25
|
+
}
|
|
26
|
+
unsafe extern "C" fn run<F: FnOnce() -> T, T>(data: *mut c_void) -> *mut c_void {
|
|
27
|
+
let call = unsafe { &mut *data.cast::<Call<F, T>>() };
|
|
28
|
+
call.result = Some(catch_unwind(AssertUnwindSafe(|| {
|
|
29
|
+
(call.operation.take().unwrap())()
|
|
30
|
+
})));
|
|
31
|
+
std::ptr::null_mut()
|
|
32
|
+
}
|
|
33
|
+
let mut call = Call {
|
|
34
|
+
operation: Some(operation),
|
|
35
|
+
result: None,
|
|
36
|
+
};
|
|
37
|
+
magnus::rb_sys::protect(|| {
|
|
38
|
+
unsafe {
|
|
39
|
+
rb_sys::rb_thread_call_without_gvl(
|
|
40
|
+
Some(run::<F, T>),
|
|
41
|
+
(&mut call as *mut Call<F, T>).cast(),
|
|
42
|
+
None,
|
|
43
|
+
std::ptr::null_mut(),
|
|
44
|
+
);
|
|
45
|
+
}
|
|
46
|
+
0
|
|
47
|
+
})?;
|
|
48
|
+
match call.result {
|
|
49
|
+
Some(Ok(value)) => Ok(value),
|
|
50
|
+
_ => Err(Error::new(
|
|
51
|
+
ruby.exception_runtime_error(),
|
|
52
|
+
"native Quack operation panicked",
|
|
53
|
+
)),
|
|
54
|
+
}
|
|
55
|
+
}
|
|
56
|
+
fn error(ruby: &Ruby, name: &str, message: impl Into<String>) -> Error {
|
|
57
|
+
let class = ruby
|
|
58
|
+
.class_object()
|
|
59
|
+
.const_get::<_, RModule>("Quack")
|
|
60
|
+
.unwrap()
|
|
61
|
+
.const_get(name)
|
|
62
|
+
.unwrap();
|
|
63
|
+
Error::new(class, message.into())
|
|
64
|
+
}
|
|
65
|
+
fn protocol_error(ruby: &Ruby, err: QuackError) -> Error {
|
|
66
|
+
let name = match &err {
|
|
67
|
+
QuackError::Http(e) if e.is_timeout() => "TimeoutError",
|
|
68
|
+
QuackError::Http(_) | QuackError::Url(_) => "ConnectionError",
|
|
69
|
+
QuackError::Server(_) => "ServerError",
|
|
70
|
+
QuackError::UnsupportedType(_) => "UnsupportedTypeError",
|
|
71
|
+
QuackError::Proxy(_) => "ProxyError",
|
|
72
|
+
_ => "ProtocolError",
|
|
73
|
+
};
|
|
74
|
+
error(ruby, name, err.to_string())
|
|
75
|
+
}
|
|
76
|
+
struct Session {
|
|
77
|
+
client: QuackClient,
|
|
78
|
+
runtime: Option<Runtime>,
|
|
79
|
+
active: AtomicBool,
|
|
80
|
+
fatal: AtomicBool,
|
|
81
|
+
}
|
|
82
|
+
struct Admission<'a> {
|
|
83
|
+
active: &'a AtomicBool,
|
|
84
|
+
release: bool,
|
|
85
|
+
}
|
|
86
|
+
impl Drop for Admission<'_> {
|
|
87
|
+
fn drop(&mut self) {
|
|
88
|
+
if self.release {
|
|
89
|
+
self.active.store(false, Ordering::Release);
|
|
90
|
+
}
|
|
91
|
+
}
|
|
92
|
+
}
|
|
93
|
+
impl Session {
|
|
94
|
+
fn runtime(&self) -> &Runtime {
|
|
95
|
+
self.runtime.as_ref().unwrap()
|
|
96
|
+
}
|
|
97
|
+
fn check(&self, ruby: &Ruby) -> Result<(), Error> {
|
|
98
|
+
if self.fatal.load(Ordering::Acquire) || !self.client.is_connected() {
|
|
99
|
+
return Err(error(
|
|
100
|
+
ruby,
|
|
101
|
+
"ConnectionError",
|
|
102
|
+
"connection is closed or failed; reconnect explicitly",
|
|
103
|
+
));
|
|
104
|
+
}
|
|
105
|
+
Ok(())
|
|
106
|
+
}
|
|
107
|
+
fn admit(&self, ruby: &Ruby) -> Result<Admission<'_>, Error> {
|
|
108
|
+
if self.active.swap(true, Ordering::AcqRel) {
|
|
109
|
+
return Err(error(
|
|
110
|
+
ruby,
|
|
111
|
+
"BusyError",
|
|
112
|
+
"consume or close the active result before another operation",
|
|
113
|
+
));
|
|
114
|
+
}
|
|
115
|
+
Ok(Admission {
|
|
116
|
+
active: &self.active,
|
|
117
|
+
release: true,
|
|
118
|
+
})
|
|
119
|
+
}
|
|
120
|
+
fn map<T>(&self, ruby: &Ruby, result: Result<T, QuackError>) -> Result<T, Error> {
|
|
121
|
+
result.map_err(|err| {
|
|
122
|
+
if err.is_connection_fatal() {
|
|
123
|
+
self.fatal.store(true, Ordering::Release);
|
|
124
|
+
}
|
|
125
|
+
protocol_error(ruby, err)
|
|
126
|
+
})
|
|
127
|
+
}
|
|
128
|
+
}
|
|
129
|
+
impl Drop for Session {
|
|
130
|
+
fn drop(&mut self) {
|
|
131
|
+
if let Some(runtime) = self.runtime.take() {
|
|
132
|
+
// GC must not block Ruby. The cleanup thread owns the runtime until
|
|
133
|
+
// a bounded DISCONNECT completes. Explicit #close reports errors.
|
|
134
|
+
struct CleanupRuntime(Option<Runtime>);
|
|
135
|
+
impl Drop for CleanupRuntime {
|
|
136
|
+
fn drop(&mut self) {
|
|
137
|
+
if let Some(runtime) = self.0.take() {
|
|
138
|
+
runtime.shutdown_background();
|
|
139
|
+
}
|
|
140
|
+
}
|
|
141
|
+
}
|
|
142
|
+
let runtime = CleanupRuntime(Some(runtime));
|
|
143
|
+
let client = self.client.clone();
|
|
144
|
+
let _ = std::thread::Builder::new()
|
|
145
|
+
.name("quack-cleanup".into())
|
|
146
|
+
.spawn(move || {
|
|
147
|
+
runtime.0.as_ref().unwrap().block_on(async {
|
|
148
|
+
let _ =
|
|
149
|
+
tokio::time::timeout(Duration::from_secs(2), client.disconnect()).await;
|
|
150
|
+
});
|
|
151
|
+
});
|
|
152
|
+
}
|
|
153
|
+
}
|
|
154
|
+
}
|
|
155
|
+
#[magnus::wrap(class = "Quack::NativeConnection", free_immediately, size)]
|
|
156
|
+
struct NativeConnection {
|
|
157
|
+
session: ManuallyDrop<Arc<Session>>,
|
|
158
|
+
pid: u32,
|
|
159
|
+
}
|
|
160
|
+
impl Drop for NativeConnection {
|
|
161
|
+
fn drop(&mut self) {
|
|
162
|
+
// A fork inherits Tokio's locks but not its worker threads. Never touch
|
|
163
|
+
// or destroy inherited runtime state in the child process.
|
|
164
|
+
if self.pid == std::process::id() {
|
|
165
|
+
unsafe {
|
|
166
|
+
ManuallyDrop::drop(&mut self.session);
|
|
167
|
+
}
|
|
168
|
+
}
|
|
169
|
+
}
|
|
170
|
+
}
|
|
171
|
+
impl NativeConnection {
|
|
172
|
+
fn check_pid(&self, ruby: &Ruby) -> Result<(), Error> {
|
|
173
|
+
if self.pid != std::process::id() {
|
|
174
|
+
return Err(error(
|
|
175
|
+
ruby,
|
|
176
|
+
"ForkError",
|
|
177
|
+
"Quack connections cannot cross fork; reconnect in the child",
|
|
178
|
+
));
|
|
179
|
+
}
|
|
180
|
+
Ok(())
|
|
181
|
+
}
|
|
182
|
+
fn connect(ruby: &Ruby, uri: String, options: RHash) -> Result<Self, Error> {
|
|
183
|
+
let mut opts = QuackClientOptions {
|
|
184
|
+
auth_token: options.lookup("auth_token")?,
|
|
185
|
+
ssl_fingerprint: options.lookup("ssl_fingerprint")?,
|
|
186
|
+
ssl: options.lookup("ssl")?,
|
|
187
|
+
client_id: options.lookup("client_id")?,
|
|
188
|
+
client_duckdb_version: options.lookup("client_duckdb_version")?,
|
|
189
|
+
client_platform: options.lookup("client_platform")?,
|
|
190
|
+
min_supported_quack_version: options.lookup("min_supported_quack_version")?,
|
|
191
|
+
max_supported_quack_version: options.lookup("max_supported_quack_version")?,
|
|
192
|
+
..Default::default()
|
|
193
|
+
};
|
|
194
|
+
for key in ["timeout", "heartbeat_timeout"] {
|
|
195
|
+
if let Some(seconds) = options.lookup::<_, Option<f64>>(key)? {
|
|
196
|
+
if !seconds.is_finite() || seconds <= 0.0 {
|
|
197
|
+
return Err(Error::new(
|
|
198
|
+
ruby.exception_arg_error(),
|
|
199
|
+
format!("{key} must be a positive finite number"),
|
|
200
|
+
));
|
|
201
|
+
}
|
|
202
|
+
let duration = Duration::try_from_secs_f64(seconds)
|
|
203
|
+
.map_err(|_| Error::new(ruby.exception_arg_error(), "duration out of range"))?;
|
|
204
|
+
if key == "timeout" {
|
|
205
|
+
opts.timeout = Some(duration);
|
|
206
|
+
} else {
|
|
207
|
+
opts.heartbeat_timeout = Some(duration);
|
|
208
|
+
}
|
|
209
|
+
}
|
|
210
|
+
}
|
|
211
|
+
if let Some(headers) = options.lookup::<_, Option<RHash>>("headers")? {
|
|
212
|
+
headers.foreach(|key: String, value: String| {
|
|
213
|
+
opts.headers.push((key, value));
|
|
214
|
+
Ok(magnus::r_hash::ForEach::Continue)
|
|
215
|
+
})?;
|
|
216
|
+
}
|
|
217
|
+
let (runtime, client) = without_gvl(ruby, move || {
|
|
218
|
+
let runtime = tokio::runtime::Builder::new_multi_thread()
|
|
219
|
+
.worker_threads(1)
|
|
220
|
+
.enable_all()
|
|
221
|
+
.build()
|
|
222
|
+
.map_err(|e| QuackError::Protocol(e.to_string()))?;
|
|
223
|
+
let client = runtime.block_on(QuackClient::connect(&uri, opts))?;
|
|
224
|
+
Ok::<_, QuackError>((runtime, client))
|
|
225
|
+
})?
|
|
226
|
+
.map_err(|e| protocol_error(ruby, e))?;
|
|
227
|
+
Ok(Self {
|
|
228
|
+
session: ManuallyDrop::new(Arc::new(Session {
|
|
229
|
+
client,
|
|
230
|
+
runtime: Some(runtime),
|
|
231
|
+
active: AtomicBool::new(false),
|
|
232
|
+
fatal: AtomicBool::new(false),
|
|
233
|
+
})),
|
|
234
|
+
pid: std::process::id(),
|
|
235
|
+
})
|
|
236
|
+
}
|
|
237
|
+
fn connected(&self) -> bool {
|
|
238
|
+
self.pid == std::process::id()
|
|
239
|
+
&& !self.session.fatal.load(Ordering::Acquire)
|
|
240
|
+
&& self.session.client.is_connected()
|
|
241
|
+
}
|
|
242
|
+
fn info(&self) -> Result<RHash, Error> {
|
|
243
|
+
let ruby = &Ruby::get().unwrap();
|
|
244
|
+
self.check_pid(ruby)?;
|
|
245
|
+
let hash = ruby.hash_new();
|
|
246
|
+
if let Some(info) = &self.session.client.info {
|
|
247
|
+
hash.aset("quack_version", info.quack_version)?;
|
|
248
|
+
hash.aset("server_version", info.server_duckdb_version.clone())?;
|
|
249
|
+
hash.aset("server_platform", info.server_platform.clone())?;
|
|
250
|
+
hash.aset(
|
|
251
|
+
"heartbeat_timeout",
|
|
252
|
+
info.heartbeat_timeout.map(|v| v.as_secs_f64()),
|
|
253
|
+
)?;
|
|
254
|
+
}
|
|
255
|
+
Ok(hash)
|
|
256
|
+
}
|
|
257
|
+
fn query(&self, sql: String) -> Result<NativeStream, Error> {
|
|
258
|
+
let ruby = &Ruby::get().unwrap();
|
|
259
|
+
self.check_pid(ruby)?;
|
|
260
|
+
self.session.check(ruby)?;
|
|
261
|
+
let mut admission = self.session.admit(ruby)?;
|
|
262
|
+
let result = without_gvl(ruby, || {
|
|
263
|
+
self.session
|
|
264
|
+
.runtime()
|
|
265
|
+
.block_on(self.session.client.query(&sql, None))
|
|
266
|
+
})?;
|
|
267
|
+
let (columns, stream) = self.session.map(ruby, result)?.into_chunks();
|
|
268
|
+
let stream = NativeStream {
|
|
269
|
+
session: ManuallyDrop::new(Arc::clone(&self.session)),
|
|
270
|
+
stream: ManuallyDrop::new(Mutex::new(Some(stream))),
|
|
271
|
+
columns,
|
|
272
|
+
pid: self.pid,
|
|
273
|
+
closed: AtomicBool::new(false),
|
|
274
|
+
};
|
|
275
|
+
admission.release = false;
|
|
276
|
+
Ok(stream)
|
|
277
|
+
}
|
|
278
|
+
fn append_rows(
|
|
279
|
+
&self,
|
|
280
|
+
table: String,
|
|
281
|
+
schema: Option<String>,
|
|
282
|
+
rows: RArray,
|
|
283
|
+
schema_sql: String,
|
|
284
|
+
batch_size: usize,
|
|
285
|
+
) -> Result<(), Error> {
|
|
286
|
+
let ruby = &Ruby::get().unwrap();
|
|
287
|
+
self.check_pid(ruby)?;
|
|
288
|
+
self.session.check(ruby)?;
|
|
289
|
+
if self
|
|
290
|
+
.session
|
|
291
|
+
.client
|
|
292
|
+
.info
|
|
293
|
+
.as_ref()
|
|
294
|
+
.and_then(|info| info.quack_version)
|
|
295
|
+
!= Some(1)
|
|
296
|
+
{
|
|
297
|
+
return Err(error(
|
|
298
|
+
ruby,
|
|
299
|
+
"UnsupportedTypeError",
|
|
300
|
+
"Quack v3 SEND_DATA is not implemented by quack_protocol; use SQL INSERT or insert_all",
|
|
301
|
+
));
|
|
302
|
+
}
|
|
303
|
+
if batch_size == 0 {
|
|
304
|
+
return Err(Error::new(
|
|
305
|
+
ruby.exception_arg_error(),
|
|
306
|
+
"batch_size must be positive",
|
|
307
|
+
));
|
|
308
|
+
}
|
|
309
|
+
let mut owned_rows = Vec::with_capacity(rows.len());
|
|
310
|
+
for index in 0..rows.len() {
|
|
311
|
+
let row: RHash = rows.entry(index as isize)?;
|
|
312
|
+
let mut owned = Row::new();
|
|
313
|
+
row.foreach(|key: String, value: RArray| {
|
|
314
|
+
owned.insert(key, from_packed(ruby, value)?);
|
|
315
|
+
Ok(magnus::r_hash::ForEach::Continue)
|
|
316
|
+
})?;
|
|
317
|
+
owned_rows.push(owned);
|
|
318
|
+
}
|
|
319
|
+
let _admission = self.session.admit(ruby)?;
|
|
320
|
+
let result = without_gvl(ruby, || {
|
|
321
|
+
self.session.runtime().block_on(async {
|
|
322
|
+
// Resolve complete nested/ENUM/UNION metadata using DuckDB's
|
|
323
|
+
// parser, including user-defined types and case-sensitive names.
|
|
324
|
+
let (columns, stream) = self
|
|
325
|
+
.session
|
|
326
|
+
.client
|
|
327
|
+
.query(&schema_sql, None)
|
|
328
|
+
.await?
|
|
329
|
+
.into_chunks();
|
|
330
|
+
drop(stream);
|
|
331
|
+
self.session
|
|
332
|
+
.client
|
|
333
|
+
.append_rows(table, schema, &owned_rows, Some(columns), Some(batch_size))
|
|
334
|
+
.await
|
|
335
|
+
})
|
|
336
|
+
})?;
|
|
337
|
+
self.session.map(ruby, result)
|
|
338
|
+
}
|
|
339
|
+
fn close(&self) -> Result<(), Error> {
|
|
340
|
+
let ruby = &Ruby::get().unwrap();
|
|
341
|
+
self.check_pid(ruby)?;
|
|
342
|
+
let _admission = self.session.admit(ruby)?;
|
|
343
|
+
let result = without_gvl(ruby, || {
|
|
344
|
+
self.session
|
|
345
|
+
.runtime()
|
|
346
|
+
.block_on(self.session.client.disconnect())
|
|
347
|
+
})?;
|
|
348
|
+
self.session.map(ruby, result)
|
|
349
|
+
}
|
|
350
|
+
}
|
|
351
|
+
type ChunkStream = BoxStream<'static, Result<DataChunk, QuackError>>;
|
|
352
|
+
|
|
353
|
+
#[magnus::wrap(class = "Quack::NativeStream", free_immediately, size)]
|
|
354
|
+
struct NativeStream {
|
|
355
|
+
session: ManuallyDrop<Arc<Session>>,
|
|
356
|
+
stream: ManuallyDrop<Mutex<Option<ChunkStream>>>,
|
|
357
|
+
columns: Vec<ColumnDefinition>,
|
|
358
|
+
pid: u32,
|
|
359
|
+
closed: AtomicBool,
|
|
360
|
+
}
|
|
361
|
+
impl NativeStream {
|
|
362
|
+
fn columns(&self) -> Result<RArray, Error> {
|
|
363
|
+
let ruby = &Ruby::get().unwrap();
|
|
364
|
+
let array = ruby.ary_new();
|
|
365
|
+
for col in &self.columns {
|
|
366
|
+
let hash = ruby.hash_new();
|
|
367
|
+
hash.aset("name", col.name.as_str())?;
|
|
368
|
+
hash.aset("type", format!("{:?}", col.logical_type.id))?;
|
|
369
|
+
hash.aset("logical_type", format!("{:?}", col.logical_type))?;
|
|
370
|
+
hash.aset("duckdb_type", type_metadata(ruby, &col.logical_type)?)?;
|
|
371
|
+
if let Some(info) = &col.logical_type.type_info {
|
|
372
|
+
hash.aset("alias", info.alias())?;
|
|
373
|
+
if let ExtraTypeInfo::Decimal { width, scale, .. } = info {
|
|
374
|
+
hash.aset("precision", *width)?;
|
|
375
|
+
hash.aset("scale", *scale)?;
|
|
376
|
+
}
|
|
377
|
+
}
|
|
378
|
+
array.push(hash)?;
|
|
379
|
+
}
|
|
380
|
+
Ok(array)
|
|
381
|
+
}
|
|
382
|
+
fn close(&self) -> Result<(), Error> {
|
|
383
|
+
let ruby = &Ruby::get().unwrap();
|
|
384
|
+
if self.pid != std::process::id() {
|
|
385
|
+
return Err(error(
|
|
386
|
+
ruby,
|
|
387
|
+
"ForkError",
|
|
388
|
+
"result belongs to the parent process",
|
|
389
|
+
));
|
|
390
|
+
}
|
|
391
|
+
if self.closed.load(Ordering::Acquire) {
|
|
392
|
+
return Ok(());
|
|
393
|
+
}
|
|
394
|
+
let mut stream = self
|
|
395
|
+
.stream
|
|
396
|
+
.try_lock()
|
|
397
|
+
.map_err(|_| error(ruby, "BusyError", "result is being read by another thread"))?;
|
|
398
|
+
stream.take();
|
|
399
|
+
self.closed.store(true, Ordering::Release);
|
|
400
|
+
self.session.active.store(false, Ordering::Release);
|
|
401
|
+
Ok(())
|
|
402
|
+
}
|
|
403
|
+
fn next_batch(&self) -> Result<Option<RArray>, Error> {
|
|
404
|
+
let ruby = &Ruby::get().unwrap();
|
|
405
|
+
if self.pid != std::process::id() {
|
|
406
|
+
return Err(error(
|
|
407
|
+
ruby,
|
|
408
|
+
"ForkError",
|
|
409
|
+
"result belongs to the parent process",
|
|
410
|
+
));
|
|
411
|
+
}
|
|
412
|
+
if self.closed.load(Ordering::Acquire) {
|
|
413
|
+
return Ok(None);
|
|
414
|
+
}
|
|
415
|
+
let result = without_gvl(ruby, || {
|
|
416
|
+
let mut guard = self.stream.try_lock().map_err(|_| ())?;
|
|
417
|
+
let next = match guard.as_mut() {
|
|
418
|
+
Some(stream) => self.session.runtime().block_on(stream.next()),
|
|
419
|
+
None => None,
|
|
420
|
+
};
|
|
421
|
+
if next.is_none() || next.as_ref().is_some_and(Result::is_err) {
|
|
422
|
+
guard.take();
|
|
423
|
+
self.closed.store(true, Ordering::Release);
|
|
424
|
+
self.session.active.store(false, Ordering::Release);
|
|
425
|
+
}
|
|
426
|
+
Ok::<_, ()>(next)
|
|
427
|
+
})?
|
|
428
|
+
.map_err(|_| error(ruby, "BusyError", "result is being read by another thread"))?;
|
|
429
|
+
match result {
|
|
430
|
+
None => Ok(None),
|
|
431
|
+
Some(result) => {
|
|
432
|
+
let chunk = self.session.map(ruby, result)?;
|
|
433
|
+
let rows = ruby.ary_new_capa(chunk.row_count);
|
|
434
|
+
for row_index in 0..chunk.row_count {
|
|
435
|
+
let row = ruby.ary_new_capa(self.columns.len());
|
|
436
|
+
for col_index in 0..self.columns.len() {
|
|
437
|
+
let value = &chunk
|
|
438
|
+
.column_values(col_index)
|
|
439
|
+
.ok_or_else(|| error(ruby, "ProtocolError", "missing chunk column"))?
|
|
440
|
+
[row_index];
|
|
441
|
+
row.push(to_ruby(ruby, value)?)?;
|
|
442
|
+
}
|
|
443
|
+
rows.push(row)?;
|
|
444
|
+
}
|
|
445
|
+
Ok(Some(rows))
|
|
446
|
+
}
|
|
447
|
+
}
|
|
448
|
+
}
|
|
449
|
+
}
|
|
450
|
+
impl Drop for NativeStream {
|
|
451
|
+
fn drop(&mut self) {
|
|
452
|
+
if self.pid == std::process::id() {
|
|
453
|
+
// An old finished stream must not clear the next query's flag.
|
|
454
|
+
if !self.closed.load(Ordering::Acquire) {
|
|
455
|
+
self.session.active.store(false, Ordering::Release);
|
|
456
|
+
}
|
|
457
|
+
unsafe {
|
|
458
|
+
ManuallyDrop::drop(&mut self.stream);
|
|
459
|
+
ManuallyDrop::drop(&mut self.session);
|
|
460
|
+
}
|
|
461
|
+
}
|
|
462
|
+
}
|
|
463
|
+
}
|
|
464
|
+
fn type_metadata(ruby: &Ruby, logical_type: &LogicalType) -> Result<RHash, Error> {
|
|
465
|
+
let metadata = ruby.hash_new();
|
|
466
|
+
metadata.aset("id", format!("{:?}", logical_type.id))?;
|
|
467
|
+
if let Some(info) = &logical_type.type_info {
|
|
468
|
+
metadata.aset("alias", info.alias())?;
|
|
469
|
+
match info {
|
|
470
|
+
ExtraTypeInfo::Decimal { width, scale, .. } => {
|
|
471
|
+
metadata.aset("width", *width)?;
|
|
472
|
+
metadata.aset("scale", *scale)?;
|
|
473
|
+
}
|
|
474
|
+
ExtraTypeInfo::List { child_type, .. } | ExtraTypeInfo::Array { child_type, .. } => {
|
|
475
|
+
metadata.aset("child", type_metadata(ruby, child_type)?)?;
|
|
476
|
+
if let ExtraTypeInfo::Array { size, .. } = info {
|
|
477
|
+
metadata.aset("size", *size)?;
|
|
478
|
+
}
|
|
479
|
+
}
|
|
480
|
+
ExtraTypeInfo::Struct { child_types, .. } => {
|
|
481
|
+
let children = ruby.ary_new();
|
|
482
|
+
for child in child_types {
|
|
483
|
+
let entry = ruby.hash_new();
|
|
484
|
+
entry.aset("name", child.name.as_str())?;
|
|
485
|
+
entry.aset("type", type_metadata(ruby, &child.logical_type)?)?;
|
|
486
|
+
children.push(entry)?;
|
|
487
|
+
}
|
|
488
|
+
metadata.aset("children", children)?;
|
|
489
|
+
}
|
|
490
|
+
ExtraTypeInfo::Enum { values, .. } => {
|
|
491
|
+
metadata.aset("values", values.clone())?;
|
|
492
|
+
}
|
|
493
|
+
ExtraTypeInfo::Geo { crs, .. } => {
|
|
494
|
+
metadata.aset("crs", crs.definition.as_deref())?;
|
|
495
|
+
}
|
|
496
|
+
_ => {}
|
|
497
|
+
}
|
|
498
|
+
}
|
|
499
|
+
Ok(metadata)
|
|
500
|
+
}
|
|
501
|
+
|
|
502
|
+
fn to_ruby(ruby: &Ruby, value: &QValue) -> Result<Value, Error> {
|
|
503
|
+
let module: RModule = ruby.class_object().const_get("Quack")?;
|
|
504
|
+
Ok(match value {
|
|
505
|
+
QValue::Null => ruby.qnil().as_value(),
|
|
506
|
+
QValue::Bool(v) => (*v).into_value_with(ruby),
|
|
507
|
+
QValue::Int(v) => ruby.integer_from_i64(*v).as_value(),
|
|
508
|
+
QValue::UInt(v) => ruby.integer_from_u64(*v).as_value(),
|
|
509
|
+
QValue::HugeInt(v) => ruby.integer_from_i128(*v).as_value(),
|
|
510
|
+
QValue::UHugeInt(v) => ruby.integer_from_u128(*v).as_value(),
|
|
511
|
+
QValue::Float(v) => (*v as f64).into_value_with(ruby),
|
|
512
|
+
QValue::Double(v) => (*v).into_value_with(ruby),
|
|
513
|
+
QValue::String(v) => ruby.str_new(v).as_value(),
|
|
514
|
+
QValue::Bytes(v) => ruby.str_from_slice(v).as_value(),
|
|
515
|
+
QValue::Decimal(v) => {
|
|
516
|
+
module.funcall("__decimal", (quack_protocol::decimal_to_string(v),))?
|
|
517
|
+
}
|
|
518
|
+
QValue::Date(v) => module.funcall("__date", (v.days,))?,
|
|
519
|
+
QValue::Time(v) => module.funcall("__time", (v.value, format!("{:?}", v.unit)))?,
|
|
520
|
+
QValue::TimeTz(v) => module.funcall("__time_tz", (v.bits,))?,
|
|
521
|
+
QValue::Timestamp(v) => module.funcall(
|
|
522
|
+
"__timestamp",
|
|
523
|
+
(v.value, format!("{:?}", v.unit), v.timezone_utc),
|
|
524
|
+
)?,
|
|
525
|
+
QValue::Interval(v) => module.funcall("__interval", (v.months, v.days, v.micros))?,
|
|
526
|
+
QValue::List(values) => {
|
|
527
|
+
let out = ruby.ary_new_capa(values.len());
|
|
528
|
+
for value in values {
|
|
529
|
+
out.push(to_ruby(ruby, value)?)?;
|
|
530
|
+
}
|
|
531
|
+
out.as_value()
|
|
532
|
+
}
|
|
533
|
+
QValue::Struct(values) => {
|
|
534
|
+
let out = ruby.hash_new();
|
|
535
|
+
for (key, value) in values {
|
|
536
|
+
out.aset(key.as_str(), to_ruby(ruby, value)?)?;
|
|
537
|
+
}
|
|
538
|
+
out.as_value()
|
|
539
|
+
}
|
|
540
|
+
})
|
|
541
|
+
}
|
|
542
|
+
#[magnus::init(name = "quack_native")]
|
|
543
|
+
fn init(ruby: &Ruby) -> Result<(), Error> {
|
|
544
|
+
let module = ruby.define_module("Quack")?;
|
|
545
|
+
let base = module.define_error("Error", ruby.exception_standard_error())?;
|
|
546
|
+
for name in [
|
|
547
|
+
"ConnectionError",
|
|
548
|
+
"TimeoutError",
|
|
549
|
+
"ServerError",
|
|
550
|
+
"ProtocolError",
|
|
551
|
+
"UnsupportedTypeError",
|
|
552
|
+
"ProxyError",
|
|
553
|
+
"BusyError",
|
|
554
|
+
"ForkError",
|
|
555
|
+
] {
|
|
556
|
+
module.define_error(name, base)?;
|
|
557
|
+
}
|
|
558
|
+
let connection = module.define_class("NativeConnection", ruby.class_object())?;
|
|
559
|
+
connection.define_singleton_method("new", function!(NativeConnection::connect, 2))?;
|
|
560
|
+
connection.define_method("query", method!(NativeConnection::query, 1))?;
|
|
561
|
+
connection.define_method("append_rows", method!(NativeConnection::append_rows, 5))?;
|
|
562
|
+
connection.define_method("close", method!(NativeConnection::close, 0))?;
|
|
563
|
+
connection.define_method("connected?", method!(NativeConnection::connected, 0))?;
|
|
564
|
+
connection.define_method("info", method!(NativeConnection::info, 0))?;
|
|
565
|
+
let stream = module.define_class("NativeStream", ruby.class_object())?;
|
|
566
|
+
stream.define_method("columns", method!(NativeStream::columns, 0))?;
|
|
567
|
+
stream.define_method("next_batch", method!(NativeStream::next_batch, 0))?;
|
|
568
|
+
stream.define_method("close", method!(NativeStream::close, 0))?;
|
|
569
|
+
Ok(())
|
|
570
|
+
}
|
|
571
|
+
|
|
572
|
+
fn from_packed(ruby: &Ruby, value: RArray) -> Result<QValue, Error> {
|
|
573
|
+
let tag: String = value.entry(0)?;
|
|
574
|
+
let invalid = || Error::new(ruby.exception_arg_error(), "append value is out of range");
|
|
575
|
+
Ok(match tag.as_str() {
|
|
576
|
+
"null" => QValue::Null,
|
|
577
|
+
"bool" => QValue::Bool(value.entry(1)?),
|
|
578
|
+
"int" => {
|
|
579
|
+
let text: String = value.entry(1)?;
|
|
580
|
+
if let Ok(number) = text.parse::<i64>() {
|
|
581
|
+
QValue::Int(number)
|
|
582
|
+
} else if let Ok(number) = text.parse::<i128>() {
|
|
583
|
+
QValue::HugeInt(number)
|
|
584
|
+
} else {
|
|
585
|
+
QValue::UHugeInt(text.parse().map_err(|_| invalid())?)
|
|
586
|
+
}
|
|
587
|
+
}
|
|
588
|
+
"double" => QValue::Double(value.entry(1)?),
|
|
589
|
+
"string" => QValue::String(value.entry(1)?),
|
|
590
|
+
"bytes" => {
|
|
591
|
+
let bytes: magnus::RString = value.entry(1)?;
|
|
592
|
+
QValue::Bytes(unsafe { bytes.as_slice() }.to_vec())
|
|
593
|
+
}
|
|
594
|
+
"decimal" => quack_protocol::decimal_value(
|
|
595
|
+
value.entry::<String>(1)?,
|
|
596
|
+
value.entry(2)?,
|
|
597
|
+
value.entry(3)?,
|
|
598
|
+
)
|
|
599
|
+
.map_err(|e| protocol_error(ruby, e))?,
|
|
600
|
+
"date" => quack_protocol::date_value(value.entry(1)?),
|
|
601
|
+
"time" => quack_protocol::time_value(
|
|
602
|
+
value.entry(1)?,
|
|
603
|
+
match value.entry::<String>(2)?.as_str() {
|
|
604
|
+
"Micros" => quack_protocol::TimeUnit::Micros,
|
|
605
|
+
_ => quack_protocol::TimeUnit::Nanos,
|
|
606
|
+
},
|
|
607
|
+
),
|
|
608
|
+
"timetz" => quack_protocol::time_tz_value(value.entry(1)?),
|
|
609
|
+
"timestamp" => quack_protocol::timestamp_value(
|
|
610
|
+
value.entry(1)?,
|
|
611
|
+
match value.entry::<String>(2)?.as_str() {
|
|
612
|
+
"Seconds" => quack_protocol::TimestampUnit::Seconds,
|
|
613
|
+
"Millis" => quack_protocol::TimestampUnit::Millis,
|
|
614
|
+
"Micros" => quack_protocol::TimestampUnit::Micros,
|
|
615
|
+
_ => quack_protocol::TimestampUnit::Nanos,
|
|
616
|
+
},
|
|
617
|
+
value.entry(3)?,
|
|
618
|
+
),
|
|
619
|
+
"interval" => {
|
|
620
|
+
quack_protocol::interval_value(value.entry(1)?, value.entry(2)?, value.entry(3)?)
|
|
621
|
+
}
|
|
622
|
+
"list" => {
|
|
623
|
+
let array: RArray = value.entry(1)?;
|
|
624
|
+
let mut values = Vec::with_capacity(array.len());
|
|
625
|
+
for i in 0..array.len() {
|
|
626
|
+
values.push(from_packed(ruby, array.entry(i as isize)?)?);
|
|
627
|
+
}
|
|
628
|
+
QValue::List(values)
|
|
629
|
+
}
|
|
630
|
+
"struct" => {
|
|
631
|
+
let hash: RHash = value.entry(1)?;
|
|
632
|
+
let mut values = Row::new();
|
|
633
|
+
hash.foreach(|key: String, value: RArray| {
|
|
634
|
+
values.insert(key, from_packed(ruby, value)?);
|
|
635
|
+
Ok(magnus::r_hash::ForEach::Continue)
|
|
636
|
+
})?;
|
|
637
|
+
QValue::Struct(values)
|
|
638
|
+
}
|
|
639
|
+
_ => {
|
|
640
|
+
return Err(error(
|
|
641
|
+
ruby,
|
|
642
|
+
"UnsupportedTypeError",
|
|
643
|
+
format!("unsupported append value {tag}"),
|
|
644
|
+
));
|
|
645
|
+
}
|
|
646
|
+
})
|
|
647
|
+
}
|