laminar_connectors/postgres/cdc/schema/
mod.rs1use std::sync::Arc;
8
9use arrow_schema::{DataType, Field, Schema, SchemaRef};
10
11use crate::error::ConnectorError;
12
13use super::types::PgColumn;
14
15#[derive(Debug, Clone)]
17pub struct RelationInfo {
18 pub relation_id: u32,
20
21 pub namespace: String,
23
24 pub name: String,
26
27 pub replica_identity: char,
30
31 pub columns: Vec<PgColumn>,
33}
34
35impl RelationInfo {
36 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#[derive(Debug, Clone, Default)]
100pub struct RelationCache {
101 relations: Vec<RelationInfo>,
102 variable_retained_bytes: usize,
103}
104
105impl RelationCache {
106 #[must_use]
108 pub fn new() -> Self {
109 Self::default()
110 }
111
112 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 #[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 #[must_use]
219 pub fn len(&self) -> usize {
220 self.relations.len()
221 }
222
223 #[must_use]
225 pub fn is_empty(&self) -> bool {
226 self.relations.is_empty()
227 }
228
229 pub fn clear(&mut self) {
231 self.relations = Vec::new();
232 self.variable_retained_bytes = 0;
233 }
234}
235
236#[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;