swiftide_integrations/lancedb/
connection_pool.rs

1use anyhow::Context as _;
2use anyhow::Result;
3use deadpool::managed::Manager;
4use derive_builder::Builder;
5use lancedb::connection::ConnectBuilder;
6
7#[derive(Builder, Debug, Clone)]
8#[builder(setter(into), build_fn(error = "anyhow::Error"))]
9pub struct LanceDBPoolManager {
10    uri: String,
11    #[builder(default)]
12    api_key: Option<String>,
13    #[builder(default)]
14    region: Option<String>,
15    #[builder(default)]
16    storage_options: Vec<(String, String)>,
17}
18
19pub type LanceDBConnectionPool = deadpool::managed::Pool<LanceDBPoolManager>;
20
21impl LanceDBPoolManager {
22    pub fn builder() -> LanceDBPoolManagerBuilder {
23        LanceDBPoolManagerBuilder::default()
24    }
25}
26
27impl Manager for LanceDBPoolManager {
28    type Type = lancedb::Connection;
29    type Error = anyhow::Error;
30
31    async fn create(&self) -> Result<Self::Type, Self::Error> {
32        let mut builder = ConnectBuilder::new(&self.uri);
33
34        if let Some(api_key) = &self.api_key {
35            builder = builder.api_key(api_key);
36        }
37
38        if let Some(region) = &self.region {
39            builder = builder.region(region);
40        }
41
42        for (key, value) in &self.storage_options {
43            builder = builder.storage_option(key, value);
44        }
45
46        builder
47            .execute()
48            .await
49            .context("Failed to create LanceDB connection")
50    }
51
52    async fn recycle(
53        &self,
54        _obj: &mut Self::Type,
55        _metrics: &deadpool::managed::Metrics,
56    ) -> deadpool::managed::RecycleResult<Self::Error> {
57        // NOTE: Should work fine with drop
58        Ok(())
59    }
60}
swiftide_integrations/lancedb/connection_pool.rs

swiftide_integrations/lancedb/
connection_pool.rs
