we_trust_postgres/
message.rs

1use bytes::{BufMut, BytesMut};
2use yykv_types::DsValue;
3
4#[derive(Debug, Clone)]
5pub enum Message {
6    Startup {
7        params: Vec<(String, String)>,
8    },
9    SslRequest,
10    Password(String),
11    Query(String),
12    Parse {
13        name: String,
14        query: String,
15        param_types: Vec<u32>,
16    },
17    Bind {
18        portal: String,
19        statement: String,
20        params: Vec<DsValue>,
21    },
22    Describe {
23        target_type: u8, // 'S' for statement, 'P' for portal
24        name: String,
25    },
26    Execute {
27        portal: String,
28        max_rows: i32,
29    },
30    Sync,
31    Terminate,
32    Close {
33        target_type: u8, // 'S' for statement, 'P' for portal
34        name: String,
35    },
36    Flush,
37}
38
39impl Message {
40    pub fn encode(&self, dst: &mut BytesMut) {
41        match self {
42            Message::SslRequest => {
43                dst.extend_from_slice(&8i32.to_be_bytes());
44                dst.extend_from_slice(&80877103i32.to_be_bytes());
45            }
46            Message::Startup { params } => {
47                let mut payload = BytesMut::new();
48                payload.extend_from_slice(&196608i32.to_be_bytes());
49                for (k, v) in params {
50                    payload.extend_from_slice(k.as_bytes());
51                    payload.extend_from_slice(&[0]);
52                    payload.extend_from_slice(v.as_bytes());
53                    payload.extend_from_slice(&[0]);
54                }
55                payload.extend_from_slice(&[0]);
56
57                dst.extend_from_slice(&(payload.len() as i32 + 4).to_be_bytes());
58                dst.extend_from_slice(&payload);
59            }
60            Message::Password(pass) => {
61                dst.extend_from_slice(b"p");
62                dst.extend_from_slice(&(pass.len() as i32 + 5).to_be_bytes());
63                dst.extend_from_slice(pass.as_bytes());
64                dst.extend_from_slice(&[0]);
65            }
66            Message::Query(q) => {
67                dst.extend_from_slice(b"Q");
68                dst.extend_from_slice(&(q.len() as i32 + 5).to_be_bytes());
69                dst.extend_from_slice(q.as_bytes());
70                dst.extend_from_slice(&[0]);
71            }
72            Message::Parse {
73                name,
74                query,
75                param_types,
76            } => {
77                dst.extend_from_slice(b"P");
78                let mut payload = BytesMut::new();
79                payload.extend_from_slice(name.as_bytes());
80                payload.extend_from_slice(&[0]);
81                payload.extend_from_slice(query.as_bytes());
82                payload.extend_from_slice(&[0]);
83                payload.extend_from_slice(&(param_types.len() as i16).to_be_bytes());
84                for &oid in param_types {
85                    payload.extend_from_slice(&oid.to_be_bytes());
86                }
87                dst.extend_from_slice(&(payload.len() as i32 + 4).to_be_bytes());
88                dst.extend_from_slice(&payload);
89            }
90            Message::Bind {
91                portal,
92                statement,
93                params,
94            } => {
95                dst.extend_from_slice(b"B");
96                let mut payload = BytesMut::new();
97                payload.extend_from_slice(portal.as_bytes());
98                payload.extend_from_slice(&[0]);
99                payload.extend_from_slice(statement.as_bytes());
100                payload.extend_from_slice(&[0]);
101
102                // Parameter format codes (0 for text)
103                payload.extend_from_slice(&0i16.to_be_bytes());
104
105                // Parameter values
106                payload.extend_from_slice(&(params.len() as i16).to_be_bytes());
107                for param in params {
108                    match param {
109                        DsValue::Null => {
110                            payload.extend_from_slice(&(-1i32).to_be_bytes());
111                        }
112                        _ => {
113                            // Simplified for now: encode everything as string/text
114                            let s = format!("{:?}", param);
115                            payload.extend_from_slice(&(s.len() as i32).to_be_bytes());
116                            payload.extend_from_slice(s.as_bytes());
117                        }
118                    }
119                }
120
121                // Result format codes (0 for text)
122                payload.extend_from_slice(&0i16.to_be_bytes());
123
124                dst.extend_from_slice(&(payload.len() as i32 + 4).to_be_bytes());
125                dst.extend_from_slice(&payload);
126            }
127            Message::Describe { target_type, name } => {
128                dst.extend_from_slice(b"D");
129                let mut payload = BytesMut::new();
130                payload.put_u8(*target_type);
131                payload.extend_from_slice(name.as_bytes());
132                payload.extend_from_slice(&[0]);
133                dst.extend_from_slice(&(payload.len() as i32 + 4).to_be_bytes());
134                dst.extend_from_slice(&payload);
135            }
136            Message::Execute { portal, max_rows } => {
137                dst.extend_from_slice(b"E");
138                let mut payload = BytesMut::new();
139                payload.extend_from_slice(portal.as_bytes());
140                payload.extend_from_slice(&[0]);
141                payload.extend_from_slice(&max_rows.to_be_bytes());
142                dst.extend_from_slice(&(payload.len() as i32 + 4).to_be_bytes());
143                dst.extend_from_slice(&payload);
144            }
145            Message::Sync => {
146                dst.extend_from_slice(b"S");
147                dst.extend_from_slice(&4i32.to_be_bytes());
148            }
149            Message::Terminate => {
150                dst.extend_from_slice(b"X");
151                dst.extend_from_slice(&4i32.to_be_bytes());
152            }
153            Message::Close { target_type, name } => {
154                dst.extend_from_slice(b"C");
155                let mut payload = BytesMut::new();
156                payload.put_u8(*target_type);
157                payload.extend_from_slice(name.as_bytes());
158                payload.extend_from_slice(&[0]);
159                dst.extend_from_slice(&(payload.len() as i32 + 4).to_be_bytes());
160                dst.extend_from_slice(&payload);
161            }
162            Message::Flush => {
163                dst.extend_from_slice(b"H");
164                dst.extend_from_slice(&4i32.to_be_bytes());
165            }
166        }
167    }
168}
we_trust_postgres/message.rs

we_trust_postgres/
message.rs
