use std::borrow::Cow;
use std::sync::{Arc, OnceLock};
use autumn_web::app::AppBuilder;
use autumn_web::plugin::Plugin;
use autumn_web::search::{IndexDefinition, SearchIndexed};
use crate::authz::SearchVisibility;
use crate::backend::SearchBackend;
use crate::client::{BackfillOptions, SearchClient, SearchClientBuilder};
use crate::config::SearchConfig;
use crate::embedding::Embedder;
use crate::jobs::search_job_infos;
use crate::postgres::PostgresSearchStore;
use crate::source::DocumentSource;
pub const BACKFILL_ENV: &str = "AUTUMN_SEARCH_BACKFILL";
pub const BACKFILL_PURGE_ENV: &str = "AUTUMN_SEARCH_BACKFILL_PURGE";
pub const BACKFILL_STARTED_MARKER: &str = "autumn-search: backfill starting";
pub struct SearchPlugin {
config: SearchConfig,
backend: Option<Arc<dyn SearchBackend>>,
embedder: Option<Arc<dyn Embedder>>,
source: Option<Arc<dyn DocumentSource>>,
visibility: Option<Arc<dyn SearchVisibility>>,
indexes: Vec<IndexDefinition>,
use_postgres: bool,
queue_override: Option<String>,
postgres_store: OnceLock<Arc<PostgresSearchStore>>,
default_backend: OnceLock<Arc<dyn SearchBackend>>,
resolved_config: OnceLock<Result<SearchConfig, String>>,
config_explicit: bool,
}
impl Default for SearchPlugin {
fn default() -> Self {
Self::new()
}
}
impl SearchPlugin {
#[must_use]
pub fn new() -> Self {
Self {
config: SearchConfig::default(),
backend: None,
embedder: None,
source: None,
visibility: None,
indexes: Vec::new(),
use_postgres: false,
queue_override: None,
postgres_store: OnceLock::new(),
default_backend: OnceLock::new(),
resolved_config: OnceLock::new(),
config_explicit: false,
}
}
#[must_use]
pub fn config(mut self, config: SearchConfig) -> Self {
self.config = config;
self.config_explicit = true;
self
}
#[must_use]
pub fn queue(mut self, queue: impl Into<String>) -> Self {
self.queue_override = Some(queue.into());
self
}
#[must_use]
pub fn backend(mut self, backend: Arc<dyn SearchBackend>) -> Self {
self.backend = Some(backend);
if std::mem::take(&mut self.use_postgres) {
self.postgres_store = OnceLock::new();
}
self
}
#[must_use]
pub const fn postgres(mut self) -> Self {
self.use_postgres = true;
self
}
#[must_use]
pub fn embedder(mut self, embedder: Arc<dyn Embedder>) -> Self {
self.embedder = Some(embedder);
self
}
#[must_use]
pub fn source(mut self, source: Arc<dyn DocumentSource>) -> Self {
self.source = Some(source);
self
}
#[must_use]
pub fn visibility(mut self, visibility: Arc<dyn SearchVisibility>) -> Self {
self.visibility = Some(visibility);
self
}
#[must_use]
pub fn index<M: SearchIndexed>(mut self) -> Self {
self.indexes.push(M::index_definition());
self
}
#[must_use]
pub fn index_definition(mut self, definition: IndexDefinition) -> Self {
self.indexes.push(definition);
self
}
#[must_use]
pub const fn uses_postgres(&self) -> bool {
self.use_postgres
}
#[must_use]
pub fn search_config(&self) -> SearchConfig {
let mut config = self.config.clone();
if let Some(queue) = &self.queue_override {
config.queue.clone_from(queue);
}
config
}
fn resolved_config(&self) -> &Result<SearchConfig, String> {
self.resolved_config.get_or_init(|| {
if self.config_explicit {
let config = self.search_config();
return config
.validate()
.map(|()| config)
.map_err(|error| error.to_string());
}
SearchConfig::resolve()
.map(|mut config| {
if let Some(queue) = &self.queue_override {
config.queue.clone_from(queue);
}
config
})
.map_err(|error| error.to_string())
})
}
fn effective_config(&self) -> SearchConfig {
self.resolved_config()
.as_ref()
.map_or_else(|_| self.search_config(), Clone::clone)
}
#[must_use]
pub fn client(&self) -> SearchClient {
self.client_builder().build()
}
fn postgres_store(&self) -> Option<&Arc<PostgresSearchStore>> {
if !self.use_postgres {
return None;
}
Some(self.postgres_store.get_or_init(|| {
Arc::new(PostgresSearchStore::new(
self.effective_config().embedding_dimensions,
))
}))
}
fn default_backend(&self) -> &Arc<dyn SearchBackend> {
self.default_backend
.get_or_init(|| Arc::new(crate::MemorySearchBackend::new()) as Arc<dyn SearchBackend>)
}
fn client_builder(&self) -> SearchClientBuilder {
let config = self.effective_config();
let mut builder = SearchClient::builder()
.enabled(config.enabled)
.batch_size(config.batch_size);
if let Some(store) = self.postgres_store() {
builder = builder
.backend(Arc::clone(store) as Arc<dyn SearchBackend>)
.source(Arc::clone(store) as Arc<dyn DocumentSource>);
}
if let Some(backend) = &self.backend {
builder = builder.backend(Arc::clone(backend));
} else if self.postgres_store().is_none() {
builder = builder.backend(Arc::clone(self.default_backend()));
}
if let Some(embedder) = &self.embedder {
builder = builder.embedder(Arc::clone(embedder));
}
if let Some(source) = &self.source {
builder = builder.source(Arc::clone(source));
}
if let Some(visibility) = &self.visibility {
builder = builder.visibility(Arc::clone(visibility));
}
for definition in &self.indexes {
builder = builder.index_definition(definition.clone());
}
builder
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum BackfillTarget {
AllIndexes,
Index(String),
}
impl BackfillTarget {
#[must_use]
pub fn name(&self) -> Option<&str> {
match self {
Self::AllIndexes => None,
Self::Index(name) => Some(name),
}
}
}
#[must_use]
pub fn backfill_request_from_env(
config: &SearchConfig,
) -> Option<(BackfillTarget, BackfillOptions)> {
let raw = std::env::var(BACKFILL_ENV).ok()?;
let target = parse_backfill_target(&raw)?;
let purge_raw = std::env::var(BACKFILL_PURGE_ENV).ok();
Some((
target,
BackfillOptions::default()
.batch_size(config.batch_size)
.purge(parse_purge_flag(purge_raw.as_deref())),
))
}
#[allow(
clippy::unnecessary_wraps,
clippy::missing_const_for_fn,
reason = "the SQLite arm of `backend_select!` returns an error"
)]
fn require_postgres_backend() -> Result<(), autumn_web::AutumnError> {
autumn_web::backend_select! {
pg => { Ok(()) },
sqlite => {
Err(autumn_web::AutumnError::internal_server_error_msg(
"SearchPlugin::postgres() needs the Postgres backend, but this build of \
autumn-web uses SQLite (`--features sqlite`); install a different search \
backend with `SearchPlugin::backend`",
))
},
}
}
fn parse_backfill_target(raw: &str) -> Option<BackfillTarget> {
let target = raw.trim();
if target.is_empty() {
return None;
}
if target.eq_ignore_ascii_case("all") {
Some(BackfillTarget::AllIndexes)
} else {
Some(BackfillTarget::Index(target.to_owned()))
}
}
fn parse_purge_flag(raw: Option<&str>) -> bool {
raw.is_some_and(|value| matches!(value.trim(), "1" | "true" | "yes"))
}
impl Plugin for SearchPlugin {
fn contract(&self) -> Option<autumn_web::plugin_contract::PluginContract> {
Some(autumn_web::plugin_contract::lockstep_contract(
env!("CARGO_PKG_NAME"),
env!("CARGO_PKG_VERSION"),
))
}
fn name(&self) -> Cow<'static, str> {
Cow::Borrowed("autumn-search")
}
fn build(self, app: AppBuilder) -> AppBuilder {
let app = app.config_section("search");
let config_error = self.resolved_config().as_ref().err().cloned();
let config = self.effective_config();
let postgres = self.postgres_store().map(Arc::clone);
let jobs = search_job_infos(&config.queue);
let client = self.client();
app.jobs(jobs).on_startup(move |state| {
let client = client.clone();
let postgres = postgres.clone();
let config = config.clone();
let config_error = config_error.clone();
async move {
if let Some(message) = config_error {
return Err(autumn_web::AutumnError::internal_server_error_msg(format!(
"autumn-search: {message}"
)));
}
if postgres.is_some() {
require_postgres_backend()?;
}
state.insert_extension(client.clone());
if config.enabled {
let embedder_width = client.embedding_dimensions();
if let Some(configured) = config.embedding_dimensions
&& embedder_width != 0
&& embedder_width != configured
{
return Err(autumn_web::AutumnError::internal_server_error_msg(format!(
"autumn-search: `search.embedding_dimensions = {configured}` does not \
match the installed embedder, which produces \
{embedder_width}-dimension vectors; set them to the same value (and \
re-run `autumn search reindex --purge` if the width really changed)"
)));
}
if let Some(store) = &postgres {
let pool = state.pool().cloned().ok_or_else(|| {
autumn_web::AutumnError::internal_server_error_msg(
"SearchPlugin::postgres() needs a database pool; configure \
`database.primary_url` or install a different backend",
)
})?;
store.install_pool(pool);
}
client.ensure_indexes().await.map_err(|error| {
autumn_web::AutumnError::internal_server_error_msg(format!(
"search index setup failed: {error}"
))
})?;
} else {
tracing::warn!(
"autumn-search is disabled (`[search] enabled = false`): index writes and \
queries are no-ops, and no backend setup was performed"
);
}
if let Some((target, options)) = backfill_request_from_env(&config) {
run_backfill_and_exit(&client, &target, &options).await;
}
Ok(())
}
})
}
}
async fn run_backfill_and_exit(
client: &SearchClient,
target: &BackfillTarget,
options: &BackfillOptions,
) -> ! {
println!(
"{BACKFILL_STARTED_MARKER} {}",
target.name().unwrap_or("all")
);
let _ = std::io::Write::flush(&mut std::io::stdout());
if !client.is_enabled() {
eprintln!(
"\u{2717} autumn-search: `[search] enabled = false` for this profile — nothing was \
rebuilt.\n Re-enable search (or pick the right profile with `--profile`) and run \
the reindex again."
);
std::process::exit(1);
}
let result = match target.name() {
Some(index) => client.backfill(index, options).await.map(|r| vec![r]),
None => client.backfill_all(options).await,
};
match result {
Ok(reports) => {
for report in &reports {
println!(
"autumn-search: reindexed {} ({} documents in {} batches{})",
report.index,
report.indexed,
report.batches,
if report.purged { ", purged first" } else { "" }
);
}
if reports.is_empty() {
println!("autumn-search: no indexes are registered; nothing to reindex");
}
std::process::exit(0)
}
Err(error) => {
eprintln!("autumn-search: reindex failed: {error}");
std::process::exit(1)
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn the_postgres_store_needs_the_postgres_backend() {
let result = require_postgres_backend();
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}");
},
}
}
#[test]
fn the_plugin_name_is_the_crate_name() {
assert_eq!(SearchPlugin::new().name(), "autumn-search");
}
#[test]
fn contract_declares_lockstep_with_own_crate() {
let contract = SearchPlugin::new().contract().expect("a contract");
assert_eq!(contract.plugin, env!("CARGO_PKG_NAME"));
assert_eq!(
contract.plugin_version.as_deref(),
Some(env!("CARGO_PKG_VERSION"))
);
assert_eq!(
contract.autumn_web.as_deref(),
Some(autumn_web::plugin_contract::lockstep_range(env!("CARGO_PKG_VERSION")).as_str())
);
assert!(contract.experimental_surfaces.is_empty());
}
#[test]
fn the_queue_override_reaches_the_registered_jobs() {
let plugin = SearchPlugin::new().queue("indexing");
assert_eq!(plugin.search_config().queue, "indexing");
let infos = search_job_infos(&plugin.search_config().queue);
assert!(infos.iter().all(|info| info.queue == "indexing"));
}
#[test]
fn a_queue_override_does_not_suppress_the_rest_of_the_file_config() {
let plugin = SearchPlugin::new().queue("indexing");
assert!(
!plugin.config_explicit,
"a queue override must leave the file config in play"
);
let plugin = SearchPlugin::new().config(SearchConfig::default());
assert!(plugin.config_explicit);
}
#[test]
fn an_explicit_config_still_honours_a_later_queue_override() {
let config = SearchConfig::from_toml_str("[search]\nbatch_size = 7\n").expect("parse");
let plugin = SearchPlugin::new().config(config).queue("indexing");
let effective = plugin.search_config();
assert_eq!(effective.queue, "indexing");
assert_eq!(effective.batch_size, 7, "the rest of the config survives");
}
#[test]
fn an_explicit_backend_turns_off_the_postgres_default() {
let plugin = SearchPlugin::new()
.postgres()
.backend(Arc::new(crate::memory::MemorySearchBackend::new()));
assert!(!plugin.uses_postgres());
}
#[test]
fn an_explicit_backend_keeps_an_explicit_document_source() {
let source = Arc::new(crate::source::MemoryDocumentSource::new());
let plugin = SearchPlugin::new()
.postgres()
.source(Arc::clone(&source) as Arc<dyn DocumentSource>)
.backend(Arc::new(crate::memory::MemorySearchBackend::new()));
assert!(
plugin.source.is_some(),
"overriding the backend must not discard the app's own source"
);
}
#[test]
fn an_explicit_config_is_validated_like_a_resolved_one() {
let config = SearchConfig {
batch_size: 0,
..SearchConfig::default()
};
let plugin = SearchPlugin::new().config(config);
let error = plugin
.resolved_config()
.as_ref()
.expect_err("a zero batch size must not resolve");
assert!(error.contains("batch_size"), "unexpected error: {error}");
}
#[test]
fn an_unset_or_blank_target_means_a_normal_boot() {
assert!(parse_backfill_target("").is_none());
assert!(parse_backfill_target(" ").is_none());
}
#[test]
fn all_means_every_index_and_a_name_means_one() {
assert_eq!(
parse_backfill_target("all"),
Some(BackfillTarget::AllIndexes)
);
assert_eq!(
parse_backfill_target("ALL"),
Some(BackfillTarget::AllIndexes)
);
assert_eq!(
parse_backfill_target(" articles "),
Some(BackfillTarget::Index("articles".to_owned()))
);
assert_eq!(BackfillTarget::AllIndexes.name(), None);
assert_eq!(
BackfillTarget::Index("articles".to_owned()).name(),
Some("articles")
);
}
#[test]
fn purge_is_opt_in_and_only_for_affirmative_values() {
assert!(parse_purge_flag(Some("1")));
assert!(parse_purge_flag(Some("true")));
assert!(parse_purge_flag(Some(" yes ")));
assert!(!parse_purge_flag(Some("0")));
assert!(!parse_purge_flag(Some("no")));
assert!(!parse_purge_flag(Some("")));
assert!(!parse_purge_flag(None));
}
}