1use std::collections::{BTreeMap, BTreeSet};
13
14use mz_proto::{IntoRustIfSome, RustType, TryFromProtoError};
15use proptest::prelude::any;
16use proptest_derive::Arbitrary;
17use serde::{Deserialize, Serialize};
18use tokio_postgres::types::Oid;
19use tracing::warn;
20
21use crate::schema_change::{KeyRef, SchemaChange, SchemaChangeError};
22
23include!(concat!(env!("OUT_DIR"), "/mz_postgres_util.desc.rs"));
24
25#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
29pub struct PostgresSchemaDesc {
30 pub oid: Oid,
32 pub name: String,
34 pub owner: Oid,
36}
37
38#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize, Arbitrary)]
40pub struct PostgresTableDesc {
41 pub oid: Oid,
43 pub namespace: String,
45 pub name: String,
47 #[proptest(strategy = "proptest::collection::vec(any::<PostgresColumnDesc>(), 1..4)")]
49 pub columns: Vec<PostgresColumnDesc>,
50 #[proptest(strategy = "proptest::collection::btree_set(any::<PostgresKeyDesc>(), 1..4)")]
53 pub keys: BTreeSet<PostgresKeyDesc>,
54}
55
56impl PostgresTableDesc {
57 pub fn determine_compatibility(
70 &self,
71 other: &PostgresTableDesc,
72 allow_type_to_change_by_col_num: &BTreeSet<u16>,
73 ) -> Result<(), SchemaChangeError> {
74 if self == other {
75 return Ok(());
76 }
77
78 if self.oid != other.oid {
79 warn!(
80 "table {}.{} changed oid from {} to {} during schema verification",
81 self.namespace, self.name, self.oid, other.oid
82 );
83 return Err(
84 self.build_schema_change_error(SchemaChange::TableDropped { oid: other.oid })
85 );
86 }
87
88 if self.namespace != other.namespace || self.name != other.name {
89 return Err(self.build_schema_change_error(SchemaChange::TableRenamed {
90 namespace: other.namespace.clone(),
91 name: other.name.clone(),
92 oid: other.oid,
93 }));
94 }
95
96 let other_cols_by_name = BTreeMap::from_iter(other.columns.iter().map(|c| (&c.name, c)));
97 for column in &self.columns {
98 let allow_type_change = allow_type_to_change_by_col_num.contains(&column.col_num);
99 let other_column = other_cols_by_name.get(&column.name).copied();
100 if let Some(change) =
101 column.get_incompatible_schema_change(other_column, allow_type_change)
102 {
103 return Err(self.build_schema_change_error(change));
104 }
105 }
106
107 if let Some(key) = self.keys.difference(&other.keys).next() {
108 return Err(self.build_schema_change_error(self.key_change(key, other)));
109 }
110
111 Ok(())
112 }
113
114 fn build_schema_change_error(&self, change: SchemaChange) -> SchemaChangeError {
115 SchemaChangeError {
116 namespace: self.namespace.clone(),
117 name: self.name.clone(),
118 oid: self.oid,
119 change,
120 }
121 }
122
123 fn key_change(&self, key: &PostgresKeyDesc, other: &PostgresTableDesc) -> SchemaChange {
124 let key_ref = KeyRef {
125 name: key.name.clone(),
126 is_primary: key.is_primary,
127 columns: key
128 .cols
129 .iter()
130 .map(|attnum| {
131 self.columns
132 .iter()
133 .find(|c| c.col_num == *attnum)
134 .map_or_else(|| format!("attnum {}", attnum), |c| c.name.clone())
135 })
136 .collect(),
137 };
138 let still_exists = other
139 .keys
140 .iter()
141 .any(|k| k.oid == key.oid || k.name == key.name);
142 if still_exists {
143 SchemaChange::KeyAltered { key: key_ref }
144 } else {
145 SchemaChange::KeyDropped { key: key_ref }
146 }
147 }
148}
149
150impl RustType<ProtoPostgresTableDesc> for PostgresTableDesc {
151 fn into_proto(&self) -> ProtoPostgresTableDesc {
152 ProtoPostgresTableDesc {
153 oid: self.oid,
154 namespace: self.namespace.clone(),
155 name: self.name.clone(),
156 columns: self.columns.iter().map(|c| c.into_proto()).collect(),
157 keys: self.keys.iter().map(PostgresKeyDesc::into_proto).collect(),
158 }
159 }
160
161 fn from_proto(proto: ProtoPostgresTableDesc) -> Result<Self, TryFromProtoError> {
162 Ok(PostgresTableDesc {
163 oid: proto.oid,
164 namespace: proto.namespace.clone(),
165 name: proto.name.clone(),
166 columns: proto
167 .columns
168 .into_iter()
169 .map(PostgresColumnDesc::from_proto)
170 .collect::<Result<_, _>>()?,
171 keys: proto
172 .keys
173 .into_iter()
174 .map(PostgresKeyDesc::from_proto)
175 .collect::<Result<_, _>>()?,
176 })
177 }
178}
179
180#[derive(
182 Debug,
183 Clone,
184 Eq,
185 PartialEq,
186 Ord,
187 PartialOrd,
188 Serialize,
189 Deserialize,
190 Arbitrary
191)]
192pub struct PostgresColumnDesc {
193 pub name: String,
195 pub col_num: u16,
198 pub type_oid: Oid,
200 pub type_mod: i32,
202 pub nullable: bool,
204}
205
206impl PostgresColumnDesc {
207 fn get_incompatible_schema_change(
214 &self,
215 other: Option<&PostgresColumnDesc>,
216 allow_type_change: bool,
217 ) -> Option<SchemaChange> {
218 let column = self.name.clone();
219 let Some(other) = other else {
220 return Some(SchemaChange::ColumnDropped { column });
221 };
222 if self.name != other.name {
223 return Some(SchemaChange::ColumnDropped { column });
224 }
225 if self.col_num != other.col_num {
226 return Some(SchemaChange::ColumnMoved { column });
227 }
228 if !allow_type_change
229 && (self.type_oid != other.type_oid || self.type_mod != other.type_mod)
230 {
231 return Some(SchemaChange::ColumnTypeChanged { column });
232 }
233 if !self.nullable && other.nullable {
238 return Some(SchemaChange::NotNullDropped { column });
239 }
240 None
241 }
242}
243
244impl RustType<ProtoPostgresColumnDesc> for PostgresColumnDesc {
245 fn into_proto(&self) -> ProtoPostgresColumnDesc {
246 ProtoPostgresColumnDesc {
247 name: self.name.clone(),
248 col_num: Some(self.col_num.into()),
249 type_oid: self.type_oid,
250 type_mod: self.type_mod,
251 nullable: self.nullable,
252 }
253 }
254
255 fn from_proto(proto: ProtoPostgresColumnDesc) -> Result<Self, TryFromProtoError> {
256 let col_num_u32: u32 = proto
257 .col_num
258 .into_rust_if_some("ProtoPostgresColumnDesc::col_num")?;
259 let col_num = u16::try_from(col_num_u32)
262 .map_err(|e| TryFromProtoError::InvalidFieldError(e.to_string()))?;
263 Ok(PostgresColumnDesc {
264 name: proto.name,
265 col_num,
266 type_oid: proto.type_oid,
267 type_mod: proto.type_mod,
268 nullable: proto.nullable,
269 })
270 }
271}
272
273#[derive(
275 Debug,
276 Clone,
277 Eq,
278 PartialEq,
279 Serialize,
280 Deserialize,
281 PartialOrd,
282 Ord,
283 Arbitrary
284)]
285pub struct PostgresKeyDesc {
286 pub oid: Oid,
288 pub name: String,
290 #[proptest(strategy = "proptest::collection::vec(any::<u16>(), 0..4)")]
293 pub cols: Vec<u16>,
294 pub is_primary: bool,
296 pub nulls_not_distinct: bool,
299}
300
301impl RustType<ProtoPostgresKeyDesc> for PostgresKeyDesc {
302 fn into_proto(&self) -> ProtoPostgresKeyDesc {
303 ProtoPostgresKeyDesc {
304 oid: self.oid,
305 name: self.name.clone(),
306 cols: self.cols.clone().into_iter().map(u32::from).collect(),
307 is_primary: self.is_primary,
308 nulls_not_distinct: self.nulls_not_distinct,
309 }
310 }
311
312 fn from_proto(proto: ProtoPostgresKeyDesc) -> Result<Self, TryFromProtoError> {
313 let cols = proto
317 .cols
318 .into_iter()
319 .map(|c| {
320 u16::try_from(c).map_err(|e| TryFromProtoError::InvalidFieldError(e.to_string()))
321 })
322 .collect::<Result<Vec<_>, _>>()?;
323 Ok(PostgresKeyDesc {
324 oid: proto.oid,
325 name: proto.name,
326 cols,
327 is_primary: proto.is_primary,
328 nulls_not_distinct: proto.nulls_not_distinct,
329 })
330 }
331}