zero_mysql/
handler.rs

1use crate::error::Result;
2use crate::protocol::command::ColumnDefinition;
3use crate::protocol::response::{OkPayload, OkPayloadBytes};
4use crate::protocol::r#trait::{BinaryResultSetHandler, TextResultSetHandler};
5use crate::protocol::{BinaryRowPayload, TextRowPayload};
6use crate::raw::FromRawRow;
7use smart_default::SmartDefault;
8
9/// A handler that ignores all result set data but captures affected_rows and last_insert_id
10///
11/// Useful for `exec_drop()` and `query_drop()` methods that discard results but need metadata.
12#[derive(Default)]
13pub struct DropHandler {
14    affected_rows: u64,
15    last_insert_id: u64,
16}
17
18impl DropHandler {
19    /// Get the number of affected rows from the last operation
20    pub fn affected_rows(&self) -> u64 {
21        self.affected_rows
22    }
23
24    /// Get the last insert ID from the last operation
25    pub fn last_insert_id(&self) -> u64 {
26        self.last_insert_id
27    }
28}
29
30impl BinaryResultSetHandler for DropHandler {
31    fn no_result_set(&mut self, ok: OkPayloadBytes) -> Result<()> {
32        let payload = OkPayload::try_from(ok)?;
33        self.affected_rows = payload.affected_rows;
34        self.last_insert_id = payload.last_insert_id;
35        Ok(())
36    }
37
38    fn resultset_start(&mut self, _: &[ColumnDefinition<'_>]) -> Result<()> {
39        Ok(())
40    }
41
42    fn row(&mut self, _: &[ColumnDefinition<'_>], _: BinaryRowPayload<'_>) -> Result<()> {
43        Ok(())
44    }
45
46    fn resultset_end(&mut self, eof: OkPayloadBytes) -> Result<()> {
47        let payload = OkPayload::try_from(eof)?;
48        self.affected_rows = payload.affected_rows;
49        self.last_insert_id = payload.last_insert_id;
50        Ok(())
51    }
52}
53
54impl TextResultSetHandler for DropHandler {
55    fn no_result_set(&mut self, ok: OkPayloadBytes) -> Result<()> {
56        let payload = OkPayload::try_from(ok)?;
57        self.affected_rows = payload.affected_rows;
58        self.last_insert_id = payload.last_insert_id;
59        Ok(())
60    }
61
62    fn resultset_start(&mut self, _: &[ColumnDefinition<'_>]) -> Result<()> {
63        Ok(())
64    }
65
66    fn row(&mut self, _: &[ColumnDefinition<'_>], _: TextRowPayload<'_>) -> Result<()> {
67        Ok(())
68    }
69
70    fn resultset_end(&mut self, eof: OkPayloadBytes) -> Result<()> {
71        let payload = OkPayload::try_from(eof)?;
72        self.affected_rows = payload.affected_rows;
73        self.last_insert_id = payload.last_insert_id;
74        Ok(())
75    }
76}
77
78/// A wrapper handler that forwards calls to an inner handler but stops after the first row
79///
80/// Useful for `exec_first()` methods that only process the first row.
81pub struct FirstRowHandler<'a, H> {
82    pub inner: &'a mut H,
83    pub found_row: bool,
84}
85
86impl<'a, H> FirstRowHandler<'a, H> {
87    pub fn new(inner: &'a mut H) -> Self {
88        Self {
89            inner,
90            found_row: false,
91        }
92    }
93}
94
95impl<'a, H: BinaryResultSetHandler> BinaryResultSetHandler for FirstRowHandler<'a, H> {
96    fn no_result_set(&mut self, ok: OkPayloadBytes) -> Result<()> {
97        self.inner.no_result_set(ok)
98    }
99
100    fn resultset_start(&mut self, cols: &[ColumnDefinition<'_>]) -> Result<()> {
101        self.inner.resultset_start(cols)
102    }
103
104    fn row(&mut self, cols: &[ColumnDefinition<'_>], row: BinaryRowPayload<'_>) -> Result<()> {
105        if !self.found_row {
106            self.found_row = true;
107            self.inner.row(cols, row)
108        } else {
109            Ok(()) // Ignore subsequent rows
110        }
111    }
112
113    fn resultset_end(&mut self, eof: OkPayloadBytes) -> Result<()> {
114        self.inner.resultset_end(eof)
115    }
116}
117
118/// A handler that collects all rows into a Vec<Row>
119///
120/// Useful for `exec()` methods that need to return all rows as a collection.
121#[derive(SmartDefault)]
122pub struct CollectHandler<Row> {
123    rows: Vec<Row>,
124    affected_rows: u64,
125    last_insert_id: u64,
126}
127
128impl<Row> CollectHandler<Row> {
129    pub fn take_rows(&mut self) -> Vec<Row> {
130        std::mem::take(&mut self.rows)
131    }
132    pub fn into_rows(self) -> Vec<Row> {
133        self.rows
134    }
135    pub fn affected_rows(&self) -> u64 {
136        self.affected_rows
137    }
138    pub fn last_insert_id(&self) -> u64 {
139        self.last_insert_id
140    }
141}
142
143impl<Row: for<'buf> FromRawRow<'buf>> BinaryResultSetHandler for CollectHandler<Row> {
144    fn no_result_set(&mut self, ok: OkPayloadBytes) -> Result<()> {
145        let payload = OkPayload::try_from(ok)?;
146        self.affected_rows = payload.affected_rows;
147        self.last_insert_id = payload.last_insert_id;
148        Ok(())
149    }
150
151    fn resultset_start(&mut self, _cols: &[ColumnDefinition<'_>]) -> Result<()> {
152        Ok(())
153    }
154
155    fn row(&mut self, cols: &[ColumnDefinition], row: BinaryRowPayload) -> Result<()> {
156        self.rows.push(Row::from_raw_row(cols, row)?);
157        Ok(())
158    }
159
160    fn resultset_end(&mut self, eof: OkPayloadBytes) -> Result<()> {
161        let payload = OkPayload::try_from(eof)?;
162        self.affected_rows = payload.affected_rows;
163        self.last_insert_id = payload.last_insert_id;
164        Ok(())
165    }
166}