reflex/vectordb/bq/
backend.rs1use 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)]
11pub enum BqBackend {
13 Real(BqClient),
15 #[cfg(any(test, feature = "mock"))]
16 Mock(MockBqClient),
18}
19
20impl BqBackend {
21 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}