Skip to main content

drasi_source_postgres/
types.rs

1// Copyright 2025 The Drasi Authors.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use chrono::{DateTime, NaiveDateTime, Utc};
16use drasi_core::models::ElementValue;
17use ordered_float::OrderedFloat;
18use postgres_types::Oid;
19use rust_decimal::prelude::ToPrimitive;
20use rust_decimal::Decimal;
21use serde_json::Value as JsonValue;
22use std::collections::HashMap;
23use std::sync::Arc;
24use uuid::Uuid;
25
26#[derive(Debug, Clone)]
27pub enum PostgresValue {
28    Null,
29    Bool(bool),
30    Int2(i16),
31    Int4(i32),
32    Int8(i64),
33    Float4(f32),
34    Float8(f64),
35    Numeric(Decimal),
36    Text(String),
37    Varchar(String),
38    Char(String),
39    Uuid(Uuid),
40    Timestamp(NaiveDateTime),
41    TimestampTz(DateTime<Utc>),
42    Date(chrono::NaiveDate),
43    Time(chrono::NaiveTime),
44    Json(JsonValue),
45    Jsonb(JsonValue),
46    Array(Vec<PostgresValue>),
47    Composite(HashMap<String, PostgresValue>),
48    Bytea(Vec<u8>),
49}
50
51#[derive(Debug, Clone)]
52pub struct ColumnInfo {
53    pub name: String,
54    pub type_oid: Oid,
55    pub type_modifier: i32,
56    pub is_key: bool,
57}
58
59#[derive(Debug, Clone)]
60pub struct RelationInfo {
61    pub id: u32,
62    pub namespace: String,
63    pub name: String,
64    pub replica_identity: ReplicaIdentity,
65    pub columns: Vec<ColumnInfo>,
66}
67
68#[derive(Debug, Clone, Copy, PartialEq)]
69pub enum ReplicaIdentity {
70    Default,
71    Nothing,
72    Full,
73    Index,
74}
75
76#[derive(Debug, Clone)]
77pub struct TransactionInfo {
78    pub xid: u32,
79    pub commit_lsn: u64,
80    pub commit_timestamp: DateTime<Utc>,
81}
82
83#[derive(Debug, Clone)]
84pub enum WalMessage {
85    Begin(TransactionInfo),
86    Commit(TransactionInfo),
87    Relation(RelationInfo),
88    Insert {
89        relation_id: u32,
90        tuple: Vec<PostgresValue>,
91    },
92    Update {
93        relation_id: u32,
94        old_tuple: Option<Vec<PostgresValue>>,
95        new_tuple: Vec<PostgresValue>,
96    },
97    Delete {
98        relation_id: u32,
99        old_tuple: Vec<PostgresValue>,
100    },
101    Truncate {
102        relation_ids: Vec<u32>,
103    },
104}
105
106#[derive(Debug, Clone)]
107pub struct ReplicationSlotInfo {
108    pub slot_name: String,
109    pub consistent_point: String,
110    pub snapshot_name: Option<String>,
111    pub output_plugin: String,
112    /// The earliest WAL position retained by this slot (`restart_lsn` from pg_replication_slots).
113    /// Only populated when querying an existing slot via `get_replication_slot_info`.
114    pub restart_lsn: Option<String>,
115}
116
117#[derive(Debug, Clone)]
118pub struct StandbyStatusUpdate {
119    pub write_lsn: u64,
120    pub flush_lsn: u64,
121    pub apply_lsn: u64,
122    pub reply_requested: bool,
123}
124
125impl PostgresValue {
126    pub fn to_json(&self) -> JsonValue {
127        match self {
128            PostgresValue::Null => JsonValue::Null,
129            PostgresValue::Bool(b) => JsonValue::Bool(*b),
130            PostgresValue::Int2(i) => JsonValue::Number((*i).into()),
131            PostgresValue::Int4(i) => JsonValue::Number((*i).into()),
132            PostgresValue::Int8(i) => JsonValue::Number((*i).into()),
133            PostgresValue::Float4(f) => {
134                if let Some(n) = serde_json::Number::from_f64(*f as f64) {
135                    JsonValue::Number(n)
136                } else {
137                    JsonValue::Null
138                }
139            }
140            PostgresValue::Float8(f) => {
141                if let Some(n) = serde_json::Number::from_f64(*f) {
142                    JsonValue::Number(n)
143                } else {
144                    JsonValue::Null
145                }
146            }
147            PostgresValue::Numeric(d) => {
148                // Convert Decimal to Number via string parsing
149                // This ensures precision is maintained and the value is valid JSON number
150                d.to_string()
151                    .parse::<serde_json::Number>()
152                    .map(JsonValue::Number)
153                    .unwrap_or(JsonValue::Null)
154            }
155            PostgresValue::Text(s) | PostgresValue::Varchar(s) | PostgresValue::Char(s) => {
156                JsonValue::String(s.clone())
157            }
158            PostgresValue::Uuid(u) => JsonValue::String(u.to_string()),
159            PostgresValue::Timestamp(ts) => JsonValue::String(ts.to_string()),
160            PostgresValue::TimestampTz(ts) => JsonValue::String(ts.to_rfc3339()),
161            PostgresValue::Date(d) => JsonValue::String(d.to_string()),
162            PostgresValue::Time(t) => JsonValue::String(t.to_string()),
163            PostgresValue::Json(j) | PostgresValue::Jsonb(j) => j.clone(),
164            PostgresValue::Array(arr) => {
165                JsonValue::Array(arr.iter().map(|v| v.to_json()).collect())
166            }
167            PostgresValue::Composite(map) => {
168                let obj: serde_json::Map<String, JsonValue> =
169                    map.iter().map(|(k, v)| (k.clone(), v.to_json())).collect();
170                JsonValue::Object(obj)
171            }
172            PostgresValue::Bytea(bytes) => JsonValue::String(base64::encode(bytes)),
173        }
174    }
175
176    /// Converts this `PostgresValue` to an `ElementValue`.
177    ///
178    /// Type mapping:
179    /// - `Timestamp` -> `LocalDateTime`, `TimestampTz` -> `ZonedDateTime`
180    ///   (datetime types preserved).
181    /// - `Numeric` -> `Float` (converted through `f64`; precision may be lost
182    ///   for very large or high-precision values).
183    /// - `Date`, `Time`, `Uuid` -> `String`.
184    /// - `Json`, `Jsonb` -> `String` (JSON serialized as text).
185    /// - `Composite`, `Bytea` -> `String` (via JSON representation fallback).
186    /// - Other primitive scalar types map directly to their
187    ///   `ElementValue` counterparts.
188    pub fn to_element_value(&self) -> ElementValue {
189        match self {
190            PostgresValue::Null => ElementValue::Null,
191            PostgresValue::Bool(b) => ElementValue::Bool(*b),
192            PostgresValue::Int2(i) => ElementValue::Integer(*i as i64),
193            PostgresValue::Int4(i) => ElementValue::Integer(*i as i64),
194            PostgresValue::Int8(i) => ElementValue::Integer(*i),
195            PostgresValue::Float4(f) => ElementValue::Float(OrderedFloat(*f as f64)),
196            PostgresValue::Float8(f) => ElementValue::Float(OrderedFloat(*f)),
197            PostgresValue::Numeric(d) => {
198                // Convert Decimal directly to f64; use NaN sentinel if out of range.
199                ElementValue::Float(OrderedFloat(d.to_f64().unwrap_or(f64::NAN)))
200            }
201            PostgresValue::Text(s) | PostgresValue::Varchar(s) | PostgresValue::Char(s) => {
202                ElementValue::String(Arc::from(s.as_str()))
203            }
204            PostgresValue::Uuid(u) => ElementValue::String(Arc::from(u.to_string())),
205            PostgresValue::Timestamp(ts) => ElementValue::LocalDateTime(*ts),
206            PostgresValue::TimestampTz(ts) => ElementValue::ZonedDateTime(ts.fixed_offset()),
207            PostgresValue::Date(d) => ElementValue::String(Arc::from(d.to_string())),
208            PostgresValue::Time(t) => ElementValue::String(Arc::from(t.to_string())),
209            PostgresValue::Json(j) | PostgresValue::Jsonb(j) => {
210                ElementValue::String(Arc::from(j.to_string()))
211            }
212            PostgresValue::Array(arr) => {
213                ElementValue::List(arr.iter().map(|v| v.to_element_value()).collect())
214            }
215            PostgresValue::Composite(_) | PostgresValue::Bytea(_) => {
216                // Fall back to JSON string representation for complex types
217                ElementValue::String(Arc::from(self.to_json().to_string()))
218            }
219        }
220    }
221
222    /// Returns true if this value is null
223    pub fn is_null(&self) -> bool {
224        matches!(self, PostgresValue::Null)
225    }
226}
227
228// Add base64 encoding support
229mod base64 {
230    pub fn encode(input: &[u8]) -> String {
231        const TABLE: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
232        let mut result = String::new();
233        let mut i = 0;
234
235        while i < input.len() {
236            let b1 = input[i];
237            let b2 = if i + 1 < input.len() { input[i + 1] } else { 0 };
238            let b3 = if i + 2 < input.len() { input[i + 2] } else { 0 };
239
240            result.push(TABLE[(b1 >> 2) as usize] as char);
241            result.push(TABLE[(((b1 & 0x03) << 4) | (b2 >> 4)) as usize] as char);
242
243            if i + 1 < input.len() {
244                result.push(TABLE[(((b2 & 0x0f) << 2) | (b3 >> 6)) as usize] as char);
245            } else {
246                result.push('=');
247            }
248
249            if i + 2 < input.len() {
250                result.push(TABLE[(b3 & 0x3f) as usize] as char);
251            } else {
252                result.push('=');
253            }
254
255            i += 3;
256        }
257
258        result
259    }
260}
261
262#[cfg(test)]
263mod tests {
264    use super::*;
265    use crate::decoder::decode_column_value_text;
266    use chrono::{FixedOffset, NaiveDate, TimeZone, Utc};
267
268    fn sample_naive_datetime() -> NaiveDateTime {
269        NaiveDate::from_ymd_opt(2024, 6, 15)
270            .unwrap()
271            .and_hms_opt(10, 30, 45)
272            .unwrap()
273    }
274
275    fn sample_utc_datetime() -> DateTime<Utc> {
276        Utc.with_ymd_and_hms(2024, 6, 15, 10, 30, 45).unwrap()
277    }
278
279    #[test]
280    fn timestamp_to_element_value_is_local_datetime() {
281        let ts = sample_naive_datetime();
282        let pv = PostgresValue::Timestamp(ts);
283        let ev = pv.to_element_value();
284        assert_eq!(ev, ElementValue::LocalDateTime(ts));
285    }
286
287    #[test]
288    fn timestamptz_to_element_value_is_zoned_datetime() {
289        let ts = sample_utc_datetime();
290        let pv = PostgresValue::TimestampTz(ts);
291        let ev = pv.to_element_value();
292        assert_eq!(ev, ElementValue::ZonedDateTime(ts.fixed_offset()));
293    }
294
295    #[test]
296    fn null_to_element_value() {
297        let pv = PostgresValue::Null;
298        assert_eq!(pv.to_element_value(), ElementValue::Null);
299        assert!(pv.is_null());
300    }
301
302    #[test]
303    fn bool_to_element_value() {
304        let pv = PostgresValue::Bool(true);
305        assert_eq!(pv.to_element_value(), ElementValue::Bool(true));
306        assert!(!pv.is_null());
307    }
308
309    #[test]
310    fn int_types_to_element_value() {
311        assert_eq!(
312            PostgresValue::Int2(42).to_element_value(),
313            ElementValue::Integer(42)
314        );
315        assert_eq!(
316            PostgresValue::Int4(100_000).to_element_value(),
317            ElementValue::Integer(100_000)
318        );
319        assert_eq!(
320            PostgresValue::Int8(9_000_000_000).to_element_value(),
321            ElementValue::Integer(9_000_000_000)
322        );
323    }
324
325    #[test]
326    fn float_types_to_element_value() {
327        let ev = PostgresValue::Float4(1.23).to_element_value();
328        match ev {
329            ElementValue::Float(f) => assert!((f.into_inner() - 1.23f64).abs() < 0.001),
330            other => panic!("Expected Float, got {other:?}"),
331        }
332
333        let ev = PostgresValue::Float8(9.876543210).to_element_value();
334        match ev {
335            ElementValue::Float(f) => {
336                assert!((f.into_inner() - 9.876543210).abs() < 1e-9)
337            }
338            other => panic!("Expected Float, got {other:?}"),
339        }
340    }
341
342    #[test]
343    fn text_to_element_value() {
344        let pv = PostgresValue::Text("hello".to_string());
345        assert_eq!(
346            pv.to_element_value(),
347            ElementValue::String(Arc::from("hello"))
348        );
349    }
350
351    #[test]
352    fn array_with_timestamps_to_element_value() {
353        let ts = sample_naive_datetime();
354        let pv = PostgresValue::Array(vec![PostgresValue::Timestamp(ts), PostgresValue::Int4(42)]);
355        let ev = pv.to_element_value();
356        match ev {
357            ElementValue::List(items) => {
358                assert_eq!(items.len(), 2);
359                assert_eq!(items[0], ElementValue::LocalDateTime(ts));
360                assert_eq!(items[1], ElementValue::Integer(42));
361            }
362            other => panic!("Expected List, got {other:?}"),
363        }
364    }
365
366    #[test]
367    fn timestamp_to_json_is_string() {
368        // to_json() should still produce a string (backwards compat)
369        let ts = sample_naive_datetime();
370        let pv = PostgresValue::Timestamp(ts);
371        let json = pv.to_json();
372        assert!(json.is_string());
373    }
374
375    #[test]
376    fn timestamptz_to_json_is_string() {
377        let ts = sample_utc_datetime();
378        let pv = PostgresValue::TimestampTz(ts);
379        let json = pv.to_json();
380        assert!(json.is_string());
381    }
382
383    // ── Bootstrap / CDC parity tests ───────────────────────────────────
384    //
385    // These tests assert that `decode_column_value_text` (bootstrap path)
386    // produces the same `ElementValue` as `PostgresValue::to_element_value`
387    // (CDC path) for equivalent input data.
388
389    #[test]
390    fn parity_bool_true() {
391        let cdc = PostgresValue::Bool(true).to_element_value();
392        let bootstrap = decode_column_value_text("true", 16).unwrap();
393        assert_eq!(cdc, bootstrap, "bool true: CDC vs bootstrap mismatch");
394    }
395
396    #[test]
397    fn parity_bool_false() {
398        let cdc = PostgresValue::Bool(false).to_element_value();
399        let bootstrap = decode_column_value_text("f", 16).unwrap();
400        assert_eq!(cdc, bootstrap, "bool false: CDC vs bootstrap mismatch");
401    }
402
403    #[test]
404    fn parity_int2() {
405        let cdc = PostgresValue::Int2(42).to_element_value();
406        let bootstrap = decode_column_value_text("42", 21).unwrap();
407        assert_eq!(cdc, bootstrap, "int2: CDC vs bootstrap mismatch");
408    }
409
410    #[test]
411    fn parity_int4() {
412        let cdc = PostgresValue::Int4(100_000).to_element_value();
413        let bootstrap = decode_column_value_text("100000", 23).unwrap();
414        assert_eq!(cdc, bootstrap, "int4: CDC vs bootstrap mismatch");
415    }
416
417    #[test]
418    fn parity_int8() {
419        let cdc = PostgresValue::Int8(9_000_000_000).to_element_value();
420        let bootstrap = decode_column_value_text("9000000000", 20).unwrap();
421        assert_eq!(cdc, bootstrap, "int8: CDC vs bootstrap mismatch");
422    }
423
424    #[test]
425    fn parity_float4() {
426        let cdc = PostgresValue::Float4(1.5).to_element_value();
427        let bootstrap = decode_column_value_text("1.5", 700).unwrap();
428        assert_eq!(cdc, bootstrap, "float4: CDC vs bootstrap mismatch");
429    }
430
431    #[test]
432    fn parity_float8() {
433        let cdc = PostgresValue::Float8(9.876543210).to_element_value();
434        let bootstrap = decode_column_value_text("9.87654321", 701).unwrap();
435        assert_eq!(cdc, bootstrap, "float8: CDC vs bootstrap mismatch");
436    }
437
438    #[test]
439    fn parity_numeric() {
440        let dec = Decimal::from_str_exact("123.45").unwrap();
441        let cdc = PostgresValue::Numeric(dec).to_element_value();
442        let bootstrap = decode_column_value_text("123.45", 1700).unwrap();
443        assert_eq!(cdc, bootstrap, "numeric: CDC vs bootstrap mismatch");
444    }
445
446    #[test]
447    fn parity_text() {
448        let cdc = PostgresValue::Text("hello world".to_string()).to_element_value();
449        let bootstrap = decode_column_value_text("hello world", 25).unwrap();
450        assert_eq!(cdc, bootstrap, "text: CDC vs bootstrap mismatch");
451    }
452
453    #[test]
454    fn parity_varchar() {
455        let cdc = PostgresValue::Varchar("drasi".to_string()).to_element_value();
456        let bootstrap = decode_column_value_text("drasi", 1043).unwrap();
457        assert_eq!(cdc, bootstrap, "varchar: CDC vs bootstrap mismatch");
458    }
459
460    #[test]
461    fn parity_timestamp() {
462        // timestamp without timezone (OID 1114)
463        let dt = NaiveDate::from_ymd_opt(2024, 6, 15)
464            .unwrap()
465            .and_hms_micro_opt(10, 30, 45, 123456)
466            .unwrap();
467        let cdc = PostgresValue::Timestamp(dt).to_element_value();
468        let bootstrap = decode_column_value_text("2024-06-15 10:30:45.123456", 1114).unwrap();
469        assert_eq!(cdc, bootstrap, "timestamp: CDC vs bootstrap mismatch");
470        // Both must be LocalDateTime, not String
471        assert!(
472            matches!(cdc, ElementValue::LocalDateTime(_)),
473            "CDC timestamp should be LocalDateTime, got {cdc:?}"
474        );
475    }
476
477    #[test]
478    fn parity_timestamp_no_fractional() {
479        let dt = NaiveDate::from_ymd_opt(2024, 6, 15)
480            .unwrap()
481            .and_hms_opt(10, 30, 45)
482            .unwrap();
483        let cdc = PostgresValue::Timestamp(dt).to_element_value();
484        let bootstrap = decode_column_value_text("2024-06-15 10:30:45", 1114).unwrap();
485        assert_eq!(
486            cdc, bootstrap,
487            "timestamp (no frac): CDC vs bootstrap mismatch"
488        );
489    }
490
491    #[test]
492    fn parity_timestamptz_utc() {
493        // timestamptz (OID 1184) with UTC
494        let utc_dt = Utc.with_ymd_and_hms(2024, 6, 15, 10, 30, 45).unwrap();
495        let cdc = PostgresValue::TimestampTz(utc_dt).to_element_value();
496        let bootstrap = decode_column_value_text("2024-06-15T10:30:45+00:00", 1184).unwrap();
497        assert_eq!(cdc, bootstrap, "timestamptz UTC: CDC vs bootstrap mismatch");
498        // Both must be ZonedDateTime, not String
499        assert!(
500            matches!(cdc, ElementValue::ZonedDateTime(_)),
501            "CDC timestamptz should be ZonedDateTime, got {cdc:?}"
502        );
503    }
504
505    #[test]
506    fn parity_timestamptz_with_offset() {
507        // CDC always stores timestamptz as UTC, but the logical instant is the same.
508        // Bootstrap parses the offset from the string.
509        // Both should represent the same instant in time even if offsets differ.
510        let utc_dt = Utc.with_ymd_and_hms(2024, 6, 15, 8, 30, 45).unwrap(); // 08:30:45 UTC
511        let cdc = PostgresValue::TimestampTz(utc_dt).to_element_value();
512
513        // Same instant expressed as +02:00 → 10:30:45+02:00
514        let bootstrap = decode_column_value_text("2024-06-15T10:30:45+02:00", 1184).unwrap();
515
516        // Both are ZonedDateTime
517        match (&cdc, &bootstrap) {
518            (ElementValue::ZonedDateTime(cdc_dt), ElementValue::ZonedDateTime(bs_dt)) => {
519                // Same instant in UTC
520                assert_eq!(
521                    cdc_dt.timestamp(),
522                    bs_dt.timestamp(),
523                    "timestamptz with offset: same UTC instant expected"
524                );
525            }
526            _ => panic!(
527                "Expected ZonedDateTime from both paths, got CDC={cdc:?}, bootstrap={bootstrap:?}"
528            ),
529        }
530    }
531
532    #[test]
533    fn parity_date() {
534        let cdc =
535            PostgresValue::Date(NaiveDate::from_ymd_opt(2024, 6, 15).unwrap()).to_element_value();
536        let bootstrap = decode_column_value_text("2024-06-15", 1082).unwrap();
537        assert_eq!(cdc, bootstrap, "date: CDC vs bootstrap mismatch");
538    }
539
540    #[test]
541    fn parity_uuid() {
542        let uuid = Uuid::parse_str("550e8400-e29b-41d4-a716-446655440000").unwrap();
543        let cdc = PostgresValue::Uuid(uuid).to_element_value();
544        let bootstrap =
545            decode_column_value_text("550e8400-e29b-41d4-a716-446655440000", 2950).unwrap();
546        assert_eq!(cdc, bootstrap, "uuid: CDC vs bootstrap mismatch");
547    }
548
549    use std::str::FromStr;
550
551    #[test]
552    fn test_decimal_to_json_as_number() {
553        // Test that decimal values are serialized as numbers, not strings
554        let decimal = Decimal::from_str("123.45").unwrap();
555        let pg_value = PostgresValue::Numeric(decimal);
556        let json = pg_value.to_json();
557
558        // Should be a Number, not a String
559        assert!(json.is_number(), "Decimal should be serialized as a number");
560
561        // Verify the value is correct
562        let num = json.as_f64().unwrap();
563        assert_eq!(num, 123.45);
564    }
565
566    #[test]
567    fn test_decimal_integer_to_json() {
568        let decimal = Decimal::from_str("100").unwrap();
569        let pg_value = PostgresValue::Numeric(decimal);
570        let json = pg_value.to_json();
571
572        assert!(
573            json.is_number(),
574            "Integer decimal should be serialized as a number"
575        );
576
577        let num = json.as_f64().unwrap();
578        assert_eq!(num, 100.0);
579    }
580
581    #[test]
582    fn test_decimal_small_value_to_json() {
583        let decimal = Decimal::from_str("0.00001").unwrap();
584        let pg_value = PostgresValue::Numeric(decimal);
585        let json = pg_value.to_json();
586
587        assert!(
588            json.is_number(),
589            "Small decimal should be serialized as a number"
590        );
591
592        let num = json.as_f64().unwrap();
593        assert_eq!(num, 0.00001);
594    }
595
596    #[test]
597    fn test_decimal_negative_to_json() {
598        let decimal = Decimal::from_str("-999.99").unwrap();
599        let pg_value = PostgresValue::Numeric(decimal);
600        let json = pg_value.to_json();
601
602        assert!(
603            json.is_number(),
604            "Negative decimal should be serialized as a number"
605        );
606
607        let num = json.as_f64().unwrap();
608        assert_eq!(num, -999.99);
609    }
610
611    #[test]
612    fn test_decimal_zero_to_json() {
613        let decimal = Decimal::from_str("0").unwrap();
614        let pg_value = PostgresValue::Numeric(decimal);
615        let json = pg_value.to_json();
616
617        assert!(
618            json.is_number(),
619            "Zero decimal should be serialized as a number"
620        );
621
622        let num = json.as_f64().unwrap();
623        assert_eq!(num, 0.0);
624    }
625}