fix(dameng): read result metadata from the frame header, not the first column

Author: datlechin

dameng-protocol/src/frame.rs

@@ -11,9 +11,20 @@
 //! 14      4     AffectedRows (i32 LE) - rows affected (for DML responses)
 //! 18      1     CompressFlag (u8)
 //! 19      1     Checksum (u8) - XOR of bytes 0-18
-//! 20      44    Reserved (zeros)
+//! 20      44    Reserved for the message type (zeros on client messages)
 //! 64      var   Payload body
 //! ```
+//!
+//! The server reuses the reserved area, and each response type spells it differently.
+//! The fields below overlap on purpose, so read the one that belongs to the message:
+//! ```text
+//! Offset  Size  Response      Field
+//! 20      8     FETCH         fetch_total (i64 LE)
+//! 22      2     EXEC/OPE ACK  column_count (u16 LE)
+//! 24      8     EXEC/OPE ACK  update_count (i64 LE)
+//! 28      4     FETCH         batch_row_count (i32 LE)
+//! 28      4     STARTUP       server_encoding (i32 LE)
+//! ```
 
 use bytes::{Buf, BufMut, BytesMut};
 
@@ -22,6 +33,10 @@ use crate::error::{Error, Result};
 /// The size of the frame header in bytes.
 pub const FRAME_HEADER_SIZE: usize = 64;
 
+/// What the server writes in a row total it cannot state yet, because rows are still
+/// queued behind the cursor. A total that is not this value is the final one.
+pub const ROW_TOTAL_UNKNOWN: i64 = i64::MAX;
+
 /// DM protocol frame header (64 bytes).
 #[derive(Debug, Clone, PartialEq)]
 pub struct Frame {
@@ -38,9 +53,17 @@ pub struct Frame {
     pub affected_rows: i32,
     /// Compression flag (0=none, 1=snappy, 2=zlib).
     pub compress_flag: u8,
-    /// Update count for DML operations, stored in the reserved area
-    /// at header offset 24 (int64 LE). Always 0 for non-DML.
+    /// Update count at header offset 24 (int64 LE): the affected count of a DML
+    /// statement, and the row total of a result set. It is
+    /// [`ROW_TOTAL_UNKNOWN`] while the server still holds rows for the cursor.
     pub update_count: u64,
+    /// Number of columns in a statement response, header offset 22 (uint16 LE).
+    pub column_count: u16,
+    /// Row total a FETCH reply reports, header offset 20 (int64 LE).
+    /// [`ROW_TOTAL_UNKNOWN`] until the batch that drains the cursor.
+    pub fetch_total: i64,
+    /// Number of rows a FETCH reply carries in its payload, header offset 28 (int32 LE).
+    pub batch_row_count: i32,
     /// Server encoding from header offset 28 (int32 LE).
     /// Used in STARTUP_RESPONSE. 0=GB18030, 1=UTF-8, 2=EUC-KR.
     pub server_encoding: u8,
@@ -57,6 +80,9 @@ impl Frame {
             affected_rows: 0,
             compress_flag: 0,
             update_count: 0,
+            column_count: 0,
+            fetch_total: 0,
+            batch_row_count: 0,
             server_encoding: 0,
         }
     }
@@ -91,25 +117,40 @@ impl Frame {
             return Err(Error::ChecksumMismatch);
         }
 
-        // Parse update_count from the reserved area at header offset 24 (int64 LE).
-        // Need to read from the raw buffer before advancing. buf currently has
-        // cursor at byte 20 (after consuming 20 bytes).
-        // Bytes [4..12] of the remaining 44-byte reserved area = absolute offset 24.
-        let update_count = if buf.remaining() >= 12 {
-            let raw = &buf.chunk()[4..12]; // offset 24-31 in absolute header
-            u64::from_le_bytes([
-                raw[0], raw[1], raw[2], raw[3], raw[4], raw[5], raw[6], raw[7],
-            ])
+        // The cursor sits at byte 20 now, so the reserved area starts at chunk()[0]
+        // and every field below is read before the buffer advances past it.
+        let reserved = if buf.remaining() >= 12 {
+            let raw = &buf.chunk()[..12];
+            let mut bytes = [0u8; 12];
+            bytes.copy_from_slice(raw);
+            bytes
         } else {
-            0
-        };
-        // Server encoding at header offset 28 (int32 LE)
-        let server_encoding = if buf.remaining() >= 12 {
-            let raw = &buf.chunk()[8..12]; // offset 28-31 in absolute header
-            u8::from_le_bytes([raw[0]])
-        } else {
-            0
+            [0u8; 12]
         };
+        let fetch_total = i64::from_le_bytes([
+            reserved[0],
+            reserved[1],
+            reserved[2],
+            reserved[3],
+            reserved[4],
+            reserved[5],
+            reserved[6],
+            reserved[7],
+        ]);
+        let column_count = u16::from_le_bytes([reserved[2], reserved[3]]);
+        let update_count = u64::from_le_bytes([
+            reserved[4],
+            reserved[5],
+            reserved[6],
+            reserved[7],
+            reserved[8],
+            reserved[9],
+            reserved[10],
+            reserved[11],
+        ]);
+        let batch_row_count =
+            i32::from_le_bytes([reserved[8], reserved[9], reserved[10], reserved[11]]);
+        let server_encoding = reserved[8];
 
         // Skip remaining 44 bytes of reserved
         buf.advance(44);
@@ -122,6 +163,9 @@ impl Frame {
             affected_rows,
             compress_flag,
             update_count,
+            column_count,
+            fetch_total,
+            batch_row_count,
             server_encoding,
         })
     }
