Skip to main content

we_trust_sqlite/
connection.rs

1use crate::executor::SqliteExecutor;
2use crate::transaction::SqliteTransaction;
3use async_trait::async_trait;
4use std::sync::Arc;
5use yykv_types::{ColumnInfo, DsError, DsResult, DsValue, EnumInfo, SchemaInspector, TableInfo};
6
7/// Limbo connection wrapper
8type LimboConn = limbo::Connection;
9
10#[async_trait]
11impl SchemaInspector for SqliteConnection {
12    async fn introspect(&self, _schema: Option<&str>) -> DsResult<(Vec<TableInfo>, Vec<EnumInfo>)> {
13        // SQLite doesn't have multiple schemas in the same way as Postgres (ignoring ATTACH DATABASE for now)
14
15        let sql = "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'";
16        let table_rows = self.query(sql, &[]).await?;
17
18        let mut tables = Vec::new();
19        for row in table_rows {
20            if let DsValue::List(fields) = row {
21                let table_name = match fields.first() {
22                    Some(DsValue::Text(s)) => s.clone(),
23                    _ => continue,
24                };
25
26                let mut columns = Vec::new();
27                let col_sql = format!("PRAGMA table_info({})", table_name);
28                let col_rows = self.query(&col_sql, &[]).await?;
29
30                for col_row in col_rows {
31                    if let DsValue::List(cfields) = col_row {
32                        let name = match cfields.get(1) {
33                            Some(DsValue::Text(s)) => s.clone(),
34                            _ => continue,
35                        };
36                        let data_type = match cfields.get(2) {
37                            Some(DsValue::Text(s)) => s.clone(),
38                            _ => continue,
39                        };
40                        let is_nullable = match cfields.get(3) {
41                            Some(DsValue::Int(v)) => *v == 0,
42                            _ => true,
43                        };
44                        let is_primary_key = match cfields.get(5) {
45                            Some(DsValue::Int(v)) => *v > 0,
46                            _ => false,
47                        };
48                        let default = match cfields.get(4) {
49                            Some(DsValue::Text(s)) => Some(s.clone()),
50                            _ => None,
51                        };
52
53                        columns.push(ColumnInfo {
54                            name,
55                            data_type,
56                            is_nullable,
57                            is_primary_key,
58                            is_enum: false,
59                            foreign_key: None,
60                            default,
61                            description: None,
62                        });
63                    }
64                }
65
66                tables.push(TableInfo {
67                    name: table_name,
68                    columns,
69                    description: None,
70                });
71            }
72        }
73
74        Ok((tables, Vec::new()))
75    }
76}
77
78type Result<T> = std::result::Result<T, DsError>;
79
80/// SQLite connection using Limbo (pure Rust implementation)
81pub struct SqliteConnection {
82    conn: Arc<LimboConn>,
83    executor: Arc<SqliteExecutor>,
84}
85
86impl SqliteConnection {
87    /// Create a new SQLite connection using Limbo
88    pub fn new(
89        conn: Arc<LimboConn>,
90        executor: Arc<SqliteExecutor>,
91    ) -> Self {
92        Self {
93            conn,
94            executor,
95        }
96    }
97
98    /// Execute a SQL statement and return the number of affected rows
99    pub async fn execute(&self, sql: &str, params: &[DsValue]) -> Result<u64> {
100        self.executor.execute(sql, params).await
101    }
102
103    /// Execute a SQL query and return the results
104    pub async fn query(&self, sql: &str, params: &[DsValue]) -> Result<Vec<DsValue>> {
105        self.executor.query(sql, params).await
106    }
107
108    /// Begin a new transaction
109    pub async fn begin_transaction(self) -> Result<SqliteTransaction> {
110        Ok(SqliteTransaction::new(self))
111    }
112}