Skip to main content

we_trust_sqlite/
connection.rs

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