Skip to main content

mz_postgres_util/
desc.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Descriptions of PostgreSQL objects.
11
12use 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/// Describes a schema in a PostgreSQL database.
26///
27/// <https://www.postgresql.org/docs/current/catalog-pg-namespace.html>
28#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
29pub struct PostgresSchemaDesc {
30    /// The OID of the schema.
31    pub oid: Oid,
32    /// The name of the schema.
33    pub name: String,
34    /// Owner of the namespace
35    pub owner: Oid,
36}
37
38/// Describes a table in a PostgreSQL database.
39#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize, Arbitrary)]
40pub struct PostgresTableDesc {
41    /// The OID of the table.
42    pub oid: Oid,
43    /// The name of the schema that the table belongs to.
44    pub namespace: String,
45    /// The name of the table.
46    pub name: String,
47    /// The description of each column, in order of their position in the table.
48    #[proptest(strategy = "proptest::collection::vec(any::<PostgresColumnDesc>(), 1..4)")]
49    pub columns: Vec<PostgresColumnDesc>,
50    /// Applicable keys for this table (i.e. primary key and unique
51    /// constraints).
52    #[proptest(strategy = "proptest::collection::btree_set(any::<PostgresKeyDesc>(), 1..4)")]
53    pub keys: BTreeSet<PostgresKeyDesc>,
54}
55
56impl PostgresTableDesc {
57    /// Determines if two `PostgresTableDesc` are compatible with one another in
58    /// a way that Materialize can handle.
59    ///
60    /// Currently this means that the values are equal except for the following
61    /// exceptions:
62    /// - `self`'s columns are a compatible prefix of `other`'s columns.
63    ///   Compatibility is defined by `PostgresColumnDesc::get_incompatible_schema_change`.
64    /// - `self`'s keys are all present in `other`
65    ///
66    /// On incompatibility, the error describes the first mismatch found and
67    /// how to recover from it. The error becomes the permanent, user-visible
68    /// error for the stalled table, so it must stand on its own.
69    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/// Describes a column in a [`PostgresTableDesc`].
181#[derive(
182    Debug,
183    Clone,
184    Eq,
185    PartialEq,
186    Ord,
187    PartialOrd,
188    Serialize,
189    Deserialize,
190    Arbitrary
191)]
192pub struct PostgresColumnDesc {
193    /// The name of the column.
194    pub name: String,
195    /// The column's monotonic position in its table, i.e. "this was the _i_th
196    /// column created" irrespective of the current number of columns.
197    pub col_num: u16,
198    /// The OID of the column's type.
199    pub type_oid: Oid,
200    /// The modifier for the column's type.
201    pub type_mod: i32,
202    /// True if the column lacks a `NOT NULL` constraint.
203    pub nullable: bool,
204}
205
206impl PostgresColumnDesc {
207    /// Determines if data a relation with a structure of `other` can be treated
208    /// the same as `self`.
209    ///
210    /// Note that this function somewhat unnecessarily errors if the names
211    /// differ; this is negotiable but we want users to understand the fixedness
212    /// of names in our schemas.
213    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        // Columns are compatible if:
234        // - self is nullable; introducing a not null constraint doesn't
235        //   change this column's behavior.
236        // - self and other are both not nullable
237        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        // `col_num` is `u16` on the Rust side. Reject u32 values that don't fit
260        // instead of panicking. This is reachable from untrusted proto bytes.
261        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/// Describes a key in a [`PostgresTableDesc`].
274#[derive(
275    Debug,
276    Clone,
277    Eq,
278    PartialEq,
279    Serialize,
280    Deserialize,
281    PartialOrd,
282    Ord,
283    Arbitrary
284)]
285pub struct PostgresKeyDesc {
286    /// This key is derived from the `pg_constraint` with this OID.
287    pub oid: Oid,
288    /// The name of the constraints.
289    pub name: String,
290    /// The `attnum` of the columns comprising the key. `attnum` is a unique identifier for a column
291    /// in a PG table; see <https://www.postgresql.org/docs/current/catalog-pg-attribute.html>
292    #[proptest(strategy = "proptest::collection::vec(any::<u16>(), 0..4)")]
293    pub cols: Vec<u16>,
294    /// Whether or not this key is the primary key.
295    pub is_primary: bool,
296    /// If this constraint was generated with NULLS NOT DISTINCT; see
297    /// <https://www.postgresql.org/about/featurematrix/detail/392/>
298    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        // `cols` is `Vec<u16>` on the Rust side but `Vec<u32>` on the wire;
314        // a u32 value above 65535 used to panic via `.expect`, which is
315        // reachable from untrusted proto bytes.
316        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}