qail_pg/protocol/
encoder.rs

1//! PostgreSQL Encoder (Visitor Pattern)
2//!
3//! Compiles Qail AST into PostgreSQL wire protocol bytes.
4//! This is pure, synchronous computation - no I/O, no async.
5//!
6//! # Architecture
7//!
8//! Layer 2 of the QAIL architecture:
9//! - Input: Qail (AST)
10//! - Output: BytesMut (ready to send over the wire)
11//!
12//! The async I/O layer (Layer 3) consumes these bytes.
13
14use super::EncodeError;
15use bytes::BytesMut;
16
17/// Takes a Qail and produces wire protocol bytes.
18/// This is the "Visitor" in the visitor pattern.
19pub struct PgEncoder;
20
21impl PgEncoder {
22    /// Wire format code for text columns.
23    pub const FORMAT_TEXT: i16 = 0;
24    /// Wire format code for binary columns.
25    pub const FORMAT_BINARY: i16 = 1;
26
27    #[inline(always)]
28    fn result_format_wire_len(result_format: i16) -> usize {
29        if result_format == Self::FORMAT_TEXT {
30            2 // result format count = 0
31        } else {
32            4 // result format count = 1 + one format code
33        }
34    }
35
36    #[inline(always)]
37    fn encode_result_formats_vec(content: &mut Vec<u8>, result_format: i16) {
38        if result_format == Self::FORMAT_TEXT {
39            content.extend_from_slice(&0i16.to_be_bytes());
40        } else {
41            content.extend_from_slice(&1i16.to_be_bytes());
42            content.extend_from_slice(&result_format.to_be_bytes());
43        }
44    }
45
46    #[inline(always)]
47    fn encode_result_formats_bytesmut(buf: &mut BytesMut, result_format: i16) {
48        if result_format == Self::FORMAT_TEXT {
49            buf.extend_from_slice(&0i16.to_be_bytes());
50        } else {
51            buf.extend_from_slice(&1i16.to_be_bytes());
52            buf.extend_from_slice(&result_format.to_be_bytes());
53        }
54    }
55
56    /// Encode a raw SQL string as a Simple Query message.
57    /// Wire format:
58    /// - 'Q' (1 byte) - message type
59    /// - length (4 bytes, big-endian, includes self)
60    /// - query string (null-terminated)
61    pub fn encode_query_string(sql: &str) -> BytesMut {
62        let mut buf = BytesMut::new();
63
64        // Bounds check: SQL + null terminator + 4 bytes length must fit in i32
65        let content_len = sql.len() + 1; // +1 for null terminator
66        if content_len > (i32::MAX as usize) - 4 {
67            // Return empty buffer — write will fail safely rather than
68            // producing a malformed message with overflowed length.
69            return buf;
70        }
71
72        // Message type 'Q' for Query
73        buf.extend_from_slice(b"Q");
74
75        let total_len = (content_len + 4) as i32; // +4 for length field itself
76
77        // Length (4 bytes, big-endian)
78        buf.extend_from_slice(&total_len.to_be_bytes());
79
80        // Query string
81        buf.extend_from_slice(sql.as_bytes());
82
83        // Null terminator
84        buf.extend_from_slice(&[0]);
85
86        buf
87    }
88
89    /// Encode a Terminate message to close the connection.
90    pub fn encode_terminate() -> BytesMut {
91        let mut buf = BytesMut::new();
92        buf.extend_from_slice(&[b'X', 0, 0, 0, 4]);
93        buf
94    }
95
96    /// Encode a Sync message (end of pipeline in extended query protocol).
97    pub fn encode_sync() -> BytesMut {
98        let mut buf = BytesMut::new();
99        buf.extend_from_slice(&[b'S', 0, 0, 0, 4]);
100        buf
101    }
102
103    // ==================== Extended Query Protocol ====================
104
105    /// Encode a Parse message (prepare a statement).
106    /// Wire format:
107    /// - 'P' (1 byte) - message type
108    /// - length (4 bytes)
109    /// - statement name (null-terminated, "" for unnamed)
110    /// - query string (null-terminated)
111    /// - parameter count (2 bytes)
112    /// - parameter OIDs (4 bytes each, 0 = infer type)
113    pub fn encode_parse(name: &str, sql: &str, param_types: &[u32]) -> BytesMut {
114        let mut buf = BytesMut::new();
115
116        // Message type 'P'
117        buf.extend_from_slice(b"P");
118
119        let mut content = Vec::new();
120
121        // Statement name (null-terminated)
122        content.extend_from_slice(name.as_bytes());
123        content.push(0);
124
125        // Query string (null-terminated)
126        content.extend_from_slice(sql.as_bytes());
127        content.push(0);
128
129        // Parameter count
130        content.extend_from_slice(&(param_types.len() as i16).to_be_bytes());
131
132        // Parameter OIDs
133        for &oid in param_types {
134            content.extend_from_slice(&oid.to_be_bytes());
135        }
136
137        // Length (includes length field itself)
138        let len = (content.len() + 4) as i32;
139        buf.extend_from_slice(&len.to_be_bytes());
140        buf.extend_from_slice(&content);
141
142        buf
143    }
144
145    /// Encode a Bind message (bind parameters to a prepared statement).
146    /// Wire format:
147    /// - 'B' (1 byte) - message type
148    /// - length (4 bytes)
149    /// - portal name (null-terminated)
150    /// - statement name (null-terminated)
151    /// - format code count (2 bytes) - we use 0 (all text)
152    /// - parameter count (2 bytes)
153    /// - for each parameter: length (4 bytes, -1 for NULL), data
154    /// - result format count + codes
155    ///
156    /// # Arguments
157    ///
158    /// * `portal` — Destination portal name (empty string for unnamed).
159    /// * `statement` — Source prepared statement name (empty string for unnamed).
160    /// * `params` — Parameter values; `None` entries encode as SQL NULL.
161    pub fn encode_bind(
162        portal: &str,
163        statement: &str,
164        params: &[Option<Vec<u8>>],
165    ) -> Result<BytesMut, EncodeError> {
166        Self::encode_bind_with_result_format(portal, statement, params, Self::FORMAT_TEXT)
167    }
168
169    /// Encode a Bind message with explicit result-column format.
170    ///
171    /// `result_format` is PostgreSQL wire format code: `0 = text`, `1 = binary`.
172    /// For `0`, this encodes "result format count = 0" (server default text).
173    /// For non-zero codes, this encodes one explicit result format code.
174    pub fn encode_bind_with_result_format(
175        portal: &str,
176        statement: &str,
177        params: &[Option<Vec<u8>>],
178        result_format: i16,
179    ) -> Result<BytesMut, EncodeError> {
180        if params.len() > i16::MAX as usize {
181            return Err(EncodeError::TooManyParameters(params.len()));
182        }
183
184        let mut buf = BytesMut::new();
185
186        // Message type 'B'
187        buf.extend_from_slice(b"B");
188
189        let mut content = Vec::new();
190
191        // Portal name (null-terminated)
192        content.extend_from_slice(portal.as_bytes());
193        content.push(0);
194
195        // Statement name (null-terminated)
196        content.extend_from_slice(statement.as_bytes());
197        content.push(0);
198
199        // Format codes count (0 = use default text format)
200        content.extend_from_slice(&0i16.to_be_bytes());
201
202        // Parameter count
203        content.extend_from_slice(&(params.len() as i16).to_be_bytes());
204
205        // Parameters
206        for param in params {
207            match param {
208                None => {
209                    // NULL: length = -1
210                    content.extend_from_slice(&(-1i32).to_be_bytes());
211                }
212                Some(data) => {
213                    if data.len() > i32::MAX as usize {
214                        return Err(EncodeError::MessageTooLarge(data.len()));
215                    }
216                    content.extend_from_slice(&(data.len() as i32).to_be_bytes());
217                    content.extend_from_slice(data);
218                }
219            }
220        }
221
222        // Result format codes: default text (count=0) or explicit code.
223        Self::encode_result_formats_vec(&mut content, result_format);
224
225        // Length
226        let len = (content.len() + 4) as i32;
227        buf.extend_from_slice(&len.to_be_bytes());
228        buf.extend_from_slice(&content);
229
230        Ok(buf)
231    }
232
233    /// Encode an Execute message (execute a bound portal).
234    /// Wire format:
235    /// - 'E' (1 byte) - message type
236    /// - length (4 bytes)
237    /// - portal name (null-terminated)
238    /// - max rows (4 bytes, 0 = unlimited)
239    pub fn encode_execute(portal: &str, max_rows: i32) -> BytesMut {
240        let mut buf = BytesMut::new();
241
242        // Message type 'E'
243        buf.extend_from_slice(b"E");
244
245        let mut content = Vec::new();
246
247        // Portal name (null-terminated)
248        content.extend_from_slice(portal.as_bytes());
249        content.push(0);
250
251        // Max rows
252        content.extend_from_slice(&max_rows.to_be_bytes());
253
254        // Length
255        let len = (content.len() + 4) as i32;
256        buf.extend_from_slice(&len.to_be_bytes());
257        buf.extend_from_slice(&content);
258
259        buf
260    }
261
262    /// Encode a Describe message (get statement/portal metadata).
263    /// Wire format:
264    /// - 'D' (1 byte) - message type
265    /// - length (4 bytes)
266    /// - 'S' for statement or 'P' for portal
267    /// - name (null-terminated)
268    pub fn encode_describe(is_portal: bool, name: &str) -> BytesMut {
269        let mut buf = BytesMut::new();
270
271        // Message type 'D'
272        buf.extend_from_slice(b"D");
273
274        let mut content = Vec::new();
275
276        // Type: 'S' for statement, 'P' for portal
277        content.push(if is_portal { b'P' } else { b'S' });
278
279        // Name (null-terminated)
280        content.extend_from_slice(name.as_bytes());
281        content.push(0);
282
283        // Length
284        let len = (content.len() + 4) as i32;
285        buf.extend_from_slice(&len.to_be_bytes());
286        buf.extend_from_slice(&content);
287
288        buf
289    }
290
291    /// Encode a complete extended query pipeline (OPTIMIZED).
292    /// This combines Parse + Bind + Execute + Sync in a single buffer.
293    /// Zero intermediate allocations - writes directly to pre-sized BytesMut.
294    pub fn encode_extended_query(
295        sql: &str,
296        params: &[Option<Vec<u8>>],
297    ) -> Result<BytesMut, EncodeError> {
298        Self::encode_extended_query_with_result_format(sql, params, Self::FORMAT_TEXT)
299    }
300
301    /// Encode a complete extended query pipeline with explicit result format.
302    ///
303    /// `result_format` is PostgreSQL wire format code: `0 = text`, `1 = binary`.
304    pub fn encode_extended_query_with_result_format(
305        sql: &str,
306        params: &[Option<Vec<u8>>],
307        result_format: i16,
308    ) -> Result<BytesMut, EncodeError> {
309        if params.len() > i16::MAX as usize {
310            return Err(EncodeError::TooManyParameters(params.len()));
311        }
312
313        // Calculate total size upfront to avoid reallocations
314        // Bind: 1 + 4 + 1 + 1 + 2 + 2 + params_data + result_formats
315        // Execute: 1 + 4 + 1 + 4 = 10
316        // Sync: 5
317        let params_size: usize = params
318            .iter()
319            .map(|p| 4 + p.as_ref().map_or(0, |v| v.len()))
320            .sum();
321        let result_formats_size = Self::result_format_wire_len(result_format);
322        let total_size = 9 + sql.len() + (11 + params_size + result_formats_size) + 10 + 5;
323
324        let mut buf = BytesMut::with_capacity(total_size);
325
326        // ===== PARSE =====
327        buf.extend_from_slice(b"P");
328        let parse_len = (1 + sql.len() + 1 + 2 + 4) as i32; // name + sql + param_count
329        buf.extend_from_slice(&parse_len.to_be_bytes());
330        buf.extend_from_slice(&[0]); // Unnamed statement
331        buf.extend_from_slice(sql.as_bytes());
332        buf.extend_from_slice(&[0]); // Null terminator
333        buf.extend_from_slice(&0i16.to_be_bytes()); // No param types (infer)
334
335        // ===== BIND =====
336        buf.extend_from_slice(b"B");
337        let bind_len = (1 + 1 + 2 + 2 + params_size + result_formats_size + 4) as i32;
338        buf.extend_from_slice(&bind_len.to_be_bytes());
339        buf.extend_from_slice(&[0]); // Unnamed portal
340        buf.extend_from_slice(&[0]); // Unnamed statement
341        buf.extend_from_slice(&0i16.to_be_bytes()); // Format codes (default text)
342        buf.extend_from_slice(&(params.len() as i16).to_be_bytes());
343        for param in params {
344            match param {
345                None => buf.extend_from_slice(&(-1i32).to_be_bytes()),
346                Some(data) => {
347                    if data.len() > i32::MAX as usize {
348                        return Err(EncodeError::MessageTooLarge(data.len()));
349                    }
350                    buf.extend_from_slice(&(data.len() as i32).to_be_bytes());
351                    buf.extend_from_slice(data);
352                }
353            }
354        }
355        Self::encode_result_formats_bytesmut(&mut buf, result_format);
356
357        // ===== EXECUTE =====
358        buf.extend_from_slice(b"E");
359        buf.extend_from_slice(&9i32.to_be_bytes()); // len = 4 + 1 + 4
360        buf.extend_from_slice(&[0]); // Unnamed portal
361        buf.extend_from_slice(&0i32.to_be_bytes()); // Unlimited rows
362
363        // ===== SYNC =====
364        buf.extend_from_slice(&[b'S', 0, 0, 0, 4]);
365
366        Ok(buf)
367    }
368
369    /// Encode a CopyFail message to abort a COPY IN with an error.
370    /// Wire format:
371    /// - 'f' (1 byte) - message type
372    /// - length (4 bytes)
373    /// - error message (null-terminated)
374    pub fn encode_copy_fail(reason: &str) -> BytesMut {
375        let mut buf = BytesMut::new();
376        buf.extend_from_slice(b"f");
377        let content_len = reason.len() + 1; // +1 for null terminator
378        let len = (content_len + 4) as i32;
379        buf.extend_from_slice(&len.to_be_bytes());
380        buf.extend_from_slice(reason.as_bytes());
381        buf.extend_from_slice(&[0]);
382        buf
383    }
384
385    /// Encode a Close message to release a prepared statement or portal.
386    /// Wire format:
387    /// - 'C' (1 byte) - message type
388    /// - length (4 bytes)
389    /// - 'S' for statement or 'P' for portal
390    /// - name (null-terminated)
391    pub fn encode_close(is_portal: bool, name: &str) -> BytesMut {
392        let mut buf = BytesMut::new();
393        buf.extend_from_slice(b"C");
394        let content_len = 1 + name.len() + 1; // type + name + null
395        let len = (content_len + 4) as i32;
396        buf.extend_from_slice(&len.to_be_bytes());
397        buf.extend_from_slice(&[if is_portal { b'P' } else { b'S' }]);
398        buf.extend_from_slice(name.as_bytes());
399        buf.extend_from_slice(&[0]);
400        buf
401    }
402}
403
404// ==================== ULTRA-OPTIMIZED Hot Path Encoders ====================
405//
406// These encoders are designed to beat C:
407// - Direct integer writes (no temp arrays, no bounds checks)
408// - Borrowed slice params (zero-copy)
409// - Single store instructions via BufMut
410//
411
412use bytes::BufMut;
413
414/// Zero-copy parameter for ultra-fast encoding.
415/// Uses borrowed slices to avoid any allocation or copy.
416pub enum Param<'a> {
417    /// SQL NULL value.
418    Null,
419    /// Non-null parameter as a borrowed byte slice.
420    Bytes(&'a [u8]),
421}
422
423impl PgEncoder {
424    /// Direct i32 write - no temp array, no bounds check.
425    /// LLVM emits a single store instruction.
426    #[inline(always)]
427    fn put_i32_be(buf: &mut BytesMut, v: i32) {
428        buf.put_i32(v);
429    }
430
431    #[inline(always)]
432    fn put_i16_be(buf: &mut BytesMut, v: i16) {
433        buf.put_i16(v);
434    }
435
436    /// Encode Bind message - ULTRA OPTIMIZED.
437    /// - Direct integer writes (no temp arrays)
438    /// - Borrowed params (zero-copy)
439    /// - Single allocation check
440    #[inline]
441    pub fn encode_bind_ultra<'a>(
442        buf: &mut BytesMut,
443        statement: &str,
444        params: &[Param<'a>],
445    ) -> Result<(), EncodeError> {
446        Self::encode_bind_ultra_with_result_format(buf, statement, params, Self::FORMAT_TEXT)
447    }
448
449    /// Encode Bind message with explicit result-column format.
450    #[inline]
451    pub fn encode_bind_ultra_with_result_format<'a>(
452        buf: &mut BytesMut,
453        statement: &str,
454        params: &[Param<'a>],
455        result_format: i16,
456    ) -> Result<(), EncodeError> {
457        if params.len() > i16::MAX as usize {
458            return Err(EncodeError::TooManyParameters(params.len()));
459        }
460
461        // Calculate content length upfront
462        let params_size: usize = params
463            .iter()
464            .map(|p| match p {
465                Param::Null => 4,
466                Param::Bytes(b) => 4 + b.len(),
467            })
468            .sum();
469        let result_formats_size = Self::result_format_wire_len(result_format);
470        let content_len = 1 + statement.len() + 1 + 2 + 2 + params_size + result_formats_size;
471
472        // Single reserve - no more allocations
473        buf.reserve(1 + 4 + content_len);
474
475        // Message type 'B'
476        buf.put_u8(b'B');
477
478        // Length (includes itself) - DIRECT WRITE
479        Self::put_i32_be(buf, (content_len + 4) as i32);
480
481        // Portal name (empty, null-terminated)
482        buf.put_u8(0);
483
484        // Statement name (null-terminated)
485        buf.extend_from_slice(statement.as_bytes());
486        buf.put_u8(0);
487
488        // Format codes count (0 = default text)
489        Self::put_i16_be(buf, 0);
490
491        // Parameter count
492        Self::put_i16_be(buf, params.len() as i16);
493
494        // Parameters - ZERO COPY from borrowed slices
495        for param in params {
496            match param {
497                Param::Null => Self::put_i32_be(buf, -1),
498                Param::Bytes(data) => {
499                    if data.len() > i32::MAX as usize {
500                        return Err(EncodeError::MessageTooLarge(data.len()));
501                    }
502                    Self::put_i32_be(buf, data.len() as i32);
503                    buf.extend_from_slice(data);
504                }
505            }
506        }
507
508        // Result format codes
509        Self::encode_result_formats_bytesmut(buf, result_format);
510        Ok(())
511    }
512
513    /// Encode Execute message - ULTRA OPTIMIZED.
514    #[inline(always)]
515    pub fn encode_execute_ultra(buf: &mut BytesMut) {
516        // Execute: 'E' + len(9) + portal("") + max_rows(0)
517        // = 'E' 00 00 00 09 00 00 00 00 00
518        buf.extend_from_slice(&[b'E', 0, 0, 0, 9, 0, 0, 0, 0, 0]);
519    }
520
521    /// Encode Sync message - ULTRA OPTIMIZED.
522    #[inline(always)]
523    pub fn encode_sync_ultra(buf: &mut BytesMut) {
524        buf.extend_from_slice(&[b'S', 0, 0, 0, 4]);
525    }
526
527    // Keep the original methods for compatibility
528
529    /// Encode Bind message directly into existing buffer (ZERO ALLOCATION).
530    /// This is the hot path optimization - no intermediate Vec allocation.
531    #[inline]
532    pub fn encode_bind_to(
533        buf: &mut BytesMut,
534        statement: &str,
535        params: &[Option<Vec<u8>>],
536    ) -> Result<(), EncodeError> {
537        Self::encode_bind_to_with_result_format(buf, statement, params, Self::FORMAT_TEXT)
538    }
539
540    /// Encode Bind into existing buffer with explicit result-column format.
541    #[inline]
542    pub fn encode_bind_to_with_result_format(
543        buf: &mut BytesMut,
544        statement: &str,
545        params: &[Option<Vec<u8>>],
546        result_format: i16,
547    ) -> Result<(), EncodeError> {
548        if params.len() > i16::MAX as usize {
549            return Err(EncodeError::TooManyParameters(params.len()));
550        }
551
552        // Calculate content length upfront
553        // portal(1) + statement(len+1) + format_codes(2) + param_count(2)
554        // + params_data + result_formats(2 or 4)
555        let params_size: usize = params
556            .iter()
557            .map(|p| 4 + p.as_ref().map_or(0, |v| v.len()))
558            .sum();
559        let result_formats_size = Self::result_format_wire_len(result_format);
560        let content_len = 1 + statement.len() + 1 + 2 + 2 + params_size + result_formats_size;
561
562        buf.reserve(1 + 4 + content_len);
563
564        // Message type 'B'
565        buf.put_u8(b'B');
566
567        // Length (includes itself) - DIRECT WRITE
568        Self::put_i32_be(buf, (content_len + 4) as i32);
569
570        // Portal name (empty, null-terminated)
571        buf.put_u8(0);
572
573        // Statement name (null-terminated)
574        buf.extend_from_slice(statement.as_bytes());
575        buf.put_u8(0);
576
577        // Format codes count (0 = default text)
578        Self::put_i16_be(buf, 0);
579
580        // Parameter count
581        Self::put_i16_be(buf, params.len() as i16);
582
583        // Parameters
584        for param in params {
585            match param {
586                None => Self::put_i32_be(buf, -1),
587                Some(data) => {
588                    if data.len() > i32::MAX as usize {
589                        return Err(EncodeError::MessageTooLarge(data.len()));
590                    }
591                    Self::put_i32_be(buf, data.len() as i32);
592                    buf.extend_from_slice(data);
593                }
594            }
595        }
596
597        // Result format codes
598        Self::encode_result_formats_bytesmut(buf, result_format);
599        Ok(())
600    }
601
602    /// Encode Execute message directly into existing buffer (ZERO ALLOCATION).
603    #[inline]
604    pub fn encode_execute_to(buf: &mut BytesMut) {
605        // Content: portal(1) + max_rows(4) = 5 bytes
606        buf.extend_from_slice(&[b'E', 0, 0, 0, 9, 0, 0, 0, 0, 0]);
607    }
608
609    /// Encode Sync message directly into existing buffer (ZERO ALLOCATION).
610    #[inline]
611    pub fn encode_sync_to(buf: &mut BytesMut) {
612        buf.extend_from_slice(&[b'S', 0, 0, 0, 4]);
613    }
614}
615
616#[cfg(test)]
617mod tests {
618    use super::*;
619
620    // NOTE: test_encode_simple_query removed - use AstEncoder instead
621    #[test]
622    fn test_encode_query_string() {
623        let sql = "SELECT 1";
624        let bytes = PgEncoder::encode_query_string(sql);
625
626        // Message type
627        assert_eq!(bytes[0], b'Q');
628
629        // Length: 4 (length field) + 8 (query) + 1 (null) = 13
630        let len = i32::from_be_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]);
631        assert_eq!(len, 13);
632
633        // Query content
634        assert_eq!(&bytes[5..13], b"SELECT 1");
635
636        // Null terminator
637        assert_eq!(bytes[13], 0);
638    }
639
640    #[test]
641    fn test_encode_terminate() {
642        let bytes = PgEncoder::encode_terminate();
643        assert_eq!(bytes.as_ref(), &[b'X', 0, 0, 0, 4]);
644    }
645
646    #[test]
647    fn test_encode_sync() {
648        let bytes = PgEncoder::encode_sync();
649        assert_eq!(bytes.as_ref(), &[b'S', 0, 0, 0, 4]);
650    }
651
652    #[test]
653    fn test_encode_parse() {
654        let bytes = PgEncoder::encode_parse("", "SELECT $1", &[]);
655
656        // Message type 'P'
657        assert_eq!(bytes[0], b'P');
658
659        // Content should include query
660        let content = String::from_utf8_lossy(&bytes[5..]);
661        assert!(content.contains("SELECT $1"));
662    }
663
664    #[test]
665    fn test_encode_bind() {
666        let params = vec![
667            Some(b"42".to_vec()),
668            None, // NULL
669        ];
670        let bytes = PgEncoder::encode_bind("", "", &params).unwrap();
671
672        // Message type 'B'
673        assert_eq!(bytes[0], b'B');
674
675        // Should have proper length
676        let len = i32::from_be_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]);
677        assert!(len > 4); // At least header
678    }
679
680    #[test]
681    fn test_encode_bind_binary_result_format() {
682        let bytes =
683            PgEncoder::encode_bind_with_result_format("", "", &[], PgEncoder::FORMAT_BINARY)
684                .unwrap();
685
686        // B + len + portal + statement + param formats + param count + result formats
687        // Result format section for binary should be: count=1, format=1.
688        assert_eq!(&bytes[11..15], &[0, 1, 0, 1]);
689    }
690
691    #[test]
692    fn test_encode_execute() {
693        let bytes = PgEncoder::encode_execute("", 0);
694
695        // Message type 'E'
696        assert_eq!(bytes[0], b'E');
697
698        // Length: 4 + 1 (null) + 4 (max_rows) = 9
699        let len = i32::from_be_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]);
700        assert_eq!(len, 9);
701    }
702
703    #[test]
704    fn test_encode_extended_query() {
705        let params = vec![Some(b"hello".to_vec())];
706        let bytes = PgEncoder::encode_extended_query("SELECT $1", &params).unwrap();
707
708        // Should contain all 4 message types: P, B, E, S
709        assert!(bytes.windows(1).any(|w| w == [b'P']));
710        assert!(bytes.windows(1).any(|w| w == [b'B']));
711        assert!(bytes.windows(1).any(|w| w == [b'E']));
712        assert!(bytes.windows(1).any(|w| w == [b'S']));
713    }
714
715    #[test]
716    fn test_encode_extended_query_binary_result_format() {
717        let bytes = PgEncoder::encode_extended_query_with_result_format(
718            "SELECT 1",
719            &[],
720            PgEncoder::FORMAT_BINARY,
721        )
722        .unwrap();
723
724        let parse_len = i32::from_be_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]) as usize;
725        let bind_start = 1 + parse_len;
726        assert_eq!(bytes[bind_start], b'B');
727
728        let bind_len = i32::from_be_bytes([
729            bytes[bind_start + 1],
730            bytes[bind_start + 2],
731            bytes[bind_start + 3],
732            bytes[bind_start + 4],
733        ]);
734        assert_eq!(bind_len, 14);
735
736        let bind_content = &bytes[bind_start + 5..bind_start + 1 + bind_len as usize];
737        assert_eq!(&bind_content[6..10], &[0, 1, 0, 1]);
738    }
739
740    #[test]
741    fn test_encode_copy_fail() {
742        let bytes = PgEncoder::encode_copy_fail("bad data");
743        assert_eq!(bytes[0], b'f');
744        let len = i32::from_be_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]);
745        assert_eq!(len as usize, 4 + "bad data".len() + 1);
746        assert_eq!(&bytes[5..13], b"bad data");
747        assert_eq!(bytes[13], 0);
748    }
749
750    #[test]
751    fn test_encode_close_statement() {
752        let bytes = PgEncoder::encode_close(false, "my_stmt");
753        assert_eq!(bytes[0], b'C');
754        assert_eq!(bytes[5], b'S'); // Statement type
755        assert_eq!(&bytes[6..13], b"my_stmt");
756        assert_eq!(bytes[13], 0);
757    }
758
759    #[test]
760    fn test_encode_close_portal() {
761        let bytes = PgEncoder::encode_close(true, "");
762        assert_eq!(bytes[0], b'C');
763        assert_eq!(bytes[5], b'P'); // Portal type
764        assert_eq!(bytes[6], 0); // Empty name null terminator
765    }
766
767    #[test]
768    fn test_encode_bind_to_binary_result_format() {
769        let mut buf = BytesMut::new();
770        PgEncoder::encode_bind_to_with_result_format(&mut buf, "", &[], PgEncoder::FORMAT_BINARY)
771            .unwrap();
772
773        assert_eq!(&buf[11..15], &[0, 1, 0, 1]);
774    }
775
776    #[test]
777    fn test_encode_bind_ultra_binary_result_format() {
778        let mut buf = BytesMut::new();
779        PgEncoder::encode_bind_ultra_with_result_format(
780            &mut buf,
781            "",
782            &[],
783            PgEncoder::FORMAT_BINARY,
784        )
785        .unwrap();
786
787        assert_eq!(&buf[11..15], &[0, 1, 0, 1]);
788    }
789}
qail_pg/protocol/encoder.rs

qail_pg/protocol/
encoder.rs
