Skip to main content

gluesql_core/store/
alter_table.rs

1use {
2    super::{Store, StoreMut},
3    crate::{
4        ast::ColumnDef, data::Value, executor::evaluate_stateless, plan::plan_scalar_expr,
5        result::Result,
6    },
7    serde::Serialize,
8    std::fmt::Debug,
9    thiserror::Error,
10};
11
12#[derive(Error, Serialize, Debug, PartialEq, Eq)]
13pub enum AlterTableError {
14    #[error("Table not found: {0}")]
15    TableNotFound(String),
16
17    #[error("Renaming column not found")]
18    RenamingColumnNotFound,
19
20    #[error("Default value is required: {0:#?}")]
21    DefaultValueRequired(ColumnDef),
22
23    #[error("Already existing column: {0}")]
24    AlreadyExistingColumn(String),
25
26    #[error("Dropping column not found: {0}")]
27    DroppingColumnNotFound(String),
28
29    #[error("Schemaless table does not support ALTER TABLE: {0}")]
30    SchemalessTableFound(String),
31}
32
33pub trait AlterTable: Store + StoreMut {
34    fn rename_schema(&mut self, table_name: &str, new_table_name: &str) -> Result<()> {
35        let mut schema = self
36            .fetch_schema(table_name)?
37            .ok_or_else(|| AlterTableError::TableNotFound(table_name.to_owned()))?;
38        new_table_name.clone_into(&mut schema.table_name);
39        self.insert_schema(&schema)?;
40
41        let rows = self.scan_data(table_name)?.collect::<Result<Vec<_>>>()?;
42
43        self.insert_data(new_table_name, rows)?;
44        self.delete_schema(table_name)
45    }
46
47    fn rename_column(
48        &mut self,
49        table_name: &str,
50        old_column_name: &str,
51        new_column_name: &str,
52    ) -> Result<()> {
53        let mut schema = self
54            .fetch_schema(table_name)?
55            .ok_or_else(|| AlterTableError::TableNotFound(table_name.to_owned()))?;
56
57        let column_defs = schema
58            .column_defs
59            .as_mut()
60            .ok_or_else(|| AlterTableError::SchemalessTableFound(table_name.to_owned()))?;
61
62        if column_defs
63            .iter()
64            .any(|column_def| column_def.name == new_column_name)
65        {
66            return Err(AlterTableError::AlreadyExistingColumn(new_column_name.to_owned()).into());
67        }
68
69        new_column_name.clone_into(
70            &mut column_defs
71                .iter_mut()
72                .find(|column_def| column_def.name == old_column_name)
73                .ok_or(AlterTableError::RenamingColumnNotFound)?
74                .name,
75        );
76
77        let rows = self.scan_data(table_name)?.collect::<Result<Vec<_>>>()?;
78
79        self.insert_schema(&schema)?;
80        self.insert_data(table_name, rows)
81    }
82
83    fn add_column(&mut self, table_name: &str, column_def: &ColumnDef) -> Result<()> {
84        let mut schema = self
85            .fetch_schema(table_name)?
86            .ok_or_else(|| AlterTableError::TableNotFound(table_name.to_owned()))?;
87
88        let default_value = match (column_def.default.as_ref(), column_def.nullable) {
89            (Some(default), _) => {
90                let default = plan_scalar_expr(default.clone());
91
92                evaluate_stateless(None, &default)?.try_into()?
93            }
94            (None, true) => Value::Null,
95            (None, false) => {
96                return Err(AlterTableError::DefaultValueRequired(column_def.clone()).into());
97            }
98        };
99
100        let column_defs = schema
101            .column_defs
102            .as_mut()
103            .ok_or_else(|| AlterTableError::SchemalessTableFound(table_name.to_owned()))?;
104
105        if column_defs.iter().any(|def| def.name == column_def.name) {
106            return Err(AlterTableError::AlreadyExistingColumn(column_def.name.clone()).into());
107        }
108
109        column_defs.push(column_def.clone());
110
111        let rows = self
112            .scan_data(table_name)?
113            .map(|row| {
114                let (key, mut values) = row?;
115                let default_value = default_value.clone();
116
117                values.push(default_value);
118                Ok((key, values))
119            })
120            .collect::<Result<Vec<_>>>()?;
121
122        self.insert_schema(&schema)?;
123        self.insert_data(table_name, rows)
124    }
125
126    fn drop_column(&mut self, table_name: &str, column_name: &str, if_exists: bool) -> Result<()> {
127        let mut schema = self
128            .fetch_schema(table_name)?
129            .ok_or_else(|| AlterTableError::TableNotFound(table_name.to_owned()))?;
130
131        let column_defs = schema
132            .column_defs
133            .as_mut()
134            .ok_or_else(|| AlterTableError::SchemalessTableFound(table_name.to_owned()))?;
135
136        let i = match column_defs
137            .iter()
138            .position(|column_def| column_def.name == column_name)
139        {
140            Some(i) => i,
141            None if if_exists => return Ok(()),
142            None => {
143                return Err(AlterTableError::DroppingColumnNotFound(column_name.to_owned()).into());
144            }
145        };
146
147        column_defs.retain(|column_def| column_def.name != column_name);
148
149        let rows = self
150            .scan_data(table_name)?
151            .map(|row| {
152                let (key, mut values) = row?;
153                values.remove(i);
154                Ok((key, values))
155            })
156            .collect::<Result<Vec<_>>>()?;
157
158        self.insert_schema(&schema)?;
159        self.insert_data(table_name, rows)
160    }
161}