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(self.build_schema_change_error(SchemaChange::TableDropped));
84 }
85
86 if self.namespace != other.namespace || self.name != other.name {
87 return Err(self.build_schema_change_error(SchemaChange::TableRenamed));
88 }
89
90 let other_cols_by_name = BTreeMap::from_iter(other.columns.iter().map(|c| (&c.name, c)));
91 for column in &self.columns {
92 let allow_type_change = allow_type_to_change_by_col_num.contains(&column.col_num);
93 let other_column = other_cols_by_name.get(&column.name).copied();
94 if let Some(change) =
95 column.get_incompatible_schema_change(other_column, allow_type_change)
96 {
97 return Err(self.build_schema_change_error(change));
98 }
99 }
100
101 if let Some(key) = self.keys.difference(&other.keys).next() {
102 return Err(self.build_schema_change_error(self.key_change(key, other)));
103 }
104
105 Ok(())
106 }
107
108 fn build_schema_change_error(&self, change: SchemaChange) -> SchemaChangeError {
109 SchemaChangeError {
110 namespace: self.namespace.clone(),
111 name: self.name.clone(),
112 oid: self.oid,
113 change,
114 }
115 }
116
117 fn key_change(&self, key: &PostgresKeyDesc, other: &PostgresTableDesc) -> SchemaChange {
118 let key_ref = KeyRef {
119 name: key.name.clone(),
120 is_primary: key.is_primary,
121 columns: key
122 .cols
123 .iter()
124 .map(|attnum| {
125 self.columns
126 .iter()
127 .find(|c| c.col_num == *attnum)
128 .map_or_else(|| format!("attnum {}", attnum), |c| c.name.clone())
129 })
130 .collect(),
131 };
132 let still_exists = other
133 .keys
134 .iter()
135 .any(|k| k.oid == key.oid || k.name == key.name);
136 if still_exists {
137 SchemaChange::KeyAltered { key: key_ref }
138 } else {
139 SchemaChange::KeyDropped { key: key_ref }
140 }
141 }
142}
143
144impl RustType<ProtoPostgresTableDesc> for PostgresTableDesc {
145 fn into_proto(&self) -> ProtoPostgresTableDesc {
146 ProtoPostgresTableDesc {
147 oid: self.oid,
148 namespace: self.namespace.clone(),
149 name: self.name.clone(),
150 columns: self.columns.iter().map(|c| c.into_proto()).collect(),
151 keys: self.keys.iter().map(PostgresKeyDesc::into_proto).collect(),
152 }
153 }
154
155 fn from_proto(proto: ProtoPostgresTableDesc) -> Result<Self, TryFromProtoError> {
156 Ok(PostgresTableDesc {
157 oid: proto.oid,
158 namespace: proto.namespace.clone(),
159 name: proto.name.clone(),
160 columns: proto
161 .columns
162 .into_iter()
163 .map(PostgresColumnDesc::from_proto)
164 .collect::<Result<_, _>>()?,
165 keys: proto
166 .keys
167 .into_iter()
168 .map(PostgresKeyDesc::from_proto)
169 .collect::<Result<_, _>>()?,
170 })
171 }
172}
173
174#[derive(
176 Debug,
177 Clone,
178 Eq,
179 PartialEq,
180 Ord,
181 PartialOrd,
182 Serialize,
183 Deserialize,
184 Arbitrary
185)]
186pub struct PostgresColumnDesc {
187 pub name: String,
189 pub col_num: u16,
192 pub type_oid: Oid,
194 pub type_mod: i32,
196 pub nullable: bool,
198}
199
200impl PostgresColumnDesc {
201 fn get_incompatible_schema_change(
208 &self,
209 other: Option<&PostgresColumnDesc>,
210 allow_type_change: bool,
211 ) -> Option<SchemaChange> {
212 let column = self.name.clone();
213 let Some(other) = other else {
214 return Some(SchemaChange::ColumnDropped { column });
215 };
216 if self.name != other.name {
217 return Some(SchemaChange::ColumnDropped { column });
218 }
219 if self.col_num != other.col_num {
220 return Some(SchemaChange::ColumnMoved { column });
221 }
222 if !allow_type_change
223 && (self.type_oid != other.type_oid || self.type_mod != other.type_mod)
224 {
225 return Some(SchemaChange::ColumnTypeChanged { column });
226 }
227 if !self.nullable && other.nullable {
232 return Some(SchemaChange::NotNullDropped { column });
233 }
234 None
235 }
236}
237
238impl RustType<ProtoPostgresColumnDesc> for PostgresColumnDesc {
239 fn into_proto(&self) -> ProtoPostgresColumnDesc {
240 ProtoPostgresColumnDesc {
241 name: self.name.clone(),
242 col_num: Some(self.col_num.into()),
243 type_oid: self.type_oid,
244 type_mod: self.type_mod,
245 nullable: self.nullable,
246 }
247 }
248
249 fn from_proto(proto: ProtoPostgresColumnDesc) -> Result<Self, TryFromProtoError> {
250 let col_num_u32: u32 = proto
251 .col_num
252 .into_rust_if_some("ProtoPostgresColumnDesc::col_num")?;
253 let col_num = u16::try_from(col_num_u32)
256 .map_err(|e| TryFromProtoError::InvalidFieldError(e.to_string()))?;
257 Ok(PostgresColumnDesc {
258 name: proto.name,
259 col_num,
260 type_oid: proto.type_oid,
261 type_mod: proto.type_mod,
262 nullable: proto.nullable,
263 })
264 }
265}
266
267#[derive(
269 Debug,
270 Clone,
271 Eq,
272 PartialEq,
273 Serialize,
274 Deserialize,
275 PartialOrd,
276 Ord,
277 Arbitrary
278)]
279pub struct PostgresKeyDesc {
280 pub oid: Oid,
282 pub name: String,
284 #[proptest(strategy = "proptest::collection::vec(any::<u16>(), 0..4)")]
287 pub cols: Vec<u16>,
288 pub is_primary: bool,
290 pub nulls_not_distinct: bool,
293}
294
295impl RustType<ProtoPostgresKeyDesc> for PostgresKeyDesc {
296 fn into_proto(&self) -> ProtoPostgresKeyDesc {
297 ProtoPostgresKeyDesc {
298 oid: self.oid,
299 name: self.name.clone(),
300 cols: self.cols.clone().into_iter().map(u32::from).collect(),
301 is_primary: self.is_primary,
302 nulls_not_distinct: self.nulls_not_distinct,
303 }
304 }
305
306 fn from_proto(proto: ProtoPostgresKeyDesc) -> Result<Self, TryFromProtoError> {
307 let cols = proto
311 .cols
312 .into_iter()
313 .map(|c| {
314 u16::try_from(c).map_err(|e| TryFromProtoError::InvalidFieldError(e.to_string()))
315 })
316 .collect::<Result<Vec<_>, _>>()?;
317 Ok(PostgresKeyDesc {
318 oid: proto.oid,
319 name: proto.name,
320 cols,
321 is_primary: proto.is_primary,
322 nulls_not_distinct: proto.nulls_not_distinct,
323 })
324 }
325}