Skip to main content

sofie_db/search/
indexing.rs

1//! Jobs keeping a search engine in step after each commit. A job reads the record as it is
2//! then, so a retried or late job writes the latest state.
3
4use std::collections::HashMap;
5use std::sync::Arc;
6
7use serde_json::{Value, json};
8use sofie_core::{Cast, Error, Result, SqlType, SqlValue, Uuid};
9
10use super::engine::{Installed, SearchEngine};
11use super::texts::{PAGE, Record, Records, read_records, rewrite_stale};
12use super::{Document, SEARCH_COLUMN, SearchIndex};
13use crate::worker::{JobContext, Worker};
14use crate::{Db, MessageKind, Outcome, Plan, Resolve, Write, query_decoded};
15
16const INDEX_JOB: &str = "sofie.search.index";
17const REINDEX_JOB: &str = "sofie.search.reindex_page";
18const FORGET_JOB: &str = "sofie.search.forget";
19const INDEX_TRIES: usize = 3;
20
21/// On the worker of a database a search engine is installed on.
22#[must_use]
23pub fn index_searches(worker: Worker) -> Worker {
24    worker
25        .job(INDEX_JOB, index_record)
26        .job(REINDEX_JOB, reindex_page)
27        .job(FORGET_JOB, forget_documents)
28}
29
30pub(crate) fn changes_document(write: &Write<'_>) -> bool {
31    let table = write.table;
32    std::iter::once(SEARCH_COLUMN)
33        .chain(table.hiding())
34        .chain(table.owner)
35        .any(|column| write.value(column) != write.loaded(column))
36}
37
38pub(crate) fn queue(write: &Write<'_>, plan: &mut Plan<'_>) -> Result<()> {
39    let job = index_job(write.table.schema, write.table.name, write.id);
40    plan.queue(MessageKind::Job, [job])
41}
42
43fn index_job(module: &str, aggregate: &str, id: Uuid) -> (&'static str, Value) {
44    let payload = json!({ "module": module, "aggregate": aggregate, "id": id });
45    (INDEX_JOB, payload)
46}
47
48/// In the prune's own commit: the rows are gone, so no later prune or reindex would find them.
49#[cfg(feature = "retention")]
50pub(crate) fn queue_forget(index: &SearchIndex, ids: &[Uuid], plan: &mut Plan<'_>) -> Result<()> {
51    let payload = json!({ "module": index.module, "aggregate": index.name, "ids": ids });
52    plan.queue(MessageKind::Job, [(FORGET_JOB, payload)])
53}
54
55#[cfg(feature = "retention")]
56pub(crate) fn index_of(db: &Db, module: &str, name: &str) -> Option<SearchIndex> {
57    db.search.get()?.find(module, name).cloned()
58}
59
60struct Forgetting {
61    indexing: Indexing,
62    ids: Vec<Uuid>,
63}
64
65impl Resolve<JobContext> for Forgetting {
66    fn resolve(context: &mut JobContext) -> impl Future<Output = Result<Self>> + Send {
67        let ids = context
68            .payload
69            .get("ids")
70            .cloned()
71            .map(serde_json::from_value::<Vec<Uuid>>)
72            .transpose()
73            .map_err(|error| Error::internal("reading a forgetting job's ids", error))
74            .and_then(|ids| {
75                ids.ok_or_else(|| Error::internal_message("a forgetting job names no ids"))
76            });
77        let forgetting = ids
78            .and_then(|ids| Indexing::read(context, "id").map(|indexing| Self { indexing, ids }));
79        std::future::ready(forgetting)
80    }
81}
82
83async fn forget_documents(job: Forgetting) -> Result<()> {
84    let Forgetting { indexing, ids } = job;
85    indexing.engine.remove(&indexing.index, &ids).await
86}
87
88struct Indexing {
89    db: Db,
90    engine: Arc<dyn SearchEngine>,
91    index: SearchIndex,
92    id: Option<Uuid>,
93}
94
95impl Indexing {
96    fn read(context: &JobContext, id_key: &str) -> Result<Self> {
97        let installed = Installed::of(&context.db)?;
98        let text = |key: &str| {
99            context
100                .payload
101                .get(key)
102                .and_then(Value::as_str)
103                .ok_or_else(|| {
104                    Error::internal_message(format!("an indexing job's payload has no `{key}`"))
105                })
106        };
107        let index = installed
108            .index(text("module")?, text("aggregate")?)?
109            .clone();
110        let id = match context.payload.get(id_key) {
111            None | Some(Value::Null) => None,
112            Some(_) => Some(text(id_key)?.parse().map_err(|_| {
113                Error::internal_message(format!("an indexing job's `{id_key}` is not an id"))
114            })?),
115        };
116        Ok(Self {
117            db: context.db.clone(),
118            engine: Arc::clone(&installed.engine),
119            index,
120            id,
121        })
122    }
123}
124
125struct OneRecord(Indexing);
126
127impl Resolve<JobContext> for OneRecord {
128    fn resolve(context: &mut JobContext) -> impl Future<Output = Result<Self>> + Send {
129        std::future::ready(Indexing::read(context, "id").map(Self))
130    }
131}
132
133struct NextPage(Indexing);
134
135impl Resolve<JobContext> for NextPage {
136    fn resolve(context: &mut JobContext) -> impl Future<Output = Result<Self>> + Send {
137        std::future::ready(Indexing::read(context, "after").map(Self))
138    }
139}
140
141/// Read again after each write: a job for a later version may have landed first, and this one's
142/// write then undid it. A record rewritten under every try is a conflict, retried later.
143async fn index_record(job: OneRecord) -> Result<()> {
144    let OneRecord(job) = job;
145    let id = job
146        .id
147        .ok_or_else(|| Error::internal_message("an indexing job names no record"))?;
148    let mut records = read_records(&job.db, &job.index, Records::One(id)).await?;
149    for _ in 0..INDEX_TRIES {
150        rewrite_stale(&job.db, &job.index, &records).await?;
151        let indexed = records.first().map(|record| record.version.clone());
152        match records.into_iter().next() {
153            Some(record) if !record.hidden => job.engine.put(&job.index, &[record.document]).await,
154            _ => job.engine.remove(&job.index, &[id]).await,
155        }?;
156        records = read_records(&job.db, &job.index, Records::One(id)).await?;
157        if records.first().map(|record| record.version.clone()) == indexed {
158            return Ok(());
159        }
160    }
161    Err(Error::Conflict)
162}
163
164async fn reindex_page(job: NextPage) -> Result<Reindexing> {
165    let NextPage(job) = job;
166    let records = read_records(&job.db, &job.index, Records::After(job.id)).await?;
167    rewrite_stale(&job.db, &job.index, &records).await?;
168    let last = records.last().map(|record| record.document.id);
169    let full = records.len() == PAGE;
170    let written: Vec<(Uuid, SqlValue)> = records
171        .iter()
172        .map(|record| (record.document.id, record.version.clone()))
173        .collect();
174    let (hidden, shown): (Vec<Record>, Vec<Record>) =
175        records.into_iter().partition(|record| record.hidden);
176    let documents: Vec<Document> = shown.into_iter().map(|record| record.document).collect();
177    if !documents.is_empty() {
178        job.engine.put(&job.index, &documents).await?;
179    }
180    let removed: Vec<Uuid> = hidden.iter().map(|record| record.document.id).collect();
181    if !removed.is_empty() {
182        job.engine.remove(&job.index, &removed).await?;
183    }
184    let moved = moved_since(&job.db, &job.index, &written).await?;
185    let pages = match last {
186        Some(last) if full => vec![(job.index.clone(), Some(last))],
187        _ => Vec::new(),
188    };
189    Ok(Reindexing {
190        pages,
191        moved: moved
192            .into_iter()
193            .map(|id| (job.index.clone(), id))
194            .collect(),
195    })
196}
197
198/// A commit between the page's read and its write was undone by it: index those records again.
199async fn moved_since(
200    db: &Db,
201    index: &SearchIndex,
202    written: &[(Uuid, SqlValue)],
203) -> Result<Vec<Uuid>> {
204    let (Some((first, _)), Some((last, _))) = (written.first(), written.last()) else {
205        return Ok(Vec::new());
206    };
207    let dialect = db.driver().dialect();
208    let id = dialect.quote("id");
209    let sql = format!(
210        "SELECT {id}, {} FROM {} WHERE {id} >= {} AND {id} <= {}",
211        dialect.quote("version"),
212        dialect.qualified(index.module, index.name),
213        dialect.placeholder(1),
214        dialect.placeholder(2),
215    );
216    let params = [SqlValue::Uuid(*first), SqlValue::Uuid(*last)];
217    let types = vec![Some(SqlType::Uuid), Some(SqlType::Int)];
218    let now: HashMap<Uuid, SqlValue> = query_decoded(db.driver(), &sql, &params, types, |row| {
219        Ok((Uuid::decode(row)?, SqlValue::Int(i64::decode(row)?)))
220    })
221    .await?
222    .into_iter()
223    .collect();
224    Ok(written
225        .iter()
226        .filter(|(id, version)| now.get(id) != Some(version))
227        .map(|(id, _)| *id)
228        .collect())
229}
230
231pub struct Reindexing {
232    pages: Vec<(SearchIndex, Option<Uuid>)>,
233    moved: Vec<(SearchIndex, Uuid)>,
234}
235
236impl Reindexing {
237    /// Each page queues the next, as a mirror's backfill does, so a stopped worker resumes.
238    pub(crate) fn every_page(indexes: &[SearchIndex]) -> Self {
239        Self {
240            pages: indexes.iter().map(|index| (index.clone(), None)).collect(),
241            moved: Vec::new(),
242        }
243    }
244}
245
246impl Outcome for Reindexing {
247    fn plan(&self, plan: &mut Plan<'_>) -> Result<()> {
248        let pages = self.pages.iter().map(|(index, after)| {
249            let payload = json!({
250                "module": index.module,
251                "aggregate": index.name,
252                "after": after,
253            });
254            (REINDEX_JOB, payload)
255        });
256        let records = self
257            .moved
258            .iter()
259            .map(|(index, id)| index_job(index.module, index.name, *id));
260        plan.queue(MessageKind::Job, pages.chain(records))
261    }
262}