#![cfg(feature = "vector")]
use prax_sqlite::vector::prelude::*;
use prax_sqlite::{SqliteConfig, SqlitePool};
fn vector_lib_available() -> bool {
std::env::var_os("SQLITE_VECTOR_RS_LIB").is_some()
}
async fn make_pool() -> SqlitePool {
SqlitePool::new(SqliteConfig::memory()).await.unwrap()
}
#[tokio::test]
async fn test_pool_creation_succeeds_even_without_cdylib() {
let pool = make_pool().await;
assert!(pool.available_permits() > 0);
let _conn = pool.get().await.unwrap();
}
#[tokio::test]
async fn test_pool_autoregisters_vector_extension() {
if !vector_lib_available() {
eprintln!("skipping: SQLITE_VECTOR_RS_LIB not set");
return;
}
let pool = make_pool().await;
let conn = pool.get().await.unwrap();
let dims: i64 = conn
.inner()
.call(|c| {
c.query_row(
"SELECT vector_dims(vector_from_json('[1.0, 2.0, 3.0]', 'float4'), 'float4')",
[],
|row| row.get(0),
)
.map_err(tokio_rusqlite::Error::Rusqlite)
})
.await
.unwrap();
assert_eq!(dims, 3);
}
#[tokio::test]
async fn test_create_virtual_table_via_vector_index_builder() {
if !vector_lib_available() {
eprintln!("skipping: SQLITE_VECTOR_RS_LIB not set");
return;
}
let pool = make_pool().await;
let conn = pool.get().await.unwrap();
let ddl = VectorIndex::new("docs_vectors")
.rowid_column("doc_id")
.column(
"embedding",
VectorElementType::Float4,
4,
DistanceMetric::Cosine,
Some(VectorIndexKind::Hnsw),
)
.to_create_sql();
conn.execute_batch(&ddl).await.unwrap();
}
#[tokio::test]
async fn test_insert_and_topk_search() {
if !vector_lib_available() {
eprintln!("skipping: SQLITE_VECTOR_RS_LIB not set");
return;
}
let pool = make_pool().await;
let conn = pool.get().await.unwrap();
conn.execute_batch(
"CREATE TABLE documents (id INTEGER PRIMARY KEY, title TEXT NOT NULL);\n\
CREATE VIRTUAL TABLE documents_vectors USING vector(\n\
rowid_column='document_id',\n\
embedding='float4[4] cosine hnsw');",
)
.await
.unwrap();
conn.execute("INSERT INTO documents (id, title) VALUES (1, 'A'), (2, 'B'), (3, 'C')")
.await
.unwrap();
conn.execute_batch(
"INSERT INTO documents_vectors (document_id, embedding) \
VALUES (1, vector_from_json('[1.0, 0.0, 0.0, 0.0]', 'float4'));\n\
INSERT INTO documents_vectors (document_id, embedding) \
VALUES (2, vector_from_json('[0.0, 1.0, 0.0, 0.0]', 'float4'));\n\
INSERT INTO documents_vectors (document_id, embedding) \
VALUES (3, vector_from_json('[0.9, 0.1, 0.0, 0.0]', 'float4'));",
)
.await
.unwrap();
let results: Vec<(i64, f64)> = conn
.inner()
.call(|c| {
let mut stmt = c.prepare(
"SELECT documents.id, \
vector_distance(v.embedding, vector_from_json('[1.0, 0.0, 0.0, 0.0]', 'float4'), 'cosine', 'float4') AS distance \
FROM documents_vectors v \
JOIN documents ON documents.id = v.document_id \
ORDER BY distance \
LIMIT 2",
)?;
let rows = stmt
.query_map([], |row: &rusqlite::Row| {
Ok((row.get::<_, i64>(0)?, row.get::<_, f64>(1)?))
})?
.collect::<Result<Vec<_>, rusqlite::Error>>()?;
Ok(rows)
})
.await
.unwrap();
assert_eq!(results.len(), 2);
assert_eq!(results[0].0, 1);
assert_eq!(results[1].0, 3);
}