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(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/// Describes a column in a [`PostgresTableDesc`].
175#[derive(
176    Debug,
177    Clone,
178    Eq,
179    PartialEq,
180    Ord,
181    PartialOrd,
182    Serialize,
183    Deserialize,
184    Arbitrary
185)]
186pub struct PostgresColumnDesc {
187    /// The name of the column.
188    pub name: String,
189    /// The column's monotonic position in its table, i.e. "this was the _i_th
190    /// column created" irrespective of the current number of columns.
191    pub col_num: u16,
192    /// The OID of the column's type.
193    pub type_oid: Oid,
194    /// The modifier for the column's type.
195    pub type_mod: i32,
196    /// True if the column lacks a `NOT NULL` constraint.
197    pub nullable: bool,
198}
199
200impl PostgresColumnDesc {
201    /// Determines if data a relation with a structure of `other` can be treated
202    /// the same as `self`.
203    ///
204    /// Note that this function somewhat unnecessarily errors if the names
205    /// differ; this is negotiable but we want users to understand the fixedness
206    /// of names in our schemas.
207    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        // Columns are compatible if:
228        // - self is nullable; introducing a not null constraint doesn't
229        //   change this column's behavior.
230        // - self and other are both not nullable
231        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        // `col_num` is `u16` on the Rust side. Reject u32 values that don't fit
254        // instead of panicking. This is reachable from untrusted proto bytes.
255        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/// Describes a key in a [`PostgresTableDesc`].
268#[derive(
269    Debug,
270    Clone,
271    Eq,
272    PartialEq,
273    Serialize,
274    Deserialize,
275    PartialOrd,
276    Ord,
277    Arbitrary
278)]
279pub struct PostgresKeyDesc {
280    /// This key is derived from the `pg_constraint` with this OID.
281    pub oid: Oid,
282    /// The name of the constraints.
283    pub name: String,
284    /// The `attnum` of the columns comprising the key. `attnum` is a unique identifier for a column
285    /// in a PG table; see <https://www.postgresql.org/docs/current/catalog-pg-attribute.html>
286    #[proptest(strategy = "proptest::collection::vec(any::<u16>(), 0..4)")]
287    pub cols: Vec<u16>,
288    /// Whether or not this key is the primary key.
289    pub is_primary: bool,
290    /// If this constraint was generated with NULLS NOT DISTINCT; see
291    /// <https://www.postgresql.org/about/featurematrix/detail/392/>
292    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        // `cols` is `Vec<u16>` on the Rust side but `Vec<u32>` on the wire;
308        // a u32 value above 65535 used to panic via `.expect`, which is
309        // reachable from untrusted proto bytes.
310        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}