Skip to main content

we_trust_sqlite/
codec.rs

1use bytes::{Buf, BufMut, BytesMut};
2use tokio_util::codec::{Decoder, Encoder};
3use yykv_types::DsError;
4
5#[derive(Debug, Clone)]
6pub enum SqliteCommand {
7    Execute(String),
8    Query(String),
9}
10
11#[derive(Debug, Clone)]
12pub enum SqliteResponse {
13    Ok,
14    Row(Vec<u8>),
15    Error(String),
16}
17
18pub struct SqliteCodec;
19
20impl Decoder for SqliteCodec {
21    type Item = SqliteCommand;
22    type Error = DsError;
23
24    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
25        if src.len() < 5 {
26            return Ok(None);
27        }
28
29        let tag = src[0];
30        let len = u32::from_le_bytes([src[1], src[2], src[3], src[4]]) as usize;
31
32        if src.len() < 5 + len {
33            src.reserve(len);
34            return Ok(None);
35        }
36
37        src.advance(5);
38        let data = src.split_to(len);
39        let sql = String::from_utf8_lossy(&data).to_string();
40
41        match tag {
42            1 => Ok(Some(SqliteCommand::Execute(sql))),
43            2 => Ok(Some(SqliteCommand::Query(sql))),
44            _ => Err(DsError::protocol("Invalid SQLite command tag")),
45        }
46    }
47}
48
49impl Encoder<SqliteResponse> for SqliteCodec {
50    type Error = DsError;
51
52    fn encode(&mut self, item: SqliteResponse, dst: &mut BytesMut) -> Result<(), Self::Error> {
53        match item {
54            SqliteResponse::Ok => {
55                dst.put_u8(0);
56                dst.put_u32_le(0);
57            }
58            SqliteResponse::Row(data) => {
59                dst.put_u8(1);
60                dst.put_u32_le(data.len() as u32);
61                dst.put_slice(&data);
62            }
63            SqliteResponse::Error(e) => {
64                dst.put_u8(255);
65                let bytes = e.as_bytes();
66                dst.put_u32_le(bytes.len() as u32);
67                dst.put_slice(bytes);
68            }
69        }
70        Ok(())
71    }
72}