From 08bec1cba81dd4f803201d6ff2f84c60a94b9ce6 Mon Sep 17 00:00:00 2001 From: Oh-Hyun Kwon Date: Tue, 7 Jan 2025 14:17:46 -0500 Subject: [PATCH] feat: add command line args for duckdb-server-rust (#647) * feat: add command line args for duckdb-server-rust * use Option * format code --- packages/duckdb-server-rust/Cargo.lock | 95 +++++++++++++++++++++++-- packages/duckdb-server-rust/Cargo.toml | 1 + packages/duckdb-server-rust/src/app.rs | 13 +++- packages/duckdb-server-rust/src/main.rs | 31 +++++++- packages/duckdb-server-rust/src/test.rs | 6 +- 5 files changed, 132 insertions(+), 14 deletions(-) diff --git a/packages/duckdb-server-rust/Cargo.lock b/packages/duckdb-server-rust/Cargo.lock index 5f50d22d..7660a76a 100644 --- a/packages/duckdb-server-rust/Cargo.lock +++ b/packages/duckdb-server-rust/Cargo.lock @@ -93,12 +93,55 @@ version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299" +[[package]] +name = "anstream" +version = "0.6.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8acc5369981196006228e28809f761875c0327210a891e941f4c683b3a99529b" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + [[package]] name = "anstyle" version = "1.0.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "55cc3b69f167a1ef2e161439aa98aed94e6028e5f9a59be9a6ffb47aef1651f9" +[[package]] +name = "anstyle-parse" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b2d16507662817a6a20a9ea92df6652ee4f94f914589377d69f3b21bc5798a9" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "79947af37f4177cfead1110013d678905c37501914fba0efea834c3fe9a8d60c" +dependencies = [ + "windows-sys 0.59.0", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2109dbce0e72be3ec00bed26e6a7479ca384ad226efdd66db8fa2e3a38c83125" +dependencies = [ + "anstyle", + "windows-sys 0.59.0", +] + [[package]] name = "anyhow" version = "1.0.93" @@ -773,28 +816,43 @@ dependencies = [ [[package]] name = "clap" -version = "4.5.21" +version = "4.5.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fb3b4b9e5a7c7514dfa52869339ee98b3156b0bfb4e8a77c4ff4babb64b1604f" +checksum = "3135e7ec2ef7b10c6ed8950f0f792ed96ee093fa088608f1c76e569722700c84" dependencies = [ "clap_builder", + "clap_derive", ] [[package]] name = "clap_builder" -version = "4.5.21" +version = "4.5.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b17a95aa67cc7b5ebd32aa5370189aa0d79069ef1c64ce893bd30fb24bff20ec" +checksum = "30582fc632330df2bd26877bde0c1f4470d57c582bbc070376afcd04d8cb4838" dependencies = [ + "anstream", "anstyle", "clap_lex", + "strsim", +] + +[[package]] +name = "clap_derive" +version = "4.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ac6a0c7b1a9e9a5186361f67dfa1b88213572f427fb9ab038efb2bd8c582dab" +dependencies = [ + "heck 0.5.0", + "proc-macro2", + "quote", + "syn 2.0.90", ] [[package]] name = "clap_lex" -version = "0.7.3" +version = "0.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "afb84c814227b90d6895e01398aee0d8033c00e7466aca416fb6a8e0eb19d8a7" +checksum = "f46ad14479a25103f283c0f10005961cf086d8dc42205bb44c46ac563475dca6" [[package]] name = "cmake" @@ -805,6 +863,12 @@ dependencies = [ "cc", ] +[[package]] +name = "colorchoice" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b63caa9aa9397e2d9480a9b13673856c78d8ac123288526c37d7839f2a86990" + [[package]] name = "comfy-table" version = "7.1.3" @@ -1018,6 +1082,7 @@ dependencies = [ "axum", "axum-server", "axum-server-dual-protocol", + "clap", "criterion", "duckdb", "futures", @@ -1624,6 +1689,12 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "is_terminal_polyfill" +version = "1.70.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7943c866cc5cd64cbc25b2e01621d07fa8eb2a1a23160ee81ce38704e97b8ecf" + [[package]] name = "itertools" version = "0.10.5" @@ -2642,6 +2713,12 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + [[package]] name = "strum" version = "0.25.0" @@ -3109,6 +3186,12 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + [[package]] name = "uuid" version = "1.11.0" diff --git a/packages/duckdb-server-rust/Cargo.toml b/packages/duckdb-server-rust/Cargo.toml index 60a5db2b..8477eb9b 100644 --- a/packages/duckdb-server-rust/Cargo.toml +++ b/packages/duckdb-server-rust/Cargo.toml @@ -19,6 +19,7 @@ async-trait = "0.1.83" axum = { version = "0.7", features = ["http1", "http2", "ws", "json", "tokio", "tracing", "macros"] } axum-server = { version = "0.7", features = ["tls-rustls"] } axum-server-dual-protocol = "0.7" +clap = { version = "4.5.23", features = ["derive"] } duckdb = { version = "1", features = ["bundled", "csv", "json", "parquet", "url", "r2d2"] } futures = "0.3" listenfd = "1.0" diff --git a/packages/duckdb-server-rust/src/app.rs b/packages/duckdb-server-rust/src/app.rs index 086b7dbc..bfe8f8b1 100644 --- a/packages/duckdb-server-rust/src/app.rs +++ b/packages/duckdb-server-rust/src/app.rs @@ -40,10 +40,17 @@ async fn handle_post( query::handle(&state, params).await } -pub fn app() -> Result { +pub fn app( + dp_path: Option<&str>, + connection_pool_size: Option, + cache_size: Option, +) -> Result { // Database and state setup - let db = ConnectionPool::new(":memory:", 10)?; - let cache = lru::LruCache::new(1000.try_into()?); + let db = ConnectionPool::new( + dp_path.unwrap_or(":memory:"), + connection_pool_size.unwrap_or(10), + )?; + let cache = lru::LruCache::new(cache_size.unwrap_or(1000).try_into()?); let state = Arc::new(AppState { db: Box::new(db), diff --git a/packages/duckdb-server-rust/src/main.rs b/packages/duckdb-server-rust/src/main.rs index d4926ff6..3922dbf5 100644 --- a/packages/duckdb-server-rust/src/main.rs +++ b/packages/duckdb-server-rust/src/main.rs @@ -1,5 +1,6 @@ use anyhow::Result; use axum_server::tls_rustls::RustlsConfig; +use clap::Parser; use listenfd::ListenFd; use std::net::TcpListener; use std::{net::Ipv4Addr, net::SocketAddr, path::PathBuf}; @@ -14,8 +15,30 @@ mod interfaces; mod query; mod websocket; +#[derive(Parser, Debug)] +#[command(version, about, long_about = None)] +struct Args { + /// Path of database file (e.g., "database.db". ":memory:" for in-memory database) + #[arg(default_value = ":memory:")] + database: String, + + /// HTTP Port + #[arg(short, long, default_value_t = 3000)] + port: u16, + + /// Max connection pool size + #[arg(long, default_value_t = 10)] + connection_pool_size: u32, + + /// Max number of cache entries + #[arg(long, default_value_t = 1000)] + cache_size: usize, +} + #[tokio::main] async fn main() -> Result<(), Box> { + let args = Args::parse(); + // Tracing setup tracing_subscriber::registry() .with( @@ -27,7 +50,11 @@ async fn main() -> Result<(), Box> { .init(); // App setup - let app = app::app()?; + let app = app::app( + Some(&args.database), + Some(args.connection_pool_size), + Some(args.cache_size), + )?; // TLS configuration let mut config = RustlsConfig::from_pem_file( @@ -42,7 +69,7 @@ async fn main() -> Result<(), Box> { } // Listenfd setup - let addr = SocketAddr::new(Ipv4Addr::LOCALHOST.into(), 3000); + let addr = SocketAddr::new(Ipv4Addr::LOCALHOST.into(), args.port); let mut listenfd = ListenFd::from_env(); let listener = match listenfd.take_tcp_listener(0)? { // if we are given a tcp listener on listen fd 0, we use that one diff --git a/packages/duckdb-server-rust/src/test.rs b/packages/duckdb-server-rust/src/test.rs index 99d2b8a9..b825684b 100644 --- a/packages/duckdb-server-rust/src/test.rs +++ b/packages/duckdb-server-rust/src/test.rs @@ -92,7 +92,7 @@ async fn get_arrow() -> Result<()> { #[tokio::test] async fn select_1_get() -> Result<()> { - let app = app::app()?; + let app = app::app(None, None, None)?; let response = app .oneshot( @@ -112,7 +112,7 @@ async fn select_1_get() -> Result<()> { #[tokio::test] async fn select_1_post() -> Result<()> { - let app = app::app()?; + let app = app::app(None, None, None)?; let response = app .oneshot( @@ -136,7 +136,7 @@ async fn select_1_post() -> Result<()> { #[tokio::test] async fn query_arrow() -> Result<()> { - let app = app::app()?; + let app = app::app(None, None, None)?; let response = app .oneshot(