aurora_db/
client.rs

1use crate::error::{AuroraError, Result};
2use crate::network::protocol::{Request, Response};
3use crate::query::SimpleQueryBuilder;
4use crate::types::{Document, FieldType, Value};
5use std::collections::HashMap;
6use tokio::io::{AsyncReadExt, AsyncWriteExt};
7use tokio::net::TcpStream;
8
9pub struct Client {
10    stream: TcpStream,
11}
12
13impl Client {
14    /// Connect to an Aurora server at the given address.
15    pub async fn connect(addr: &str) -> Result<Self> {
16        let stream = TcpStream::connect(addr).await?;
17        Ok(Self { stream })
18    }
19
20    /// Sends a request to the server and awaits a response.
21    async fn send_request(&mut self, request: Request) -> Result<Response> {
22        let request_bytes = bincode::serialize(&request).map_err(AuroraError::Bincode)?;
23        let len_bytes = (request_bytes.len() as u32).to_le_bytes();
24
25        self.stream.write_all(&len_bytes).await?;
26        self.stream.write_all(&request_bytes).await?;
27
28        let mut len_bytes = [0u8; 4];
29        self.stream.read_exact(&mut len_bytes).await?;
30        let len = u32::from_le_bytes(len_bytes) as usize;
31
32        let mut buffer = vec![0u8; len];
33        self.stream.read_exact(&mut buffer).await?;
34
35        let response: Response = bincode::deserialize(&buffer).map_err(AuroraError::Bincode)?;
36
37        Ok(response)
38    }
39
40    pub async fn new_collection(
41        &mut self,
42        name: &str,
43        fields: Vec<(String, FieldType, bool)>,
44    ) -> Result<()> {
45        let req = Request::NewCollection {
46            name: name.to_string(),
47            fields,
48        };
49        match self.send_request(req).await? {
50            Response::Done => Ok(()),
51            Response::Error(e) => Err(AuroraError::Protocol(e)),
52            _ => Err(AuroraError::Protocol("Unexpected response".into())),
53        }
54    }
55
56    pub async fn insert(
57        &mut self,
58        collection: &str,
59        data: HashMap<String, Value>,
60    ) -> Result<String> {
61        let req = Request::Insert {
62            collection: collection.to_string(),
63            data,
64        };
65        match self.send_request(req).await? {
66            Response::Message(id) => Ok(id),
67            Response::Error(e) => Err(AuroraError::Protocol(e)),
68            _ => Err(AuroraError::Protocol("Unexpected response".into())),
69        }
70    }
71
72    pub async fn get_document(&mut self, collection: &str, id: &str) -> Result<Option<Document>> {
73        let req = Request::GetDocument {
74            collection: collection.to_string(),
75            id: id.to_string(),
76        };
77        match self.send_request(req).await? {
78            Response::Document(doc) => Ok(doc),
79            Response::Error(e) => Err(AuroraError::Protocol(e)),
80            _ => Err(AuroraError::Protocol("Unexpected response".into())),
81        }
82    }
83
84    pub async fn query(&mut self, builder: SimpleQueryBuilder) -> Result<Vec<Document>> {
85        let req = Request::Query(builder);
86        match self.send_request(req).await? {
87            Response::Documents(docs) => Ok(docs),
88            Response::Error(e) => Err(AuroraError::Protocol(e)),
89            _ => Err(AuroraError::Protocol("Unexpected response".into())),
90        }
91    }
92
93    pub async fn begin_transaction(&mut self) -> Result<()> {
94        match self.send_request(Request::BeginTransaction).await? {
95            Response::Done => Ok(()),
96            Response::Error(e) => Err(AuroraError::Protocol(e)),
97            _ => Err(AuroraError::Protocol("Unexpected response".into())),
98        }
99    }
100
101    pub async fn commit_transaction(&mut self) -> Result<()> {
102        match self.send_request(Request::CommitTransaction).await? {
103            Response::Done => Ok(()),
104            Response::Error(e) => Err(AuroraError::Protocol(e)),
105            _ => Err(AuroraError::Protocol("Unexpected response".into())),
106        }
107    }
108
109    pub async fn delete(&mut self, key: &str) -> Result<()> {
110        match self.send_request(Request::Delete(key.to_string())).await? {
111            Response::Done => Ok(()),
112            Response::Error(e) => Err(AuroraError::Protocol(e)),
113            _ => Err(AuroraError::Protocol("Unexpected response".into())),
114        }
115    }
116}