Skip to main content

reflect_db/
sqlite.rs

1use async_trait::async_trait;
2use sqlx::{SqlitePool, Row};
3use crate::error::ReflectError;
4use crate::executor::{Executor, TableInfo};
5use crate::metadata::{Column, ForeignKey, Index, PrimaryKey, SqlType};
6
7pub struct SqliteExecutor {
8    pool: SqlitePool,
9}
10
11impl SqliteExecutor {
12    pub fn new(pool: SqlitePool) -> Self {
13        Self { pool }
14    }
15}
16
17#[async_trait]
18impl Executor for SqliteExecutor {
19    async fn fetch_tables(
20        &self,
21        _schema: Option<&str>,
22        include_views: bool,
23    ) -> Result<Vec<TableInfo>, ReflectError> {
24        let mut types = vec!["'table'"];
25        if include_views {
26            types.push("'view'");
27        }
28        let types_list = types.join(", ");
29
30        let query_str = format!(
31            "SELECT name, type FROM sqlite_master WHERE type IN ({}) AND name NOT LIKE 'sqlite_%'",
32            types_list
33        );
34
35        let rows = sqlx::query(&query_str)
36            .fetch_all(&self.pool)
37            .await?;
38
39        let mut infos = Vec::new();
40        for row in rows {
41            let name: String = row.try_get("name")?;
42            let ttype: String = row.try_get("type")?;
43            infos.push(TableInfo {
44                name,
45                is_view: ttype == "view",
46            });
47        }
48        Ok(infos)
49    }
50
51    async fn fetch_columns(
52        &self,
53        table: &str,
54        _schema: Option<&str>,
55    ) -> Result<Vec<Column>, ReflectError> {
56        let query_str = format!("PRAGMA table_info('{}')", table);
57        let rows = sqlx::query(&query_str).fetch_all(&self.pool).await?;
58
59        let mut columns = Vec::new();
60        for row in rows {
61            let name: String = row.try_get("name")?;
62            let data_type_str: String = row.try_get("type")?;
63            let notnull: i64 = row.try_get("notnull")?; // i64 is generally what sqlx gives for INTEGER
64            
65            let default: Option<String> = row.try_get("dflt_value").unwrap_or(None);
66
67            let nullable = notnull == 0;
68            let dt = data_type_str.to_uppercase();
69
70            let data_type = if dt.contains("INT") {
71                SqlType::Integer
72            } else if dt.contains("CHAR") || dt.contains("TEXT") {
73                SqlType::Text
74            } else if dt.contains("REAL") || dt.contains("FLOA") || dt.contains("DOUB") {
75                SqlType::Float
76            } else if dt.contains("BOOL") {
77                SqlType::Boolean
78            } else if dt.contains("DATE") || dt.contains("TIME") {
79                SqlType::Timestamp
80            } else {
81                SqlType::Custom(dt.clone())
82            };
83
84
85            columns.push(Column {
86                name,
87                data_type,
88                nullable,
89                default,
90            });
91        }
92
93        Ok(columns)
94    }
95
96    async fn fetch_primary_key(
97        &self,
98        table: &str,
99        _schema: Option<&str>,
100    ) -> Result<Option<PrimaryKey>, ReflectError> {
101        let query_str = format!("PRAGMA table_info('{}')", table);
102        let rows = sqlx::query(&query_str).fetch_all(&self.pool).await?;
103
104        let mut pk_cols = Vec::new();
105        for row in rows {
106            let pk: i64 = row.try_get("pk")?;
107            if pk > 0 {
108                let name: String = row.try_get("name")?;
109                pk_cols.push((pk, name));
110            }
111        }
112        
113        if pk_cols.is_empty() {
114            return Ok(None);
115        }
116        
117        pk_cols.sort_by_key(|(seq, _)| *seq);
118        let columns = pk_cols.into_iter().map(|(_, name)| name).collect();
119
120        Ok(Some(PrimaryKey { columns }))
121    }
122
123    async fn fetch_foreign_keys(
124        &self,
125        table: &str,
126        _schema: Option<&str>,
127    ) -> Result<Vec<ForeignKey>, ReflectError> {
128        let query_str = format!("PRAGMA foreign_key_list('{}')", table);
129        let rows = sqlx::query(&query_str).fetch_all(&self.pool).await?;
130
131        let mut fks = Vec::new();
132        for row in rows {
133            let table: String = row.try_get("table")?;
134            let from: String = row.try_get("from")?;
135            let to: String = row.try_get("to")?;
136
137            fks.push(ForeignKey {
138                column: from,
139                referenced_table: table,
140                referenced_column: to,
141            });
142        }
143
144        Ok(fks)
145    }
146
147    async fn fetch_indexes(
148        &self,
149        table: &str,
150        _schema: Option<&str>,
151    ) -> Result<Vec<Index>, ReflectError> {
152        let query_str = format!("PRAGMA index_list('{}')", table);
153        let rows = sqlx::query(&query_str).fetch_all(&self.pool).await?;
154
155        let mut indexes = Vec::new();
156        for row in rows {
157            let name: String = row.try_get("name")?;
158            let unique: i64 = row.try_get("unique")?;
159            
160            // fetch index columns
161            let idx_info_query = format!("PRAGMA index_info('{}')", name);
162            let col_rows = sqlx::query(&idx_info_query).fetch_all(&self.pool).await?;
163            
164            let mut columns = Vec::new();
165            for col_row in col_rows {
166                let col_name: String = col_row.try_get("name")?;
167                columns.push(col_name);
168            }
169            
170            indexes.push(Index {
171                name,
172                columns,
173                unique: unique > 0,
174            });
175        }
176
177        Ok(indexes)
178    }
179}