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) => JsonValue::String(d.to_string()),
141            PostgresValue::Text(s) | PostgresValue::Varchar(s) | PostgresValue::Char(s) => {
142                JsonValue::String(s.clone())
143            }
144            PostgresValue::Uuid(u) => JsonValue::String(u.to_string()),
145            PostgresValue::Timestamp(ts) => JsonValue::String(ts.to_string()),
146            PostgresValue::TimestampTz(ts) => JsonValue::String(ts.to_rfc3339()),
147            PostgresValue::Date(d) => JsonValue::String(d.to_string()),
148            PostgresValue::Time(t) => JsonValue::String(t.to_string()),
149            PostgresValue::Json(j) | PostgresValue::Jsonb(j) => j.clone(),
150            PostgresValue::Array(arr) => {
151                JsonValue::Array(arr.iter().map(|v| v.to_json()).collect())
152            }
153            PostgresValue::Composite(map) => {
154                let obj: serde_json::Map<String, JsonValue> =
155                    map.iter().map(|(k, v)| (k.clone(), v.to_json())).collect();
156                JsonValue::Object(obj)
157            }
158            PostgresValue::Bytea(bytes) => JsonValue::String(base64::encode(bytes)),
159        }
160    }
161}
162
163// Add base64 encoding support
164mod base64 {
165    pub fn encode(input: &[u8]) -> String {
166        const TABLE: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
167        let mut result = String::new();
168        let mut i = 0;
169
170        while i < input.len() {
171            let b1 = input[i];
172            let b2 = if i + 1 < input.len() { input[i + 1] } else { 0 };
173            let b3 = if i + 2 < input.len() { input[i + 2] } else { 0 };
174
175            result.push(TABLE[(b1 >> 2) as usize] as char);
176            result.push(TABLE[(((b1 & 0x03) << 4) | (b2 >> 4)) as usize] as char);
177
178            if i + 1 < input.len() {
179                result.push(TABLE[(((b2 & 0x0f) << 2) | (b3 >> 6)) as usize] as char);
180            } else {
181                result.push('=');
182            }
183
184            if i + 2 < input.len() {
185                result.push(TABLE[(b3 & 0x3f) as usize] as char);
186            } else {
187                result.push('=');
188            }
189
190            i += 3;
191        }
192
193        result
194    }
195}