zero_trust_sdk/database.rs
1//! Database management for the Zero Trust SDK
2
3use crate::{
4 config::Config,
5 error::{Result, ZeroTrustError},
6 types::{Database, Table, QueryResult, QueryOptions},
7};
8use reqwest::Client;
9use serde_json::json;
10use std::sync::Arc;
11
12/// Database manager for handling database operations and queries
13#[derive(Debug, Clone)]
14pub struct DatabaseManager {
15 config: Arc<Config>,
16 http_client: Arc<Client>,
17}
18
19impl DatabaseManager {
20 /// Create a new database manager
21 pub(crate) fn new(config: Arc<Config>, http_client: Arc<Client>) -> Self {
22 Self {
23 config,
24 http_client,
25 }
26 }
27
28 /// List all databases
29 ///
30 /// # Examples
31 ///
32 /// ```rust,no_run
33 /// use zero_trust_sdk::{ZeroTrustClient, Config};
34 ///
35 /// #[tokio::main]
36 /// async fn main() -> Result<(), Box<dyn std::error::Error>> {
37 /// let config = Config::new("https://api.zerotrust.com")?
38 /// .with_token("your-jwt-token");
39 /// let client = ZeroTrustClient::new(config).await?;
40 ///
41 /// let databases = client.databases().list().await?;
42 /// for db in databases {
43 /// println!("Database: {} ({} tables)", db.name, db.tables.len());
44 /// }
45 ///
46 /// Ok(())
47 /// }
48 /// ```
49 pub async fn list(&self) -> Result<Vec<Database>> {
50 self.ensure_authenticated()?;
51
52 let url = format!("{}/api/v1/databases", self.config.api_url);
53 let response = self.send_authenticated_request(
54 self.http_client.get(&url)
55 ).await?;
56
57 if response.status().is_success() {
58 let databases: Vec<Database> = response.json().await?;
59 Ok(databases)
60 } else {
61 self.handle_error_response(response).await
62 }
63 }
64
65 /// Create a new database
66 ///
67 /// # Arguments
68 ///
69 /// * `name` - Database name
70 ///
71 /// # Examples
72 ///
73 /// ```rust,no_run
74 /// use zero_trust_sdk::{ZeroTrustClient, Config};
75 ///
76 /// #[tokio::main]
77 /// async fn main() -> Result<(), Box<dyn std::error::Error>> {
78 /// let config = Config::new("https://api.zerotrust.com")?
79 /// .with_token("your-jwt-token");
80 /// let client = ZeroTrustClient::new(config).await?;
81 ///
82 /// let database = client.databases().create("my-app-db").await?;
83 /// println!("Created database: {}", database.name);
84 ///
85 /// Ok(())
86 /// }
87 /// ```
88 pub async fn create<S: AsRef<str>>(&self, name: S) -> Result<Database> {
89 self.ensure_authenticated()?;
90
91 let url = format!("{}/api/v1/databases", self.config.api_url);
92 let payload = json!({
93 "name": name.as_ref()
94 });
95
96 let response = self.send_authenticated_request(
97 self.http_client
98 .post(&url)
99 .json(&payload)
100 ).await?;
101
102 if response.status().is_success() {
103 let database: Database = response.json().await?;
104 Ok(database)
105 } else {
106 self.handle_error_response(response).await
107 }
108 }
109
110 /// Get database information
111 ///
112 /// # Arguments
113 ///
114 /// * `name` - Database name
115 ///
116 /// # Examples
117 ///
118 /// ```rust,no_run
119 /// use zero_trust_sdk::{ZeroTrustClient, Config};
120 ///
121 /// #[tokio::main]
122 /// async fn main() -> Result<(), Box<dyn std::error::Error>> {
123 /// let config = Config::new("https://api.zerotrust.com")?
124 /// .with_token("your-jwt-token");
125 /// let client = ZeroTrustClient::new(config).await?;
126 ///
127 /// let database = client.databases().get("my-app-db").await?;
128 /// println!("Database {} has {} tables", database.name, database.tables.len());
129 ///
130 /// Ok(())
131 /// }
132 /// ```
133 pub async fn get<S: AsRef<str>>(&self, name: S) -> Result<Database> {
134 self.ensure_authenticated()?;
135
136 let url = format!("{}/api/v1/databases/{}", self.config.api_url, name.as_ref());
137 let response = self.send_authenticated_request(
138 self.http_client.get(&url)
139 ).await?;
140
141 if response.status().is_success() {
142 let database: Database = response.json().await?;
143 Ok(database)
144 } else {
145 self.handle_error_response(response).await
146 }
147 }
148
149 /// Delete a database
150 ///
151 /// # Arguments
152 ///
153 /// * `name` - Database name
154 ///
155 /// # Examples
156 ///
157 /// ```rust,no_run
158 /// use zero_trust_sdk::{ZeroTrustClient, Config};
159 ///
160 /// #[tokio::main]
161 /// async fn main() -> Result<(), Box<dyn std::error::Error>> {
162 /// let config = Config::new("https://api.zerotrust.com")?
163 /// .with_token("your-jwt-token");
164 /// let client = ZeroTrustClient::new(config).await?;
165 ///
166 /// client.databases().delete("old-database").await?;
167 /// println!("Database deleted successfully");
168 ///
169 /// Ok(())
170 /// }
171 /// ```
172 pub async fn delete<S: AsRef<str>>(&self, name: S) -> Result<()> {
173 self.ensure_authenticated()?;
174
175 let url = format!("{}/api/v1/databases/{}", self.config.api_url, name.as_ref());
176 let response = self.send_authenticated_request(
177 self.http_client.delete(&url)
178 ).await?;
179
180 if response.status().is_success() {
181 Ok(())
182 } else {
183 self.handle_error_response(response).await
184 }
185 }
186
187 /// Get a query builder for the specified database
188 ///
189 /// # Arguments
190 ///
191 /// * `database` - Database name
192 ///
193 /// # Examples
194 ///
195 /// ```rust,no_run
196 /// use zero_trust_sdk::{ZeroTrustClient, Config};
197 ///
198 /// #[tokio::main]
199 /// async fn main() -> Result<(), Box<dyn std::error::Error>> {
200 /// let config = Config::new("https://api.zerotrust.com")?
201 /// .with_token("your-jwt-token");
202 /// let client = ZeroTrustClient::new(config).await?;
203 ///
204 /// let result = client.databases()
205 /// .query("my-app-db")
206 /// .execute("SELECT * FROM users LIMIT 10")
207 /// .await?;
208 ///
209 /// println!("Found {} rows", result.meta.row_count);
210 /// Ok(())
211 /// }
212 /// ```
213 pub fn query<S: AsRef<str>>(&self, database: S) -> QueryBuilder {
214 QueryBuilder::new(
215 database.as_ref().to_string(),
216 self.config.clone(),
217 self.http_client.clone(),
218 )
219 }
220
221 /// List tables in a database
222 ///
223 /// # Arguments
224 ///
225 /// * `database` - Database name
226 ///
227 /// # Examples
228 ///
229 /// ```rust,no_run
230 /// use zero_trust_sdk::{ZeroTrustClient, Config};
231 ///
232 /// #[tokio::main]
233 /// async fn main() -> Result<(), Box<dyn std::error::Error>> {
234 /// let config = Config::new("https://api.zerotrust.com")?
235 /// .with_token("your-jwt-token");
236 /// let client = ZeroTrustClient::new(config).await?;
237 ///
238 /// let tables = client.databases().list_tables("my-app-db").await?;
239 /// for table in tables {
240 /// println!("Table: {} ({} columns)", table.name, table.columns.len());
241 /// }
242 ///
243 /// Ok(())
244 /// }
245 /// ```
246 pub async fn list_tables<S: AsRef<str>>(&self, database: S) -> Result<Vec<Table>> {
247 self.ensure_authenticated()?;
248
249 let url = format!(
250 "{}/api/v1/databases/{}/tables",
251 self.config.api_url,
252 database.as_ref()
253 );
254
255 let response = self.send_authenticated_request(
256 self.http_client.get(&url)
257 ).await?;
258
259 if response.status().is_success() {
260 let tables: Vec<Table> = response.json().await?;
261 Ok(tables)
262 } else {
263 self.handle_error_response(response).await
264 }
265 }
266
267 /// Get table information
268 ///
269 /// # Arguments
270 ///
271 /// * `database` - Database name
272 /// * `table` - Table name
273 ///
274 /// # Examples
275 ///
276 /// ```rust,no_run
277 /// use zero_trust_sdk::{ZeroTrustClient, Config};
278 ///
279 /// #[tokio::main]
280 /// async fn main() -> Result<(), Box<dyn std::error::Error>> {
281 /// let config = Config::new("https://api.zerotrust.com")?
282 /// .with_token("your-jwt-token");
283 /// let client = ZeroTrustClient::new(config).await?;
284 ///
285 /// let table = client.databases().get_table("my-app-db", "users").await?;
286 /// println!("Table {} has {} rows", table.name, table.row_count.unwrap_or(0));
287 ///
288 /// Ok(())
289 /// }
290 /// ```
291 pub async fn get_table<S1, S2>(&self, database: S1, table: S2) -> Result<Table>
292 where
293 S1: AsRef<str>,
294 S2: AsRef<str>,
295 {
296 self.ensure_authenticated()?;
297
298 let url = format!(
299 "{}/api/v1/databases/{}/tables/{}",
300 self.config.api_url,
301 database.as_ref(),
302 table.as_ref()
303 );
304
305 let response = self.send_authenticated_request(
306 self.http_client.get(&url)
307 ).await?;
308
309 if response.status().is_success() {
310 let table: Table = response.json().await?;
311 Ok(table)
312 } else {
313 self.handle_error_response(response).await
314 }
315 }
316
317 // Helper methods
318 fn ensure_authenticated(&self) -> Result<()> {
319 if !self.config.is_authenticated() {
320 return Err(ZeroTrustError::auth("Authentication required"));
321 }
322 Ok(())
323 }
324
325 async fn send_authenticated_request(
326 &self,
327 mut request: reqwest::RequestBuilder,
328 ) -> Result<reqwest::Response> {
329 if let Some(token) = &self.config.token {
330 request = request.header("Authorization", format!("Bearer {}", token));
331 }
332
333 let response = request
334 .header("Content-Type", "application/json")
335 .send()
336 .await?;
337
338 Ok(response)
339 }
340
341 async fn handle_error_response<T>(&self, response: reqwest::Response) -> Result<T> {
342 let status = response.status();
343 let error_text = response.text().await.unwrap_or_default();
344
345 match status.as_u16() {
346 401 => Err(ZeroTrustError::auth("Authentication failed")),
347 403 => Err(ZeroTrustError::permission_denied("Insufficient permissions")),
348 404 => Err(ZeroTrustError::not_found("Resource not found")),
349 400 => Err(ZeroTrustError::validation(error_text)),
350 500..=599 => Err(ZeroTrustError::server_error(status.as_u16(), error_text)),
351 _ => Err(ZeroTrustError::generic(format!("HTTP {}: {}", status, error_text))),
352 }
353 }
354}
355
356/// Query builder for database operations
357#[derive(Debug)]
358pub struct QueryBuilder {
359 database: String,
360 config: Arc<Config>,
361 http_client: Arc<Client>,
362 options: QueryOptions,
363}
364
365impl QueryBuilder {
366 fn new(database: String, config: Arc<Config>, http_client: Arc<Client>) -> Self {
367 Self {
368 database,
369 config,
370 http_client,
371 options: QueryOptions::default(),
372 }
373 }
374
375 /// Set maximum number of rows to return
376 ///
377 /// # Arguments
378 ///
379 /// * `max_rows` - Maximum number of rows
380 ///
381 /// # Examples
382 ///
383 /// ```rust,no_run
384 /// use zero_trust_sdk::{ZeroTrustClient, Config};
385 ///
386 /// #[tokio::main]
387 /// async fn main() -> Result<(), Box<dyn std::error::Error>> {
388 /// let config = Config::new("https://api.zerotrust.com")?
389 /// .with_token("your-jwt-token");
390 /// let client = ZeroTrustClient::new(config).await?;
391 ///
392 /// let result = client.databases()
393 /// .query("my-app-db")
394 /// .max_rows(100)
395 /// .execute("SELECT * FROM users")
396 /// .await?;
397 ///
398 /// Ok(())
399 /// }
400 /// ```
401 pub fn max_rows(mut self, max_rows: u64) -> Self {
402 self.options.max_rows = Some(max_rows);
403 self
404 }
405
406 /// Set query timeout
407 ///
408 /// # Arguments
409 ///
410 /// * `timeout_ms` - Timeout in milliseconds
411 ///
412 /// # Examples
413 ///
414 /// ```rust,no_run
415 /// use zero_trust_sdk::{ZeroTrustClient, Config};
416 ///
417 /// #[tokio::main]
418 /// async fn main() -> Result<(), Box<dyn std::error::Error>> {
419 /// let config = Config::new("https://api.zerotrust.com")?
420 /// .with_token("your-jwt-token");
421 /// let client = ZeroTrustClient::new(config).await?;
422 ///
423 /// let result = client.databases()
424 /// .query("my-app-db")
425 /// .timeout(30000) // 30 seconds
426 /// .execute("SELECT * FROM large_table")
427 /// .await?;
428 ///
429 /// Ok(())
430 /// }
431 /// ```
432 pub fn timeout(mut self, timeout_ms: u64) -> Self {
433 self.options.timeout_ms = Some(timeout_ms);
434 self
435 }
436
437 /// Include execution metadata in the response
438 ///
439 /// # Examples
440 ///
441 /// ```rust,no_run
442 /// use zero_trust_sdk::{ZeroTrustClient, Config};
443 ///
444 /// #[tokio::main]
445 /// async fn main() -> Result<(), Box<dyn std::error::Error>> {
446 /// let config = Config::new("https://api.zerotrust.com")?
447 /// .with_token("your-jwt-token");
448 /// let client = ZeroTrustClient::new(config).await?;
449 ///
450 /// let result = client.databases()
451 /// .query("my-app-db")
452 /// .include_meta()
453 /// .execute("SELECT * FROM users")
454 /// .await?;
455 ///
456 /// if let Some(execution_time) = result.meta.execution_time_ms {
457 /// println!("Query executed in {}ms", execution_time);
458 /// }
459 ///
460 /// Ok(())
461 /// }
462 /// ```
463 pub fn include_meta(mut self) -> Self {
464 self.options.include_meta = true;
465 self
466 }
467
468 /// Execute a SQL query
469 ///
470 /// # Arguments
471 ///
472 /// * `sql` - SQL query string
473 ///
474 /// # Examples
475 ///
476 /// ```rust,no_run
477 /// use zero_trust_sdk::{ZeroTrustClient, Config};
478 ///
479 /// #[tokio::main]
480 /// async fn main() -> Result<(), Box<dyn std::error::Error>> {
481 /// let config = Config::new("https://api.zerotrust.com")?
482 /// .with_token("your-jwt-token");
483 /// let client = ZeroTrustClient::new(config).await?;
484 ///
485 /// let result = client.databases()
486 /// .query("my-app-db")
487 /// .execute("SELECT name, email FROM users WHERE active = true")
488 /// .await?;
489 ///
490 /// for row in result.data.rows {
491 /// println!("User: {} - {}", row[0], row[1]);
492 /// }
493 ///
494 /// Ok(())
495 /// }
496 /// ```
497 pub async fn execute<S: AsRef<str>>(self, sql: S) -> Result<QueryResult> {
498 let url = format!("{}/api/v1/query/execute", self.config.api_url);
499
500 let mut payload = json!({
501 "database": self.database,
502 "sql": sql.as_ref()
503 });
504
505 // Add options to payload
506 if let Some(max_rows) = self.options.max_rows {
507 payload["max_rows"] = json!(max_rows);
508 }
509 if let Some(timeout_ms) = self.options.timeout_ms {
510 payload["timeout_ms"] = json!(timeout_ms);
511 }
512 if self.options.include_meta {
513 payload["include_meta"] = json!(true);
514 }
515
516 let response = self.send_authenticated_request(
517 self.http_client
518 .post(&url)
519 .json(&payload)
520 ).await?;
521
522 if response.status().is_success() {
523 let result: QueryResult = response.json().await?;
524 Ok(result)
525 } else {
526 self.handle_error_response(response).await
527 }
528 }
529
530 /// Execute multiple queries in a batch
531 ///
532 /// # Arguments
533 ///
534 /// * `queries` - Vector of SQL query strings
535 ///
536 /// # Examples
537 ///
538 /// ```rust,no_run
539 /// use zero_trust_sdk::{ZeroTrustClient, Config};
540 ///
541 /// #[tokio::main]
542 /// async fn main() -> Result<(), Box<dyn std::error::Error>> {
543 /// let config = Config::new("https://api.zerotrust.com")?
544 /// .with_token("your-jwt-token");
545 /// let client = ZeroTrustClient::new(config).await?;
546 ///
547 /// let queries = vec![
548 /// "SELECT COUNT(*) FROM users",
549 /// "SELECT COUNT(*) FROM products",
550 /// "SELECT COUNT(*) FROM orders"
551 /// ];
552 ///
553 /// let results = client.databases()
554 /// .query("my-app-db")
555 /// .batch(queries)
556 /// .await?;
557 ///
558 /// for result in results {
559 /// println!("Query returned {} rows", result.meta.row_count);
560 /// }
561 ///
562 /// Ok(())
563 /// }
564 /// ```
565 pub async fn batch<S: AsRef<str>>(self, queries: Vec<S>) -> Result<Vec<QueryResult>> {
566 let url = format!("{}/api/v1/query/batch", self.config.api_url);
567
568 let sql_queries: Vec<&str> = queries.iter().map(|q| q.as_ref()).collect();
569
570 let mut payload = json!({
571 "database": self.database,
572 "queries": sql_queries
573 });
574
575 // Add options to payload
576 if let Some(max_rows) = self.options.max_rows {
577 payload["max_rows"] = json!(max_rows);
578 }
579 if let Some(timeout_ms) = self.options.timeout_ms {
580 payload["timeout_ms"] = json!(timeout_ms);
581 }
582 if self.options.include_meta {
583 payload["include_meta"] = json!(true);
584 }
585
586 let response = self.send_authenticated_request(
587 self.http_client
588 .post(&url)
589 .json(&payload)
590 ).await?;
591
592 if response.status().is_success() {
593 let results: Vec<QueryResult> = response.json().await?;
594 Ok(results)
595 } else {
596 self.handle_error_response(response).await
597 }
598 }
599
600 // Helper methods (duplicate from DatabaseManager for QueryBuilder)
601 async fn send_authenticated_request(
602 &self,
603 mut request: reqwest::RequestBuilder,
604 ) -> Result<reqwest::Response> {
605 if let Some(token) = &self.config.token {
606 request = request.header("Authorization", format!("Bearer {}", token));
607 }
608
609 let response = request
610 .header("Content-Type", "application/json")
611 .send()
612 .await?;
613
614 Ok(response)
615 }
616
617 async fn handle_error_response<T>(&self, response: reqwest::Response) -> Result<T> {
618 let status = response.status();
619 let error_text = response.text().await.unwrap_or_default();
620
621 match status.as_u16() {
622 401 => Err(ZeroTrustError::auth("Authentication failed")),
623 403 => Err(ZeroTrustError::permission_denied("Insufficient permissions")),
624 404 => Err(ZeroTrustError::not_found("Resource not found")),
625 400 => Err(ZeroTrustError::validation(error_text)),
626 500..=599 => Err(ZeroTrustError::server_error(status.as_u16(), error_text)),
627 _ => Err(ZeroTrustError::generic(format!("HTTP {}: {}", status, error_text))),
628 }
629 }
630}
631
632#[cfg(test)]
633mod tests {
634 use super::*;
635 use mockito::{Matcher, Server};
636
637 async fn create_test_database_manager() -> (DatabaseManager, mockito::ServerGuard) {
638 let server = Server::new_async().await;
639 let url = server.url();
640
641 let config = Config::new(&url).unwrap().with_token("test-token");
642 let http_client = reqwest::Client::new();
643
644 let db_manager = DatabaseManager::new(
645 Arc::new(config),
646 Arc::new(http_client),
647 );
648
649 (db_manager, server)
650 }
651
652 #[tokio::test]
653 async fn test_list_databases() {
654 let (db_manager, mut server) = create_test_database_manager().await;
655
656 let mock = server
657 .mock("GET", "/api/v1/databases")
658 .with_status(200)
659 .with_header("content-type", "application/json")
660 .with_body(r#"[
661 {
662 "name": "test-db",
663 "tables": ["users", "products"],
664 "created_at": "2023-01-01T00:00:00Z",
665 "size": 1024,
666 "record_count": 100
667 }
668 ]"#)
669 .create_async()
670 .await;
671
672 let result = db_manager.list().await;
673 assert!(result.is_ok());
674
675 let databases = result.unwrap();
676 assert_eq!(databases.len(), 1);
677 assert_eq!(databases[0].name, "test-db");
678 assert_eq!(databases[0].tables.len(), 2);
679
680 mock.assert_async().await;
681 }
682
683 #[tokio::test]
684 async fn test_query_execute() {
685 let (db_manager, mut server) = create_test_database_manager().await;
686
687 let mock = server
688 .mock("POST", "/api/v1/query/execute")
689 .with_status(200)
690 .with_header("content-type", "application/json")
691 .with_body(r#"{
692 "data": {
693 "columns": ["id", "name"],
694 "rows": [[1, "John"], [2, "Jane"]]
695 },
696 "meta": {
697 "row_count": 2,
698 "execution_time_ms": 150,
699 "operation_type": "Read"
700 }
701 }"#)
702 .create_async()
703 .await;
704
705 let result = db_manager
706 .query("test-db")
707 .execute("SELECT * FROM users")
708 .await;
709
710 assert!(result.is_ok());
711
712 let query_result = result.unwrap();
713 assert_eq!(query_result.meta.row_count, 2);
714 assert_eq!(query_result.data.rows.len(), 2);
715
716 mock.assert_async().await;
717 }
718}