@@ -224,4 +268,31 @@ mod tests {
         assert_eq!(frame.handle, 3);
         assert_eq!(frame.body_len, 100);
     }
+
+    /// The reserved area carries several overlapping counts. Only the checksum over
+    /// bytes 0-18 is fixed, so a reply can be forged by patching the tail.
+    #[test]
+    fn test_frame_parse_reserved_counts() {
+        let mut encoded = Frame::new(187, 0, 64).encode();
+        encoded[20..28].copy_from_slice(&20_000i64.to_le_bytes());
+        encoded[28..32].copy_from_slice(&662i32.to_le_bytes());
+        let frame = Frame::parse(&mut encoded).unwrap();
+        assert_eq!(frame.fetch_total, 20_000);
+        assert_eq!(frame.batch_row_count, 662);
+
+        let mut encoded = Frame::new(187, 0, 64).encode();
+        encoded[22..24].copy_from_slice(&5u16.to_le_bytes());
+        encoded[24..32].copy_from_slice(&2i64.to_le_bytes());
+        let frame = Frame::parse(&mut encoded).unwrap();
+        assert_eq!(frame.column_count, 5);
+        assert_eq!(frame.update_count, 2);
+    }
+
+    #[test]
+    fn test_frame_parse_unknown_row_total() {
+        let mut encoded = Frame::new(187, 0, 0).encode();
+        encoded[24..32].copy_from_slice(&ROW_TOTAL_UNKNOWN.to_le_bytes());
+        let frame = Frame::parse(&mut encoded).unwrap();
+        assert_eq!(frame.update_count, ROW_TOTAL_UNKNOWN as u64);
+    }
 }

dameng-protocol/src/lib.rs

@@ -8,6 +8,6 @@ pub mod frame;
 pub mod message;
 
 pub use error::{Error, Result};
-pub use frame::Frame;
+pub use frame::{Frame, ROW_TOTAL_UNKNOWN};
 pub use message::explain::ExplainResponse;
 pub use message::response::{Column, ExecResponse, Row};

dameng-protocol/src/message/fetch.rs

@@ -14,21 +14,27 @@
 //! 38      4     prefetchBytes (i32 LE) — max bytes to fetch, clamped [32, 65536]
 //! ```
 //!
-//! Response wire format:
+//! DM 8.1.3.62 reads none of the three: it streams the next batch from wherever its
+//! cursor stands, sized to its own budget, whatever the request asked for.
+//!
+//! The reply's payload is bare inline row data from byte 0, in the same format the
+//! EXEC/OPE payload uses after its column descriptors. It carries no header and no
+//! column metadata, so the columns of the statement that opened the cursor are what
+//! the rows are parsed against. The counts travel in the frame header instead:
 //! ```text
 //! Offset  Size  Field
