From 5790f85383195141914a0e3d8bb22d85f7798137 Mon Sep 17 00:00:00 2001 From: Nathan Sobo Date: Thu, 24 Oct 2024 11:33:45 -0600 Subject: [PATCH] Experiment with duckdb for analytics --- Cargo.lock | 349 ++++++++++++++++++++++++- crates/collab/Cargo.toml | 2 + crates/collab/src/api/events.rs | 69 +++-- crates/collab/src/clickhouse.rs | 4 +- crates/collab/src/duckdb.rs | 18 ++ crates/collab/src/lib.rs | 20 ++ crates/collab/src/tests/test_server.rs | 2 + 7 files changed, 433 insertions(+), 31 deletions(-) create mode 100644 crates/collab/src/duckdb.rs diff --git a/Cargo.lock b/Cargo.lock index 4e86627d80..6bae99a69b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -71,6 +71,7 @@ checksum = "e89da841a80418a9b391ebaea17f5c112ffaaa96f621d2c285b5174da76b9011" dependencies = [ "cfg-if", "const-random", + "getrandom 0.2.15", "once_cell", "version_check", "zerocopy", @@ -306,6 +307,168 @@ dependencies = [ "serde", ] +[[package]] +name = "arrow" +version = "53.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a9ba0d7248932f4e2a12fb37f0a2e3ec82b3bdedbac2a1dce186e036843b8f8c" +dependencies = [ + "arrow-arith", + "arrow-array", + "arrow-buffer", + "arrow-cast", + "arrow-data", + "arrow-ord", + "arrow-row", + "arrow-schema", + "arrow-select", + "arrow-string", +] + +[[package]] +name = "arrow-arith" +version = "53.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d60afcdc004841a5c8d8da4f4fa22d64eb19c0c01ef4bcedd77f175a7cf6e38f" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "chrono", + "half", + "num", +] + +[[package]] +name = "arrow-array" +version = "53.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f16835e8599dbbb1659fd869d865254c4cf32c6c2bb60b6942ac9fc36bfa5da" +dependencies = [ + "ahash 0.8.11", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "chrono", + "half", + "hashbrown 0.14.5", + "num", +] + +[[package]] +name = "arrow-buffer" +version = "53.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a1f34f0faae77da6b142db61deba2cb6d60167592b178be317b341440acba80" +dependencies = [ + "bytes 1.7.2", + "half", + "num", +] + +[[package]] +name = "arrow-cast" +version = "53.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "450e4abb5775bca0740bec0bcf1b1a5ae07eff43bd625661c4436d8e8e4540c4" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "arrow-select", + "atoi", + "base64 0.22.1", + "chrono", + "comfy-table", + "half", + "lexical-core", + "num", + "ryu", +] + +[[package]] +name = "arrow-data" +version = "53.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b1e618bbf714c7a9e8d97203c806734f012ff71ae3adc8ad1b075689f540634" +dependencies = [ + "arrow-buffer", + "arrow-schema", + "half", + "num", +] + +[[package]] +name = "arrow-ord" +version = "53.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2427f37b4459a4b9e533045abe87a5183a5e0995a3fc2c2fd45027ae2cc4ef3f" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "arrow-select", + "half", + "num", +] + +[[package]] +name = "arrow-row" +version = "53.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15959657d92e2261a7a323517640af87f5afd9fd8a6492e424ebee2203c567f6" +dependencies = [ + "ahash 0.8.11", + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "half", +] + +[[package]] +name = "arrow-schema" +version = "53.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fbf0388a18fd7f7f3fe3de01852d30f54ed5182f9004db700fbe3ba843ed2794" +dependencies = [ + "bitflags 2.6.0", +] + +[[package]] +name = "arrow-select" +version = "53.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b83e5723d307a38bf00ecd2972cd078d1339c7fd3eb044f609958a9a24463f3a" +dependencies = [ + "ahash 0.8.11", + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "num", +] + +[[package]] +name = "arrow-string" +version = "53.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ab3db7c09dd826e74079661d84ed01ed06547cf75d52c2818ef776d0d852305" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "arrow-select", + "memchr", + "num", + "regex", + "regex-syntax 0.8.4", +] + [[package]] name = "as-raw-xcb-connection" version = "1.0.1" @@ -870,7 +1033,7 @@ dependencies = [ "libc", "pin-project", "redox_syscall 0.2.16", - "xattr", + "xattr 0.2.3", ] [[package]] @@ -1580,7 +1743,7 @@ dependencies = [ "bitflags 2.6.0", "cexpr", "clang-sys", - "itertools 0.10.5", + "itertools 0.12.1", "lazy_static", "lazycell", "proc-macro2", @@ -2551,6 +2714,7 @@ dependencies = [ "dashmap 6.0.1", "derive_more", "dev_server_projects", + "duckdb", "editor", "env_logger", "envy", @@ -2584,6 +2748,7 @@ dependencies = [ "project", "prometheus", "prost", + "r2d2", "rand 0.8.5", "recent_projects", "release_channel", @@ -2699,6 +2864,17 @@ dependencies = [ "memchr", ] +[[package]] +name = "comfy-table" +version = "7.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b34115915337defe99b2aff5c2ce6771e5fbc4079f4b506301f5cf394c8452f7" +dependencies = [ + "strum 0.26.3", + "strum_macros 0.26.4", + "unicode-width", +] + [[package]] name = "command_palette" version = "0.1.0" @@ -3634,6 +3810,26 @@ dependencies = [ "phf", ] +[[package]] +name = "duckdb" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "86844939330ba6ce345c4b5333d3be45c4f0c092779bf9617bba92efb8b841f5" +dependencies = [ + "arrow", + "cast", + "fallible-iterator", + "fallible-streaming-iterator", + "hashlink", + "libduckdb-sys", + "memchr", + "num-integer", + "r2d2", + "rust_decimal", + "smallvec", + "strum 0.25.0", +] + [[package]] name = "dwrote" version = "0.11.1" @@ -4189,6 +4385,12 @@ version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2acce4a10f12dc2fb14a218589d4f1f62ef011b2d0cc4b3cb1bba8e94da14649" +[[package]] +name = "fallible-streaming-iterator" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a" + [[package]] name = "fancy-regex" version = "0.12.0" @@ -5182,6 +5384,7 @@ checksum = "6dd08c532ae367adf81c312a4580bc67f1d0fe8bc9c460520283f4c0ff277888" dependencies = [ "cfg-if", "crunchy", + "num-traits", ] [[package]] @@ -5591,7 +5794,7 @@ dependencies = [ "httpdate", "itoa", "pin-project-lite", - "socket2 0.4.10", + "socket2 0.5.7", "tokio", "tower-service", "tracing", @@ -6430,6 +6633,70 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "03087c2bad5e1034e8cace5926dec053fb3790248370865f5117a7d0213354c8" +[[package]] +name = "lexical-core" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0431c65b318a590c1de6b8fd6e72798c92291d27762d94c9e6c37ed7a73d8458" +dependencies = [ + "lexical-parse-float", + "lexical-parse-integer", + "lexical-util", + "lexical-write-float", + "lexical-write-integer", +] + +[[package]] +name = "lexical-parse-float" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eb17a4bdb9b418051aa59d41d65b1c9be5affab314a872e5ad7f06231fb3b4e0" +dependencies = [ + "lexical-parse-integer", + "lexical-util", + "static_assertions", +] + +[[package]] +name = "lexical-parse-integer" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5df98f4a4ab53bf8b175b363a34c7af608fe31f93cc1fb1bf07130622ca4ef61" +dependencies = [ + "lexical-util", + "static_assertions", +] + +[[package]] +name = "lexical-util" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85314db53332e5c192b6bca611fb10c114a80d1b831ddac0af1e9be1b9232ca0" +dependencies = [ + "static_assertions", +] + +[[package]] +name = "lexical-write-float" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e7c3ad4e37db81c1cbe7cf34610340adc09c322871972f74877a712abc6c809" +dependencies = [ + "lexical-util", + "lexical-write-integer", + "static_assertions", +] + +[[package]] +name = "lexical-write-integer" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eb89e9f6958b83258afa3deed90b5de9ef68eef090ad5086c791cd2345610162" +dependencies = [ + "lexical-util", + "static_assertions", +] + [[package]] name = "libc" version = "0.2.159" @@ -6446,6 +6713,21 @@ dependencies = [ "pkg-config", ] +[[package]] +name = "libduckdb-sys" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eac2de5219db852597558df5dcd617ffccd5cbd7b9f5402ccbf899aca6cb6047" +dependencies = [ + "autocfg", + "flate2", + "pkg-config", + "serde", + "serde_json", + "tar", + "vcpkg", +] + [[package]] name = "libfuzzer-sys" version = "0.4.7" @@ -6476,7 +6758,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4979f22fdb869068da03c9f7528f8297c6fd2606bc3a4affe42e6a823fdb8da4" dependencies = [ "cfg-if", - "windows-targets 0.48.5", + "windows-targets 0.52.6", ] [[package]] @@ -8795,6 +9077,17 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "r2d2" +version = "0.8.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51de85fb3fb6524929c8a2eb85e6b6d363de4e8c48f9e2c2eac4944abc181c93" +dependencies = [ + "log", + "parking_lot", + "scheduled-thread-pool", +] + [[package]] name = "radium" version = "0.7.0" @@ -9870,6 +10163,15 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "scheduled-thread-pool" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3cbc66816425a074528352f5789333ecff06ca41b36b0b0efdfbb29edc391a19" +dependencies = [ + "parking_lot", +] + [[package]] name = "schemars" version = "0.8.21" @@ -11097,7 +11399,7 @@ version = "0.25.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "290d54ea6f91c969195bdbcd7442c8c2a2ba87da8bf60a7ee86a235d4bc1e125" dependencies = [ - "strum_macros", + "strum_macros 0.25.3", ] [[package]] @@ -11119,6 +11421,19 @@ dependencies = [ "syn 2.0.76", ] +[[package]] +name = "strum_macros" +version = "0.26.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c6bee85a5a24955dc440386795aa378cd9cf82acd5f764469152d2270e581be" +dependencies = [ + "heck 0.5.0", + "proc-macro2", + "quote", + "rustversion", + "syn 2.0.76", +] + [[package]] name = "subtle" version = "2.6.1" @@ -11484,6 +11799,17 @@ version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369" +[[package]] +name = "tar" +version = "0.4.42" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ff6c40d3aedb5e06b57c6f669ad17ab063dd1e63d977c6a88e7f4dfa4f04020" +dependencies = [ + "filetime", + "libc", + "xattr 1.3.1", +] + [[package]] name = "target-lexicon" version = "0.12.16" @@ -13747,7 +14073,7 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.59.0", ] [[package]] @@ -14422,6 +14748,17 @@ dependencies = [ "libc", ] +[[package]] +name = "xattr" +version = "1.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8da84f1a25939b27f6820d92aed108f83ff920fdf11a7b19366c27c4cda81d4f" +dependencies = [ + "libc", + "linux-raw-sys 0.4.14", + "rustix 0.38.35", +] + [[package]] name = "xcursor" version = "0.3.8" diff --git a/crates/collab/Cargo.toml b/crates/collab/Cargo.toml index ad2c013668..0f1938e7ad 100644 --- a/crates/collab/Cargo.toml +++ b/crates/collab/Cargo.toml @@ -33,6 +33,7 @@ clock.workspace = true collections.workspace = true dashmap.workspace = true derive_more.workspace = true +duckdb = { version = "1.1.1", features = ["r2d2"] } envy = "0.4.2" futures.workspace = true google_ai.workspace = true @@ -51,6 +52,7 @@ reqwest = { version = "0.11", features = ["json"] } reqwest_client.workspace = true rpc.workspace = true rustc-demangle.workspace = true +r2d2 = "0.8.9" scrypt = "0.11" sea-orm = { version = "1.1.0-rc.1", features = ["sqlx-postgres", "postgres-array", "runtime-tokio-rustls", "with-uuid"] } semantic_version.workspace = true diff --git a/crates/collab/src/api/events.rs b/crates/collab/src/api/events.rs index c1ab3d7939..f1458a9953 100644 --- a/crates/collab/src/api/events.rs +++ b/crates/collab/src/api/events.rs @@ -1,6 +1,7 @@ use super::ips_file::IpsFile; use crate::api::CloudflareIpCountryHeader; -use crate::clickhouse::write_to_table; +use crate::clickhouse; +use crate::duckdb; use crate::{api::slack, AppState, Error, Result}; use anyhow::{anyhow, Context}; use aws_sdk_s3::primitives::ByteStream; @@ -11,6 +12,7 @@ use axum::{ routing::post, Extension, Router, TypedHeader, }; +use duckdb::Connection as DuckDbConnection; use rpc::ExtensionMetadata; use semantic_version::SemanticVersion; use serde::{Serialize, Serializer}; @@ -388,13 +390,6 @@ pub async fn post_events( country_code_header: Option>, body: Bytes, ) -> Result<()> { - let Some(clickhouse_client) = app.clickhouse_client.clone() else { - Err(Error::http( - StatusCode::NOT_IMPLEMENTED, - "not supported".into(), - ))? - }; - let Some(expected) = calculate_json_checksum(app.clone(), &body) else { return Err(Error::http( StatusCode::INTERNAL_SERVER_ERROR, @@ -527,10 +522,26 @@ pub async fn post_events( } } - to_upload - .upload(&clickhouse_client) + if let Some(clickhouse_client) = app.clickhouse_client.clone() { + to_upload + .upload_to_clickhouse(&clickhouse_client) + .await + .map_err(|err| Error::Internal(anyhow!(err)))?; + } + + if let Some(pool) = app.duckdb_pool.clone() { + tokio::task::spawn_blocking(move || { + let connection = pool + .get() + .context("can't get duckdb connection from pool")?; + to_upload + .upload_to_duckdb(&connection) + .map_err(|err| Error::Internal(anyhow!(err)))?; + anyhow::Ok(()) + }) .await - .map_err(|err| Error::Internal(anyhow!(err)))?; + .context("error spawning duckdb write")??; + } Ok(()) } @@ -552,14 +563,17 @@ struct ToUpload { } impl ToUpload { - pub async fn upload(&self, clickhouse_client: &clickhouse::Client) -> anyhow::Result<()> { + pub async fn upload_to_clickhouse( + &self, + clickhouse_client: &clickhouse::Client, + ) -> anyhow::Result<()> { const EDITOR_EVENTS_TABLE: &str = "editor_events"; - write_to_table(EDITOR_EVENTS_TABLE, &self.editor_events, clickhouse_client) + clickhouse::write_to_table(EDITOR_EVENTS_TABLE, &self.editor_events, clickhouse_client) .await .with_context(|| format!("failed to upload to table '{EDITOR_EVENTS_TABLE}'"))?; const INLINE_COMPLETION_EVENTS_TABLE: &str = "inline_completion_events"; - write_to_table( + clickhouse::write_to_table( INLINE_COMPLETION_EVENTS_TABLE, &self.inline_completion_events, clickhouse_client, @@ -568,7 +582,7 @@ impl ToUpload { .with_context(|| format!("failed to upload to table '{INLINE_COMPLETION_EVENTS_TABLE}'"))?; const ASSISTANT_EVENTS_TABLE: &str = "assistant_events"; - write_to_table( + clickhouse::write_to_table( ASSISTANT_EVENTS_TABLE, &self.assistant_events, clickhouse_client, @@ -577,27 +591,27 @@ impl ToUpload { .with_context(|| format!("failed to upload to table '{ASSISTANT_EVENTS_TABLE}'"))?; const CALL_EVENTS_TABLE: &str = "call_events"; - write_to_table(CALL_EVENTS_TABLE, &self.call_events, clickhouse_client) + clickhouse::write_to_table(CALL_EVENTS_TABLE, &self.call_events, clickhouse_client) .await .with_context(|| format!("failed to upload to table '{CALL_EVENTS_TABLE}'"))?; const CPU_EVENTS_TABLE: &str = "cpu_events"; - write_to_table(CPU_EVENTS_TABLE, &self.cpu_events, clickhouse_client) + clickhouse::write_to_table(CPU_EVENTS_TABLE, &self.cpu_events, clickhouse_client) .await .with_context(|| format!("failed to upload to table '{CPU_EVENTS_TABLE}'"))?; const MEMORY_EVENTS_TABLE: &str = "memory_events"; - write_to_table(MEMORY_EVENTS_TABLE, &self.memory_events, clickhouse_client) + clickhouse::write_to_table(MEMORY_EVENTS_TABLE, &self.memory_events, clickhouse_client) .await .with_context(|| format!("failed to upload to table '{MEMORY_EVENTS_TABLE}'"))?; const APP_EVENTS_TABLE: &str = "app_events"; - write_to_table(APP_EVENTS_TABLE, &self.app_events, clickhouse_client) + clickhouse::write_to_table(APP_EVENTS_TABLE, &self.app_events, clickhouse_client) .await .with_context(|| format!("failed to upload to table '{APP_EVENTS_TABLE}'"))?; const SETTING_EVENTS_TABLE: &str = "setting_events"; - write_to_table( + clickhouse::write_to_table( SETTING_EVENTS_TABLE, &self.setting_events, clickhouse_client, @@ -606,7 +620,7 @@ impl ToUpload { .with_context(|| format!("failed to upload to table '{SETTING_EVENTS_TABLE}'"))?; const EXTENSION_EVENTS_TABLE: &str = "extension_events"; - write_to_table( + clickhouse::write_to_table( EXTENSION_EVENTS_TABLE, &self.extension_events, clickhouse_client, @@ -615,22 +629,29 @@ impl ToUpload { .with_context(|| format!("failed to upload to table '{EXTENSION_EVENTS_TABLE}'"))?; const EDIT_EVENTS_TABLE: &str = "edit_events"; - write_to_table(EDIT_EVENTS_TABLE, &self.edit_events, clickhouse_client) + clickhouse::write_to_table(EDIT_EVENTS_TABLE, &self.edit_events, clickhouse_client) .await .with_context(|| format!("failed to upload to table '{EDIT_EVENTS_TABLE}'"))?; const ACTION_EVENTS_TABLE: &str = "action_events"; - write_to_table(ACTION_EVENTS_TABLE, &self.action_events, clickhouse_client) + clickhouse::write_to_table(ACTION_EVENTS_TABLE, &self.action_events, clickhouse_client) .await .with_context(|| format!("failed to upload to table '{ACTION_EVENTS_TABLE}'"))?; const REPL_EVENTS_TABLE: &str = "repl_events"; - write_to_table(REPL_EVENTS_TABLE, &self.repl_events, clickhouse_client) + clickhouse::write_to_table(REPL_EVENTS_TABLE, &self.repl_events, clickhouse_client) .await .with_context(|| format!("failed to upload to table '{REPL_EVENTS_TABLE}'"))?; Ok(()) } + + pub fn upload_to_duckdb(&self, connection: &DuckDbConnection) -> anyhow::Result<()> { + duckdb::write_to_table("edit_events", &self.edit_events, &connection) + .with_context(|| format!("failed to upload to table 'edit_events"))?; + + Ok(()) + } } pub fn serialize_country_code(country_code: &str, serializer: S) -> Result diff --git a/crates/collab/src/clickhouse.rs b/crates/collab/src/clickhouse.rs index 2937116bad..c5fb241668 100644 --- a/crates/collab/src/clickhouse.rs +++ b/crates/collab/src/clickhouse.rs @@ -1,4 +1,6 @@ -use serde::Serialize; +pub use clickhouse::*; + +use ::serde::Serialize; /// Writes the given rows to the specified Clickhouse table. pub async fn write_to_table( diff --git a/crates/collab/src/duckdb.rs b/crates/collab/src/duckdb.rs new file mode 100644 index 0000000000..4d02dceac5 --- /dev/null +++ b/crates/collab/src/duckdb.rs @@ -0,0 +1,18 @@ +pub use duckdb::*; + +pub fn write_to_table( + table_name: &str, + rows: &[T], + connection: &duckdb::Connection, +) -> anyhow::Result<()> +where + T: serde::Serialize, +{ + let mut stmt = connection.prepare(&format!( + "INSERT INTO {} SELECT * FROM json_each(?)", + table_name + ))?; + let json = serde_json::to_string(rows)?; + stmt.execute([json])?; + Ok(()) +} diff --git a/crates/collab/src/lib.rs b/crates/collab/src/lib.rs index 78f514adb7..173d6d0c45 100644 --- a/crates/collab/src/lib.rs +++ b/crates/collab/src/lib.rs @@ -3,6 +3,7 @@ pub mod auth; mod cents; pub mod clickhouse; pub mod db; +mod duckdb; pub mod env; pub mod executor; pub mod llm; @@ -24,6 +25,7 @@ use axum::{ }; pub use cents::*; use db::{ChannelId, Database}; +use duckdb::DuckdbConnectionManager; use executor::Executor; use llm::db::LlmDatabase; pub use rate_limiter::*; @@ -155,6 +157,7 @@ pub struct Config { pub clickhouse_user: Option, pub clickhouse_password: Option, pub clickhouse_database: Option, + pub duckdb_path: Option, pub invite_link_prefix: String, pub live_kit_server: Option, pub live_kit_key: Option, @@ -230,6 +233,7 @@ impl Config { clickhouse_user: None, clickhouse_password: None, clickhouse_database: None, + duckdb_path: None, zed_client_checksum_seed: None, slack_panics_webhook: None, auto_join_channel_id: None, @@ -276,9 +280,16 @@ pub struct AppState { pub rate_limiter: Arc, pub executor: Executor, pub clickhouse_client: Option<::clickhouse::Client>, + pub duckdb_pool: Option>, pub config: Config, } +#[test] +fn test_app_state_is_send_and_sync() { + fn assert_send_sync() {} + assert_send_sync::(); +} + impl AppState { pub async fn new(config: Config, executor: Executor) -> Result> { let mut db_options = db::ConnectOptions::new(config.database_url.clone()); @@ -317,6 +328,14 @@ impl AppState { let db = Arc::new(db); let stripe_client = build_stripe_client(&config).map(Arc::new).log_err(); + + let duckdb_pool = config.duckdb_path.as_ref().and_then(|path| { + r2d2::Pool::builder() + .max_size(15) + .build(DuckdbConnectionManager::file(path).log_err()?) + .log_err() + }); + let this = Self { db: db.clone(), llm_db, @@ -332,6 +351,7 @@ impl AppState { .clickhouse_url .as_ref() .and_then(|_| build_clickhouse_client(&config).log_err()), + duckdb_pool, config, }; Ok(Arc::new(this)) diff --git a/crates/collab/src/tests/test_server.rs b/crates/collab/src/tests/test_server.rs index 210a049e0b..fe4df76aae 100644 --- a/crates/collab/src/tests/test_server.rs +++ b/crates/collab/src/tests/test_server.rs @@ -643,6 +643,7 @@ impl TestServer { rate_limiter: Arc::new(RateLimiter::new(test_db.db().clone())), executor, clickhouse_client: None, + duckdb_pool: None, config: Config { http_port: 0, database_url: "".into(), @@ -673,6 +674,7 @@ impl TestServer { clickhouse_user: None, clickhouse_password: None, clickhouse_database: None, + duckdb_path: None, zed_client_checksum_seed: None, slack_panics_webhook: None, auto_join_channel_id: None,