use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{OnceLock, RwLock};
use autumn_web::pagination::Page;
use autumn_web::search::{IndexDefinition, SearchDocument};
use diesel::sql_types::{BigInt, Double, Nullable, Text};
use diesel_async::RunQueryDsl;
use diesel_async::pooled_connection::deadpool::Pool;
use crate::backend::{
BackendCapabilities, BoxFuture, IndexedDocument, KeywordQuery, SearchBackend, SearchFilter,
SearchHit, VectorQuery, empty_page,
};
use crate::error::{SearchError, SearchResult};
use crate::source::DocumentSource;
use crate::text::query_tokens;
type RuntimePool = Pool<autumn_web::RuntimeConnection>;
type BoxedQuery<'a> = diesel::query_builder::BoxedSqlQuery<
'a,
autumn_web::RuntimeBackend,
diesel::query_builder::SqlQuery,
>;
pub const DOCUMENTS_TABLE: &str = "autumn_search_documents";
pub const DELETES_TABLE: &str = "autumn_search_deletes";
const DELETE_RETENTION: &str = "7 days";
const TS_RANK_WEIGHTS: &str = "'{0.1, 0.2, 0.5, 1.0}'::float4[]";
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum VectorMode {
PgVector {
dimensions: usize,
},
Array,
}
impl VectorMode {
#[must_use]
pub const fn is_pgvector(self) -> bool {
matches!(self, Self::PgVector { .. })
}
}
#[derive(diesel::QueryableByName)]
struct HitRow {
#[diesel(sql_type = BigInt)]
record_id: i64,
#[diesel(sql_type = Double)]
score: f64,
}
#[derive(diesel::QueryableByName)]
struct CountRow {
#[diesel(sql_type = BigInt)]
total: i64,
}
#[derive(diesel::QueryableByName)]
struct EmbeddingRow {
#[diesel(sql_type = Nullable<Text>)]
embedding: Option<String>,
}
#[derive(diesel::QueryableByName)]
struct SourceRow {
#[diesel(sql_type = BigInt)]
id: i64,
#[diesel(sql_type = Text)]
fields: String,
#[diesel(sql_type = Nullable<Text>)]
tenant_id: Option<String>,
}
#[derive(diesel::QueryableByName)]
struct WatermarkRow {
#[diesel(sql_type = Text)]
watermark: String,
}
#[derive(diesel::QueryableByName)]
struct WidthRow {
#[diesel(sql_type = BigInt)]
width: i64,
}
#[derive(diesel::QueryableByName)]
struct ExistsRow {
#[diesel(sql_type = diesel::sql_types::Bool)]
present: bool,
}
#[derive(diesel::QueryableByName)]
struct OptionalColumnsRow {
#[diesel(sql_type = Nullable<diesel::sql_types::Bool>)]
has_tenant: Option<bool>,
#[diesel(sql_type = Nullable<diesel::sql_types::Bool>)]
has_deleted_at: Option<bool>,
}
pub struct PostgresSearchStore {
pool: OnceLock<RuntimePool>,
dimensions: Option<usize>,
vector_mode: RwLock<Option<VectorMode>>,
schema_ready: AtomicBool,
optional_columns: RwLock<std::collections::HashMap<&'static str, OptionalColumns>>,
vector_column_width: RwLock<Option<usize>>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct OptionalColumns {
tenant_id: bool,
deleted_at: bool,
}
impl std::fmt::Debug for PostgresSearchStore {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PostgresSearchStore")
.field("pool_installed", &self.pool.get().is_some())
.field("dimensions", &self.dimensions)
.field("vector_mode", &self.vector_mode())
.field(
"schema_ready",
&AtomicBool::load(&self.schema_ready, Ordering::SeqCst),
)
.finish_non_exhaustive()
}
}
impl PostgresSearchStore {
#[must_use]
pub fn new(dimensions: Option<usize>) -> Self {
Self {
pool: OnceLock::new(),
dimensions,
vector_mode: RwLock::new(None),
schema_ready: AtomicBool::new(false),
optional_columns: RwLock::new(std::collections::HashMap::new()),
vector_column_width: RwLock::new(None),
}
}
pub fn install_pool(&self, pool: RuntimePool) {
let _ = self.pool.set(pool);
}
#[must_use]
pub fn vector_mode(&self) -> Option<VectorMode> {
self.vector_mode.read().ok().and_then(|guard| *guard)
}
#[must_use]
pub fn vector_column_width(&self) -> Option<usize> {
self.vector_column_width
.read()
.ok()
.and_then(|guard| *guard)
}
#[allow(clippy::too_many_lines)] async fn write_documents(
&self,
definition: &IndexDefinition,
documents: &[IndexedDocument],
watermark: Option<&str>,
) -> SearchResult<()> {
use std::fmt::Write as _;
const PARAM_BUDGET: usize = 60_000;
const MAX_ROWS_PER_CHUNK: usize = 1_000;
checked(definition)?;
if documents.is_empty() {
return Ok(());
}
let mut conn = self.conn().await?;
let vector_width = self.vector_column_width();
let (vec_column, vec_update) = if vector_width.is_some() {
(
", embedding_vec",
", embedding_vec = EXCLUDED.embedding_vec",
)
} else {
("", "")
};
validate_embedding_widths(documents, vector_width, self.vector_mode())?;
let mut kept_documents = dedupe_by_id(documents, watermark.is_some());
kept_documents.sort_unstable_by_key(|document| document.id());
let mut cursor = 0usize;
while cursor < kept_documents.len() {
let mut binds = vec![
Bound::Text(definition.name.to_owned()),
Bound::Text(definition.language.to_owned()),
];
let watermark_param = watermark.map(|value| {
binds.push(Bound::Text(value.to_owned()));
3usize
});
let mut next_param: usize = if watermark_param.is_some() { 4 } else { 3 };
let mut rows = Vec::new();
let mut record_ids = Vec::new();
while let Some(&document) = kept_documents.get(cursor) {
let doc_param_count = 5usize
.saturating_add(usize::from(vector_width.is_some()))
.saturating_add(
document
.document
.fields
.iter()
.filter(|field| definition.weight_of(field.name).is_some())
.count(),
);
if !rows.is_empty()
&& (next_param.saturating_add(doc_param_count) > PARAM_BUDGET
|| rows.len() >= MAX_ROWS_PER_CHUNK)
{
break;
}
cursor = cursor.saturating_add(1);
let record_id = document.id();
record_ids.push(record_id);
let record_id_param = format!("${next_param}");
binds.push(Bound::BigInt(record_id));
next_param = next_param.saturating_add(1);
let tenant_param = format!("${next_param}");
binds.push(Bound::NullableText(document.document.tenant_id.clone()));
next_param = next_param.saturating_add(1);
let fields_param = format!("${next_param}");
binds.push(Bound::Text(fields_json(document)?));
next_param = next_param.saturating_add(1);
let content_param = format!("${next_param}");
binds.push(Bound::Text(document.document.text()));
next_param = next_param.saturating_add(1);
let embedding_param = format!("${next_param}");
binds.push(Bound::NullableText(
document.embedding.as_deref().map(array_literal),
));
next_param = next_param.saturating_add(1);
let vec_value = if let Some(width) = vector_width {
let literal = match document.embedding.as_deref() {
Some(embedding) if embedding.len() != width => {
if self.vector_mode().is_some_and(VectorMode::is_pgvector) {
return Err(SearchError::DimensionMismatch {
expected: width,
actual: embedding.len(),
});
}
None
}
other => other.map(vector_literal),
};
let placeholder = format!(", ${next_param}::vector");
binds.push(Bound::NullableText(literal));
next_param = next_param.saturating_add(1);
placeholder
} else {
String::new()
};
let mut vector_sql = String::new();
for field in &document.document.fields {
let Some(weight) = definition.weight_of(field.name) else {
continue;
};
if !vector_sql.is_empty() {
vector_sql.push_str(" || ");
}
let field_value_param = format!("${next_param}");
let _ = write!(
vector_sql,
"setweight(to_tsvector($2::text::regconfig, {field_value_param}), '{weight}')"
);
binds.push(Bound::Text(field.value.clone()));
next_param = next_param.saturating_add(1);
}
if vector_sql.is_empty() {
"to_tsvector($2::text::regconfig, '')".clone_into(&mut vector_sql);
}
let columns = format!(
"$1, {record_id_param}, {tenant_param}, $2, {fields_param}::jsonb, \
{content_param}, {vector_sql}, {embedding_param}::double precision[]{vec_value}"
);
rows.push(UpsertRow {
columns,
record_id_param,
});
}
let ids_param = next_param;
if watermark_param.is_none() {
binds.push(Bound::Ids(record_ids));
}
let sql = upsert_sql(&rows, vec_column, vec_update, ids_param, watermark_param);
bind_all(
diesel::sql_query(sql).into_boxed::<autumn_web::RuntimeBackend>(),
binds,
)
.execute(&mut conn)
.await
.map_err(SearchError::backend)?;
}
Ok(())
}
fn pool(&self) -> SearchResult<&RuntimePool> {
require_postgres_runtime()?;
self.pool.get().ok_or_else(|| {
SearchError::Backend(
"the search store has no database pool; is `database.primary_url` configured?"
.to_owned(),
)
})
}
async fn conn(
&self,
) -> SearchResult<
diesel_async::pooled_connection::deadpool::Object<autumn_web::RuntimeConnection>,
> {
self.pool()?.get().await.map_err(SearchError::backend)
}
}
#[allow(
clippy::cast_possible_truncation,
reason = "scores are ranking-only; f32 precision is deliberate"
)]
const fn narrow_score(score: f64) -> f32 {
if score.is_finite() { score as f32 } else { 0.0 }
}
fn checked(definition: &IndexDefinition) -> SearchResult<()> {
definition
.validate()
.map_err(|e| SearchError::InvalidIndex(e.to_string()))
}
fn array_literal(vector: &[f32]) -> String {
let mut out = String::with_capacity(vector.len().saturating_mul(8).saturating_add(2));
out.push('{');
for (index, value) in vector.iter().enumerate() {
if index > 0 {
out.push(',');
}
let value = if value.is_finite() { *value } else { 0.0 };
out.push_str(&value.to_string());
}
out.push('}');
out
}
fn vector_literal(vector: &[f32]) -> String {
let mut out = String::with_capacity(vector.len().saturating_mul(8).saturating_add(2));
out.push('[');
for (index, value) in vector.iter().enumerate() {
if index > 0 {
out.push(',');
}
let value = if value.is_finite() { *value } else { 0.0 };
out.push_str(&value.to_string());
}
out.push(']');
out
}
fn parse_vector(raw: &str) -> Vec<f32> {
raw.trim_matches(['{', '}', '[', ']'])
.split(',')
.filter(|part| !part.trim().is_empty())
.filter_map(|part| part.trim().parse::<f32>().ok())
.collect()
}
fn filter_sql(
definition: &IndexDefinition,
filter: &SearchFilter,
next_param: &mut usize,
binds: &mut Vec<Bound>,
) -> String {
use std::fmt::Write as _;
let mut sql = String::new();
if let Some(tenant) = &filter.tenant_id {
let _ = write!(sql, " AND tenant_id = ${next_param}");
binds.push(Bound::Text(tenant.clone()));
*next_param = next_param.saturating_add(1);
}
if let Some(allowed) = &filter.allowed_ids {
if allowed.is_empty() {
sql.push_str(" AND FALSE");
} else {
let _ = write!(sql, " AND record_id = ANY(${next_param})");
binds.push(Bound::Ids(allowed.clone()));
*next_param = next_param.saturating_add(1);
}
}
if !filter.excluded_ids.is_empty() {
let _ = write!(sql, " AND NOT (record_id = ANY(${next_param}))");
binds.push(Bound::Ids(filter.excluded_ids.clone()));
*next_param = next_param.saturating_add(1);
}
for (field, value) in &filter.equals {
if field == crate::backend::TENANT_FILTER_KEY {
let _ = write!(sql, " AND tenant_id = ${next_param}");
} else if definition.has_field(field) {
let _ = write!(sql, " AND fields ->> '{field}' = ${next_param}");
} else {
sql.push_str(" AND FALSE");
continue;
}
binds.push(Bound::Text(value.clone()));
*next_param = next_param.saturating_add(1);
}
sql
}
fn score_threshold_sql(score_expr: &str, min_score: Option<f32>, next_param: &mut usize) -> String {
if !min_score.is_some_and(f32::is_finite) {
return String::new();
}
let sql = format!(" AND {score_expr} >= ${next_param}::double precision");
*next_param = next_param.saturating_add(1);
sql
}
fn pgvector_order_by(filtered: bool) -> String {
if filtered {
"ORDER BY (1 - (embedding_vec <=> $2::vector)) DESC, record_id ASC".to_owned()
} else {
"ORDER BY (embedding_vec <=> $2::vector) ASC, record_id ASC".to_owned()
}
}
fn pgvector_embedding_predicate() -> String {
"embedding_vec IS NOT NULL AND (embedding_vec <=> $2::vector) = (embedding_vec <=> $2::vector)"
.to_owned()
}
fn dedupe_by_id(documents: &[IndexedDocument], keep_first: bool) -> Vec<&IndexedDocument> {
let mut index_for_id: std::collections::HashMap<i64, usize> =
std::collections::HashMap::with_capacity(documents.len());
for (index, document) in documents.iter().enumerate() {
if keep_first {
index_for_id.entry(document.id()).or_insert(index);
} else {
index_for_id.insert(document.id(), index);
}
}
let mut kept_indices: Vec<usize> = index_for_id.into_values().collect();
kept_indices.sort_unstable();
kept_indices
.into_iter()
.filter_map(|index| documents.get(index))
.collect()
}
fn validate_embedding_widths(
documents: &[IndexedDocument],
vector_width: Option<usize>,
vector_mode: Option<VectorMode>,
) -> SearchResult<()> {
let Some(width) = vector_width else {
return Ok(());
};
if !vector_mode.is_some_and(VectorMode::is_pgvector) {
return Ok(());
}
for document in documents {
if let Some(embedding) = document.embedding.as_deref()
&& embedding.len() != width
{
return Err(SearchError::DimensionMismatch {
expected: width,
actual: embedding.len(),
});
}
}
Ok(())
}
fn ascending_ids(ids: &[i64]) -> Vec<i64> {
let mut sorted = ids.to_vec();
sorted.sort_unstable();
sorted
}
struct UpsertRow {
columns: String,
record_id_param: String,
}
fn upsert_sql(
rows: &[UpsertRow],
vec_column: &str,
vec_update: &str,
ids_param: usize,
watermark_param: Option<usize>,
) -> String {
let (guard, values_source) = watermark_param.map_or_else(
|| {
(
String::new(),
format!(
"VALUES {}",
rows.iter()
.map(|row| format!("({})", row.columns))
.collect::<Vec<_>>()
.join(", ")
),
)
},
|slot| {
(
format!(" WHERE {DOCUMENTS_TABLE}.updated_at <= ${slot}::timestamptz"),
rows.iter()
.map(|row| {
format!(
"SELECT {} \
WHERE NOT EXISTS (SELECT 1 FROM {DELETES_TABLE} d \
WHERE d.index_name = $1 AND d.record_id = {} \
AND d.deleted_at > ${slot}::timestamptz)",
row.columns, row.record_id_param
)
})
.collect::<Vec<_>>()
.join(" UNION ALL "),
)
},
);
let upsert = format!(
"INSERT INTO {DOCUMENTS_TABLE} \
(index_name, record_id, tenant_id, language, fields, content, \
search_vector, embedding{vec_column}) \
{values_source} \
ON CONFLICT (index_name, record_id) DO UPDATE SET \
tenant_id = EXCLUDED.tenant_id, \
language = EXCLUDED.language, \
fields = EXCLUDED.fields, \
content = EXCLUDED.content, \
search_vector = EXCLUDED.search_vector, \
embedding = EXCLUDED.embedding{vec_update}, \
updated_at = NOW(){guard}"
);
if watermark_param.is_some() {
return upsert;
}
format!(
"WITH upserted AS ( \
{upsert} \
), \
cleared AS ( \
SELECT record_id FROM {DELETES_TABLE} \
WHERE index_name = $1 AND record_id = ANY(${ids_param}) \
ORDER BY record_id \
FOR UPDATE \
) \
DELETE FROM {DELETES_TABLE} \
WHERE index_name = $1 AND record_id IN (SELECT record_id FROM cleared)"
)
}
fn score_threshold_bind(min_score: Option<f32>) -> Option<Bound> {
min_score
.filter(|min| min.is_finite())
.map(|min| Bound::Double(f64::from(min)))
}
#[derive(Debug, Clone, PartialEq)]
enum Bound {
Text(String),
NullableText(Option<String>),
BigInt(i64),
Double(f64),
Ids(Vec<i64>),
}
fn bind_all(mut query: BoxedQuery<'_>, binds: impl IntoIterator<Item = Bound>) -> BoxedQuery<'_> {
for bound in binds {
query = match bound {
Bound::Text(value) => query.bind::<Text, _>(value),
Bound::NullableText(value) => query.bind::<Nullable<Text>, _>(value),
Bound::BigInt(value) => query.bind::<BigInt, _>(value),
Bound::Double(value) => query.bind::<Double, _>(value),
Bound::Ids(values) => bind_ids(query, values),
};
}
query
}
#[allow(
clippy::unnecessary_wraps,
clippy::missing_const_for_fn,
reason = "the SQLite arm of `backend_select!` returns an error"
)]
fn require_postgres_runtime() -> SearchResult<()> {
autumn_web::backend_select! {
pg => { Ok(()) },
sqlite => {
Err(SearchError::Backend(
"PostgresSearchStore needs the Postgres backend, but this build of autumn-web \
uses SQLite (`--features sqlite`)"
.to_owned(),
))
},
}
}
fn bind_ids(query: BoxedQuery<'_>, values: Vec<i64>) -> BoxedQuery<'_> {
autumn_web::backend_select! {
pg => { query.bind::<diesel::sql_types::Array<BigInt>, _>(values) },
sqlite => { query.bind::<Text, _>(serde_json::Value::from(values).to_string()) },
}
}
impl SearchBackend for PostgresSearchStore {
fn name(&self) -> &'static str {
"postgres"
}
fn capabilities(&self) -> BackendCapabilities {
BackendCapabilities::default()
.with_vector(true)
.with_weighted_fields(true)
.with_embedding_readback(true)
.with_conditional_index(true)
}
fn write_watermark(&self) -> BoxFuture<'_, SearchResult<Option<String>>> {
Box::pin(async move {
require_postgres_runtime()?;
let mut conn = self.conn().await?;
let row = diesel::sql_query("SELECT NOW()::text AS watermark")
.get_result::<WatermarkRow>(&mut conn)
.await
.map_err(SearchError::backend)?;
Ok(Some(row.watermark))
})
}
fn index_unless_newer<'a>(
&'a self,
definition: &'a IndexDefinition,
documents: &'a [IndexedDocument],
watermark: Option<&'a str>,
) -> BoxFuture<'a, SearchResult<()>> {
Box::pin(async move {
require_postgres_runtime()?;
self.write_documents(definition, documents, watermark).await
})
}
fn ensure_index<'a>(
&'a self,
definition: &'a IndexDefinition,
) -> BoxFuture<'a, SearchResult<()>> {
Box::pin(async move {
require_postgres_runtime()?;
checked(definition)?;
let mut conn = self.conn().await?;
if !self.schema_ready.swap(true, Ordering::SeqCst)
&& let Err(error) = Self::apply_schema(&mut conn).await
{
self.schema_ready.store(false, Ordering::SeqCst);
return Err(error);
}
if self.vector_mode().is_none() {
let mode = match self.dimensions {
Some(dimensions) if self.try_enable_pgvector(&mut conn, dimensions).await? => {
VectorMode::PgVector { dimensions }
}
_ => VectorMode::Array,
};
if let Ok(mut guard) = self.vector_mode.write() {
*guard = Some(mode);
}
let width = Self::vector_column_width_of(&mut conn).await?;
if let Ok(mut guard) = self.vector_column_width.write() {
*guard = width;
}
tracing::info!(
?mode,
vector_column_width = ?width,
"autumn-search postgres vector storage resolved"
);
if let Err(error) = diesel::sql_query(format!(
"DELETE FROM {DELETES_TABLE} WHERE deleted_at < NOW() - INTERVAL '{DELETE_RETENTION}'"
))
.execute(&mut conn)
.await
{
tracing::warn!(%error, "autumn-search could not prune the delete ledger");
}
}
Ok(())
})
}
fn index<'a>(
&'a self,
definition: &'a IndexDefinition,
documents: &'a [IndexedDocument],
) -> BoxFuture<'a, SearchResult<()>> {
Box::pin(async move {
require_postgres_runtime()?;
self.write_documents(definition, documents, None).await
})
}
fn delete<'a>(
&'a self,
definition: &'a IndexDefinition,
ids: &'a [i64],
) -> BoxFuture<'a, SearchResult<()>> {
Box::pin(async move {
require_postgres_runtime()?;
checked(definition)?;
if ids.is_empty() {
return Ok(());
}
let mut conn = self.conn().await?;
let sorted_ids = ascending_ids(ids);
bind_all(
diesel::sql_query(format!(
"WITH doomed AS ( \
SELECT record_id FROM {DOCUMENTS_TABLE} \
WHERE index_name = $1 AND record_id = ANY($2) \
ORDER BY record_id \
FOR UPDATE \
), \
removed AS ( \
DELETE FROM {DOCUMENTS_TABLE} \
WHERE index_name = $1 AND record_id IN (SELECT record_id FROM doomed) \
) \
INSERT INTO {DELETES_TABLE} (index_name, record_id) \
SELECT $1, unnest($2::bigint[]) \
ON CONFLICT (index_name, record_id) DO UPDATE SET deleted_at = NOW()"
))
.into_boxed::<autumn_web::RuntimeBackend>(),
[
Bound::Text(definition.name.to_owned()),
Bound::Ids(sorted_ids),
],
)
.execute(&mut conn)
.await
.map_err(SearchError::backend)?;
Ok(())
})
}
fn clear<'a>(&'a self, definition: &'a IndexDefinition) -> BoxFuture<'a, SearchResult<()>> {
Box::pin(async move {
require_postgres_runtime()?;
checked(definition)?;
let mut conn = self.conn().await?;
diesel::sql_query(format!(
"WITH doomed AS ( \
SELECT record_id FROM {DOCUMENTS_TABLE} \
WHERE index_name = $1 \
ORDER BY record_id \
FOR UPDATE \
), \
removed AS ( \
DELETE FROM {DOCUMENTS_TABLE} \
WHERE index_name = $1 AND record_id IN (SELECT record_id FROM doomed) \
) \
DELETE FROM {DELETES_TABLE} WHERE index_name = $1"
))
.bind::<Text, _>(definition.name.to_owned())
.execute(&mut conn)
.await
.map_err(SearchError::backend)?;
Ok(())
})
}
fn keyword_search<'a>(
&'a self,
definition: &'a IndexDefinition,
query: &'a KeywordQuery,
) -> BoxFuture<'a, SearchResult<Page<SearchHit>>> {
Box::pin(async move {
require_postgres_runtime()?;
checked(definition)?;
if query_tokens(&query.text).is_none() || query.filter.matches_nothing() {
return Ok(empty_page(&query.page));
}
let mut conn = self.conn().await?;
let mut param = 4_usize;
let mut binds: Vec<Bound> = Vec::new();
let predicate = filter_sql(definition, &query.filter, &mut param, &mut binds);
let head = [
Bound::Text(definition.name.to_owned()),
Bound::Text(definition.language.to_owned()),
Bound::Text(query.text.clone()),
];
let where_clause = format!(
"index_name = $1 \
AND search_vector @@ plainto_tsquery($2::text::regconfig, $3){predicate}"
);
let total = bind_all(
diesel::sql_query(format!(
"SELECT COUNT(*)::bigint AS total FROM {DOCUMENTS_TABLE} WHERE {where_clause}"
))
.into_boxed::<autumn_web::RuntimeBackend>(),
head.iter().cloned().chain(binds.iter().cloned()),
)
.get_result::<CountRow>(&mut conn)
.await
.map_err(SearchError::backend)?
.total;
let limit_param = param;
let offset_param = param.saturating_add(1);
let size = i64::from(query.page.size());
let offset = i64::from(query.page.page().saturating_sub(1)).saturating_mul(size);
let rows = bind_all(
diesel::sql_query(format!(
"SELECT record_id, \
ts_rank_cd({TS_RANK_WEIGHTS}, search_vector, \
plainto_tsquery($2::text::regconfig, $3))::double precision \
AS score \
FROM {DOCUMENTS_TABLE} WHERE {where_clause} \
ORDER BY score DESC, record_id ASC \
LIMIT ${limit_param} OFFSET ${offset_param}"
))
.into_boxed::<autumn_web::RuntimeBackend>(),
head.into_iter()
.chain(binds)
.chain([Bound::BigInt(size), Bound::BigInt(offset)]),
)
.load::<HitRow>(&mut conn)
.await
.map_err(SearchError::backend)?;
Ok(Page::new(
rows.into_iter()
.map(|row| {
SearchHit::new(definition.name, row.record_id, narrow_score(row.score))
})
.collect(),
total,
&query.page,
))
})
}
fn vector_search<'a>(
&'a self,
definition: &'a IndexDefinition,
query: &'a VectorQuery,
) -> BoxFuture<'a, SearchResult<Vec<SearchHit>>> {
Box::pin(async move {
require_postgres_runtime()?;
checked(definition)?;
if !definition.supports_vector_search() {
return Err(SearchError::VectorUnsupported {
index: definition.name.to_owned(),
});
}
if query.filter.matches_nothing() || query.limit == 0 || query.vector.is_empty() {
return Ok(Vec::new());
}
let mut conn = self.conn().await?;
let mut param = 3_usize;
let mut binds: Vec<Bound> = Vec::new();
let predicate = filter_sql(definition, &query.filter, &mut param, &mut binds);
let pgvector = self.vector_mode().is_some_and(VectorMode::is_pgvector);
let filtered = query.filter != SearchFilter::default();
let (query_vector, distance, score_expr, order_by, embedding_predicate) = if pgvector {
(
vector_literal(&query.vector),
"(embedding_vec <=> $2::vector)".to_owned(),
"(1 - (embedding_vec <=> $2::vector))::double precision".to_owned(),
pgvector_order_by(filtered),
pgvector_embedding_predicate(),
)
} else {
let expr = "autumn_search_cosine(embedding, $2::double precision[])".to_owned();
(
array_literal(&query.vector),
String::new(),
expr.clone(),
format!("ORDER BY {expr} DESC, record_id ASC"),
{
let slot = param;
param = param.saturating_add(1);
binds.push(Bound::BigInt(
i64::try_from(query.vector.len()).unwrap_or(i64::MAX),
));
format!(
"embedding IS NOT NULL AND array_length(embedding, 1) = ${slot}::integer"
)
},
)
};
let _ = distance;
let threshold = score_threshold_sql(&score_expr, query.min_score, &mut param);
let limit_param = param;
let limit = i64::try_from(query.limit).unwrap_or(i64::MAX);
let rows = bind_all(
diesel::sql_query(format!(
"SELECT record_id, {score_expr} AS score FROM {DOCUMENTS_TABLE} \
WHERE index_name = $1 AND {embedding_predicate}{predicate}{threshold} \
{order_by} LIMIT ${limit_param}"
))
.into_boxed::<autumn_web::RuntimeBackend>(),
[
Bound::Text(definition.name.to_owned()),
Bound::Text(query_vector),
]
.into_iter()
.chain(binds)
.chain(score_threshold_bind(query.min_score))
.chain([Bound::BigInt(limit)]),
)
.load::<HitRow>(&mut conn)
.await
.map_err(SearchError::backend)?;
Ok(rows
.into_iter()
.filter(|row| row.score.is_finite())
.map(|row| SearchHit::new(definition.name, row.record_id, narrow_score(row.score)))
.collect())
})
}
fn embedding<'a>(
&'a self,
definition: &'a IndexDefinition,
id: i64,
filter: &'a SearchFilter,
) -> BoxFuture<'a, SearchResult<Option<Vec<f32>>>> {
Box::pin(async move {
require_postgres_runtime()?;
checked(definition)?;
if filter.matches_nothing() {
return Ok(None);
}
let mut conn = self.conn().await?;
let mut param = 3_usize;
let mut binds: Vec<Bound> = Vec::new();
let predicate = filter_sql(definition, filter, &mut param, &mut binds);
let rows = bind_all(
diesel::sql_query(format!(
"SELECT embedding::text AS embedding FROM {DOCUMENTS_TABLE} \
WHERE index_name = $1 AND record_id = $2{predicate}"
))
.into_boxed::<autumn_web::RuntimeBackend>(),
[Bound::Text(definition.name.to_owned()), Bound::BigInt(id)]
.into_iter()
.chain(binds),
)
.load::<EmbeddingRow>(&mut conn)
.await
.map_err(SearchError::backend)?;
Ok(rows
.into_iter()
.next()
.and_then(|row| row.embedding)
.map(|raw| parse_vector(&raw)))
})
}
}
impl PostgresSearchStore {
async fn apply_schema(
conn: &mut diesel_async::pooled_connection::deadpool::Object<autumn_web::RuntimeConnection>,
) -> SearchResult<()> {
const LOCK: &str = "hashtext('autumn_search_schema')::bigint";
diesel::sql_query(format!("SELECT pg_advisory_lock({LOCK})"))
.execute(&mut *conn)
.await
.map_err(SearchError::backend)?;
let mut outcome = Ok(());
for statement in SCHEMA_STATEMENTS {
if let Err(error) = diesel::sql_query(*statement).execute(&mut *conn).await {
outcome = Err(SearchError::backend(error));
break;
}
}
let _ = diesel::sql_query(format!("SELECT pg_advisory_unlock({LOCK})"))
.execute(&mut *conn)
.await;
outcome
}
async fn vector_column_width_of(
conn: &mut diesel_async::pooled_connection::deadpool::Object<autumn_web::RuntimeConnection>,
) -> SearchResult<Option<usize>> {
let rows = diesel::sql_query(format!(
"SELECT COALESCE(atttypmod, 0)::bigint AS width FROM pg_attribute \
WHERE attrelid = to_regclass('{DOCUMENTS_TABLE}') \
AND attname = 'embedding_vec' AND NOT attisdropped"
))
.get_results::<WidthRow>(conn)
.await
.map_err(SearchError::backend)?;
Ok(rows
.into_iter()
.next()
.and_then(|row| usize::try_from(row.width).ok()))
}
async fn try_enable_pgvector(
&self,
conn: &mut diesel_async::pooled_connection::deadpool::Object<autumn_web::RuntimeConnection>,
dimensions: usize,
) -> SearchResult<bool> {
if diesel::sql_query("CREATE EXTENSION IF NOT EXISTS vector")
.execute(&mut *conn)
.await
.is_err()
{
tracing::info!(
"pgvector is not available; autumn-search will store embeddings as \
double precision[] and rank with autumn_search_cosine()"
);
return Ok(false);
}
let present = diesel::sql_query(
"SELECT EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'vector') AS present",
)
.get_result::<ExistsRow>(&mut *conn)
.await
.map_err(SearchError::backend)?
.present;
if !present {
return Ok(false);
}
diesel::sql_query(format!(
"ALTER TABLE {DOCUMENTS_TABLE} \
ADD COLUMN IF NOT EXISTS embedding_vec vector({dimensions})"
))
.execute(&mut *conn)
.await
.map_err(SearchError::backend)?;
let existing_width = Self::vector_column_width_of(conn)
.await?
.and_then(|w| i64::try_from(w).ok());
if let Some(width) = existing_width
&& width > 0
&& usize::try_from(width).unwrap_or(0) != dimensions
{
tracing::warn!(
configured = dimensions,
existing = width,
"autumn-search: the pgvector column has a different width than \
`search.embedding_dimensions`; falling back to the portable \
double precision[] path. Drop `autumn_search_documents.embedding_vec` \
and re-run a full backfill to adopt the new width."
);
return Ok(false);
}
diesel::sql_query(format!(
"UPDATE {DOCUMENTS_TABLE} SET embedding_vec = embedding::text::vector \
WHERE embedding IS NOT NULL AND embedding_vec IS NULL \
AND array_length(embedding, 1) = {dimensions}"
))
.execute(&mut *conn)
.await
.map_err(SearchError::backend)?;
let _ = diesel::sql_query(format!(
"CREATE INDEX IF NOT EXISTS autumn_search_documents_embedding_vec_idx \
ON {DOCUMENTS_TABLE} USING ivfflat (embedding_vec vector_cosine_ops) \
WITH (lists = 100)"
))
.execute(&mut *conn)
.await;
Ok(true)
}
}
impl DocumentSource for PostgresSearchStore {
fn fetch<'a>(
&'a self,
definition: &'a IndexDefinition,
ids: &'a [i64],
) -> BoxFuture<'a, SearchResult<Vec<SearchDocument>>> {
Box::pin(async move {
require_postgres_runtime()?;
checked(definition)?;
if ids.is_empty() {
return Ok(Vec::new());
}
self.load_documents(definition, Selection::Ids(ids)).await
})
}
fn scan<'a>(
&'a self,
definition: &'a IndexDefinition,
after: Option<i64>,
limit: usize,
) -> BoxFuture<'a, SearchResult<Vec<SearchDocument>>> {
Box::pin(async move {
require_postgres_runtime()?;
checked(definition)?;
self.load_documents(definition, Selection::After { after, limit })
.await
})
}
}
enum Selection<'a> {
Ids(&'a [i64]),
After { after: Option<i64>, limit: usize },
}
impl PostgresSearchStore {
async fn optional_columns(
&self,
definition: &IndexDefinition,
conn: &mut diesel_async::pooled_connection::deadpool::Object<autumn_web::RuntimeConnection>,
) -> SearchResult<OptionalColumns> {
if let Ok(cache) = self.optional_columns.read()
&& let Some(columns) = cache.get(definition.name)
{
return Ok(*columns);
}
let row = diesel::sql_query(
"SELECT bool_or(attname = 'tenant_id') AS has_tenant, \
bool_or(attname = 'deleted_at') AS has_deleted_at \
FROM pg_attribute \
WHERE attrelid = to_regclass($1) AND attnum > 0 AND NOT attisdropped",
)
.bind::<Text, _>(definition.name.to_owned())
.get_result::<OptionalColumnsRow>(conn)
.await
.map_err(SearchError::backend)?;
let columns = OptionalColumns {
tenant_id: row.has_tenant.unwrap_or(false),
deleted_at: row.has_deleted_at.unwrap_or(false),
};
if let Ok(mut cache) = self.optional_columns.write() {
cache.insert(definition.name, columns);
}
Ok(columns)
}
async fn load_documents(
&self,
definition: &IndexDefinition,
selection: Selection<'_>,
) -> SearchResult<Vec<SearchDocument>> {
let mut conn = self.conn().await?;
let columns = self.optional_columns(definition, &mut conn).await?;
let key = format!("\"{}\"", definition.key_column);
let mut binds: Vec<Bound> = Vec::new();
let mut param = 1_usize;
let mut predicates: Vec<String> = Vec::new();
let mut limit_clause = String::new();
match selection {
Selection::Ids(ids) => {
predicates.push(format!("{key} = ANY(${param})"));
binds.push(Bound::Ids(ids.to_vec()));
}
Selection::After { after, limit } => {
if let Some(after) = after {
predicates.push(format!("{key} > ${param}"));
binds.push(Bound::BigInt(after));
param = param.saturating_add(1);
}
limit_clause = format!(" LIMIT ${param}");
binds.push(Bound::BigInt(i64::try_from(limit).unwrap_or(i64::MAX)));
}
}
if columns.deleted_at && definition.soft_delete {
predicates.push("deleted_at IS NULL".to_owned());
}
let where_clause = if predicates.is_empty() {
String::new()
} else {
format!("WHERE {}", predicates.join(" AND "))
};
let tenant = if columns.tenant_id {
"tenant_id::text"
} else {
"NULL::text"
};
let rows = bind_all(
diesel::sql_query(format!(
"SELECT {key}::bigint AS id, {} AS fields, {tenant} AS tenant_id \
FROM \"{}\" {where_clause} ORDER BY {key} ASC{limit_clause}",
fields_projection(definition),
definition.name,
))
.into_boxed::<autumn_web::RuntimeBackend>(),
binds,
)
.load::<SourceRow>(&mut conn)
.await
.map_err(SearchError::backend)?;
let mut documents = Vec::with_capacity(rows.len());
for row in rows {
let values: std::collections::BTreeMap<String, String> =
serde_json::from_str(&row.fields).map_err(SearchError::backend)?;
let mut document = SearchDocument::new(definition.name, row.id);
for field in definition.fields.iter() {
let value = values.get(field.name).cloned().unwrap_or_default();
document = document.with_field(field.name, field.weight, value);
}
if let Some(embed) = definition.embed_field
&& let Some(text) = values.get(embed)
&& !text.is_empty()
{
document.embed_text = Some(text.clone());
}
document.tenant_id = row.tenant_id;
documents.push(document);
}
Ok(documents)
}
}
fn fields_projection(definition: &IndexDefinition) -> String {
const PAIRS_PER_CHUNK: usize = 40;
let pairs: Vec<String> = definition
.fields
.iter()
.map(|field| format!("'{0}', COALESCE(\"{0}\"::text, '')", field.name))
.collect();
if pairs.is_empty() {
return "'{}'::jsonb::text".to_owned();
}
let merged = pairs
.chunks(PAIRS_PER_CHUNK)
.map(|chunk| format!("jsonb_build_object({})", chunk.join(", ")))
.collect::<Vec<_>>()
.join(" || ");
format!("({merged})::text")
}
fn fields_json(document: &IndexedDocument) -> SearchResult<String> {
serde_json::to_string(
&document
.document
.fields
.iter()
.map(|f| (f.name, f.value.clone()))
.collect::<std::collections::BTreeMap<_, _>>(),
)
.map_err(SearchError::backend)
}
const SCHEMA_STATEMENTS: &[&str] = &[
"CREATE TABLE IF NOT EXISTS autumn_search_deletes (
index_name TEXT NOT NULL,
record_id BIGINT NOT NULL,
deleted_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
PRIMARY KEY (index_name, record_id)
)",
"CREATE INDEX IF NOT EXISTS autumn_search_deletes_pruning_idx
ON autumn_search_deletes (deleted_at)",
"CREATE TABLE IF NOT EXISTS autumn_search_documents (
index_name TEXT NOT NULL,
record_id BIGINT NOT NULL,
tenant_id TEXT,
language TEXT NOT NULL DEFAULT 'simple',
fields JSONB NOT NULL DEFAULT '{}'::jsonb,
content TEXT NOT NULL DEFAULT '',
search_vector TSVECTOR NOT NULL DEFAULT to_tsvector('simple', ''),
embedding DOUBLE PRECISION[],
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
PRIMARY KEY (index_name, record_id)
)",
"CREATE INDEX IF NOT EXISTS autumn_search_documents_vector_idx
ON autumn_search_documents USING GIN (search_vector)",
"CREATE INDEX IF NOT EXISTS autumn_search_documents_tenant_idx
ON autumn_search_documents (index_name, tenant_id)",
"CREATE OR REPLACE FUNCTION autumn_search_cosine(a DOUBLE PRECISION[], b DOUBLE PRECISION[])
RETURNS DOUBLE PRECISION LANGUAGE sql IMMUTABLE AS $$
SELECT COALESCE(
SUM(t.x * t.y) / NULLIF(SQRT(SUM(t.x * t.x)) * SQRT(SUM(t.y * t.y)), 0),
0
)
FROM unnest(a, b) AS t(x, y)
$$",
];
#[cfg(test)]
mod tests {
use autumn_web::search::SearchIndexField;
use super::*;
const FIELDS: &[SearchIndexField] = &[
SearchIndexField::new("title", 'A'),
SearchIndexField::new("body", 'B'),
];
#[test]
fn the_store_needs_the_postgres_backend() {
let result = require_postgres_runtime();
autumn_web::backend_select! {
pg => {
assert!(result.is_ok(), "{result:?}");
},
sqlite => {
let message = result.expect_err("refused on SQLite").to_string();
assert!(message.contains("SQLite"), "{message}");
},
}
}
fn definition() -> IndexDefinition {
IndexDefinition::new("articles", "english", FIELDS, Some("body"), false)
}
#[tokio::test]
async fn an_early_return_does_not_skip_the_backend_check() {
let store = PostgresSearchStore::new(None);
let result = store.delete(&definition(), &[]).await;
autumn_web::backend_select! {
pg => {
assert!(result.is_ok(), "{result:?}");
},
sqlite => {
let message = result.expect_err("refused on SQLite").to_string();
assert!(message.contains("SQLite"), "{message}");
},
}
}
fn doc(id: i64) -> IndexedDocument {
IndexedDocument::new(autumn_web::search::SearchDocument::new("articles", id))
}
#[test]
fn unconditional_dedup_keeps_the_last_occurrence() {
let documents = vec![doc(1), doc(2), doc(1), doc(3), doc(2)];
let kept: Vec<i64> = dedupe_by_id(&documents, false)
.into_iter()
.map(IndexedDocument::id)
.collect();
assert_eq!(kept, vec![1, 3, 2], "{kept:?}");
}
#[test]
fn watermark_guarded_dedup_keeps_the_first_occurrence() {
let documents = vec![doc(1), doc(2), doc(1), doc(3), doc(2)];
let kept: Vec<i64> = dedupe_by_id(&documents, true)
.into_iter()
.map(IndexedDocument::id)
.collect();
assert_eq!(kept, vec![1, 2, 3], "{kept:?}");
}
#[test]
fn no_duplicates_keeps_every_index_in_order_either_direction() {
let documents = vec![doc(10), doc(20), doc(30)];
for keep_first in [true, false] {
let kept: Vec<i64> = dedupe_by_id(&documents, keep_first)
.into_iter()
.map(IndexedDocument::id)
.collect();
assert_eq!(kept, vec![10, 20, 30], "keep_first={keep_first}");
}
}
#[test]
fn deduplicated_documents_sort_by_record_id_regardless_of_input_order() {
for keep_first in [true, false] {
for documents in [
vec![doc(2), doc(1), doc(3)],
vec![doc(3), doc(1), doc(2)],
vec![doc(1), doc(3), doc(2)],
] {
let mut kept: Vec<&IndexedDocument> = dedupe_by_id(&documents, keep_first);
kept.sort_unstable_by_key(|document| document.id());
let ids: Vec<i64> = kept.into_iter().map(IndexedDocument::id).collect();
assert_eq!(
ids,
vec![1, 2, 3],
"keep_first={keep_first} input={documents:?}"
);
}
}
}
fn doc_with_embedding(id: i64, width: usize) -> IndexedDocument {
doc(id).with_embedding(vec![0.0; width])
}
const PGVECTOR_MODE: Option<VectorMode> = Some(VectorMode::PgVector { dimensions: 4 });
fn assert_dimension_mismatch(error: &SearchError, expected: usize, actual: usize) {
assert!(
matches!(
error,
SearchError::DimensionMismatch {
expected: exp,
actual: act
} if *exp == expected && *act == actual
),
"{error:?}"
);
}
#[test]
fn embedding_widths_are_validated_before_unconditional_dedup_can_hide_a_bad_duplicate() {
let documents = vec![
doc_with_embedding(1, 3), doc_with_embedding(1, 4), ];
let error = validate_embedding_widths(&documents, Some(4), PGVECTOR_MODE)
.expect_err("a wrong-width embedding must be rejected");
assert_dimension_mismatch(&error, 4, 3);
}
#[test]
fn embedding_widths_are_validated_before_guarded_dedup_can_hide_a_bad_duplicate() {
let documents = vec![
doc_with_embedding(1, 4), doc_with_embedding(1, 3), ];
let error = validate_embedding_widths(&documents, Some(4), PGVECTOR_MODE)
.expect_err("a wrong-width embedding must be rejected");
assert_dimension_mismatch(&error, 4, 3);
}
#[test]
fn embedding_width_validation_accepts_a_clean_batch() {
let documents = vec![
doc_with_embedding(1, 4),
doc_with_embedding(2, 4),
doc(3), ];
validate_embedding_widths(&documents, Some(4), PGVECTOR_MODE)
.expect("a clean batch must validate");
}
#[test]
fn embedding_width_validation_only_rejects_in_pgvector_mode() {
let documents = vec![doc_with_embedding(1, 3)];
validate_embedding_widths(&documents, Some(4), Some(VectorMode::Array))
.expect("non-pgvector mode must not reject a width mismatch");
}
#[test]
fn embedding_width_validation_skips_when_there_is_no_physical_column() {
let documents = vec![doc_with_embedding(1, 3)];
validate_embedding_widths(&documents, None, PGVECTOR_MODE)
.expect("no physical column means no validation");
}
#[test]
fn greedy_chunking_never_lets_a_statement_approach_the_bind_limit_even_with_repeated_fields() {
const PARAM_BUDGET: usize = 60_000;
let weight_of = |name: &str| if name == "title" { Some('A') } else { None };
let mut documents = Vec::new();
for id in 0..800i64 {
let mut document = autumn_web::search::SearchDocument::new("articles", id);
for _ in 0..90 {
document = document.with_field("title", 'A', "x");
}
documents.push(IndexedDocument::new(document));
}
let mut chunks: Vec<usize> = Vec::new();
let mut next_param = 3usize; let mut chunk_len = 0usize;
for document in &documents {
let doc_param_count = 5usize.saturating_add(
document
.document
.fields
.iter()
.filter(|field| weight_of(field.name).is_some())
.count(),
);
if chunk_len > 0 && next_param.saturating_add(doc_param_count) > PARAM_BUDGET {
chunks.push(next_param);
next_param = 3;
chunk_len = 0;
}
next_param = next_param.saturating_add(doc_param_count);
chunk_len = chunk_len.saturating_add(1);
}
chunks.push(next_param);
assert!(
chunks.len() > 1,
"the fixture must force more than one chunk: {chunks:?}"
);
for &total in &chunks {
assert!(
total < 65_535,
"a chunk's total bind count {total} must stay under Postgres' bind limit: \
{chunks:?}"
);
}
}
#[test]
fn a_row_count_cap_bounds_chunk_size_even_when_the_param_budget_never_would() {
const PARAM_BUDGET: usize = 60_000;
const MAX_ROWS_PER_CHUNK: usize = 1_000;
let doc_param_count = 5usize;
let mut chunks: Vec<usize> = Vec::new();
let mut next_param = 3usize;
let mut chunk_len = 0usize;
for _ in 0..5_000 {
if chunk_len > 0
&& (next_param.saturating_add(doc_param_count) > PARAM_BUDGET
|| chunk_len >= MAX_ROWS_PER_CHUNK)
{
chunks.push(chunk_len);
next_param = 3;
chunk_len = 0;
}
next_param = next_param.saturating_add(doc_param_count);
chunk_len = chunk_len.saturating_add(1);
}
chunks.push(chunk_len);
assert!(
chunks.len() > 1,
"5,000 minimal-field documents must still split into more than one chunk: {chunks:?}"
);
for &len in &chunks {
assert!(
len <= MAX_ROWS_PER_CHUNK,
"a chunk of {len} rows must never exceed MAX_ROWS_PER_CHUNK: {chunks:?}"
);
}
}
#[test]
fn a_filtered_vector_query_does_not_let_ivfflat_decide_what_is_visible() {
let open = pgvector_order_by(false);
assert!(open.contains("<=>"), "{open}");
assert!(open.contains("ASC, record_id ASC"), "{open}");
let scoped = pgvector_order_by(true);
assert!(
scoped.starts_with("ORDER BY (1 - "),
"a filtered query must not be served by the approximate index: {scoped}"
);
assert_ne!(open, scoped);
}
#[test]
fn the_pgvector_predicate_excludes_nan_scores_before_limit() {
let predicate = pgvector_embedding_predicate();
assert!(
predicate.contains("embedding_vec IS NOT NULL"),
"{predicate}"
);
assert!(
predicate.contains("(embedding_vec <=> $2::vector) = (embedding_vec <=> $2::vector)"),
"NaN must be excluded by self-inequality: {predicate}"
);
}
#[test]
fn an_unconditional_upsert_clears_the_tombstone_in_the_same_statement() {
let row = || UpsertRow {
columns: "$1, $4, $5, $2, $6::jsonb, $7, to_tsvector('simple', ''), \
$8::double precision[]"
.to_owned(),
record_id_param: "$4".to_owned(),
};
let sql = upsert_sql(&[row()], "", "", 9, None);
assert!(
sql.starts_with("WITH upserted AS ("),
"the upsert and the tombstone clear must be one statement: {sql}"
);
assert!(
sql.contains(&format!("DELETE FROM {DELETES_TABLE}")),
"{sql}"
);
assert!(sql.contains("record_id = ANY($9)"), "{sql}");
let lock_block = sql
.find(&format!("DELETE FROM {DELETES_TABLE}"))
.map(|at| &sql[..at])
.expect("the tombstone-clear DELETE exists");
assert!(lock_block.contains("ORDER BY record_id"), "{sql}");
assert!(lock_block.contains("FOR UPDATE"), "{sql}");
let guarded = upsert_sql(&[row()], "", "", 9, Some(3));
assert!(
!guarded.contains(&format!("DELETE FROM {DELETES_TABLE}")),
"a backfill write must never clear a tombstone: {guarded}"
);
assert!(
guarded.contains(&format!("SELECT 1 FROM {DELETES_TABLE}")),
"but it must still be blocked by one: {guarded}"
);
let second = UpsertRow {
columns: "$1, $9, $10, $2, $11::jsonb, $12, to_tsvector('simple', ''), \
$13::double precision[]"
.to_owned(),
record_id_param: "$9".to_owned(),
};
let multi = upsert_sql(&[row(), second], "", "", 99, Some(3));
assert_eq!(multi.matches(" UNION ALL ").count(), 1, "{multi}");
assert_eq!(multi.matches("$3::timestamptz").count(), 3, "{multi}");
assert!(multi.contains("record_id = $4"), "{multi}");
assert!(multi.contains("record_id = $9"), "{multi}");
}
#[test]
fn delete_binds_its_ledger_ids_in_ascending_order() {
assert_eq!(ascending_ids(&[3, 1, 2]), vec![1, 2, 3]);
assert_eq!(ascending_ids(&[1, 2, 3]), vec![1, 2, 3]);
assert_eq!(ascending_ids(&[5, 5, 1]), vec![1, 5, 5]);
assert_eq!(ascending_ids(&[]), Vec::<i64>::new());
}
#[test]
fn a_min_score_is_bound_so_the_statement_cache_stays_bounded() {
let expr = "(1 - (embedding_vec <=> $2::vector))::double precision";
let mut a = 5;
let mut b = 5;
let first = score_threshold_sql(expr, Some(0.25), &mut a);
let second = score_threshold_sql(expr, Some(0.999_999), &mut b);
assert_eq!(
first, second,
"two thresholds must produce ONE statement, not two"
);
assert_eq!(a, b, "and consume the same parameter slot");
assert!(first.contains("$5"), "{first}");
assert!(
!first.contains("0.25"),
"the value must not reach the SQL text: {first}"
);
let mut param = 5;
assert!(score_threshold_sql(expr, None, &mut param).is_empty());
assert_eq!(param, 5);
let mut param = 5;
assert!(score_threshold_sql(expr, Some(f32::NAN), &mut param).is_empty());
assert!(score_threshold_sql(expr, Some(f32::INFINITY), &mut param).is_empty());
assert_eq!(param, 5);
for candidate in [None, Some(0.5), Some(f32::NAN), Some(f32::INFINITY)] {
let mut param = 5;
let emitted = !score_threshold_sql(expr, candidate, &mut param).is_empty();
assert_eq!(
emitted,
score_threshold_bind(candidate).is_some(),
"fragment/bind disagreement for {candidate:?}"
);
assert_eq!(param, if emitted { 6 } else { 5 });
}
}
#[test]
fn ts_rank_weights_track_the_framework_factors() {
use autumn_web::search::weight_factor;
let declared = parse_vector(
TS_RANK_WEIGHTS
.trim_end_matches("::float4[]")
.trim_matches('\''),
);
let top = weight_factor('A');
let expected: Vec<f32> = ['D', 'C', 'B', 'A']
.into_iter()
.map(|weight| weight_factor(weight) / top)
.collect();
assert_eq!(
declared.len(),
expected.len(),
"ts_rank_cd takes exactly {{D, C, B, A}}: {TS_RANK_WEIGHTS}"
);
for (index, (got, want)) in declared.iter().zip(&expected).enumerate() {
assert!(
(got - want).abs() < 1e-6,
"weight {} is {got}, expected {want} โ the SQL array must stay in step with \
`weight_factor`",
['D', 'C', 'B', 'A'][index]
);
}
}
#[test]
fn array_literals_round_trip() {
let vector = vec![1.0_f32, -0.5, 0.25];
assert_eq!(parse_vector(&array_literal(&vector)), vector);
assert_eq!(parse_vector(&vector_literal(&vector)), vector);
}
#[test]
fn non_finite_values_never_produce_invalid_sql() {
let literal = array_literal(&[f32::NAN, f32::INFINITY, 1.0]);
assert_eq!(literal, "{0,0,1}");
assert!(!literal.contains("NaN"));
assert!(!literal.contains("inf"));
}
#[test]
fn parsing_tolerates_empty_and_malformed_elements() {
assert!(parse_vector("{}").is_empty());
assert!(parse_vector("[]").is_empty());
assert_eq!(parse_vector("{1, ,2}"), vec![1.0, 2.0]);
assert_eq!(parse_vector("{1,oops,2}"), vec![1.0, 2.0]);
}
#[test]
fn an_empty_filter_adds_no_predicate() {
let mut param = 2;
let mut binds = Vec::new();
assert_eq!(
filter_sql(
&definition(),
&SearchFilter::default(),
&mut param,
&mut binds
),
""
);
assert_eq!(param, 2);
assert!(binds.is_empty());
}
#[test]
fn a_tenant_filter_binds_its_value_rather_than_interpolating_it() {
let mut param = 2;
let mut binds = Vec::new();
let sql = filter_sql(
&definition(),
&SearchFilter::default().tenant("acme'; DROP TABLE users; --"),
&mut param,
&mut binds,
);
assert_eq!(sql, " AND tenant_id = $2");
assert_eq!(
binds,
vec![Bound::Text("acme'; DROP TABLE users; --".to_owned())]
);
assert_eq!(param, 3);
}
#[test]
fn an_empty_allowlist_renders_a_fail_closed_predicate() {
let mut param = 2;
let mut binds = Vec::new();
let sql = filter_sql(
&definition(),
&SearchFilter::default().allow_ids(Vec::<i64>::new()),
&mut param,
&mut binds,
);
assert!(sql.contains("AND FALSE"), "{sql}");
}
#[test]
fn id_lists_are_bound_as_arrays_never_interpolated() {
let mut param = 2;
let mut binds = Vec::new();
let sql = filter_sql(
&definition(),
&SearchFilter::default()
.allow_ids([1_i64, 2])
.exclude_ids([3_i64]),
&mut param,
&mut binds,
);
assert!(sql.contains("record_id = ANY($2)"), "{sql}");
assert!(sql.contains("NOT (record_id = ANY($3))"), "{sql}");
assert!(
!sql.contains("IN ("),
"ids must never be interpolated: {sql}"
);
assert_eq!(binds, vec![Bound::Ids(vec![1, 2]), Bound::Ids(vec![3])]);
}
#[test]
fn equals_predicates_bind_values_and_number_parameters_in_order() {
let mut param = 4;
let mut binds = Vec::new();
let sql = filter_sql(
&definition(),
&SearchFilter::default()
.equals("title", "Hello")
.equals("body", "World"),
&mut param,
&mut binds,
);
assert!(sql.contains("fields ->> 'body' = $4"), "{sql}");
assert!(sql.contains("fields ->> 'title' = $5"), "{sql}");
assert_eq!(
binds,
vec![
Bound::Text("World".to_owned()),
Bound::Text("Hello".to_owned())
]
);
assert_eq!(param, 6);
}
#[test]
fn an_undeclared_equals_key_never_reaches_sql() {
let mut param = 4;
let mut binds = Vec::new();
let sql = filter_sql(
&definition(),
&SearchFilter::default().equals("title' = 'x' OR '1'='1", "boom"),
&mut param,
&mut binds,
);
assert_eq!(sql, " AND FALSE", "{sql}");
assert!(!sql.contains("OR"), "{sql}");
assert!(binds.is_empty(), "nothing may be bound for a dropped key");
assert_eq!(param, 4, "the parameter counter must not advance");
}
#[test]
fn an_undeclared_key_does_not_disturb_the_numbering_of_declared_ones() {
let mut param = 4;
let mut binds = Vec::new();
let sql = filter_sql(
&definition(),
&SearchFilter::default()
.equals("aaa_not_a_field", "x")
.equals("title", "Hello"),
&mut param,
&mut binds,
);
assert!(sql.contains("fields ->> 'title' = $4"), "{sql}");
assert!(sql.contains("AND FALSE"), "{sql}");
assert_eq!(binds, vec![Bound::Text("Hello".to_owned())]);
assert_eq!(param, 5);
}
#[test]
fn the_reserved_tenant_key_is_still_honoured() {
let mut param = 2;
let mut binds = Vec::new();
let sql = filter_sql(
&definition(),
&SearchFilter::default().equals(crate::backend::TENANT_FILTER_KEY, "acme"),
&mut param,
&mut binds,
);
assert_eq!(sql, " AND tenant_id = $2");
assert_eq!(binds, vec![Bound::Text("acme".to_owned())]);
}
#[test]
fn the_fields_projection_casts_the_whole_merge_not_just_the_last_chunk() {
const PAIRS_PER_CHUNK: usize = 40;
let single = fields_projection(&definition());
assert!(single.starts_with("(jsonb_build_object("), "{single}");
assert!(single.ends_with(")::text"), "{single}");
assert_eq!(single.matches("jsonb_build_object(").count(), 1);
let fields: Vec<SearchIndexField> = (0..=(PAIRS_PER_CHUNK * 2))
.map(|i| SearchIndexField::new(format!("field_{i}").leak(), 'D'))
.collect();
let mut wide = definition();
wide.fields = std::borrow::Cow::Owned(fields);
wide.embed_field = None;
let projection = fields_projection(&wide);
assert_eq!(
projection.matches("jsonb_build_object(").count(),
3,
"81 pairs at 40/chunk is 3 calls: {projection}"
);
assert!(projection.contains(" || "), "{projection}");
assert!(projection.starts_with('('), "{projection}");
assert!(projection.ends_with(")::text"), "{projection}");
assert!(
!projection.contains(")::text ||"),
"the cast must wrap the merge, not a chunk: {projection}"
);
}
#[test]
fn an_empty_field_list_still_yields_valid_json() {
let mut empty = definition();
empty.fields = std::borrow::Cow::Borrowed(&[]);
assert_eq!(fields_projection(&empty), "'{}'::jsonb::text");
}
#[test]
fn an_invalid_definition_is_rejected_before_any_sql_is_built() {
let mut bad = definition();
bad.name = "articles\"; DROP TABLE users; --";
assert!(matches!(checked(&bad), Err(SearchError::InvalidIndex(_))));
assert!(checked(&definition()).is_ok());
}
#[test]
fn the_schema_is_idempotent_by_construction() {
for statement in SCHEMA_STATEMENTS {
let statement = statement.trim_start();
assert!(
statement.contains("IF NOT EXISTS") || statement.starts_with("CREATE OR REPLACE"),
"not idempotent: {statement}"
);
}
}
#[test]
fn a_store_without_a_pool_reports_the_missing_configuration() {
let store = PostgresSearchStore::new(None);
let Err(error) = store.pool() else {
panic!("a store with no installed pool must not report one");
};
let expected = autumn_web::backend_select! {
pg => { "database.primary_url" },
sqlite => { "SQLite" },
};
assert!(error.to_string().contains(expected), "{error}");
assert!(store.vector_mode().is_none());
}
#[test]
fn vector_mode_reports_whether_pgvector_is_in_use() {
assert!(VectorMode::PgVector { dimensions: 4 }.is_pgvector());
assert!(!VectorMode::Array.is_pgvector());
}
}