-//! 0       20    Reserved
-//! 20      8     updateCount (i64 LE) — total row count in result set
-//! 28      4     rsSizeof    (i32 LE) — byte size of row data
-//! 32      N     row data (same format as EXEC_RESPONSE inline rows)
+//! 20      8     fetch_total     (i64 LE) — rows in the whole result set,
+//!                                          ROW_TOTAL_UNKNOWN until the last batch
+//! 28      4     batch_row_count (i32 LE) — rows in this reply
 //! ```
+//! A reply with an empty body means the cursor is drained.
 
 use bytes::{BufMut, BytesMut};
 
-use crate::error::Result;
+use crate::frame::Frame;
 use dameng_types::encoding::ServerEncoding;
 
-use super::response::{Column, ExecResponse, Row};
+use super::response::{parse_inline_rows, Column, Row};
 
 /// Default prefetch byte budget for FETCH requests.
 pub const DEFAULT_PREFETCH_BYTES: i32 = 8192;
@@ -107,85 +113,27 @@ impl FetchMessage {
 /// Response from a FETCH request (msg_type=7).
 #[derive(Debug, Clone)]
 pub struct FetchResponse {
-    /// Total number of rows in the entire result set.
+    /// Rows in the whole result set, or [`crate::frame::ROW_TOTAL_UNKNOWN`] while the
+    /// server still holds rows for the cursor.
     pub total_row_count: i64,
-    /// Column metadata (may be empty if already known from initial query).
-    pub columns: Vec<Column>,
-    /// Row data fetched in this batch.
+    /// Rows this batch carried.
     pub rows: Vec<Row>,
 }
 
 impl FetchResponse {
-    /// Parse a FETCH response from raw payload bytes.
+    /// Parse a FETCH reply: the payload is bare inline rows, the counts are in `frame`.
     ///
-    /// Response format:
-    /// - Offset 0-19: reserved
-    /// - Offset 20-27: updateCount (i64 LE) — total row count
-    /// - Offset 28-31: rsSizeof (i32 LE) — byte size of row data
-    /// - Offset 32+: row data (same format as EXEC_RESPONSE inline rows)
-    pub fn from_bytes(data: &[u8], server_encoding: ServerEncoding) -> Result<Self> {
-        if data.len() < 32 {
-            return Err(crate::error::Error::Incomplete);
-        }
-
-        // updateCount at offset 20
-        let total_row_count = i64::from_le_bytes([
-            data[20], data[21], data[22], data[23], data[24], data[25], data[26], data[27],
-        ]);
-
-        // rsSizeof at offset 28
-        let rs_sizeof = if data.len() >= 32 {
-            i32::from_le_bytes([data[28], data[29], data[30], data[31]]) as usize
-        } else {
-            0
-        };
-
-        // Row data starts at offset 32
-        let row_data_start = 32;
-        let row_data_end = (row_data_start + rs_sizeof).min(data.len());
-
-        if row_data_start >= data.len() || rs_sizeof == 0 {
-            return Ok(FetchResponse {
-                total_row_count,
-                columns: vec![],
-                rows: vec![],
-            });
-        }
-
-        let row_data = &data[row_data_start..row_data_end];
-
-        // The row data follows the same inline format as EXEC_RESPONSE.
-        // Parse it using the ExecResponse parser.
-        // Guard against parsing garbage: if row_data is all zeros or too short,
-        // the server returned metadata only (no inline data).
-        let has_real_data =
-            row_data.len() > 16 && !row_data.iter().all(|&b| b == 0) && row_data[0] != 0;
-
-        if !has_real_data {
-            // Server returned a cursor/total count but no inline row data.
-            // This happens when the cursor_id is invalid or the result is empty.
-            return Ok(FetchResponse {
-                total_row_count,
-                columns: vec![],
-                rows: vec![],
-            });
-        }
-
-        match ExecResponse::from_bytes(row_data, server_encoding) {
-            Ok(resp) => Ok(FetchResponse {
-                total_row_count,
-                columns: resp.columns,
-                rows: resp.rows,
-            }),
-            Err(_) => {
-                // If we can't parse the row data as EXEC_RESPONSE format,
-                // return what we have with empty rows.
-                Ok(FetchResponse {
-                    total_row_count,
-                    columns: vec![],
-                    rows: vec![],
-                })
-            }
+    /// `columns` describes the result set the cursor belongs to. The reply repeats no
+    /// metadata, so nothing else can say how wide a row is.
+    pub fn from_frame(
+        frame: &Frame,
+        data: &[u8],
+        columns: &[Column],
+        server_encoding: ServerEncoding,
+    ) -> Self {
+        Self {
+            total_row_count: frame.fetch_total,
+            rows: parse_inline_rows(data, 0, columns, server_encoding),
         }
     }
 
@@ -285,19 +233,66 @@ mod tests {
         assert_eq!(fetch.prefetch_bytes, DEFAULT_PREFETCH_BYTES);
     }
 
-    #[test]
-    fn test_fetch_response_incomplete() {
-        let data = [0u8; 10];
-        let result = FetchResponse::from_bytes(&data, ServerEncoding::Utf8);
-        assert!(result.is_err());
+    fn varchar_column() -> Column {
+        Column {
+            name: "NAME".to_string(),
+            type_code: 3,
+            type_name: "VARCHAR".to_string(),
+            precision: 100,
+            scale: 0,
+            nullable: true,
+            display_size: 0,
+            table_name: "CUSTOMER".to_string(),
+            schema_name: "APP".to_string(),
+            lob_tab_id: 0,
+            lob_col_id: 0,
+        }
     }
 
+    /// A drained cursor answers with an empty body. That is the end of the result set,
+    /// not a truncated message.
     #[test]
-    fn test_fetch_response_empty_data() {
-        let data = vec![0u8; 42];
-        let resp = FetchResponse::from_bytes(&data, ServerEncoding::Utf8).unwrap();
-        assert_eq!(resp.total_row_count, 0);
+    fn test_fetch_response_empty_body_ends_the_cursor() {
+        let mut frame = Frame::new(7, 0, 0);
+        frame.fetch_total = 2;
+        let resp =
+            FetchResponse::from_frame(&frame, &[], &[varchar_column()], ServerEncoding::Utf8);
+        assert_eq!(resp.total_row_count, 2);
         assert!(resp.rows.is_empty());
+        assert!(!resp.has_more(2));
+    }
+
+    /// Captured from DM 8.1.3.62: the payload is rows from byte 0, no header in front.
+    #[test]
+    fn test_fetch_response_parses_bare_rows() {
+        let data: Vec<u8> = vec![
+            0x13, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x0c, 0x00, 0x05, 0x00,
+            0x41, 0x6c, 0x69, 0x63, 0x65, // "Alice"
+            0x11, 0x00, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x0c, 0x00, 0x03, 0x00,
+            0x42, 0x6f, 0x62, // "Bob"
+        ];
+        let mut frame = Frame::new(7, 0, data.len() as i32);
+        frame.fetch_total = 2;
+        frame.batch_row_count = 2;
+
+        let resp =
+            FetchResponse::from_frame(&frame, &data, &[varchar_column()], ServerEncoding::Utf8);
+
+        assert_eq!(resp.rows.len(), frame.batch_row_count as usize);
+        assert_eq!(resp.total_row_count, 2);
+        assert_eq!(resp.rows[0].get_str(0).unwrap(), "Alice");
+        assert_eq!(resp.rows[1].get_str(0).unwrap(), "Bob");
+    }
+
+    /// Every batch before the last one leaves the total unknown.
+    #[test]
+    fn test_fetch_response_unknown_total() {
+        let mut frame = Frame::new(7, 0, 0);
+        frame.fetch_total = crate::frame::ROW_TOTAL_UNKNOWN;
+        let resp =
+            FetchResponse::from_frame(&frame, &[], &[varchar_column()], ServerEncoding::Utf8);
+        assert_eq!(resp.total_row_count, crate::frame::ROW_TOTAL_UNKNOWN);
+        assert!(resp.has_more(662));
     }
 
     #[test]
@@ -313,7 +308,6 @@ mod tests {
     fn test_has_more() {
         let resp = FetchResponse {
             total_row_count: 1000,
-            columns: vec![],
             rows: vec![],
         };
         assert!(resp.has_more(0));