Skip to main content

gluesql_core/executor/alter/
table.rs

1use {
2    super::{AlterError, validate, validate_column_names},
3    crate::{
4        ast::{ColumnDef, ColumnUniqueOption, ForeignKey, ToSql},
5        data::{Row, Schema},
6        executor::{
7            evaluate_stateless,
8            query::{self, OutputBody},
9        },
10        plan::{
11            FilterInputPlan, ProjectInputPlan, ProjectPlan, ProjectionPlan, QueryPlan,
12            SelectItemPlan, SourcePlan, ValuesPlan,
13        },
14        prelude::{DataType, Value},
15        result::Result,
16        store::{GStore, GStoreMut},
17    },
18    serde::Serialize,
19    std::fmt,
20};
21
22pub struct CreateTableOptions<'a> {
23    pub target_table_name: &'a str,
24    pub column_defs: Option<&'a [ColumnDef]>,
25    pub if_not_exists: bool,
26    pub source: &'a Option<Box<QueryPlan>>,
27    pub engine: &'a Option<String>,
28    pub foreign_keys: &'a Vec<ForeignKey>,
29    pub comment: &'a Option<String>,
30}
31
32pub fn create_table<T: GStore + GStoreMut>(
33    storage: &mut T,
34    CreateTableOptions {
35        target_table_name,
36        column_defs,
37        if_not_exists,
38        source,
39        engine,
40        foreign_keys,
41        comment,
42    }: CreateTableOptions<'_>,
43) -> Result<()> {
44    let mut selected_source_rows = None;
45    let target_columns_defs = match source.as_deref() {
46        Some(source_query) => match query::output_body(source_query) {
47            OutputBody::Project(project) => match source_for_schema_copy(project) {
48                Some(SourcePlan::Table(table)) => {
49                    let schema = storage.fetch_schema(&table.name)?;
50                    let Schema {
51                        column_defs: source_column_defs,
52                        ..
53                    } = schema
54                        .ok_or_else(|| AlterError::CtasSourceTableNotFound(table.name.clone()))?;
55
56                    source_column_defs
57                }
58                Some(SourcePlan::Series(_)) => {
59                    let column_def = ColumnDef {
60                        name: "N".into(),
61                        data_type: DataType::Int,
62                        nullable: false,
63                        default: None,
64                        unique: None,
65                        comment: None,
66                    };
67
68                    Some(vec![column_def])
69                }
70                _ => {
71                    let (labels, rows) = query::execute_with_labels(storage, source_query, None)?;
72                    let rows = rows
73                        .map(|row| row.map(Row::into_values))
74                        .collect::<Result<Vec<_>>>()?;
75                    let column_defs = column_defs_from_rows(labels, &rows);
76                    selected_source_rows = Some(rows);
77
78                    Some(column_defs)
79                }
80            },
81            OutputBody::Values(ValuesPlan(values_list)) => {
82                let first_len = values_list[0].len();
83                let mut column_types = vec![None; first_len];
84
85                for exprs in values_list {
86                    for (i, expr) in exprs.iter().enumerate() {
87                        if column_types[i].is_some() {
88                            continue;
89                        }
90
91                        column_types[i] = evaluate_stateless(None, expr)
92                            .and_then(Value::try_from)
93                            .map(|value| value.get_type())?;
94                    }
95
96                    if column_types.iter().all(Option::is_some) {
97                        break;
98                    }
99                }
100
101                let column_defs = column_types
102                    .iter()
103                    .map(|column_type| match column_type {
104                        Some(column_type) => column_type.to_owned(),
105                        None => DataType::Text,
106                    })
107                    .enumerate()
108                    .map(|(i, data_type)| ColumnDef {
109                        name: format!("column{}", i + 1),
110                        data_type,
111                        nullable: true,
112                        default: None,
113                        unique: None,
114                        comment: None,
115                    })
116                    .collect::<Vec<_>>();
117
118                Some(column_defs)
119            }
120        },
121        None if column_defs.is_some() => column_defs.map(<[ColumnDef]>::to_vec),
122        None => None,
123    };
124
125    if let Some(column_defs) = target_columns_defs.as_deref() {
126        validate_column_names(column_defs)?;
127
128        for column_def in column_defs {
129            validate(column_def)?;
130        }
131    }
132
133    for foreign_key in foreign_keys {
134        let ForeignKey {
135            referencing_column_name,
136            referenced_table_name,
137            referenced_column_name,
138            ..
139        } = foreign_key;
140
141        let column_defs = if referenced_table_name == target_table_name {
142            target_columns_defs.clone()
143        } else {
144            let referenced_schema =
145                storage
146                    .fetch_schema(referenced_table_name)?
147                    .ok_or_else(|| {
148                        AlterError::ReferencedTableNotFound(referenced_table_name.to_owned())
149                    })?;
150
151            referenced_schema.column_defs
152        };
153
154        let referenced_column_def = column_defs
155            .and_then(|column_defs| {
156                column_defs
157                    .into_iter()
158                    .find(|column_def| column_def.name == *referenced_column_name)
159            })
160            .ok_or_else(|| AlterError::ReferencedColumnNotFound(referenced_column_name.to_owned()))?
161            .clone();
162
163        let referencing_column_def = target_columns_defs
164            .as_deref()
165            .and_then(|column_defs| {
166                column_defs
167                    .iter()
168                    .find(|column_def| column_def.name == *referencing_column_name)
169            })
170            .ok_or_else(|| {
171                AlterError::ReferencingColumnNotFound(referencing_column_name.to_owned())
172            })?;
173
174        if referencing_column_def.data_type != referenced_column_def.data_type {
175            return Err(AlterError::ForeignKeyDataTypeMismatch {
176                referencing_column: referencing_column_name.to_owned(),
177                referencing_column_type: referencing_column_def.data_type.clone(),
178                referenced_column: referenced_column_name.to_owned(),
179                referenced_column_type: referenced_column_def.data_type.clone(),
180            }
181            .into());
182        }
183
184        if referenced_column_def.unique != Some(ColumnUniqueOption { is_primary: true }) {
185            return Err(AlterError::ReferencingNonPKColumn {
186                referenced_table: referenced_table_name.to_owned(),
187                referenced_column: referenced_column_name.to_owned(),
188            }
189            .into());
190        }
191    }
192
193    if storage.fetch_schema(target_table_name)?.is_none() {
194        let schema = Schema {
195            table_name: target_table_name.to_owned(),
196            column_defs: target_columns_defs,
197            indexes: vec![],
198            engine: engine.clone(),
199            foreign_keys: foreign_keys.clone(),
200            comment: comment.clone(),
201        };
202
203        storage.insert_schema(&schema)?;
204    } else if !if_not_exists {
205        return Err(AlterError::TableAlreadyExists(target_table_name.to_owned()).into());
206    }
207
208    match source {
209        Some(query) => {
210            let rows = match selected_source_rows {
211                Some(rows) => rows,
212                None => query::execute(storage, query, None)?
213                    .map(|row| row.map(Row::into_values))
214                    .collect::<Result<Vec<_>>>()?,
215            };
216
217            storage.append_data(target_table_name, rows)
218        }
219        None => Ok(()),
220    }
221}
222
223fn source_for_schema_copy(project: &ProjectPlan) -> Option<&SourcePlan> {
224    let source = match &project.input {
225        ProjectInputPlan::Source(relation) => relation,
226        ProjectInputPlan::Filter(filter) => match &filter.input {
227            FilterInputPlan::Source(relation) => relation,
228            FilterInputPlan::InnerJoin(_) | FilterInputPlan::LeftOuterJoin(_) => return None,
229        },
230        ProjectInputPlan::InnerJoin(_)
231        | ProjectInputPlan::LeftOuterJoin(_)
232        | ProjectInputPlan::Aggregation(_)
233        | ProjectInputPlan::Having(_) => return None,
234    };
235
236    match &project.projection {
237        ProjectionPlan::SchemalessMap => true,
238        ProjectionPlan::SelectItems(items) => items.iter().all(|item| {
239            matches!(
240                item,
241                SelectItemPlan::Wildcard | SelectItemPlan::QualifiedWildcard(_)
242            )
243        }),
244    }
245    .then_some(source)
246}
247
248fn column_defs_from_rows(labels: Vec<String>, rows: &[Vec<Value>]) -> Vec<ColumnDef> {
249    labels
250        .into_iter()
251        .enumerate()
252        .map(|(index, name)| {
253            let data_type = rows
254                .iter()
255                .filter_map(|row| row.get(index))
256                .find_map(Value::get_type)
257                .unwrap_or(DataType::Text);
258
259            ColumnDef {
260                name,
261                data_type,
262                nullable: true,
263                default: None,
264                unique: None,
265                comment: None,
266            }
267        })
268        .collect()
269}
270
271pub fn drop_table<T: GStore + GStoreMut>(
272    storage: &mut T,
273    table_names: &[String],
274    if_exists: bool,
275    cascade: bool,
276) -> Result<usize> {
277    let mut n = 0;
278
279    for table_name in table_names {
280        let schema = storage.fetch_schema(table_name)?;
281
282        match (schema, if_exists) {
283            (None, true) => {
284                continue;
285            }
286            (None, false) => {
287                return Err(AlterError::TableNotFound(table_name.to_owned()).into());
288            }
289            _ => {}
290        }
291
292        let referencings = storage.fetch_referencings(table_name)?;
293
294        if !referencings.is_empty() && !cascade {
295            return Err(AlterError::CannotDropTableWithReferencing {
296                referenced_table_name: table_name.into(),
297                referencings,
298            }
299            .into());
300        }
301
302        for Referencing {
303            table_name,
304            foreign_key: ForeignKey { name, .. },
305        } in referencings
306        {
307            let mut schema = storage
308                .fetch_schema(&table_name)?
309                .ok_or_else(|| AlterError::TableNotFound(table_name.clone()))?;
310            schema
311                .foreign_keys
312                .retain(|foreign_key| foreign_key.name != name);
313            storage.insert_schema(&schema)?;
314        }
315        storage.delete_schema(table_name)?;
316
317        n += 1;
318    }
319
320    Ok(n)
321}
322
323#[derive(Debug, PartialEq, Eq, Serialize)]
324pub struct Referencing {
325    pub table_name: String,
326    pub foreign_key: ForeignKey,
327}
328
329impl fmt::Display for Referencing {
330    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
331        write!(
332            f,
333            r#"{} on table "{}""#,
334            self.foreign_key.to_sql(),
335            self.table_name
336        )
337    }
338}
339
340#[cfg(test)]
341mod tests {
342    use {super::*, crate::ast::ReferentialAction};
343
344    #[test]
345    fn test_referencing_display() {
346        let referencing = Referencing {
347            table_name: "Referencing".to_owned(),
348            foreign_key: ForeignKey {
349                name: "FK_referenced_id-Referenced_id".to_owned(),
350                referencing_column_name: "referenced_id".to_owned(),
351                referenced_table_name: "Referenced".to_owned(),
352                referenced_column_name: "id".to_owned(),
353                on_delete: ReferentialAction::NoAction,
354                on_update: ReferentialAction::NoAction,
355            },
356        };
357
358        assert_eq!(
359            format!("{referencing}"),
360            r#"CONSTRAINT "FK_referenced_id-Referenced_id" FOREIGN KEY ("referenced_id") REFERENCES "Referenced" ("id") ON DELETE NO ACTION ON UPDATE NO ACTION on table "Referencing""#
361        );
362    }
363}