crabka_connect_postgres/
schema.rs

1use bytes::Bytes;
2use crabka_schema_serde::wire::encode_protobuf;
3use prost::Message as _;
4use prost_reflect::prost_types::field_descriptor_proto::{Label, Type};
5use prost_reflect::prost_types::{
6    DescriptorProto, FieldDescriptorProto, FileDescriptorProto, FileDescriptorSet,
7};
8use prost_reflect::{DescriptorPool, DynamicMessage, MessageDescriptor, Value};
9
10use crate::PostgresConnectError;
11use crate::model::{ColumnValue, EntityDifference, EntityKey, Operation, ScalarValue};
12
13// Local placeholder IDs until schema registry allocation is wired in.
14const KEY_SCHEMA_ID: u32 = 1;
15const VALUE_SCHEMA_ID: u32 = 2;
16
17const KEY_MESSAGE_INDEX: &[i32] = &[1];
18const VALUE_MESSAGE_INDEX: &[i32] = &[2];
19const PACKAGE: &str = "crabka.connect.postgres";
20const COLUMN_VALUE: &str = "ColumnValue";
21const ENTITY_KEY: &str = "EntityKey";
22const ENTITY_DIFFERENCE: &str = "EntityDifference";
23
24#[derive(Debug, Clone)]
25pub struct PostgresProtoEncoder {
26    key: MessageDescriptor,
27    value: MessageDescriptor,
28    column_value: MessageDescriptor,
29}
30
31impl PostgresProtoEncoder {
32    pub fn new() -> Result<Self, PostgresConnectError> {
33        let pool = DescriptorPool::from_file_descriptor_set(schema_descriptor_set())
34            .map_err(convert_error)?;
35
36        Ok(Self {
37            key: message_descriptor(&pool, ENTITY_KEY)?,
38            value: message_descriptor(&pool, ENTITY_DIFFERENCE)?,
39            column_value: message_descriptor(&pool, COLUMN_VALUE)?,
40        })
41    }
42
43    pub fn encode_key(&self, key: &EntityKey) -> Result<Bytes, PostgresConnectError> {
44        let message = self.key_to_message(key)?;
45        Ok(encode_protobuf(
46            KEY_SCHEMA_ID,
47            KEY_MESSAGE_INDEX,
48            &message.encode_to_vec(),
49        ))
50    }
51
52    pub fn encode_value(&self, value: &EntityDifference) -> Result<Bytes, PostgresConnectError> {
53        let message = self.difference_to_message(value)?;
54        Ok(encode_protobuf(
55            VALUE_SCHEMA_ID,
56            VALUE_MESSAGE_INDEX,
57            &message.encode_to_vec(),
58        ))
59    }
60
61    fn key_to_message(&self, key: &EntityKey) -> Result<DynamicMessage, PostgresConnectError> {
62        let mut message = DynamicMessage::new(self.key.clone());
63        set_field(&mut message, "table", Value::String(key.table.clone()))?;
64        set_field(
65            &mut message,
66            "columns",
67            Value::List(self.columns_to_values(&key.columns)?),
68        )?;
69        Ok(message)
70    }
71
72    fn difference_to_message(
73        &self,
74        difference: &EntityDifference,
75    ) -> Result<DynamicMessage, PostgresConnectError> {
76        let mut message = DynamicMessage::new(self.value.clone());
77        set_field(
78            &mut message,
79            "table",
80            Value::String(difference.table.clone()),
81        )?;
82        set_field(
83            &mut message,
84            "operation",
85            Value::String(operation_name(difference.op).to_owned()),
86        )?;
87        set_field(
88            &mut message,
89            "lsn",
90            Value::String(difference.lsn.to_string()),
91        )?;
92        set_field(
93            &mut message,
94            "before",
95            Value::List(self.columns_to_values(&difference.before)?),
96        )?;
97        set_field(
98            &mut message,
99            "after",
100            Value::List(self.columns_to_values(&difference.after)?),
101        )?;
102        set_field(
103            &mut message,
104            "key",
105            Value::Message(self.key_to_message(&difference.key)?),
106        )?;
107        if let Some(txid) = difference.txid {
108            set_field(&mut message, "txid", Value::I64(txid))?;
109        }
110        if let Some(commit_timestamp_ms) = difference.commit_timestamp_ms {
111            set_field(
112                &mut message,
113                "commit_timestamp_ms",
114                Value::I64(commit_timestamp_ms),
115            )?;
116        }
117        Ok(message)
118    }
119
120    fn columns_to_values(
121        &self,
122        columns: &[ColumnValue],
123    ) -> Result<Vec<Value>, PostgresConnectError> {
124        columns
125            .iter()
126            .map(|column| self.column_to_value(column).map(Value::Message))
127            .collect()
128    }
129
130    fn column_to_value(
131        &self,
132        column: &ColumnValue,
133    ) -> Result<DynamicMessage, PostgresConnectError> {
134        let mut message = DynamicMessage::new(self.column_value.clone());
135        set_field(&mut message, "name", Value::String(column.name.clone()))?;
136        set_scalar_fields(&mut message, &column.value)?;
137        Ok(message)
138    }
139}
140
141fn schema_descriptor_set() -> FileDescriptorSet {
142    FileDescriptorSet {
143        file: vec![FileDescriptorProto {
144            name: Some("crabka/connect/postgres/cdc.proto".to_owned()),
145            package: Some(PACKAGE.to_owned()),
146            syntax: Some("proto3".to_owned()),
147            message_type: vec![
148                DescriptorProto {
149                    name: Some(COLUMN_VALUE.to_owned()),
150                    field: vec![
151                        field("name", 1, Type::String),
152                        field("kind", 2, Type::String),
153                        field("string_value", 3, Type::String),
154                        field("bool_value", 4, Type::Bool),
155                        field("int_value", 5, Type::Int64),
156                        field("bytes_value", 6, Type::Bytes),
157                        field("is_null", 7, Type::Bool),
158                    ],
159                    ..DescriptorProto::default()
160                },
161                DescriptorProto {
162                    name: Some(ENTITY_KEY.to_owned()),
163                    field: vec![
164                        field("table", 1, Type::String),
165                        repeated_message_field("columns", 2, COLUMN_VALUE),
166                    ],
167                    ..DescriptorProto::default()
168                },
169                DescriptorProto {
170                    name: Some(ENTITY_DIFFERENCE.to_owned()),
171                    field: vec![
172                        field("table", 1, Type::String),
173                        field("operation", 2, Type::String),
174                        field("lsn", 3, Type::String),
175                        repeated_message_field("before", 4, COLUMN_VALUE),
176                        repeated_message_field("after", 5, COLUMN_VALUE),
177                        message_field("key", 6, ENTITY_KEY),
178                        field("txid", 7, Type::Int64),
179                        field("commit_timestamp_ms", 8, Type::Int64),
180                    ],
181                    ..DescriptorProto::default()
182                },
183            ],
184            ..FileDescriptorProto::default()
185        }],
186    }
187}
188
189fn field(name: &str, number: i32, field_type: Type) -> FieldDescriptorProto {
190    FieldDescriptorProto {
191        name: Some(name.to_owned()),
192        number: Some(number),
193        label: Some(Label::Optional as i32),
194        r#type: Some(field_type as i32),
195        ..FieldDescriptorProto::default()
196    }
197}
198
199fn message_field(name: &str, number: i32, message_name: &str) -> FieldDescriptorProto {
200    FieldDescriptorProto {
201        type_name: Some(full_type_name(message_name)),
202        ..field(name, number, Type::Message)
203    }
204}
205
206fn repeated_message_field(name: &str, number: i32, message_name: &str) -> FieldDescriptorProto {
207    FieldDescriptorProto {
208        label: Some(Label::Repeated as i32),
209        ..message_field(name, number, message_name)
210    }
211}
212
213fn full_type_name(message_name: &str) -> String {
214    format!(".{PACKAGE}.{message_name}")
215}
216
217fn message_descriptor(
218    pool: &DescriptorPool,
219    name: &str,
220) -> Result<MessageDescriptor, PostgresConnectError> {
221    pool.get_message_by_name(&format!("{PACKAGE}.{name}"))
222        .ok_or_else(|| PostgresConnectError::Convert(format!("protobuf message {name} not found")))
223}
224
225fn set_field(
226    message: &mut DynamicMessage,
227    name: &str,
228    value: Value,
229) -> Result<(), PostgresConnectError> {
230    message
231        .try_set_field_by_name(name, value)
232        .map_err(convert_error)
233}
234
235fn operation_name(operation: Operation) -> &'static str {
236    match operation {
237        Operation::Insert => "insert",
238        Operation::Update => "update",
239        Operation::Delete => "delete",
240    }
241}
242
243fn set_scalar_fields(
244    message: &mut DynamicMessage,
245    value: &ScalarValue,
246) -> Result<(), PostgresConnectError> {
247    match value {
248        ScalarValue::Null => {
249            set_field(message, "kind", Value::String("null".to_owned()))?;
250            set_field(message, "is_null", Value::Bool(true))?;
251        }
252        ScalarValue::UnchangedToast => {
253            set_field(message, "kind", Value::String("unchanged_toast".to_owned()))?;
254        }
255        ScalarValue::Bool(value) => {
256            set_field(message, "kind", Value::String("bool".to_owned()))?;
257            set_field(message, "bool_value", Value::Bool(*value))?;
258        }
259        ScalarValue::Int(value) => {
260            set_field(message, "kind", Value::String("int".to_owned()))?;
261            set_field(message, "int_value", Value::I64(*value))?;
262        }
263        ScalarValue::Float(value) => {
264            set_field(message, "kind", Value::String("float".to_owned()))?;
265            set_field(message, "string_value", Value::String(value.clone()))?;
266        }
267        ScalarValue::Text(value) => {
268            set_field(message, "kind", Value::String("text".to_owned()))?;
269            set_field(message, "string_value", Value::String(value.clone()))?;
270        }
271        ScalarValue::Bytes(value) => {
272            set_field(message, "kind", Value::String("bytes".to_owned()))?;
273            set_field(
274                message,
275                "bytes_value",
276                Value::Bytes(Bytes::copy_from_slice(value)),
277            )?;
278        }
279    }
280
281    Ok(())
282}
283
284fn convert_error(error: impl std::fmt::Display) -> PostgresConnectError {
285    PostgresConnectError::Convert(error.to_string())
286}
287
288#[cfg(test)]
289mod tests {
290    use bytes::Bytes;
291    use prost_reflect::{DescriptorPool, DynamicMessage, Value};
292
293    use crabka_schema_serde::wire::decode_protobuf;
294
295    use crate::model::{ColumnSchema, ScalarValue};
296    use crate::pgoutput::{RelationCache, RelationEvent, RowEvent, RowEventKind};
297    use crate::{ColumnValue, EntityDifference, EntityKey, Operation, PgLsn, TableSchema};
298
299    use super::{
300        COLUMN_VALUE, ENTITY_DIFFERENCE, ENTITY_KEY, KEY_SCHEMA_ID, PostgresProtoEncoder,
301        VALUE_SCHEMA_ID, message_descriptor, schema_descriptor_set,
302    };
303
304    #[test]
305    fn encoder_frames_key_and_value_as_protobuf() {
306        let encoder = PostgresProtoEncoder::new().expect("encoder builds descriptors");
307        let diff = sample_difference();
308        let pool = DescriptorPool::from_file_descriptor_set(schema_descriptor_set())
309            .expect("descriptor pool builds");
310
311        let key = encoder.encode_key(&diff.key).expect("key encodes");
312        let value = encoder.encode_value(&diff).expect("value encodes");
313
314        let (key_id, key_index, key_body) = decode_protobuf(&key).expect("key frame decodes");
315        assert_eq!(key_id, KEY_SCHEMA_ID);
316        assert_eq!(key_index, vec![1]);
317        assert!(!key_body.is_empty());
318
319        let (value_id, value_index, value_body) =
320            decode_protobuf(&value).expect("value frame decodes");
321        assert_eq!(value_id, VALUE_SCHEMA_ID);
322        assert_eq!(value_index, vec![2]);
323        assert!(!value_body.is_empty());
324
325        let key_message = DynamicMessage::decode(
326            message_descriptor(&pool, ENTITY_KEY).expect("key descriptor"),
327            key_body,
328        )
329        .expect("key body decodes");
330        assert_eq!(string_field(&key_message, "table"), "public.accounts");
331        let key_columns = list_field(&key_message, "columns");
332        let id_column = message_value(&key_columns[0]);
333        assert_eq!(string_field(id_column, "name"), "id");
334        assert_eq!(string_field(id_column, "kind"), "int");
335        assert_eq!(i64_field(id_column, "int_value"), 42);
336
337        let value_message = DynamicMessage::decode(
338            message_descriptor(&pool, ENTITY_DIFFERENCE).expect("value descriptor"),
339            value_body,
340        )
341        .expect("value body decodes");
342        assert_eq!(string_field(&value_message, "table"), "public.accounts");
343        assert_eq!(string_field(&value_message, "operation"), "update");
344        assert_eq!(string_field(&value_message, "lsn"), "0/2A");
345
346        let after = list_field(&value_message, "after");
347        let name_column = message_value(&after[0]);
348        assert_eq!(string_field(name_column, "name"), "name");
349        assert_eq!(string_field(name_column, "kind"), "text");
350        assert_eq!(string_field(name_column, "string_value"), "new");
351
352        let before = list_field(&value_message, "before");
353        let null_column = message_value(&before[1]);
354        assert_eq!(string_field(null_column, "name"), "nickname");
355        assert_eq!(string_field(null_column, "kind"), "null");
356        assert!(bool_field(null_column, "is_null"));
357
358        let avatar_column = message_value(&after[1]);
359        assert_eq!(string_field(avatar_column, "name"), "avatar");
360        assert_eq!(string_field(avatar_column, "kind"), "bytes");
361        assert_eq!(bytes_field(avatar_column, "bytes_value").as_ref(), b"abc");
362
363        let unchanged_toast_column = message_value(&after[2]);
364        assert_eq!(string_field(unchanged_toast_column, "name"), "details");
365        assert_eq!(
366            string_field(unchanged_toast_column, "kind"),
367            "unchanged_toast"
368        );
369    }
370
371    #[test]
372    fn decoded_int8_key_encodes_as_int_scalar_kind() {
373        let mut cache = RelationCache::default();
374        cache.apply_relation(RelationEvent {
375            relation_id: 7,
376            schema: "public".to_owned(),
377            table: "orders".to_owned(),
378            columns: vec![ColumnSchema {
379                name: "id".to_owned(),
380                type_name: "int8".to_owned(),
381                key: true,
382            }],
383        });
384        let difference = cache
385            .translate(RowEvent {
386                relation_id: 7,
387                lsn: PgLsn(0x2a),
388                commit_lsn: None,
389                txid: None,
390                commit_timestamp_ms: None,
391                kind: RowEventKind::Insert,
392                values: vec![ColumnValue {
393                    name: "col0".to_owned(),
394                    value: ScalarValue::Text("42".to_owned()),
395                }],
396            })
397            .expect("decoded row should translate");
398        assert_eq!(difference.key.columns[0].value, ScalarValue::Int(42));
399
400        let encoder = PostgresProtoEncoder::new().expect("encoder builds descriptors");
401        let key = encoder.encode_key(&difference.key).expect("key encodes");
402        let pool = DescriptorPool::from_file_descriptor_set(schema_descriptor_set())
403            .expect("descriptor pool builds");
404        let (_, _, key_body) = decode_protobuf(&key).expect("key frame decodes");
405        let key_message = DynamicMessage::decode(
406            message_descriptor(&pool, ENTITY_KEY).expect("key descriptor"),
407            key_body,
408        )
409        .expect("key body decodes");
410        let key_columns = list_field(&key_message, "columns");
411        let id_column = message_value(&key_columns[0]);
412
413        assert_eq!(string_field(id_column, "kind"), "int");
414        assert_eq!(i64_field(id_column, "int_value"), 42);
415    }
416
417    #[test]
418    fn descriptor_fields_have_expected_proto3_labels() {
419        let descriptor_set = schema_descriptor_set();
420        let file = descriptor_set.file.first().expect("descriptor file");
421        let column_value = file
422            .message_type
423            .iter()
424            .find(|message| message.name.as_deref() == Some(COLUMN_VALUE))
425            .expect("column value message");
426        let entity_key = file
427            .message_type
428            .iter()
429            .find(|message| message.name.as_deref() == Some(ENTITY_KEY))
430            .expect("entity key message");
431        let entity_difference = file
432            .message_type
433            .iter()
434            .find(|message| message.name.as_deref() == Some(ENTITY_DIFFERENCE))
435            .expect("entity difference message");
436
437        assert_eq!(
438            column_value
439                .field
440                .iter()
441                .find(|field| field.name.as_deref() == Some("name"))
442                .and_then(|field| field.label),
443            Some(prost_reflect::prost_types::field_descriptor_proto::Label::Optional as i32)
444        );
445        assert_eq!(
446            entity_key
447                .field
448                .iter()
449                .find(|field| field.name.as_deref() == Some("columns"))
450                .and_then(|field| field.label),
451            Some(prost_reflect::prost_types::field_descriptor_proto::Label::Repeated as i32)
452        );
453        assert_eq!(
454            entity_difference
455                .field
456                .iter()
457                .find(|field| field.name.as_deref() == Some("key"))
458                .and_then(|field| field.label),
459            Some(prost_reflect::prost_types::field_descriptor_proto::Label::Optional as i32)
460        );
461    }
462
463    fn sample_difference() -> EntityDifference {
464        let key = EntityKey {
465            table: "public.accounts".to_owned(),
466            columns: vec![ColumnValue {
467                name: "id".to_owned(),
468                value: ScalarValue::Int(42),
469            }],
470        };
471
472        EntityDifference {
473            table: "public.accounts".to_owned(),
474            key,
475            op: Operation::Update,
476            before: vec![
477                ColumnValue {
478                    name: "name".to_owned(),
479                    value: ScalarValue::Text("old".to_owned()),
480                },
481                ColumnValue {
482                    name: "nickname".to_owned(),
483                    value: ScalarValue::Null,
484                },
485            ],
486            after: vec![
487                ColumnValue {
488                    name: "name".to_owned(),
489                    value: ScalarValue::Text("new".to_owned()),
490                },
491                ColumnValue {
492                    name: "avatar".to_owned(),
493                    value: ScalarValue::Bytes(b"abc".to_vec()),
494                },
495                ColumnValue {
496                    name: "details".to_owned(),
497                    value: ScalarValue::UnchangedToast,
498                },
499            ],
500            lsn: PgLsn(42),
501            txid: Some(7),
502            commit_timestamp_ms: Some(1_700_000_000_000),
503            schema: TableSchema {
504                schema: "public".to_owned(),
505                table: "accounts".to_owned(),
506                columns: vec![
507                    ColumnSchema {
508                        name: "id".to_owned(),
509                        type_name: "int8".to_owned(),
510                        key: true,
511                    },
512                    ColumnSchema {
513                        name: "name".to_owned(),
514                        type_name: "text".to_owned(),
515                        key: false,
516                    },
517                ],
518            },
519        }
520    }
521
522    fn string_field(message: &DynamicMessage, name: &str) -> String {
523        match message
524            .get_field_by_name(name)
525            .expect("field exists")
526            .as_ref()
527        {
528            Value::String(value) => value.clone(),
529            other => panic!("field {name} was not a string: {other:?}"),
530        }
531    }
532
533    fn bool_field(message: &DynamicMessage, name: &str) -> bool {
534        match message
535            .get_field_by_name(name)
536            .expect("field exists")
537            .as_ref()
538        {
539            Value::Bool(value) => *value,
540            other => panic!("field {name} was not a bool: {other:?}"),
541        }
542    }
543
544    fn i64_field(message: &DynamicMessage, name: &str) -> i64 {
545        match message
546            .get_field_by_name(name)
547            .expect("field exists")
548            .as_ref()
549        {
550            Value::I64(value) => *value,
551            other => panic!("field {name} was not an int64: {other:?}"),
552        }
553    }
554
555    fn bytes_field(message: &DynamicMessage, name: &str) -> Bytes {
556        match message
557            .get_field_by_name(name)
558            .expect("field exists")
559            .as_ref()
560        {
561            Value::Bytes(value) => value.clone(),
562            other => panic!("field {name} was not bytes: {other:?}"),
563        }
564    }
565
566    fn list_field(message: &DynamicMessage, name: &str) -> Vec<Value> {
567        match message
568            .get_field_by_name(name)
569            .expect("field exists")
570            .as_ref()
571        {
572            Value::List(value) => value.clone(),
573            other => panic!("field {name} was not a list: {other:?}"),
574        }
575    }
576
577    fn message_value(value: &Value) -> &DynamicMessage {
578        match value {
579            Value::Message(message) => message,
580            other => panic!("value was not a message: {other:?}"),
581        }
582    }
583}
crabka_connect_postgres/schema.rs

crabka_connect_postgres/
schema.rs
