tank-chdb 0.44.0

chDB driver implementation for Tank: the Rust data layer
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130

use anyhow::anyhow;
use chdb_rust::connection::Connection as ChConnection;
use std::{ffi::c_char, slice};
use tank_core::{Result, error_message_from_ptr};

enum ChdbConnection {}
enum ChdbResult {}

unsafe extern "C" {
    fn chdb_stream_query_n(
        connection: *mut ChdbConnection,
        query: *const c_char,
        query_len: usize,
        format: *const c_char,
        format_len: usize,
    ) -> *mut ChdbResult;
    fn chdb_stream_fetch_result(
        connection: *mut ChdbConnection,
        result: *mut ChdbResult,
    ) -> *mut ChdbResult;
    fn chdb_stream_cancel_query(connection: *mut ChdbConnection, result: *mut ChdbResult);
    fn chdb_result_buffer(result: *mut ChdbResult) -> *mut c_char;
    fn chdb_result_length(result: *mut ChdbResult) -> usize;
    fn chdb_result_error(result: *mut ChdbResult) -> *const c_char;
    fn chdb_destroy_query_result(result: *mut ChdbResult);
}

pub(crate) struct ChDBStream {
    connection: *mut ChdbConnection,
    result: *mut ChdbResult,
    finished: bool,
}

unsafe impl Send for ChDBStream {}

impl ChDBStream {
    pub(crate) fn start(connection: &ChConnection, sql: &str) -> Result<Self> {
        // `chdb_rust::Connection` wraps `*mut *mut chdb_connection_` in a private field.
        let connection: *mut ChdbConnection =
            unsafe { **(connection as *const ChConnection).cast::<*mut *mut ChdbConnection>() };
        let sql = sql.trim().trim_end_matches(';').trim_end();
        let format = b"JSONEachRow";
        let result = unsafe {
            chdb_stream_query_n(
                connection,
                sql.as_ptr().cast(),
                sql.len(),
                format.as_ptr().cast(),
                format.len(),
            )
        };
        if result.is_null() {
            return Err(anyhow!("chDB streaming query returned no result"));
        }
        let error = unsafe { chdb_result_error(result) };
        if !error.is_null() {
            let error = anyhow!(
                "chDB streaming query failed: {}",
                error_message_from_ptr(&error)
            );
            unsafe { chdb_destroy_query_result(result) };
            return Err(error);
        }
        Ok(Self {
            connection,
            result,
            finished: false,
        })
    }

    pub(crate) fn next(&mut self) -> Result<Option<ChDBChunk>> {
        if self.finished {
            return Ok(None);
        }
        let chunk = unsafe { chdb_stream_fetch_result(self.connection, self.result) };
        if chunk.is_null() {
            self.finished = true;
            return Ok(None);
        }
        let error = unsafe { chdb_result_error(chunk) };
        if !error.is_null() {
            let error = anyhow!(
                "chDB streaming fetch failed: {}",
                error_message_from_ptr(&error)
            );
            unsafe { chdb_destroy_query_result(chunk) };
            self.finished = true;
            return Err(error);
        }
        if unsafe { chdb_result_length(chunk) } == 0 {
            unsafe { chdb_destroy_query_result(chunk) };
            self.finished = true;
            return Ok(None);
        }
        Ok(Some(ChDBChunk { result: chunk }))
    }
}

impl Drop for ChDBStream {
    fn drop(&mut self) {
        if self.result.is_null() {
            return;
        }
        if !self.finished {
            unsafe { chdb_stream_cancel_query(self.connection, self.result) };
        }
        unsafe { chdb_destroy_query_result(self.result) };
    }
}

pub(crate) struct ChDBChunk {
    result: *mut ChdbResult,
}

impl ChDBChunk {
    pub(crate) fn data(&self) -> &[u8] {
        let buffer = unsafe { chdb_result_buffer(self.result) };
        let length = unsafe { chdb_result_length(self.result) };
        if buffer.is_null() || length == 0 {
            return &[];
        }
        unsafe { slice::from_raw_parts(buffer.cast(), length) }
    }
}

impl Drop for ChDBChunk {
    fn drop(&mut self) {
        unsafe { chdb_destroy_query_result(self.result) };
    }
}