connectorx/sources/sqlite/
mod.rs

1//! Source implementation for SQLite embedded database.
2
3mod errors;
4mod typesystem;
5
6pub use self::errors::SQLiteSourceError;
7use crate::{
8    data_order::DataOrder,
9    errors::ConnectorXError,
10    sources::{PartitionParser, Produce, Source, SourcePartition},
11    sql::{count_query, limit1_query, CXQuery},
12    utils::DummyBox,
13};
14use anyhow::anyhow;
15use chrono::{NaiveDate, NaiveDateTime, NaiveTime};
16use fallible_streaming_iterator::FallibleStreamingIterator;
17use fehler::{throw, throws};
18use log::debug;
19use owning_ref::OwningHandle;
20use r2d2::{Pool, PooledConnection};
21use r2d2_sqlite::SqliteConnectionManager;
22use rusqlite::{Row, Rows, Statement};
23use sqlparser::dialect::SQLiteDialect;
24use std::convert::TryFrom;
25pub use typesystem::SQLiteTypeSystem;
26use urlencoding::decode;
27
28pub struct SQLiteSource {
29    pool: Pool<SqliteConnectionManager>,
30    origin_query: Option<String>,
31    queries: Vec<CXQuery<String>>,
32    names: Vec<String>,
33    schema: Vec<SQLiteTypeSystem>,
34}
35
36impl SQLiteSource {
37    #[throws(SQLiteSourceError)]
38    pub fn new(conn: &str, nconn: usize) -> Self {
39        let decoded_conn = decode(conn)?.into_owned();
40        debug!("decoded conn: {}", decoded_conn);
41        let manager = SqliteConnectionManager::file(decoded_conn);
42        let pool = r2d2::Pool::builder()
43            .max_size(nconn as u32)
44            .build(manager)?;
45
46        Self {
47            pool,
48            origin_query: None,
49            queries: vec![],
50            names: vec![],
51            schema: vec![],
52        }
53    }
54}
55
56impl Source for SQLiteSource
57where
58    SQLiteSourcePartition: SourcePartition<TypeSystem = SQLiteTypeSystem>,
59{
60    const DATA_ORDERS: &'static [DataOrder] = &[DataOrder::RowMajor];
61    type Partition = SQLiteSourcePartition;
62    type TypeSystem = SQLiteTypeSystem;
63    type Error = SQLiteSourceError;
64
65    #[throws(SQLiteSourceError)]
66    fn set_data_order(&mut self, data_order: DataOrder) {
67        if !matches!(data_order, DataOrder::RowMajor) {
68            throw!(ConnectorXError::UnsupportedDataOrder(data_order));
69        }
70    }
71
72    fn set_queries<Q: ToString>(&mut self, queries: &[CXQuery<Q>]) {
73        self.queries = queries.iter().map(|q| q.map(Q::to_string)).collect();
74    }
75
76    fn set_origin_query(&mut self, query: Option<String>) {
77        self.origin_query = query;
78    }
79
80    #[throws(SQLiteSourceError)]
81    fn fetch_metadata(&mut self) {
82        assert!(!self.queries.is_empty());
83        let conn = self.pool.get()?;
84
85        // Use prepare + stmt.columns() to get schema without executing the query.
86        // This gives us column names and decl_type from the DDL, which is sufficient
87        // for most cases and avoids running the query at all (issue #768).
88        let stmt = conn.prepare(self.queries[0].as_str())?;
89        let columns = stmt.columns();
90
91        let mut names: Vec<String> = Vec::with_capacity(columns.len());
92        let mut types: Vec<Option<SQLiteTypeSystem>> = Vec::with_capacity(columns.len());
93
94        for col in &columns {
95            names.push(col.name().to_string());
96            let decl_type = col.decl_type();
97            match SQLiteTypeSystem::try_from((decl_type, rusqlite::types::Type::Null)) {
98                Ok(t) => types.push(Some(t)),
99                Err(_) => types.push(None),
100            }
101        }
102
103        // If all types were resolved from decl_type, we're done
104        if !types.contains(&None) {
105            self.names = names;
106            self.schema = types.into_iter().map(|t| t.unwrap()).collect();
107            return;
108        }
109
110        // Some columns lack decl_type (computed expressions, aggregates, etc.).
111        // Fall back to executing with LIMIT 1 to infer types from actual values.
112        drop(columns);
113        drop(stmt);
114
115        let mut num_empty = 0;
116        for (i, query) in self.queries.iter().enumerate() {
117            let l1query = limit1_query(query, &SQLiteDialect {})?;
118
119            let is_success = conn.query_row(l1query.as_str(), [], |row| {
120                for (j, col) in row.as_ref().columns().iter().enumerate() {
121                    if types[j].is_none() {
122                        let vr = row.get_ref(j)?;
123                        if let Ok(t) = SQLiteTypeSystem::try_from((col.decl_type(), vr.data_type()))
124                        {
125                            types[j] = Some(t);
126                        }
127                    }
128                }
129                Ok(())
130            });
131
132            match is_success {
133                Ok(()) => {
134                    if !types.contains(&None) {
135                        self.names = names;
136                        self.schema = types.into_iter().map(|t| t.unwrap()).collect();
137                        return;
138                    } else if i == self.queries.len() - 1 {
139                        debug!(
140                            "cannot get metadata for '{}' due to null value: {:?}",
141                            query, types
142                        );
143                        throw!(SQLiteSourceError::InferTypeFromNull);
144                    }
145                }
146                Err(e) => {
147                    if let rusqlite::Error::QueryReturnedNoRows = e {
148                        num_empty += 1;
149                    }
150                    if i == self.queries.len() - 1 && num_empty < self.queries.len() {
151                        debug!("cannot get metadata for '{}': {}", query, e);
152                        throw!(e)
153                    }
154                }
155            }
156        }
157
158        // All partitions returned empty results - use decl_type where available,
159        // fall back to Text for columns without type info
160        self.names = names;
161        self.schema = types
162            .into_iter()
163            .map(|t| t.unwrap_or(SQLiteTypeSystem::Text(false)))
164            .collect();
165    }
166
167    #[throws(SQLiteSourceError)]
168    fn result_rows(&mut self) -> Option<usize> {
169        match &self.origin_query {
170            Some(q) => {
171                let cxq = CXQuery::Naked(q.clone());
172                let conn = self.pool.get()?;
173                let nrows =
174                    conn.query_row(count_query(&cxq, &SQLiteDialect {})?.as_str(), [], |row| {
175                        Ok(row.get::<_, i64>(0)? as usize)
176                    })?;
177                Some(nrows)
178            }
179            None => None,
180        }
181    }
182
183    fn names(&self) -> Vec<String> {
184        self.names.clone()
185    }
186
187    fn schema(&self) -> Vec<Self::TypeSystem> {
188        self.schema.clone()
189    }
190
191    #[throws(SQLiteSourceError)]
192    fn partition(self) -> Vec<Self::Partition> {
193        let mut ret = vec![];
194        for query in self.queries {
195            let conn = self.pool.get()?;
196
197            ret.push(SQLiteSourcePartition::new(conn, &query, &self.schema));
198        }
199        ret
200    }
201}
202
203pub struct SQLiteSourcePartition {
204    conn: PooledConnection<SqliteConnectionManager>,
205    query: CXQuery<String>,
206    schema: Vec<SQLiteTypeSystem>,
207    nrows: usize,
208    ncols: usize,
209}
210
211impl SQLiteSourcePartition {
212    pub fn new(
213        conn: PooledConnection<SqliteConnectionManager>,
214        query: &CXQuery<String>,
215        schema: &[SQLiteTypeSystem],
216    ) -> Self {
217        Self {
218            conn,
219            query: query.clone(),
220            schema: schema.to_vec(),
221            nrows: 0,
222            ncols: schema.len(),
223        }
224    }
225}
226
227impl SourcePartition for SQLiteSourcePartition {
228    type TypeSystem = SQLiteTypeSystem;
229    type Parser<'a> = SQLiteSourcePartitionParser<'a>;
230    type Error = SQLiteSourceError;
231
232    #[throws(SQLiteSourceError)]
233    fn result_rows(&mut self) {
234        self.nrows = self.conn.query_row(
235            count_query(&self.query, &SQLiteDialect {})?.as_str(),
236            [],
237            |row| Ok(row.get::<_, i64>(0)? as usize),
238        )?;
239    }
240
241    #[throws(SQLiteSourceError)]
242    fn parser(&mut self) -> Self::Parser<'_> {
243        SQLiteSourcePartitionParser::new(&self.conn, self.query.as_str(), &self.schema)?
244    }
245
246    fn nrows(&self) -> usize {
247        self.nrows
248    }
249
250    fn ncols(&self) -> usize {
251        self.ncols
252    }
253}
254
255unsafe impl<'a> Send for SQLiteSourcePartitionParser<'a> {}
256
257pub struct SQLiteSourcePartitionParser<'a> {
258    rows: OwningHandle<Box<Statement<'a>>, DummyBox<Rows<'a>>>,
259    ncols: usize,
260    current_col: usize,
261    current_consumed: bool,
262    is_finished: bool,
263}
264
265impl<'a> SQLiteSourcePartitionParser<'a> {
266    #[throws(SQLiteSourceError)]
267    pub fn new(
268        conn: &'a PooledConnection<SqliteConnectionManager>,
269        query: &str,
270        schema: &[SQLiteTypeSystem],
271    ) -> Self {
272        let stmt: Statement<'a> = conn.prepare(query)?;
273
274        // Safety: DummyBox borrows the on-heap stmt, which is owned by the OwningHandle.
275        // No matter how we move the owning handle (thus the Box<Statment>), the Statement
276        // keeps its address static on the heap, thus the borrow of MyRows keeps valid.
277        let rows: OwningHandle<Box<Statement<'a>>, DummyBox<Rows<'a>>> =
278            OwningHandle::new_with_fn(Box::new(stmt), |stmt: *const Statement<'a>| unsafe {
279                DummyBox((*(stmt as *mut Statement<'_>)).query([]).unwrap())
280            });
281        Self {
282            rows,
283            ncols: schema.len(),
284            current_col: 0,
285            current_consumed: true,
286            is_finished: false,
287        }
288    }
289
290    #[throws(SQLiteSourceError)]
291    fn next_loc(&mut self) -> (&Row<'_>, usize) {
292        self.current_consumed = true;
293        let row: &Row = (*self.rows)
294            .get()
295            .ok_or_else(|| anyhow!("Sqlite empty current row"))?;
296        let col = self.current_col;
297        self.current_col = (self.current_col + 1) % self.ncols;
298        (row, col)
299    }
300}
301
302impl<'a> PartitionParser<'a> for SQLiteSourcePartitionParser<'a> {
303    type TypeSystem = SQLiteTypeSystem;
304    type Error = SQLiteSourceError;
305
306    #[throws(SQLiteSourceError)]
307    fn fetch_next(&mut self) -> (usize, bool) {
308        assert!(self.current_col == 0);
309
310        if !self.current_consumed {
311            return (1, false);
312        } else if self.is_finished {
313            return (0, true);
314        }
315
316        match (*self.rows).next()? {
317            Some(_) => {
318                self.current_consumed = false;
319                (1, false)
320            }
321            None => {
322                self.is_finished = true;
323                (0, true)
324            }
325        }
326    }
327}
328
329macro_rules! impl_produce {
330    ($($t: ty,)+) => {
331        $(
332            impl<'r, 'a> Produce<'r, $t> for SQLiteSourcePartitionParser<'a> {
333                type Error = SQLiteSourceError;
334
335                #[throws(SQLiteSourceError)]
336                fn produce(&'r mut self) -> $t {
337                    let (row, col) = self.next_loc()?;
338                    let val = row.get(col)?;
339                    val
340                }
341            }
342
343            impl<'r, 'a> Produce<'r, Option<$t>> for SQLiteSourcePartitionParser<'a> {
344                type Error = SQLiteSourceError;
345
346                #[throws(SQLiteSourceError)]
347                fn produce(&'r mut self) -> Option<$t> {
348                    let (row, col) = self.next_loc()?;
349                    let val = row.get(col)?;
350                    val
351                }
352            }
353        )+
354    };
355}
356
357impl_produce!(
358    bool,
359    i64,
360    i32,
361    i16,
362    f64,
363    Box<str>,
364    NaiveDate,
365    NaiveTime,
366    NaiveDateTime,
367    Vec<u8>,
368);
connectorx/sources/sqlite/mod.rs

connectorx/sources/sqlite/
mod.rs
