Skip to main content

safe_migrate/model/
relation.rs

1use crate::ast::identifiers::ObjectId;
2use crate::model::column::Column;
3use serde::{Deserialize, Serialize};
4use std::collections::{HashMap, HashSet};
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
7pub enum Privilege {
8    Select,
9    Insert,
10    Update,
11    Delete,
12    Truncate,
13    References,
14    Trigger,
15    All,
16}
17
18#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
19pub struct PrivilegeMatrix {
20    /// Maps role identity to the set of privileges they possess on this relation
21    pub grants: HashMap<ObjectId, HashSet<Privilege>>,
22}
23
24impl PrivilegeMatrix {
25    pub fn grant(&mut self, role: ObjectId, privileges: HashSet<Privilege>) {
26        self.grants.entry(role).or_default().extend(privileges);
27    }
28
29    pub fn revoke(&mut self, role: &ObjectId, privileges: &HashSet<Privilege>) {
30        if let Some(owned) = self.grants.get_mut(role) {
31            if privileges.contains(&Privilege::All) {
32                owned.clear();
33            } else {
34                for p in privileges {
35                    owned.remove(p);
36                }
37            }
38        }
39    }
40
41    pub fn has_privilege(&self, role: &ObjectId, privilege: Privilege) -> bool {
42        self.grants.get(role).is_some_and(|set| {
43            set.contains(&privilege)
44                || (privilege != Privilege::All && set.contains(&Privilege::All))
45        })
46    }
47}
48
49#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
50pub enum RelationKind {
51    Table,
52    View,
53    MaterializedView,
54}
55
56#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
57pub enum Persistence {
58    Permanent,
59    Temporary,
60    Unlogged,
61}
62
63#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
64pub struct RelationState {
65    pub id: ObjectId,
66    pub owner: ObjectId,
67    pub columns: Vec<Column>,
68    pub generation: u64,
69    pub estimated_rows: Option<u64>,
70    pub relpages: Option<u64>,
71    pub kind: RelationKind,
72    pub persistence: Persistence,
73    pub triggers: HashSet<String>,
74    pub policies: HashSet<String>,
75    pub last_analyze: Option<String>,
76    pub last_autoanalyze: Option<String>,
77    /// Transaction depth at creation, used for same-transaction index checks.
78    pub created_at_tx_depth: usize,
79    pub privileges: PrivilegeMatrix,
80    pub partition_type: Option<String>, // e.g., "RANGE", "LIST", "HASH"
81    pub partition_by: Option<String>,   // The partition key expression
82    #[serde(default)]
83    pub is_fk_dependency: bool,
84}
85
86impl Default for RelationState {
87    fn default() -> Self {
88        Self {
89            id: ObjectId::new("public", "dummy"),
90            owner: ObjectId::new("public", "postgres"),
91            columns: Vec::new(),
92            generation: 0,
93            estimated_rows: Some(0),
94            relpages: None,
95            kind: RelationKind::Table,
96            persistence: Persistence::Permanent,
97            triggers: HashSet::new(),
98            policies: HashSet::new(),
99            last_analyze: None,
100            last_autoanalyze: None,
101            created_at_tx_depth: 0,
102            privileges: PrivilegeMatrix::default(),
103            partition_type: None,
104            partition_by: None,
105            is_fk_dependency: false,
106        }
107    }
108}
109
110impl RelationState {
111    pub fn new(
112        id: ObjectId,
113        owner: ObjectId,
114        generation: u64,
115        estimated_rows: Option<u64>,
116        kind: RelationKind,
117        persistence: Persistence,
118        created_at_tx_depth: usize,
119    ) -> Self {
120        Self {
121            id,
122            owner,
123            columns: Vec::new(),
124            generation,
125            estimated_rows,
126            relpages: None,
127            kind,
128            persistence,
129            triggers: HashSet::new(),
130            policies: HashSet::new(),
131            last_analyze: None,
132            last_autoanalyze: None,
133            created_at_tx_depth,
134            privileges: PrivilegeMatrix::default(),
135            partition_type: None,
136            partition_by: None,
137            is_fk_dependency: false,
138        }
139    }
140
141    pub fn mark_fk_dependency(&mut self) {
142        self.is_fk_dependency = true;
143    }
144
145    pub fn apply_column_action(&mut self, action: &ColumnAction) {
146        match action {
147            ColumnAction::Add {
148                name,
149                data_type,
150                not_null,
151                default,
152            } => {
153                if !self.columns.iter().any(|c| c.name == *name) {
154                    let serial_type = data_type
155                        .as_deref()
156                        .map(str::trim)
157                        .map(str::to_ascii_lowercase)
158                        .and_then(|ty| match ty.as_str() {
159                            "smallserial" | "serial2" => Some("smallint"),
160                            "serial" | "serial4" => Some("integer"),
161                            "bigserial" | "serial8" => Some("bigint"),
162                            _ => None,
163                        });
164                    let is_serial = serial_type.is_some();
165                    let normalized_default = if is_serial {
166                        Some(crate::analysis::expr_ir::ExprIr::FunctionCall {
167                            name: "nextval".to_string(),
168                            args: Vec::new(),
169                        })
170                    } else if matches!(
171                        default,
172                        Some(crate::analysis::expr_ir::ExprIr::Literal(value))
173                            if value.trim().eq_ignore_ascii_case("null")
174                    ) {
175                        None
176                    } else {
177                        default.clone()
178                    };
179                    self.columns.push(Column {
180                        name: name.clone(),
181                        data_type: serial_type
182                            .map(str::to_string)
183                            .or_else(|| data_type.clone()),
184                        type_id: None,
185                        default: normalized_default,
186                        is_nullable: !(*not_null || is_serial),
187                        avg_width: None,
188                        default_expr_text: None,
189                        type_modifier: None,
190                    });
191                }
192            }
193            ColumnAction::Drop { name } => {
194                self.columns.retain(|c| c.name != *name);
195            }
196            ColumnAction::Rename { from, to } => {
197                if let Some(pos) = self.columns.iter().position(|c| c.name == *from)
198                    && !self.columns.iter().any(|c| c.name == *to)
199                {
200                    self.columns[pos].name = to.clone();
201                }
202            }
203            ColumnAction::SetNotNull { name } => {
204                if let Some(col) = self.columns.iter_mut().find(|c| c.name == *name) {
205                    col.is_nullable = false;
206                }
207            }
208            ColumnAction::DropNotNull { name } => {
209                if let Some(col) = self.columns.iter_mut().find(|c| c.name == *name) {
210                    col.is_nullable = true;
211                }
212            }
213            ColumnAction::SetType { name, data_type } => {
214                if let Some(col) = self.columns.iter_mut().find(|c| c.name == *name) {
215                    col.data_type = Some(data_type.clone());
216                }
217            }
218            ColumnAction::SetDefault { name, default } => {
219                if let Some(col) = self.columns.iter_mut().find(|c| c.name == *name) {
220                    col.default = if matches!(
221                        default,
222                        Some(crate::analysis::expr_ir::ExprIr::Literal(value))
223                            if value.trim().eq_ignore_ascii_case("null")
224                    ) {
225                        None
226                    } else {
227                        default.clone()
228                    };
229                    // A migration mutation supersedes raw baseline catalog text.
230                    col.default_expr_text = None;
231                }
232            }
233        }
234    }
235
236    pub fn has_column(&self, name: &str) -> bool {
237        self.columns.iter().any(|c| c.name == name)
238    }
239
240    pub fn get_column(&self, name: &str) -> Option<&Column> {
241        self.columns.iter().find(|c| c.name == name)
242    }
243
244    pub fn is_stale(&self) -> bool {
245        self.last_analyze.is_none() && self.last_autoanalyze.is_none()
246    }
247}
248
249#[derive(Debug, Clone, PartialEq)]
250pub enum ColumnAction {
251    Add {
252        name: String,
253        data_type: Option<String>,
254        not_null: bool,
255        default: Option<crate::analysis::expr_ir::ExprIr>,
256    },
257    Drop {
258        name: String,
259    },
260    Rename {
261        from: String,
262        to: String,
263    },
264    SetNotNull {
265        name: String,
266    },
267    DropNotNull {
268        name: String,
269    },
270    SetType {
271        name: String,
272        data_type: String,
273    },
274    SetDefault {
275        name: String,
276        default: Option<crate::analysis::expr_ir::ExprIr>,
277    },
278}
279
280#[allow(clippy::large_enum_variant)]
281#[derive(Debug, Clone, PartialEq)]
282pub enum RelationOverlay {
283    Present(RelationState),
284    Dropped,
285}