duckdb_server/
query.rs

1use anyhow::Result;
2use std::path::{Path, PathBuf};
3
4use crate::bundle::{create, load};
5use crate::cache::retrieve;
6use crate::interfaces::{AppError, AppState, Command, QueryParams, QueryResponse};
7
8fn create_bundle_path(bundle_name: &str) -> PathBuf {
9    Path::new(".mosaic").join("bundle").join(bundle_name)
10}
11
12pub async fn handle(state: &AppState, params: QueryParams) -> Result<QueryResponse, AppError> {
13    let command = &params.query_type;
14    tracing::info!("Command: '{:?}', Params: '{:?}'", command, params);
15    match command {
16        Some(Command::Arrow) => {
17            if let Some(sql) = params.sql.as_deref() {
18                let persist = params.persist.unwrap_or(true);
19                let buffer = retrieve(&state.cache, sql, &Command::Arrow, persist, || {
20                    state.db.get_arrow(sql)
21                })
22                .await?;
23                Ok(QueryResponse::Arrow(buffer))
24            } else {
25                Err(AppError::BadRequest)
26            }
27        }
28        Some(Command::Exec) => {
29            if let Some(sql) = params.sql.as_deref() {
30                state.db.execute(sql).await?;
31                Ok(QueryResponse::Empty)
32            } else {
33                Err(AppError::BadRequest)
34            }
35        }
36        Some(Command::Json) => {
37            if let Some(sql) = params.sql.as_deref() {
38                let persist = params.persist.unwrap_or(true);
39                let json: Vec<u8> = retrieve(&state.cache, sql, &Command::Json, persist, || {
40                    state.db.get_json(sql)
41                })
42                .await?;
43                let string = String::from_utf8(json)?;
44                Ok(QueryResponse::Json(string))
45            } else {
46                Err(AppError::BadRequest)
47            }
48        }
49        Some(Command::CreateBundle) => {
50            if let Some(queries) = params.queries {
51                let bundle_name = params.name.unwrap_or_else(|| "default".to_string());
52                let bundle_path = create_bundle_path(&bundle_name);
53                create(state.db.as_ref(), &state.cache, queries, &bundle_path).await?;
54                Ok(QueryResponse::Empty)
55            } else {
56                Err(AppError::BadRequest)
57            }
58        }
59        Some(Command::LoadBundle) => {
60            if let Some(bundle_name) = params.name {
61                let bundle_path = create_bundle_path(&bundle_name);
62                load(state.db.as_ref(), &state.cache, &bundle_path).await?;
63                Ok(QueryResponse::Empty)
64            } else {
65                Err(AppError::BadRequest)
66            }
67        }
68        None => Err(AppError::BadRequest),
69    }
70}
duckdb_server/query.rs

duckdb_server/
query.rs
