Skip to main content

reddb_grpc_proto/
client.rs

1use std::error::Error;
2
3use reddb_wire::auth::{bearer_authorization_value, login_payload_json, AUTHORIZATION_HEADER};
4use tonic::transport::Channel;
5use tonic::Request;
6
7use crate::red_db_client::RedDbClient;
8use crate::{
9    BulkEntityReply, DeleteEntityRequest, Empty, EntityReply, HealthReply, JsonBulkCreateRequest,
10    JsonCreateRequest, JsonPayloadRequest, OperationReply, PayloadReply, QueryReply, QueryRequest,
11    QueryValue, ScanReply, ScanRequest, StatsReply, TopologyRequest,
12};
13
14/// Thin authenticated client over the generated tonic stub.
15#[derive(Clone)]
16pub struct RedDBClient {
17    inner: RedDbClient<Channel>,
18    token: Option<String>,
19    pub addr: String,
20}
21
22impl RedDBClient {
23    pub async fn connect(addr: &str, token: Option<String>) -> Result<Self, Box<dyn Error>> {
24        let endpoint = if addr.starts_with("http") {
25            addr.to_string()
26        } else {
27            format!("http://{addr}")
28        };
29        let inner = RedDbClient::connect(endpoint.clone()).await?;
30        Ok(Self {
31            inner,
32            token,
33            addr: endpoint,
34        })
35    }
36
37    fn auth_request<T>(&self, inner: T) -> Request<T> {
38        let mut request = Request::new(inner);
39        if let Some(token) = &self.token {
40            if let Ok(value) = bearer_authorization_value(token).parse() {
41                request.metadata_mut().insert(AUTHORIZATION_HEADER, value);
42            }
43        }
44        request
45    }
46
47    /// Update the bearer token, for example after a successful login.
48    pub fn set_token(&mut self, token: String) {
49        self.token = Some(token);
50    }
51
52    pub async fn health_status(&mut self) -> Result<HealthReply, Box<dyn Error>> {
53        let request = self.auth_request(Empty {});
54        Ok(self.inner.health(request).await?.into_inner())
55    }
56
57    pub async fn query_reply(&mut self, sql: &str) -> Result<QueryReply, Box<dyn Error>> {
58        self.query_reply_with_params(sql, Vec::new()).await
59    }
60
61    pub async fn query_reply_with_params(
62        &mut self,
63        sql: &str,
64        params: Vec<QueryValue>,
65    ) -> Result<QueryReply, Box<dyn Error>> {
66        let request = self.auth_request(QueryRequest {
67            query: sql.to_string(),
68            entity_types: vec![],
69            capabilities: vec![],
70            params,
71        });
72        Ok(self.inner.query(request).await?.into_inner())
73    }
74
75    pub async fn query(&mut self, sql: &str) -> Result<String, Box<dyn Error>> {
76        Ok(self.query_reply(sql).await?.result_json)
77    }
78
79    pub async fn collections(&mut self) -> Result<Vec<String>, Box<dyn Error>> {
80        let request = self.auth_request(Empty {});
81        Ok(self
82            .inner
83            .collections(request)
84            .await?
85            .into_inner()
86            .collections)
87    }
88
89    pub async fn scan(
90        &mut self,
91        collection: &str,
92        limit: u64,
93    ) -> Result<ScanReply, Box<dyn Error>> {
94        let request = self.auth_request(ScanRequest {
95            collection: collection.to_string(),
96            offset: 0,
97            limit,
98        });
99        Ok(self.inner.scan(request).await?.into_inner())
100    }
101
102    pub async fn stats(&mut self) -> Result<StatsReply, Box<dyn Error>> {
103        let request = self.auth_request(Empty {});
104        Ok(self.inner.stats(request).await?.into_inner())
105    }
106
107    pub async fn create_row_entity(
108        &mut self,
109        collection: &str,
110        json: &str,
111    ) -> Result<EntityReply, Box<dyn Error>> {
112        let request = self.auth_request(JsonCreateRequest {
113            collection: collection.to_string(),
114            payload_json: json.to_string(),
115        });
116        Ok(self.inner.create_row(request).await?.into_inner())
117    }
118
119    pub async fn bulk_create_rows(
120        &mut self,
121        collection: &str,
122        payload_json: Vec<String>,
123    ) -> Result<BulkEntityReply, Box<dyn Error>> {
124        let request = self.auth_request(JsonBulkCreateRequest {
125            collection: collection.to_string(),
126            payload_json,
127        });
128        Ok(self.inner.bulk_create_rows(request).await?.into_inner())
129    }
130
131    pub async fn explain(&mut self, sql: &str) -> Result<String, Box<dyn Error>> {
132        let request = self.auth_request(QueryRequest {
133            query: sql.to_string(),
134            entity_types: vec![],
135            capabilities: vec![],
136            params: vec![],
137        });
138        Ok(self
139            .inner
140            .explain_query(request)
141            .await?
142            .into_inner()
143            .payload)
144    }
145
146    pub async fn login(
147        &mut self,
148        username: &str,
149        password: &str,
150    ) -> Result<String, Box<dyn Error>> {
151        let request = self.auth_request(JsonPayloadRequest {
152            payload_json: login_payload_json(username, password),
153        });
154        let reply: PayloadReply = self.inner.auth_login(request).await?.into_inner();
155        Ok(reply.payload)
156    }
157
158    pub async fn replication_status(&mut self) -> Result<String, Box<dyn Error>> {
159        let request = self.auth_request(Empty {});
160        Ok(self
161            .inner
162            .replication_status(request)
163            .await?
164            .into_inner()
165            .payload)
166    }
167
168    /// Fetch the canonical topology envelope bytes.
169    pub async fn topology(&mut self) -> Result<Vec<u8>, Box<dyn Error>> {
170        let request = self.auth_request(TopologyRequest {});
171        Ok(self
172            .inner
173            .topology(request)
174            .await?
175            .into_inner()
176            .topology_bytes)
177    }
178
179    pub async fn delete_entity(
180        &mut self,
181        collection: &str,
182        id: u64,
183    ) -> Result<OperationReply, Box<dyn Error>> {
184        let request = self.auth_request(DeleteEntityRequest {
185            collection: collection.to_string(),
186            id,
187        });
188        Ok(self.inner.delete_entity(request).await?.into_inner())
189    }
190}