sea_schema/sqlite/def/
table.rs

1use sea_query::{
2    Alias, ColumnDef, Expr, ForeignKey, Index, Keyword, Query, Table, TableCreateStatement, Value,
3};
4
5use super::{
6    ColumnInfo, DefaultType, ForeignKeysInfo, IndexInfo, IndexedColumns, PartialIndexInfo,
7};
8use crate::sqlite::query::SqliteMaster;
9use crate::sqlite::{error::DiscoveryResult, executor::Executor};
10
11#[allow(unused_imports)]
12use crate::sqlx_types::{sqlite::SqliteRow, Row};
13
14/// Defines a table for SQLite
15#[derive(Debug, Default, Clone)]
16pub struct TableDef {
17    /// The table name
18    pub name: String,
19    /// A list of foreign keys in the table
20    pub foreign_keys: Vec<ForeignKeysInfo>,
21    /// A list of the indexes in the table
22    pub indexes: Vec<IndexInfo>,
23    /// A list of UNIQUE and PRIMARY KEY constraints on the table
24    pub constraints: Vec<IndexInfo>,
25    /// A list of all the columns and their types
26    pub columns: Vec<ColumnInfo>,
27    /// Whether the primary key should autoincrement
28    pub auto_increment: bool,
29}
30
31#[cfg(feature = "sqlx-sqlite")]
32/// Gets the table name from a `SqliteRow` and maps it to the [TableDef]
33impl From<&SqliteRow> for TableDef {
34    fn from(row: &SqliteRow) -> Self {
35        let row: String = row.get(0);
36        TableDef {
37            name: row,
38            foreign_keys: Vec::default(),
39            indexes: Vec::default(),
40            constraints: Vec::default(),
41            columns: Vec::default(),
42            auto_increment: bool::default(),
43        }
44    }
45}
46
47#[cfg(not(feature = "sqlx-sqlite"))]
48/// Gets the table name from a `SqliteRow` and maps it to the [TableDef]
49impl From<&SqliteRow> for TableDef {
50    fn from(_: &SqliteRow) -> Self {
51        Self::default()
52    }
53}
54
55impl TableDef {
56    /// Check if the primary key in the table is set to autoincrement as a result of using query
57    /// `SELECT COUNT(*) from sqlite_sequence where name = 'table_name';
58    pub async fn pk_is_autoincrement(&mut self, executor: &Executor) -> DiscoveryResult<&mut Self> {
59        let check_autoincrement = Query::select()
60            .expr(Expr::val(1))
61            .from(SqliteMaster)
62            .and_where(Expr::col(Alias::new("type")).eq("table"))
63            .and_where(Expr::col(Alias::new("name")).eq(self.name.as_str()))
64            .and_where(Expr::col(Alias::new("sql")).like("%AUTOINCREMENT%"))
65            .to_owned();
66
67        if !executor.fetch_all(check_autoincrement).await?.is_empty() {
68            self.auto_increment = true;
69        }
70
71        Ok(self)
72    }
73
74    /// Get a list of most of the UNIQUE and PRIMARY KEY constraints on the table.
75    /// These are implemented by indexes in most cases. These indexes have type "u" or "pk".
76    /// Note that this does not get the column name mapped by the index.
77    /// To get the column name mapped by the index, the `self.get_single_indexinfo` method is invoked
78    pub async fn get_constraints(&mut self, executor: &Executor) -> DiscoveryResult<()> {
79        let mut index_query = String::default();
80        index_query.push_str("PRAGMA index_list('");
81        index_query.push_str(&self.name);
82        index_query.push_str("')");
83
84        let partial_index_info_rows = executor.fetch_all_raw(index_query).await?;
85        let mut partial_indexes: Vec<PartialIndexInfo> = Vec::default();
86
87        partial_index_info_rows.iter().for_each(|info| {
88            let partial_index_info: PartialIndexInfo = info.into();
89
90            if partial_index_info.origin.as_str() == "u" {
91                partial_indexes.push(partial_index_info);
92            }
93        });
94
95        for partial_index in partial_indexes {
96            let partial_index_column: IndexedColumns = self
97                .get_single_indexinfo(executor, &partial_index.name)
98                .await?;
99
100            self.constraints.push(IndexInfo {
101                r#type: partial_index_column.r#type,
102                index_name: partial_index_column.name,
103                table_name: partial_index_column.table,
104                unique: partial_index.unique,
105                origin: partial_index.origin,
106                partial: partial_index.partial,
107                columns: partial_index_column.indexed_columns,
108            });
109        }
110
111        Ok(())
112    }
113
114    /// Get a list of all the indexes in the table.
115    /// Note that this does not get the column name mapped by the index.
116    /// To get the column name mapped by the index, the `self.get_single_indexinfo` method is invoked
117    pub async fn get_indexes(&mut self, executor: &Executor) -> DiscoveryResult<()> {
118        let mut index_query = String::default();
119        index_query.push_str("PRAGMA index_list('");
120        index_query.push_str(&self.name);
121        index_query.push_str("')");
122
123        let partial_index_info_rows = executor.fetch_all_raw(index_query).await?;
124        let mut partial_indexes: Vec<PartialIndexInfo> = Vec::default();
125
126        partial_index_info_rows.iter().for_each(|info| {
127            let partial_index_info: PartialIndexInfo = info.into();
128
129            if partial_index_info.origin.as_str() == "c" {
130                partial_indexes.push(partial_index_info);
131            }
132        });
133
134        for partial_index in partial_indexes {
135            let partial_index_column: IndexedColumns = self
136                .get_single_indexinfo(executor, &partial_index.name)
137                .await?;
138
139            self.indexes.push(IndexInfo {
140                r#type: partial_index_column.r#type,
141                index_name: partial_index_column.name,
142                table_name: partial_index_column.table,
143                unique: partial_index.unique,
144                origin: partial_index.origin,
145                partial: partial_index.partial,
146                columns: partial_index_column.indexed_columns,
147            });
148        }
149
150        Ok(())
151    }
152
153    /// Get a list of all the foreign keys in the table
154    pub async fn get_foreign_keys(&mut self, executor: &Executor) -> DiscoveryResult<&mut Self> {
155        let mut index_query = String::default();
156        index_query.push_str("PRAGMA foreign_key_list('");
157        index_query.push_str(&self.name);
158        index_query.push_str("')");
159
160        let index_info_rows = executor.fetch_all_raw(index_query).await?;
161
162        let mut last_fk_id = None;
163        index_info_rows.iter().for_each(|info| {
164            let mut index_info: ForeignKeysInfo = info.into();
165            let fk_id = index_info.id;
166            if last_fk_id == Some(fk_id) {
167                let last_fk = self.foreign_keys.last_mut().unwrap();
168                last_fk.from.push(index_info.from.pop().unwrap());
169                last_fk.to.push(index_info.to.pop().unwrap());
170            } else {
171                self.foreign_keys.push(index_info);
172            }
173            last_fk_id = Some(fk_id);
174        });
175
176        Ok(self)
177    }
178
179    /// Get a list of all the columns in the table mapped as [ColumnInfo]
180    pub async fn get_column_info(&mut self, executor: &Executor) -> DiscoveryResult<&TableDef> {
181        let mut index_query = String::default();
182        index_query.push_str("PRAGMA table_info('");
183        index_query.push_str(&self.name);
184        index_query.push_str("')");
185
186        let index_info_rows = executor.fetch_all_raw(index_query).await?;
187
188        for info in index_info_rows {
189            let column = ColumnInfo::to_column_def(&info)?;
190            self.columns.push(column);
191        }
192
193        Ok(self)
194    }
195
196    /// Gets the columns that are mapped to an index
197    pub(crate) async fn get_single_indexinfo(
198        &mut self,
199        executor: &Executor,
200        index_name: &str,
201    ) -> DiscoveryResult<IndexedColumns> {
202        let index_query = Query::select()
203            .expr(Expr::cust("*"))
204            .from(SqliteMaster)
205            .and_where(Expr::col(Alias::new("name")).eq(index_name))
206            .to_owned();
207
208        let index_info = executor.fetch_one(index_query).await?;
209
210        let mut index_column_query = String::default();
211        index_column_query.push_str("PRAGMA index_info('");
212        index_column_query.push_str(index_name);
213        index_column_query.push_str("')");
214
215        let index_column_info_rows = executor.fetch_all_raw(index_column_query).await?;
216
217        Ok((&index_info, index_column_info_rows.as_slice()).into())
218    }
219
220    pub fn write(&self) -> TableCreateStatement {
221        let mut primary_keys = Vec::new();
222
223        let mut new_table = Table::create();
224        new_table.table(Alias::new(&self.name));
225
226        self.columns.iter().for_each(|column_info| {
227            let mut new_column =
228                ColumnDef::new_with_type(Alias::new(&column_info.name), column_info.r#type.clone());
229            if column_info.not_null {
230                new_column.not_null();
231            }
232
233            if self.auto_increment && column_info.primary_key {
234                new_column.primary_key().auto_increment();
235            } else if column_info.primary_key {
236                primary_keys.push(column_info.name.clone());
237            }
238
239            match &column_info.default_value {
240                DefaultType::Integer(integer_value) => {
241                    new_column.default(Value::Int(Some(*integer_value)));
242                }
243                DefaultType::Float(float_value) => {
244                    new_column.default(Value::Float(Some(*float_value)));
245                }
246                DefaultType::String(string_value) => {
247                    new_column.default(Value::String(Some(Box::new(string_value.to_string()))));
248                }
249                DefaultType::Null => (),
250                DefaultType::Unspecified => (),
251                DefaultType::CurrentTimestamp => {
252                    new_column.default(Keyword::CurrentTimestamp);
253                }
254            }
255
256            new_table.col(&mut new_column);
257        });
258
259        self.foreign_keys.iter().for_each(|foreign_key| {
260            let mut fk = ForeignKey::create();
261            for from in foreign_key.from.iter() {
262                fk.from(Alias::new(&self.name), Alias::new(from));
263            }
264            for to in foreign_key.to.iter() {
265                fk.to(Alias::new(&foreign_key.table), Alias::new(to));
266            }
267            fk.on_delete(foreign_key.on_delete.to_seaquery_foreign_key_action())
268                .on_update(foreign_key.on_update.to_seaquery_foreign_key_action());
269            new_table.foreign_key(&mut fk);
270        });
271
272        self.constraints.iter().for_each(|index| {
273            new_table.index(&mut index.write());
274        });
275
276        if !primary_keys.is_empty() {
277            let mut primary_key_stmt = Index::create();
278            for primary_key in primary_keys.iter() {
279                primary_key_stmt.col(Alias::new(primary_key));
280            }
281            new_table.primary_key(&mut primary_key_stmt);
282        }
283
284        new_table
285    }
286}