laminar_connectors/postgres/cdc/
schema.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 {
259    use super::*;
260    use crate::postgres::cdc::types::{INT4_OID, INT8_OID, TEXT_OID};
261
262    fn sample_relation() -> RelationInfo {
263        RelationInfo {
264            relation_id: 16384,
265            namespace: "public".to_string(),
266            name: "users".to_string(),
267            replica_identity: 'd',
268            columns: vec![
269                PgColumn::new("id".to_string(), INT8_OID, -1, true),
270                PgColumn::new("name".to_string(), TEXT_OID, -1, false),
271                PgColumn::new("age".to_string(), INT4_OID, -1, false),
272            ],
273        }
274    }
275
276    #[test]
277    fn test_relation_full_name_public() {
278        let rel = sample_relation();
279        assert_eq!(rel.full_name().unwrap(), "public.users");
280    }
281
282    #[test]
283    fn test_relation_full_name_custom_schema() {
284        let mut rel = sample_relation();
285        rel.namespace = "app".to_string();
286        assert_eq!(rel.full_name().unwrap(), "app.users");
287    }
288
289    #[test]
290    fn test_relation_full_name_rejects_empty_components() {
291        let mut rel = sample_relation();
292        rel.namespace.clear();
293        assert!(rel.full_name().is_err());
294        rel.namespace = "public".into();
295        rel.name.clear();
296        assert!(rel.full_name().is_err());
297    }
298
299    #[test]
300    fn test_relation_cache() {
301        let mut cache = RelationCache::new();
302        assert!(cache.is_empty());
303
304        cache.insert(sample_relation()).unwrap();
305        assert_eq!(cache.len(), 1);
306        assert!(cache.get(16384).is_some());
307        assert!(cache.get(99999).is_none());
308    }
309
310    #[test]
311    fn test_cache_replace() {
312        let mut cache = RelationCache::new();
313        cache.insert(sample_relation()).unwrap();
314
315        let mut updated = sample_relation();
316        updated
317            .columns
318            .push(PgColumn::new("email".to_string(), TEXT_OID, -1, false));
319        cache.insert(updated).unwrap();
320
321        assert_eq!(cache.len(), 1);
322        assert_eq!(cache.get(16384).unwrap().columns.len(), 4);
323    }
324
325    #[test]
326    fn test_cache_clear() {
327        let mut cache = RelationCache::new();
328        cache.insert(sample_relation()).unwrap();
329        cache.clear();
330        assert!(cache.is_empty());
331    }
332
333    #[test]
334    fn test_cdc_envelope_schema() {
335        let schema = cdc_envelope_schema();
336        assert_eq!(schema.fields().len(), 6);
337        assert_eq!(schema.field(0).name(), "_table");
338        assert_eq!(schema.field(1).name(), "_op");
339        assert_eq!(schema.field(2).name(), "_lsn");
340        assert_eq!(schema.field(3).name(), "_ts_ms");
341        assert_eq!(schema.field(4).name(), "_before");
342        assert_eq!(schema.field(5).name(), "_after");
343        // _before and _after are nullable
344        assert!(schema.field(4).is_nullable());
345        assert!(schema.field(5).is_nullable());
346    }
347}
laminar_connectors/postgres/cdc/schema.rs

laminar_connectors/postgres/cdc/
schema.rs
