zero_mysql/protocol/command/
utility.rs

1use crate::constant::CommandByte;
2use crate::error::Result;
3use crate::protocol::command::ColumnDefinition;
4use crate::protocol::primitive::*;
5use crate::protocol::response::{OkPayload, OkPayloadBytes};
6use crate::protocol::r#trait::{BinaryResultSetHandler, TextResultSetHandler};
7use crate::protocol::{BinaryRowPayload, TextRowPayload};
8
9/// Write COM_QUIT command
10pub fn write_quit(out: &mut Vec<u8>) {
11    write_int_1(out, CommandByte::Quit as u8);
12}
13
14/// Write COM_PING command
15pub fn write_ping(out: &mut Vec<u8>) {
16    write_int_1(out, CommandByte::Ping as u8);
17}
18
19/// Write COM_INIT_DB command
20pub fn write_init_db(out: &mut Vec<u8>, database: &str) {
21    write_int_1(out, CommandByte::InitDb as u8);
22    out.extend_from_slice(database.as_bytes());
23}
24
25/// Write COM_RESET_CONNECTION command
26pub fn write_reset_connection(out: &mut Vec<u8>) {
27    write_int_1(out, CommandByte::ResetConnection as u8);
28}
29
30// ============================================================================
31// Handler Utilities
32// ============================================================================
33
34/// A handler that ignores all result set data but captures affected_rows and last_insert_id
35///
36/// Useful for `exec_drop()` and `query_drop()` methods that discard results but need metadata.
37#[derive(Default)]
38pub struct DropHandler {
39    affected_rows: u64,
40    last_insert_id: u64,
41}
42
43impl DropHandler {
44    /// Get the number of affected rows from the last operation
45    pub fn affected_rows(&self) -> u64 {
46        self.affected_rows
47    }
48
49    /// Get the last insert ID from the last operation
50    pub fn last_insert_id(&self) -> u64 {
51        self.last_insert_id
52    }
53}
54
55impl BinaryResultSetHandler for DropHandler {
56    fn no_result_set(&mut self, ok: OkPayloadBytes) -> Result<()> {
57        let payload = OkPayload::try_from(ok)?;
58        self.affected_rows = payload.affected_rows;
59        self.last_insert_id = payload.last_insert_id;
60        Ok(())
61    }
62
63    fn resultset_start(&mut self, _: &[ColumnDefinition<'_>]) -> Result<()> {
64        Ok(())
65    }
66
67    fn row(&mut self, _: &[ColumnDefinition<'_>], _: BinaryRowPayload<'_>) -> Result<()> {
68        Ok(())
69    }
70
71    fn resultset_end(&mut self, eof: OkPayloadBytes) -> Result<()> {
72        let payload = OkPayload::try_from(eof)?;
73        self.affected_rows = payload.affected_rows;
74        self.last_insert_id = payload.last_insert_id;
75        Ok(())
76    }
77}
78
79impl TextResultSetHandler for DropHandler {
80    fn no_result_set(&mut self, ok: OkPayloadBytes) -> Result<()> {
81        let payload = OkPayload::try_from(ok)?;
82        self.affected_rows = payload.affected_rows;
83        self.last_insert_id = payload.last_insert_id;
84        Ok(())
85    }
86
87    fn resultset_start(&mut self, _: &[ColumnDefinition<'_>]) -> Result<()> {
88        Ok(())
89    }
90
91    fn row(&mut self, _: &[ColumnDefinition<'_>], _: TextRowPayload<'_>) -> Result<()> {
92        Ok(())
93    }
94
95    fn resultset_end(&mut self, eof: OkPayloadBytes) -> Result<()> {
96        let payload = OkPayload::try_from(eof)?;
97        self.affected_rows = payload.affected_rows;
98        self.last_insert_id = payload.last_insert_id;
99        Ok(())
100    }
101}
102
103/// A wrapper handler that forwards calls to an inner handler but stops after the first row
104///
105/// Useful for `exec_first()` methods that only process the first row.
106pub struct FirstRowHandler<'a, H> {
107    pub inner: &'a mut H,
108    pub found_row: bool,
109}
110
111impl<'a, H> FirstRowHandler<'a, H> {
112    pub fn new(inner: &'a mut H) -> Self {
113        Self {
114            inner,
115            found_row: false,
116        }
117    }
118}
119
120impl<'a, H: BinaryResultSetHandler> BinaryResultSetHandler for FirstRowHandler<'a, H> {
121    fn no_result_set(&mut self, ok: OkPayloadBytes) -> Result<()> {
122        self.inner.no_result_set(ok)
123    }
124
125    fn resultset_start(&mut self, cols: &[ColumnDefinition<'_>]) -> Result<()> {
126        self.inner.resultset_start(cols)
127    }
128
129    fn row(&mut self, cols: &[ColumnDefinition<'_>], row: BinaryRowPayload<'_>) -> Result<()> {
130        if !self.found_row {
131            self.found_row = true;
132            self.inner.row(cols, row)
133        } else {
134            Ok(()) // Ignore subsequent rows
135        }
136    }
137
138    fn resultset_end(&mut self, eof: OkPayloadBytes) -> Result<()> {
139        self.inner.resultset_end(eof)
140    }
141}