Skip to main content

sea_schema/sqlite/def/
table.rs

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