Skip to main content

laminar_connectors/postgres/cdc/schema/
mod.rs

1//! `PostgreSQL` relation schema cache.
2//!
3//! Caches the schema (column metadata) for each relation received in
4//! `pgoutput` Relation messages. Required because DML messages only
5//! reference relations by OID --- the schema must be looked up from this cache.
6
7use std::sync::Arc;
8
9use arrow_schema::{DataType, Field, Schema, SchemaRef};
10
11use crate::error::ConnectorError;
12
13use super::types::PgColumn;
14
15/// Cached information about a `PostgreSQL` relation (table).
16#[derive(Debug, Clone)]
17pub struct RelationInfo {
18    /// The relation OID from `pgoutput`.
19    pub relation_id: u32,
20
21    /// Schema (namespace) name.
22    pub namespace: String,
23
24    /// Table name.
25    pub name: String,
26
27    /// Replica identity setting: 'd' (default), 'n' (nothing),
28    /// 'f' (full), 'i' (index).
29    pub replica_identity: char,
30
31    /// Column descriptors in ordinal order.
32    pub columns: Vec<PgColumn>,
33}
34
35impl RelationInfo {
36    /// Returns the fully qualified table name: `namespace.name`.
37    pub(crate) fn full_name(&self) -> Result<String, ConnectorError> {
38        if self.namespace.is_empty() || self.name.is_empty() {
39            return Err(ConnectorError::ReadError(
40                "PostgreSQL CDC relation has an empty schema or table name".into(),
41            ));
42        }
43        let length = self
44            .namespace
45            .len()
46            .checked_add(1)
47            .and_then(|length| length.checked_add(self.name.len()))
48            .ok_or_else(|| {
49                ConnectorError::ReadError(
50                    "PostgreSQL CDC schema-qualified table name size overflow".into(),
51                )
52            })?;
53        let mut table = String::new();
54        table.try_reserve_exact(length).map_err(|error| {
55            ConnectorError::ReadError(format!(
56                "PostgreSQL CDC could not reserve {length} table-name bytes: {error}"
57            ))
58        })?;
59        table.push_str(&self.namespace);
60        table.push('.');
61        table.push_str(&self.name);
62        debug_assert_eq!(table.len(), length);
63        Ok(table)
64    }
65
66    pub(crate) fn variable_retained_bytes(&self) -> Result<usize, ConnectorError> {
67        let mut retained = self
68            .namespace
69            .capacity()
70            .checked_add(self.name.capacity())
71            .and_then(|bytes| {
72                self.columns
73                    .capacity()
74                    .checked_mul(std::mem::size_of::<PgColumn>())
75                    .and_then(|column_bytes| bytes.checked_add(column_bytes))
76            })
77            .ok_or_else(|| {
78                ConnectorError::ReadError(
79                    "PostgreSQL CDC relation-cache retained-byte size overflow".into(),
80                )
81            })?;
82        for column in &self.columns {
83            retained = retained
84                .checked_add(column.name.capacity())
85                .ok_or_else(|| {
86                    ConnectorError::ReadError(
87                        "PostgreSQL CDC relation-cache retained-byte size overflow".into(),
88                    )
89                })?;
90        }
91        Ok(retained)
92    }
93}
94
95/// Cache of relation schemas received from `pgoutput` Relation messages.
96///
97/// The decoder populates this cache as it encounters Relation messages.
98/// DML decoders look up column metadata by relation ID.
99#[derive(Debug, Clone, Default)]
100pub struct RelationCache {
101    relations: Vec<RelationInfo>,
102    variable_retained_bytes: usize,
103}
104
105impl RelationCache {
106    /// Creates an empty relation cache.
107    #[must_use]
108    pub fn new() -> Self {
109        Self::default()
110    }
111
112    /// Adds or replaces a relation in the cache.
113    pub(crate) fn try_reserve_for(&mut self, relation_id: u32) -> Result<(), ConnectorError> {
114        if self
115            .relations
116            .binary_search_by_key(&relation_id, |relation| relation.relation_id)
117            .is_err()
118        {
119            self.relations.try_reserve_exact(1).map_err(|error| {
120                ConnectorError::ReadError(format!(
121                    "PostgreSQL CDC could not reserve relation-cache storage: {error}"
122                ))
123            })?;
124        }
125        Ok(())
126    }
127
128    pub(crate) fn reservation_growth_bytes(
129        &self,
130        relation_id: u32,
131    ) -> Result<usize, ConnectorError> {
132        if self
133            .relations
134            .binary_search_by_key(&relation_id, |relation| relation.relation_id)
135            .is_ok()
136            || self.relations.len() < self.relations.capacity()
137        {
138            return Ok(0);
139        }
140        self.relations
141            .capacity()
142            .max(1)
143            .checked_mul(std::mem::size_of::<RelationInfo>())
144            .ok_or_else(|| {
145                ConnectorError::ReadError(
146                    "PostgreSQL CDC relation-cache growth size overflow".into(),
147                )
148            })
149    }
150
151    pub(crate) fn insert(&mut self, info: RelationInfo) -> Result<(), ConnectorError> {
152        self.try_reserve_for(info.relation_id)?;
153        let new_bytes = info.variable_retained_bytes()?;
154        match self
155            .relations
156            .binary_search_by_key(&info.relation_id, |relation| relation.relation_id)
157        {
158            Ok(index) => {
159                let old_bytes = self.relations[index].variable_retained_bytes()?;
160                self.variable_retained_bytes = self
161                    .variable_retained_bytes
162                    .checked_sub(old_bytes)
163                    .and_then(|bytes| bytes.checked_add(new_bytes))
164                    .ok_or_else(|| {
165                        ConnectorError::Internal(
166                            "PostgreSQL CDC relation-cache retained-byte invariant failed".into(),
167                        )
168                    })?;
169                self.relations[index] = info;
170            }
171            Err(index) => {
172                self.variable_retained_bytes = self
173                    .variable_retained_bytes
174                    .checked_add(new_bytes)
175                    .ok_or_else(|| {
176                        ConnectorError::ReadError(
177                            "PostgreSQL CDC relation-cache retained-byte accounting overflow"
178                                .into(),
179                        )
180                    })?;
181                self.relations.insert(index, info);
182            }
183        }
184        Ok(())
185    }
186
187    pub(crate) fn retained_bytes(&self) -> Result<usize, ConnectorError> {
188        self.container_retained_bytes()?
189            .checked_add(self.variable_retained_bytes)
190            .ok_or_else(|| {
191                ConnectorError::ReadError(
192                    "PostgreSQL CDC relation-cache retained-byte accounting overflow".into(),
193                )
194            })
195    }
196
197    fn container_retained_bytes(&self) -> Result<usize, ConnectorError> {
198        self.relations
199            .capacity()
200            .checked_mul(std::mem::size_of::<RelationInfo>())
201            .ok_or_else(|| {
202                ConnectorError::ReadError(
203                    "PostgreSQL CDC relation-cache container size overflow".into(),
204                )
205            })
206    }
207
208    /// Looks up a relation by its OID.
209    #[must_use]
210    pub fn get(&self, relation_id: u32) -> Option<&RelationInfo> {
211        self.relations
212            .binary_search_by_key(&relation_id, |relation| relation.relation_id)
213            .ok()
214            .map(|index| &self.relations[index])
215    }
216
217    /// Returns the number of cached relations.
218    #[must_use]
219    pub fn len(&self) -> usize {
220        self.relations.len()
221    }
222
223    /// Returns `true` if no relations are cached.
224    #[must_use]
225    pub fn is_empty(&self) -> bool {
226        self.relations.is_empty()
227    }
228
229    /// Clears the cache.
230    pub fn clear(&mut self) {
231        self.relations = Vec::new();
232        self.variable_retained_bytes = 0;
233    }
234}
235
236/// Builds the CDC envelope schema used by [`PostgresCdcSource`](super::source::PostgresCdcSource).
237///
238/// This schema wraps change events in a uniform envelope with metadata
239/// columns, making it compatible with any table structure.
240#[must_use]
241pub fn cdc_envelope_schema() -> SchemaRef {
242    use arrow_schema::TimeUnit;
243    Arc::new(Schema::new(vec![
244        Field::new("_table", DataType::Utf8, false),
245        Field::new("_op", DataType::Utf8, false),
246        Field::new("_lsn", DataType::UInt64, false),
247        Field::new(
248            "_ts_ms",
249            DataType::Timestamp(TimeUnit::Millisecond, None),
250            false,
251        ),
252        Field::new("_before", DataType::Utf8, true),
253        Field::new("_after", DataType::Utf8, true),
254    ]))
255}
256
257#[cfg(test)]
258mod tests;