indexlake_catalog_sqlite/
catalog.rs

1use std::path::PathBuf;
2use std::time::Duration;
3
4use arrow::array::RecordBatch;
5use arrow::datatypes::Schema;
6use futures::StreamExt;
7use indexlake::catalog::{
8    Catalog, CatalogDataType, CatalogDatabase, CatalogSchemaRef, Row, RowStream, Scalar,
9    Transaction,
10};
11use indexlake::expr::{BinaryExpr, Expr};
12use indexlake::{ILError, ILResult};
13use log::{error, trace};
14use rusqlite::{OpenFlags, ToSql};
15use uuid::Uuid;
16
17use crate::types::scalar_to_sqlite_param;
18
19#[derive(Debug)]
20pub struct SqliteCatalog {
21    path: PathBuf,
22}
23
24impl SqliteCatalog {
25    pub fn try_new(path: impl Into<String>) -> ILResult<Self> {
26        let path = PathBuf::from(path.into());
27        if !path.exists() {
28            return Err(ILError::catalog(format!(
29                "sqlite path {} does not exist",
30                path.display()
31            )));
32        }
33        Ok(SqliteCatalog { path })
34    }
35}
36
37#[async_trait::async_trait]
38impl Catalog for SqliteCatalog {
39    fn database(&self) -> CatalogDatabase {
40        CatalogDatabase::Sqlite
41    }
42
43    async fn query(&self, sql: &str, schema: CatalogSchemaRef) -> ILResult<RowStream<'static>> {
44        trace!("sqlite query: {sql}");
45        let conn = rusqlite::Connection::open_with_flags(
46            &self.path,
47            OpenFlags::SQLITE_OPEN_READ_ONLY
48                | OpenFlags::SQLITE_OPEN_NO_MUTEX
49                | OpenFlags::SQLITE_OPEN_URI,
50        )
51        .map_err(|e| ILError::catalog(format!("failed to open sqlite db: {e}")))?;
52        let mut stmt = conn
53            .prepare(sql)
54            .map_err(|e| ILError::catalog(format!("failed to prepare sqlite stmt: {e}")))?;
55        let mut sqlite_rows = stmt
56            .query([])
57            .map_err(|e| ILError::catalog(format!("failed to query sqlite stmt: {e}")))?;
58
59        let mut rows: Vec<Row> = Vec::new();
60        while let Some(sqlite_row) = sqlite_rows
61            .next()
62            .map_err(|e| ILError::catalog(format!("failed to get next sqlite row: {e}")))?
63        {
64            let row = sqlite_row_to_row(sqlite_row, &schema)?;
65            rows.push(row);
66        }
67        Ok(Box::pin(futures::stream::iter(rows).map(Ok)))
68    }
69
70    async fn transaction(&self) -> ILResult<Box<dyn Transaction>> {
71        let conn = rusqlite::Connection::open_with_flags(
72            &self.path,
73            OpenFlags::SQLITE_OPEN_READ_WRITE
74                | OpenFlags::SQLITE_OPEN_NO_MUTEX
75                | OpenFlags::SQLITE_OPEN_URI,
76        )
77        .map_err(|e| ILError::catalog(format!("failed to open sqlite db: {e}")))?;
78        conn.busy_timeout(Duration::from_secs(30))
79            .map_err(|e| ILError::catalog(format!("failed to set sqlite busy timeout: {e}")))?;
80        conn.execute_batch("BEGIN IMMEDIATE")
81            .map_err(|e| ILError::catalog(format!("failed to begin sqlite txn: {e}")))?;
82        Ok(Box::new(SqliteTransaction { conn, done: false }))
83    }
84
85    async fn truncate(&self, table_name: &str) -> ILResult<()> {
86        let conn = rusqlite::Connection::open_with_flags(
87            &self.path,
88            OpenFlags::SQLITE_OPEN_READ_WRITE
89                | OpenFlags::SQLITE_OPEN_NO_MUTEX
90                | OpenFlags::SQLITE_OPEN_URI,
91        )
92        .map_err(|e| ILError::catalog(format!("failed to open sqlite db: {e}")))?;
93        conn.busy_timeout(Duration::from_secs(30))
94            .map_err(|e| ILError::catalog(format!("failed to set sqlite busy timeout: {e}")))?;
95        conn.execute(
96            &format!("DELETE FROM {}", self.sql_identifier(table_name)),
97            [],
98        )
99        .map_err(|e| {
100            ILError::catalog(format!(
101                "failed to truncate table {table_name} on sqlite: {e}"
102            ))
103        })?;
104        Ok(())
105    }
106
107    async fn size(&self, table_name: &str) -> ILResult<usize> {
108        let conn = rusqlite::Connection::open_with_flags(
109            &self.path,
110            OpenFlags::SQLITE_OPEN_READ_WRITE
111                | OpenFlags::SQLITE_OPEN_NO_MUTEX
112                | OpenFlags::SQLITE_OPEN_URI,
113        )
114        .map_err(|e| ILError::catalog(format!("failed to open sqlite db: {e}")))?;
115        conn.busy_timeout(Duration::from_secs(30))
116            .map_err(|e| ILError::catalog(format!("failed to set sqlite busy timeout: {e}")))?;
117        let size: usize = conn
118            .query_row(
119                &format!("SELECT SUM(pgsize) FROM dbstat WHERE name='{table_name}'"),
120                [],
121                |row| row.get(0),
122            )
123            .map_err(|e| {
124                ILError::catalog(format!(
125                    "failed to get size of table {table_name} on sqlite: {e}"
126                ))
127            })?;
128        Ok(size)
129    }
130
131    fn sql_identifier(&self, ident: &str) -> String {
132        format!("`{ident}`")
133    }
134
135    fn sql_binary_literal(&self, value: &[u8]) -> String {
136        format!("X'{}'", hex::encode(value))
137    }
138
139    fn sql_uuid_literal(&self, value: &Uuid) -> String {
140        self.sql_binary_literal(value.as_bytes())
141    }
142
143    fn sql_string_literal(&self, value: &str) -> String {
144        let value = value.replace("'", "''");
145        format!("'{value}'")
146    }
147
148    // TODO impl this
149    fn supports_filter(&self, filter: &Expr, _schema: &Schema) -> ILResult<bool> {
150        match filter {
151            Expr::Function(_) => Ok(false),
152            Expr::BinaryExpr(BinaryExpr { left, right, .. }) => {
153                if let Expr::Column(_) = left.as_ref()
154                    && let Expr::Literal(lit) = right.as_ref()
155                    && matches!(lit.value, Scalar::List(_))
156                {
157                    Ok(false)
158                } else if let Expr::Literal(lit) = left.as_ref()
159                    && let Expr::Column(_) = right.as_ref()
160                    && matches!(lit.value, Scalar::List(_))
161                {
162                    Ok(false)
163                } else {
164                    Ok(true)
165                }
166            }
167            _ => Ok(true),
168        }
169    }
170
171    fn unparse_expr(&self, expr: &Expr) -> ILResult<String> {
172        match expr {
173            Expr::Column(name) => Ok(self.sql_identifier(name)),
174            Expr::Literal(literal) => literal.value.to_sql(self),
175            Expr::BinaryExpr(binary_expr) => {
176                let left = self.unparse_expr(&binary_expr.left)?;
177                let right = self.unparse_expr(&binary_expr.right)?;
178                Ok(format!("({} {} {})", left, binary_expr.op, right))
179            }
180            Expr::Not(expr) => Ok(format!("NOT {}", self.unparse_expr(expr)?)),
181            Expr::IsNull(expr) => Ok(format!("{} IS NULL", self.unparse_expr(expr)?)),
182            Expr::IsNotNull(expr) => Ok(format!("{} IS NOT NULL", self.unparse_expr(expr)?)),
183            Expr::InList(in_list) => {
184                let list = in_list
185                    .list
186                    .iter()
187                    .map(|expr| self.unparse_expr(expr))
188                    .collect::<ILResult<Vec<_>>>()?
189                    .join(", ");
190                Ok(format!(
191                    "{} IN ({})",
192                    self.unparse_expr(&in_list.expr)?,
193                    list
194                ))
195            }
196            Expr::Function(_) => Err(ILError::invalid_input(
197                "Function can only be used for index",
198            )),
199            Expr::Like(like) => {
200                let expr = self.unparse_expr(&like.expr)?;
201                let pattern = self.unparse_expr(&like.pattern)?;
202                // For case-sensitive LIKE, SQLite requires `PRAGMA case_sensitive_like = ON;`
203                // to be set on the connection. This function only generates the SQL string
204                // and does not set the PRAGMA.
205                // For case-insensitive ILIKE, we use the `UPPER()` function on both
206                // the expression and the pattern to ensure case-insensitivity.
207                match (like.negated, like.case_insensitive) {
208                    (false, false) => Ok(format!("{expr} LIKE {pattern}")),
209                    (true, false) => Ok(format!("{expr} NOT LIKE {pattern}")),
210                    (false, true) => Ok(format!("UPPER({expr}) LIKE UPPER({pattern})")),
211                    (true, true) => Ok(format!("UPPER({expr}) NOT LIKE UPPER({pattern})")),
212                }
213            }
214            Expr::Cast(cast) => {
215                let catalog_datatype = CatalogDataType::from_arrow(&cast.cast_type)?;
216                let expr_sql = self.unparse_expr(&cast.expr)?;
217                Ok(format!(
218                    "CAST({} AS {})",
219                    expr_sql,
220                    self.unparse_catalog_data_type(catalog_datatype),
221                ))
222            }
223            Expr::TryCast(_) => Err(ILError::invalid_input("TRY_CAST is not supported in SQL")),
224            Expr::Negative(expr) => Ok(format!("-{}", self.unparse_expr(expr)?)),
225            Expr::Case(case) => {
226                let mut sql = String::new();
227                sql.push_str("CASE");
228                for (when, then) in &case.when_then {
229                    sql.push_str(&format!(
230                        " WHEN {} THEN {}",
231                        self.unparse_expr(when)?,
232                        self.unparse_expr(then)?,
233                    ));
234                }
235                if let Some(else_expr) = &case.else_expr {
236                    sql.push_str(&format!(" ELSE {}", self.unparse_expr(else_expr)?));
237                }
238                sql.push_str(" END");
239                Ok(sql)
240            }
241        }
242    }
243
244    fn unparse_catalog_data_type(&self, data_type: CatalogDataType) -> String {
245        match data_type {
246            CatalogDataType::Boolean => "BOOLEAN".to_string(),
247            CatalogDataType::Int8 => "TINYINT".to_string(),
248            CatalogDataType::Int16 => "SMALLINT".to_string(),
249            CatalogDataType::Int32 => "INTEGER".to_string(),
250            CatalogDataType::Int64 => "BIGINT".to_string(),
251            CatalogDataType::UInt8 => "TINYINT UNSIGNED".to_string(),
252            CatalogDataType::UInt16 => "SMALLINT UNSIGNED".to_string(),
253            CatalogDataType::UInt32 => "INTEGER UNSIGNED".to_string(),
254            CatalogDataType::UInt64 => "BIGINT UNSIGNED".to_string(),
255            CatalogDataType::Float32 => "FLOAT".to_string(),
256            CatalogDataType::Float64 => "DOUBLE".to_string(),
257            CatalogDataType::Utf8 => "VARCHAR".to_string(),
258            CatalogDataType::Binary => "BLOB".to_string(),
259            CatalogDataType::Uuid => "BLOB".to_string(),
260        }
261    }
262}
263
264#[derive(Debug)]
265pub struct SqliteTransaction {
266    conn: rusqlite::Connection,
267    done: bool,
268}
269
270impl SqliteTransaction {
271    fn check_done(&self) -> ILResult<()> {
272        if self.done {
273            return Err(ILError::catalog(
274                "Transaction already committed or rolled back",
275            ));
276        }
277        Ok(())
278    }
279}
280
281#[async_trait::async_trait]
282impl Transaction for SqliteTransaction {
283    async fn query(&mut self, sql: &str, schema: CatalogSchemaRef) -> ILResult<RowStream> {
284        trace!("sqlite txn query: {sql}");
285        self.check_done()?;
286        let mut stmt = self
287            .conn
288            .prepare(sql)
289            .map_err(|e| ILError::catalog(format!("failed to prepare sqlite stmt: {sql} {e}")))?;
290        let mut sqlite_rows = stmt
291            .query([])
292            .map_err(|e| ILError::catalog(format!("failed to query sqlite stmt: {sql} {e}")))?;
293
294        let mut rows: Vec<Row> = Vec::new();
295        while let Some(sqlite_row) = sqlite_rows
296            .next()
297            .map_err(|e| ILError::catalog(format!("failed to get next sqlite row: {e}")))?
298        {
299            let row = sqlite_row_to_row(sqlite_row, &schema)?;
300            rows.push(row);
301        }
302        Ok(Box::pin(futures::stream::iter(rows).map(Ok)))
303    }
304
305    async fn execute(&mut self, sql: &str) -> ILResult<usize> {
306        trace!("sqlite txn execute: {sql}");
307        self.check_done()?;
308        self.conn
309            .execute(sql, [])
310            .map_err(|e| ILError::catalog(format!("failed to execute sqlite stmt: {sql} {e}")))
311    }
312
313    async fn execute_batch(&mut self, sqls: &[String]) -> ILResult<()> {
314        trace!("sqlite txn execute batch: {:?}", sqls);
315        self.check_done()?;
316        let sql = sqls.join(";");
317        self.conn
318            .execute_batch(&sql)
319            .map_err(|e| ILError::catalog(format!("failed to execute sqlite batch: {sql} {e}")))
320    }
321
322    async fn insert_rows(
323        &mut self,
324        table_name: &str,
325        field_names: &[String],
326        batches: &[RecordBatch],
327    ) -> ILResult<()> {
328        trace!("sqlite txn insert rows: {table_name}");
329        self.check_done()?;
330
331        let num_columns = field_names.len();
332        let columns = field_names.join(", ");
333        let single_row_placeholders = format!(
334            "({})",
335            std::iter::repeat_n("?", num_columns)
336                .collect::<Vec<_>>()
337                .join(",")
338        );
339
340        const CHUNK_SIZE: usize = 256;
341
342        for batch in batches {
343            let num_rows = batch.num_rows();
344            for chunk_start in (0..num_rows).step_by(CHUNK_SIZE) {
345                let chunk_end = (chunk_start + CHUNK_SIZE).min(num_rows);
346                let chunk_rows = chunk_end - chunk_start;
347
348                let sql = format!(
349                    "INSERT INTO {} ({}) VALUES {}",
350                    table_name,
351                    columns,
352                    std::iter::repeat_n(single_row_placeholders.as_str(), chunk_rows)
353                        .collect::<Vec<_>>()
354                        .join(",")
355                );
356                let mut stmt = self
357                    .conn
358                    .prepare_cached(&sql)
359                    .map_err(|e| ILError::catalog(format!("failed to prepare sqlite stmt: {e}")))?;
360
361                let mut params = Vec::with_capacity(chunk_rows * num_columns);
362                for row_idx in chunk_start..chunk_end {
363                    for col in batch.columns() {
364                        let scalar = Scalar::try_from_array(col.as_ref(), row_idx)?;
365                        params.push(scalar_to_sqlite_param(scalar)?);
366                    }
367                }
368                let refs: Vec<&dyn ToSql> = params.iter().map(|p| p as &dyn ToSql).collect();
369                stmt.execute(refs.as_slice()).map_err(|e| {
370                    ILError::catalog(format!("failed to execute sqlite insert: {e}"))
371                })?;
372            }
373        }
374        Ok(())
375    }
376
377    async fn commit(&mut self) -> ILResult<()> {
378        trace!("sqlite txn commit");
379        self.check_done()?;
380        self.conn
381            .execute_batch("COMMIT")
382            .map_err(|e| ILError::catalog(format!("failed to commit sqlite txn: {e}")))?;
383        self.done = true;
384        Ok(())
385    }
386
387    async fn rollback(&mut self) -> ILResult<()> {
388        trace!("sqlite txn rollback");
389        self.check_done()?;
390        self.conn
391            .execute_batch("ROLLBACK")
392            .map_err(|e| ILError::catalog(format!("failed to rollback sqlite txn: {e}")))?;
393        self.done = true;
394        Ok(())
395    }
396}
397
398impl Drop for SqliteTransaction {
399    fn drop(&mut self) {
400        if self.done {
401            return;
402        }
403        if let Err(e) = self.conn.execute_batch("ROLLBACK") {
404            error!("[indexlake] failed to rollback sqlite txn: {e}");
405        }
406    }
407}
408
409fn sqlite_row_to_row(sqlite_row: &rusqlite::Row, schema: &CatalogSchemaRef) -> ILResult<Row> {
410    let mut row_values = Vec::new();
411    let err_mapping =
412        |e: rusqlite::Error| ILError::catalog(format!("failed to get row value: {e}"));
413    for (idx, field) in schema.columns.iter().enumerate() {
414        let scalar = match field.data_type {
415            CatalogDataType::Boolean => {
416                let v: Option<bool> = sqlite_row.get(idx).map_err(err_mapping)?;
417                Scalar::Boolean(v)
418            }
419            CatalogDataType::Int8 => {
420                let v: Option<i8> = sqlite_row.get(idx).map_err(err_mapping)?;
421                Scalar::Int8(v)
422            }
423            CatalogDataType::Int16 => {
424                let v: Option<i16> = sqlite_row.get(idx).map_err(err_mapping)?;
425                Scalar::Int16(v)
426            }
427            CatalogDataType::Int32 => {
428                let v: Option<i32> = sqlite_row.get(idx).map_err(err_mapping)?;
429                Scalar::Int32(v)
430            }
431            CatalogDataType::Int64 => {
432                let v: Option<i64> = sqlite_row.get(idx).map_err(err_mapping)?;
433                Scalar::Int64(v)
434            }
435            CatalogDataType::UInt8 => {
436                let v: Option<u8> = sqlite_row.get(idx).map_err(err_mapping)?;
437                Scalar::UInt8(v)
438            }
439            CatalogDataType::UInt16 => {
440                let v: Option<u16> = sqlite_row.get(idx).map_err(err_mapping)?;
441                Scalar::UInt16(v)
442            }
443            CatalogDataType::UInt32 => {
444                let v: Option<u32> = sqlite_row.get(idx).map_err(err_mapping)?;
445                Scalar::UInt32(v)
446            }
447            CatalogDataType::UInt64 => {
448                let v: Option<f64> = sqlite_row.get(idx).map_err(err_mapping)?;
449                Scalar::UInt64(v.map(|v| v as u64))
450            }
451            CatalogDataType::Float32 => {
452                let v: Option<f32> = sqlite_row.get(idx).map_err(err_mapping)?;
453                Scalar::Float32(v)
454            }
455            CatalogDataType::Float64 => {
456                let v: Option<f64> = sqlite_row.get(idx).map_err(err_mapping)?;
457                Scalar::Float64(v)
458            }
459            CatalogDataType::Utf8 => {
460                let v: Option<String> = sqlite_row.get(idx).map_err(err_mapping)?;
461                Scalar::Utf8(v)
462            }
463            CatalogDataType::Binary => {
464                let v: Option<Vec<u8>> = sqlite_row.get(idx).map_err(err_mapping)?;
465                Scalar::Binary(v)
466            }
467            CatalogDataType::Uuid => {
468                let v: Option<Vec<u8>> = sqlite_row.get(idx).map_err(err_mapping)?;
469                Scalar::Binary(v)
470            }
471        };
472        if !field.nullable && scalar.is_null() {
473            return Err(ILError::catalog(format!(
474                "column {} is not nullable but got null value",
475                field.name
476            )));
477        }
478        row_values.push(scalar);
479    }
480    Ok(Row::new(schema.clone(), row_values))
481}
indexlake_catalog_sqlite/catalog.rs

indexlake_catalog_sqlite/
catalog.rs
