reflex/vectordb/bq/
backend.rs

1use crate::cache::BqSearchBackend;
2use crate::vectordb::{SearchResult, VectorDbError, VectorPoint, WriteConsistency};
3
4use super::client::BqClient;
5use super::config::BqConfig;
6
7#[cfg(any(test, feature = "mock"))]
8use super::mock::MockBqClient;
9
10#[derive(Clone)]
11/// Binary-quantized backend wrapper (real or mock).
12pub enum BqBackend {
13    /// Real Qdrant-backed client.
14    Real(BqClient),
15    #[cfg(any(test, feature = "mock"))]
16    /// In-memory mock backend.
17    Mock(MockBqClient),
18}
19
20impl BqBackend {
21    /// Builds a backend from a URL and config (`mock:` URLs require `mock` feature).
22    pub async fn from_config(url: &str, config: BqConfig) -> Result<Self, VectorDbError> {
23        if url.starts_with("mock:") {
24            #[cfg(any(test, feature = "mock"))]
25            {
26                Ok(Self::Mock(MockBqClient::with_config(config)))
27            }
28            #[cfg(not(any(test, feature = "mock")))]
29            {
30                Err(VectorDbError::ConnectionFailed {
31                    url: url.to_string(),
32                    message: "Mock backend not enabled. Compile with --features mock".to_string(),
33                })
34            }
35        } else {
36            let client = BqClient::with_config(url, config).await?;
37            Ok(Self::Real(client))
38        }
39    }
40}
41
42impl BqSearchBackend for BqBackend {
43    async fn is_ready(&self) -> bool {
44        match self {
45            BqBackend::Real(c) => c.health_check().await.is_ok(),
46            #[cfg(any(test, feature = "mock"))]
47            BqBackend::Mock(_) => true,
48        }
49    }
50
51    async fn ensure_collection(&self, name: &str, vector_size: u64) -> Result<(), VectorDbError> {
52        match self {
53            BqBackend::Real(c) => c.ensure_bq_collection(name, vector_size).await,
54            #[cfg(any(test, feature = "mock"))]
55            BqBackend::Mock(c) => c.ensure_bq_collection(name, vector_size).await,
56        }
57    }
58
59    async fn search_bq(
60        &self,
61        collection: &str,
62        query: Vec<f32>,
63        limit: u64,
64        tenant_filter: Option<u64>,
65    ) -> Result<Vec<SearchResult>, VectorDbError> {
66        match self {
67            BqBackend::Real(c) => c.search_bq(collection, query, limit, tenant_filter).await,
68            #[cfg(any(test, feature = "mock"))]
69            BqBackend::Mock(c) => c.search_bq(collection, query, limit, tenant_filter).await,
70        }
71    }
72
73    async fn upsert_points(
74        &self,
75        collection: &str,
76        points: Vec<VectorPoint>,
77        consistency: WriteConsistency,
78    ) -> Result<(), VectorDbError> {
79        match self {
80            BqBackend::Real(c) => c.upsert_points(collection, points, consistency).await,
81            #[cfg(any(test, feature = "mock"))]
82            BqBackend::Mock(c) => c.upsert_points(collection, points, consistency).await,
83        }
84    }
85}