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    current_transaction: Option<u64>,
12}
13
14impl Client {
15    /// Connect to an Aurora server at the given address.
16    pub async fn connect(addr: &str) -> Result<Self> {
17        let stream = TcpStream::connect(addr).await?;
18        Ok(Self {
19            stream,
20            current_transaction: None,
21        })
22    }
23
24    /// Sends a request to the server and awaits a response.
25    async fn send_request(&mut self, request: Request) -> Result<Response> {
26        let request_bytes = bincode::serialize(&request).map_err(AuroraError::Bincode)?;
27        let len_bytes = (request_bytes.len() as u32).to_le_bytes();
28
29        self.stream.write_all(&len_bytes).await?;
30        self.stream.write_all(&request_bytes).await?;
31
32        let mut len_bytes = [0u8; 4];
33        self.stream.read_exact(&mut len_bytes).await?;
34        let len = u32::from_le_bytes(len_bytes) as usize;
35
36        let mut buffer = vec![0u8; len];
37        self.stream.read_exact(&mut buffer).await?;
38
39        let response: Response = bincode::deserialize(&buffer).map_err(AuroraError::Bincode)?;
40
41        Ok(response)
42    }
43
44    pub async fn new_collection(
45        &mut self,
46        name: &str,
47        fields: Vec<(String, FieldType, bool)>,
48    ) -> Result<()> {
49        let req = Request::NewCollection {
50            name: name.to_string(),
51            fields,
52        };
53        match self.send_request(req).await? {
54            Response::Done => Ok(()),
55            Response::Error(e) => Err(AuroraError::Protocol(e)),
56            _ => Err(AuroraError::Protocol("Unexpected response".into())),
57        }
58    }
59
60    pub async fn insert(
61        &mut self,
62        collection: &str,
63        data: HashMap<String, Value>,
64    ) -> Result<String> {
65        let req = Request::Insert {
66            collection: collection.to_string(),
67            data,
68        };
69        match self.send_request(req).await? {
70            Response::Message(id) => Ok(id),
71            Response::Error(e) => Err(AuroraError::Protocol(e)),
72            _ => Err(AuroraError::Protocol("Unexpected response".into())),
73        }
74    }
75
76    pub async fn get_document(&mut self, collection: &str, id: &str) -> Result<Option<Document>> {
77        let req = Request::GetDocument {
78            collection: collection.to_string(),
79            id: id.to_string(),
80        };
81        match self.send_request(req).await? {
82            Response::Document(doc) => Ok(doc),
83            Response::Error(e) => Err(AuroraError::Protocol(e)),
84            _ => Err(AuroraError::Protocol("Unexpected response".into())),
85        }
86    }
87
88    pub async fn query(&mut self, builder: SimpleQueryBuilder) -> Result<Vec<Document>> {
89        let req = Request::Query(builder);
90        match self.send_request(req).await? {
91            Response::Documents(docs) => Ok(docs),
92            Response::Error(e) => Err(AuroraError::Protocol(e)),
93            _ => Err(AuroraError::Protocol("Unexpected response".into())),
94        }
95    }
96
97    pub async fn begin_transaction(&mut self) -> Result<u64> {
98        match self.send_request(Request::BeginTransaction).await? {
99            Response::TransactionId(tx_id) => {
100                self.current_transaction = Some(tx_id);
101                Ok(tx_id)
102            }
103            Response::Error(e) => Err(AuroraError::Protocol(e)),
104            _ => Err(AuroraError::Protocol("Unexpected response".into())),
105        }
106    }
107
108    pub async fn commit_transaction(&mut self) -> Result<()> {
109        let tx_id = self
110            .current_transaction
111            .ok_or_else(|| AuroraError::InvalidOperation("No active transaction".into()))?;
112
113        match self.send_request(Request::CommitTransaction(tx_id)).await? {
114            Response::Done => {
115                self.current_transaction = None;
116                Ok(())
117            }
118            Response::Error(e) => Err(AuroraError::Protocol(e)),
119            _ => Err(AuroraError::Protocol("Unexpected response".into())),
120        }
121    }
122
123    pub async fn rollback_transaction(&mut self) -> Result<()> {
124        let tx_id = self
125            .current_transaction
126            .ok_or_else(|| AuroraError::InvalidOperation("No active transaction".into()))?;
127
128        match self
129            .send_request(Request::RollbackTransaction(tx_id))
130            .await?
131        {
132            Response::Done => {
133                self.current_transaction = None;
134                Ok(())
135            }
136            Response::Error(e) => Err(AuroraError::Protocol(e)),
137            _ => Err(AuroraError::Protocol("Unexpected response".into())),
138        }
139    }
140
141    pub async fn delete(&mut self, key: &str) -> Result<()> {
142        match self.send_request(Request::Delete(key.to_string())).await? {
143            Response::Done => Ok(()),
144            Response::Error(e) => Err(AuroraError::Protocol(e)),
145            _ => Err(AuroraError::Protocol("Unexpected response".into())),
146        }
147    }
148}
aurora_db/client.rs

aurora_db/
client.rs
