use anyhow::{Context, Result};
use qdrant_client::qdrant::{
CreateCollectionBuilder, CreateFieldIndexCollectionBuilder, Distance, FieldType,
VectorParamsBuilder,
};
use tracing::info;
use super::VectorDb;
#[expect(
async_fn_in_trait,
reason = "async trait method is required for the db interfaces"
)]
pub trait VectorConnectExt {
async fn connect(url: &str, collection: &str, embed_dim: u64) -> Result<Self>
where
Self: Sized;
async fn ensure_collection(&self) -> Result<()>;
}
impl VectorConnectExt for VectorDb {
async fn connect(url: &str, collection: &str, embed_dim: u64) -> Result<VectorDb> {
let client = qdrant_client::Qdrant::from_url(url)
.timeout(std::time::Duration::from_secs(300))
.connect_timeout(std::time::Duration::from_secs(10))
.build()
.context("Failed to build Qdrant client")?;
Ok(VectorDb {
client,
collection: collection.to_owned(),
embed_dim,
})
}
#[expect(
clippy::cognitive_complexity,
reason = "Collection checking logic is sequential"
)]
async fn ensure_collection(&self) -> Result<()> {
let exists = self
.client
.collection_exists(&self.collection)
.await
.context("Failed to check collection existence")?;
if !exists {
info!(
"Creating Qdrant collection '{}' (dim={}, distance=Cosine)",
self.collection, self.embed_dim
);
self.client
.create_collection(
CreateCollectionBuilder::new(&self.collection)
.vectors_config(VectorParamsBuilder::new(self.embed_dim, Distance::Cosine)),
)
.await
.context("Failed to create Qdrant collection")?;
} else {
info!("Qdrant collection '{}' already exists", self.collection);
}
info!(
"Ensuring Keyword payload index on 'repo_name' for collection '{}'",
self.collection
);
self.client
.create_field_index(CreateFieldIndexCollectionBuilder::new(
&self.collection,
"repo_name",
FieldType::Keyword,
))
.await
.context("Failed to create payload index on 'repo_name'")?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::super::VectorDb;
use super::VectorConnectExt;
#[ignore = "requires local Qdrant instance running on http://localhost:6334"]
#[tokio::test]
async fn test_vector_db_connection() {
let result = VectorDb::connect("http://localhost:6334", "test_collection", 384).await;
assert!(result.is_ok(), "Should be able to connect to Qdrant");
}
#[ignore = "requires local Qdrant instance running on http://localhost:6334"]
#[tokio::test]
async fn test_vector_db_ensure_collection() {
let vector_db = VectorDb::connect("http://localhost:6334", "test_collection_ensure", 384)
.await
.expect("Failed to connect to Qdrant");
let result = vector_db.ensure_collection().await;
assert!(result.is_ok(), "Should be able to ensure collection exists");
}
#[ignore = "requires local Qdrant instance running on http://localhost:6334"]
#[tokio::test]
async fn test_vector_db_ensure_collection_idempotent() {
let vector_db =
VectorDb::connect("http://localhost:6334", "test_collection_idempotent", 384)
.await
.expect("Failed to connect to Qdrant");
let result1 = vector_db.ensure_collection().await;
assert!(result1.is_ok());
let result2 = vector_db.ensure_collection().await;
assert!(result2.is_ok());
}
}