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