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 postgres_types::Oid;
17use rust_decimal::Decimal;
18use serde_json::Value as JsonValue;
19use std::collections::HashMap;
20use uuid::Uuid;
21
22#[derive(Debug, Clone)]
23pub enum PostgresValue {
24    Null,
25    Bool(bool),
26    Int2(i16),
27    Int4(i32),
28    Int8(i64),
29    Float4(f32),
30    Float8(f64),
31    Numeric(Decimal),
32    Text(String),
33    Varchar(String),
34    Char(String),
35    Uuid(Uuid),
36    Timestamp(NaiveDateTime),
37    TimestampTz(DateTime<Utc>),
38    Date(chrono::NaiveDate),
39    Time(chrono::NaiveTime),
40    Json(JsonValue),
41    Jsonb(JsonValue),
42    Array(Vec<PostgresValue>),
43    Composite(HashMap<String, PostgresValue>),
44    Bytea(Vec<u8>),
45}
46
47#[derive(Debug, Clone)]
48pub struct ColumnInfo {
49    pub name: String,
50    pub type_oid: Oid,
51    pub type_modifier: i32,
52    pub is_key: bool,
53}
54
55#[derive(Debug, Clone)]
56pub struct RelationInfo {
57    pub id: u32,
58    pub namespace: String,
59    pub name: String,
60    pub replica_identity: ReplicaIdentity,
61    pub columns: Vec<ColumnInfo>,
62}
63
64#[derive(Debug, Clone, Copy, PartialEq)]
65pub enum ReplicaIdentity {
66    Default,
67    Nothing,
68    Full,
69    Index,
70}
71
72#[derive(Debug, Clone)]
73pub struct TransactionInfo {
74    pub xid: u32,
75    pub commit_lsn: u64,
76    pub commit_timestamp: DateTime<Utc>,
77}
78
79#[derive(Debug, Clone)]
80pub enum WalMessage {
81    Begin(TransactionInfo),
82    Commit(TransactionInfo),
83    Relation(RelationInfo),
84    Insert {
85        relation_id: u32,
86        tuple: Vec<PostgresValue>,
87    },
88    Update {
89        relation_id: u32,
90        old_tuple: Option<Vec<PostgresValue>>,
91        new_tuple: Vec<PostgresValue>,
92    },
93    Delete {
94        relation_id: u32,
95        old_tuple: Vec<PostgresValue>,
96    },
97    Truncate {
98        relation_ids: Vec<u32>,
99    },
100}
101
102#[derive(Debug, Clone)]
103pub struct ReplicationSlotInfo {
104    pub slot_name: String,
105    pub consistent_point: String,
106    pub snapshot_name: Option<String>,
107    pub output_plugin: String,
108}
109
110#[derive(Debug, Clone)]
111pub struct StandbyStatusUpdate {
112    pub write_lsn: u64,
113    pub flush_lsn: u64,
114    pub apply_lsn: u64,
115    pub reply_requested: bool,
116}
117
118impl PostgresValue {
119    pub fn to_json(&self) -> JsonValue {
120        match self {
121            PostgresValue::Null => JsonValue::Null,
122            PostgresValue::Bool(b) => JsonValue::Bool(*b),
123            PostgresValue::Int2(i) => JsonValue::Number((*i).into()),
124            PostgresValue::Int4(i) => JsonValue::Number((*i).into()),
125            PostgresValue::Int8(i) => JsonValue::Number((*i).into()),
126            PostgresValue::Float4(f) => {
127                if let Some(n) = serde_json::Number::from_f64(*f as f64) {
128                    JsonValue::Number(n)
129                } else {
130                    JsonValue::Null
131                }
132            }
133            PostgresValue::Float8(f) => {
134                if let Some(n) = serde_json::Number::from_f64(*f) {
135                    JsonValue::Number(n)
136                } else {
137                    JsonValue::Null
138                }
139            }
140            PostgresValue::Numeric(d) => {
141                // Convert Decimal to Number via string parsing
142                // This ensures precision is maintained and the value is valid JSON number
143                d.to_string()
144                    .parse::<serde_json::Number>()
145                    .map(JsonValue::Number)
146                    .unwrap_or(JsonValue::Null)
147            }
148            PostgresValue::Text(s) | PostgresValue::Varchar(s) | PostgresValue::Char(s) => {
149                JsonValue::String(s.clone())
150            }
151            PostgresValue::Uuid(u) => JsonValue::String(u.to_string()),
152            PostgresValue::Timestamp(ts) => JsonValue::String(ts.to_string()),
153            PostgresValue::TimestampTz(ts) => JsonValue::String(ts.to_rfc3339()),
154            PostgresValue::Date(d) => JsonValue::String(d.to_string()),
155            PostgresValue::Time(t) => JsonValue::String(t.to_string()),
156            PostgresValue::Json(j) | PostgresValue::Jsonb(j) => j.clone(),
157            PostgresValue::Array(arr) => {
158                JsonValue::Array(arr.iter().map(|v| v.to_json()).collect())
159            }
160            PostgresValue::Composite(map) => {
161                let obj: serde_json::Map<String, JsonValue> =
162                    map.iter().map(|(k, v)| (k.clone(), v.to_json())).collect();
163                JsonValue::Object(obj)
164            }
165            PostgresValue::Bytea(bytes) => JsonValue::String(base64::encode(bytes)),
166        }
167    }
168}
169
170// Add base64 encoding support
171mod base64 {
172    pub fn encode(input: &[u8]) -> String {
173        const TABLE: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
174        let mut result = String::new();
175        let mut i = 0;
176
177        while i < input.len() {
178            let b1 = input[i];
179            let b2 = if i + 1 < input.len() { input[i + 1] } else { 0 };
180            let b3 = if i + 2 < input.len() { input[i + 2] } else { 0 };
181
182            result.push(TABLE[(b1 >> 2) as usize] as char);
183            result.push(TABLE[(((b1 & 0x03) << 4) | (b2 >> 4)) as usize] as char);
184
185            if i + 1 < input.len() {
186                result.push(TABLE[(((b2 & 0x0f) << 2) | (b3 >> 6)) as usize] as char);
187            } else {
188                result.push('=');
189            }
190
191            if i + 2 < input.len() {
192                result.push(TABLE[(b3 & 0x3f) as usize] as char);
193            } else {
194                result.push('=');
195            }
196
197            i += 3;
198        }
199
200        result
201    }
202}
203
204#[cfg(test)]
205mod tests {
206    use super::*;
207    use std::str::FromStr;
208
209    #[test]
210    fn test_decimal_to_json_as_number() {
211        // Test that decimal values are serialized as numbers, not strings
212        let decimal = Decimal::from_str("123.45").unwrap();
213        let pg_value = PostgresValue::Numeric(decimal);
214        let json = pg_value.to_json();
215
216        // Should be a Number, not a String
217        assert!(json.is_number(), "Decimal should be serialized as a number");
218
219        // Verify the value is correct
220        let num = json.as_f64().unwrap();
221        assert_eq!(num, 123.45);
222    }
223
224    #[test]
225    fn test_decimal_integer_to_json() {
226        let decimal = Decimal::from_str("100").unwrap();
227        let pg_value = PostgresValue::Numeric(decimal);
228        let json = pg_value.to_json();
229
230        assert!(
231            json.is_number(),
232            "Integer decimal should be serialized as a number"
233        );
234
235        let num = json.as_f64().unwrap();
236        assert_eq!(num, 100.0);
237    }
238
239    #[test]
240    fn test_decimal_small_value_to_json() {
241        let decimal = Decimal::from_str("0.00001").unwrap();
242        let pg_value = PostgresValue::Numeric(decimal);
243        let json = pg_value.to_json();
244
245        assert!(
246            json.is_number(),
247            "Small decimal should be serialized as a number"
248        );
249
250        let num = json.as_f64().unwrap();
251        assert_eq!(num, 0.00001);
252    }
253
254    #[test]
255    fn test_decimal_negative_to_json() {
256        let decimal = Decimal::from_str("-999.99").unwrap();
257        let pg_value = PostgresValue::Numeric(decimal);
258        let json = pg_value.to_json();
259
260        assert!(
261            json.is_number(),
262            "Negative decimal should be serialized as a number"
263        );
264
265        let num = json.as_f64().unwrap();
266        assert_eq!(num, -999.99);
267    }
268
269    #[test]
270    fn test_decimal_zero_to_json() {
271        let decimal = Decimal::from_str("0").unwrap();
272        let pg_value = PostgresValue::Numeric(decimal);
273        let json = pg_value.to_json();
274
275        assert!(
276            json.is_number(),
277            "Zero decimal should be serialized as a number"
278        );
279
280        let num = json.as_f64().unwrap();
281        assert_eq!(num, 0.0);
282    }
283}