sift_query_sqlx/enrichment/
label.rs

1//! PostgreSQL label enrichment.
2//!
3//! This module deliberately keeps PostgreSQL's `ExtensionError`, `Page`, and
4//! extension registry contract. The portable implementation in
5//! `sift-query-sqlx-common` serves the MySQL/SQLite backends, whose errors and
6//! page types are separate public contracts.
7
8use crate::{
9    EnrichContext, ExtensionError, ExtensionFuture, Page, PageEnricher, RowMap, TypedPageEnricher,
10};
11use sift_query_algorithm::TtlLruMap;
12use sift_query_core::RowValue;
13use std::collections::{BTreeMap, BTreeSet, HashMap};
14use std::sync::{Arc, Mutex};
15use std::time::{Duration, Instant};
16use tokio::sync::Semaphore;
17use tokio::time::timeout;
18
19#[derive(Debug, Clone, PartialEq, Eq)]
20pub struct Label {
21    pub id: String,
22    pub label: String,
23}
24pub trait LabelLoader: Send + Sync {
25    fn supports(&self, key: &str) -> bool;
26    fn load(&self, key: &str, ids: Vec<String>) -> ExtensionFuture<Vec<Label>>;
27}
28
29#[derive(Debug, Clone, Copy, Default)]
30pub enum MissingLabelPolicy {
31    #[default]
32    Omit,
33    Null,
34    Error,
35}
36#[derive(Debug, Clone, Copy, Default)]
37pub enum LabelFailurePolicy {
38    #[default]
39    Error,
40    KeepRows,
41}
42
43/// Typed opt-in: implementations explicitly name the source fields and writable
44/// label fields. `set_label` must not modify database fields or ordering keys.
45pub trait LabelTarget: Send + 'static {
46    fn label_id(&self, source: &str) -> Option<String>;
47    fn set_label(&mut self, source: &str, label: Option<&str>) -> Result<(), ExtensionError>;
48}
49
50type LabelKey = (String, String, String);
51type Cache = Arc<LabelCache>;
52
53struct LabelCache {
54    origin: Instant,
55    entries: Mutex<TtlLruMap<LabelKey, String>>,
56}
57
58#[derive(Clone)]
59pub struct LabelEnricher {
60    bindings: Vec<(String, String)>,
61    loaders: Vec<Arc<dyn LabelLoader>>,
62    max_ids: usize,
63    timeout: Duration,
64    gate: Arc<Semaphore>,
65    missing: MissingLabelPolicy,
66    failure: LabelFailurePolicy,
67    cache: Option<Cache>,
68}
69impl std::fmt::Debug for LabelEnricher {
70    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
71        f.debug_struct("LabelEnricher")
72            .field("bindings", &self.bindings)
73            .field("max_ids", &self.max_ids)
74            .field("timeout", &self.timeout)
75            .field("cached", &self.cache.is_some())
76            .finish_non_exhaustive()
77    }
78}
79impl Default for LabelEnricher {
80    fn default() -> Self {
81        Self::new()
82    }
83}
84impl LabelEnricher {
85    #[must_use]
86    pub fn new() -> Self {
87        Self {
88            bindings: Vec::new(),
89            loaders: Vec::new(),
90            max_ids: 512,
91            timeout: Duration::from_secs(2),
92            gate: Arc::new(Semaphore::new(4)),
93            missing: MissingLabelPolicy::Omit,
94            failure: LabelFailurePolicy::Error,
95            cache: None,
96        }
97    }
98    #[must_use]
99    pub fn bind(mut self, source: impl Into<String>, key: impl Into<String>) -> Self {
100        self.bindings.push((source.into(), key.into()));
101        self
102    }
103
104    #[must_use]
105    pub fn from_entity<E: sift_query_core::SearchEntity>(self) -> Self {
106        E::label_bindings()
107            .iter()
108            .fold(self, |enricher, (source, key)| enricher.bind(*source, *key))
109    }
110    #[must_use]
111    pub fn with_loader<L: LabelLoader + 'static>(mut self, loader: L) -> Self {
112        self.loaders.push(Arc::new(loader));
113        self
114    }
115    #[must_use]
116    pub fn with_budget(mut self, max_ids: usize, timeout: Duration) -> Self {
117        self.max_ids = max_ids;
118        self.timeout = timeout;
119        self
120    }
121    #[must_use]
122    pub fn with_missing_policy(mut self, policy: MissingLabelPolicy) -> Self {
123        self.missing = policy;
124        self
125    }
126    #[must_use]
127    pub fn with_failure_policy(mut self, policy: LabelFailurePolicy) -> Self {
128        self.failure = policy;
129        self
130    }
131    pub fn with_concurrency(
132        mut self,
133        maximum: usize,
134    ) -> Result<Self, sift_query_core::ConfigError> {
135        if maximum == 0 || maximum > 1024 {
136            return Err(sift_query_core::ConfigError::Invalid(
137                "invalid label concurrency".to_owned(),
138            ));
139        }
140        self.gate = Arc::new(Semaphore::new(maximum));
141        Ok(self)
142    }
143    /// Optional positive-result cache. Cache ownership is this enricher and its
144    /// clones, never global. Use distinct instances for different tenant/locale
145    /// dictionaries; remote request-dependent loaders should leave it disabled.
146    pub fn with_cache(
147        mut self,
148        capacity: usize,
149        ttl: Duration,
150    ) -> Result<Self, sift_query_core::ConfigError> {
151        if capacity == 0 || capacity > 16_384 || ttl.is_zero() {
152            return Err(sift_query_core::ConfigError::Invalid(
153                "invalid label cache budget".to_owned(),
154            ));
155        }
156        self.cache = Some(Arc::new(LabelCache {
157            origin: Instant::now(),
158            entries: Mutex::new(TtlLruMap::new(capacity, ttl).map_err(|_| {
159                sift_query_core::ConfigError::Invalid("invalid label cache budget".to_owned())
160            })?),
161        }));
162        Ok(self)
163    }
164
165    async fn apply(
166        &self,
167        mut page: Page<RowMap>,
168        context: EnrichContext,
169    ) -> Result<Page<RowMap>, ExtensionError> {
170        if self.timeout.is_zero() || self.max_ids == 0 {
171            return Err(ExtensionError::rejected("label", "invalid label budget"));
172        }
173        if self.bindings.len() > 32 {
174            return Err(ExtensionError::rejected("label", "too many label bindings"));
175        }
176        if page.data.is_empty() {
177            return Ok(page);
178        }
179        let mut groups = BTreeMap::<String, BTreeSet<String>>::new();
180        let mut seen_sources = BTreeSet::new();
181        for (source, key) in &self.bindings {
182            if !seen_sources.insert(source) {
183                return Err(ExtensionError::rejected(
184                    "label",
185                    "duplicate source binding",
186                ));
187            }
188            groups.entry(key.clone()).or_default();
189            for row in &page.data {
190                if let Some(id) = row.get(source).and_then(row_key) {
191                    groups.entry(key.clone()).or_default().insert(id);
192                }
193            }
194        }
195        if groups.values().map(BTreeSet::len).sum::<usize>() > self.max_ids {
196            return Err(ExtensionError::LabelIdLimit(self.max_ids));
197        }
198        for (key, ids) in groups {
199            let loaders = self
200                .loaders
201                .iter()
202                .filter(|loader| loader.supports(&key))
203                .collect::<Vec<_>>();
204            if loaders.len() != 1 {
205                return Err(ExtensionError::rejected(
206                    "label",
207                    "missing or ambiguous loader",
208                ));
209            }
210            let mut labels = HashMap::new();
211            let mut misses = Vec::new();
212            for id in &ids {
213                let cached = self.cache.as_ref().and_then(|cache| {
214                    cache
215                        .entries
216                        .lock()
217                        .unwrap_or_else(std::sync::PoisonError::into_inner)
218                        .get(
219                            &(context.entity_name.to_owned(), key.clone(), id.clone()),
220                            cache.origin.elapsed(),
221                        )
222                        .cloned()
223                });
224                if let Some(label) = cached {
225                    labels.insert(id.clone(), label);
226                } else {
227                    misses.push(id.clone());
228                }
229            }
230            if !misses.is_empty() {
231                let loaded = timeout(self.timeout, async {
232                    let _permit = self
233                        .gate
234                        .acquire()
235                        .await
236                        .map_err(|_| ExtensionError::rejected("label", "loader gate closed"))?;
237                    loaders[0].load(&key, misses.clone()).await
238                })
239                .await
240                .map_err(|_| ExtensionError::LabelTimeout)
241                .and_then(std::convert::identity);
242                let loaded = match loaded {
243                    Ok(loaded) => loaded,
244                    Err(_) if matches!(self.failure, LabelFailurePolicy::KeepRows) => {
245                        tracing::warn!(
246                            outcome = "label_unavailable",
247                            "label loading failed; keeping rows"
248                        );
249                        continue;
250                    }
251                    Err(error) => return Err(error),
252                };
253                let mut fresh = HashMap::new();
254                for label in loaded {
255                    if !misses.contains(&label.id) || fresh.insert(label.id, label.label).is_some()
256                    {
257                        return Err(ExtensionError::rejected(
258                            "label",
259                            "loader returned duplicate or unrequested IDs",
260                        ));
261                    }
262                }
263                // Cache only after the complete batch is validated. Failed or
264                // malicious responses cannot poison the shared cache.
265                if let Some(cache) = &self.cache {
266                    let now = cache.origin.elapsed();
267                    let mut cache = cache
268                        .entries
269                        .lock()
270                        .unwrap_or_else(std::sync::PoisonError::into_inner);
271                    for (id, label) in &fresh {
272                        cache.insert(
273                            (context.entity_name.to_owned(), key.clone(), id.clone()),
274                            label.clone(),
275                            now,
276                        );
277                    }
278                }
279                labels.extend(fresh);
280            }
281            for (source, binding_key) in &self.bindings {
282                if binding_key != &key {
283                    continue;
284                }
285                let target = format!("{source}Label");
286                for row in &mut page.data {
287                    if row.contains_key(&target) {
288                        return Err(ExtensionError::rejected(
289                            "label",
290                            "label target already exists",
291                        ));
292                    }
293                    match row
294                        .get(source)
295                        .and_then(row_key)
296                        .and_then(|id| labels.get(&id))
297                    {
298                        Some(label) => {
299                            row.insert(target.clone(), RowValue::Text(label.clone()));
300                        }
301                        None => match self.missing {
302                            MissingLabelPolicy::Omit => {}
303                            MissingLabelPolicy::Null => {
304                                row.insert(target.clone(), RowValue::Null);
305                            }
306                            MissingLabelPolicy::Error => {
307                                return Err(ExtensionError::rejected("label", "missing label"));
308                            }
309                        },
310                    }
311                }
312            }
313        }
314        Ok(page)
315    }
316}
317
318impl PageEnricher for LabelEnricher {
319    fn enrich(&self, page: Page<RowMap>, context: EnrichContext) -> ExtensionFuture<Page<RowMap>> {
320        let enricher = self.clone();
321        Box::pin(async move { enricher.apply(page, context).await })
322    }
323}
324impl<T: LabelTarget> TypedPageEnricher<T> for LabelEnricher {
325    fn enrich(&self, mut rows: Vec<T>, context: EnrichContext) -> ExtensionFuture<Vec<T>> {
326        let enricher = self.clone();
327        Box::pin(async move {
328            let maps = rows
329                .iter()
330                .map(|row| {
331                    enricher
332                        .bindings
333                        .iter()
334                        .filter_map(|(source, _)| {
335                            row.label_id(source)
336                                .map(|id| (source.clone(), RowValue::Text(id)))
337                        })
338                        .collect()
339                })
340                .collect();
341            let page = enricher
342                .apply(
343                    Page {
344                        data: maps,
345                        total_count: 0,
346                        summaries: Vec::new(),
347                        next_cursor: None,
348                    },
349                    context,
350                )
351                .await?;
352            for (row, map) in rows.iter_mut().zip(&page.data) {
353                for (source, _) in &enricher.bindings {
354                    if let Some(value) = map.get(&format!("{source}Label")) {
355                        row.set_label(
356                            source,
357                            match value {
358                                RowValue::Text(label) => Some(label),
359                                _ => None,
360                            },
361                        )?;
362                    }
363                }
364            }
365            Ok(rows)
366        })
367    }
368}
369
370/// Enum/local dictionaries share the same one-batch loader contract as remote services.
371#[derive(Debug, Clone, Default)]
372pub struct DictionaryLabelLoader {
373    dictionaries: HashMap<String, HashMap<String, String>>,
374}
375impl DictionaryLabelLoader {
376    #[must_use]
377    pub fn with_labels(
378        mut self,
379        key: impl Into<String>,
380        labels: impl IntoIterator<Item = (String, String)>,
381    ) -> Self {
382        self.dictionaries
383            .insert(key.into(), labels.into_iter().collect());
384        self
385    }
386}
387impl LabelLoader for DictionaryLabelLoader {
388    fn supports(&self, key: &str) -> bool {
389        self.dictionaries.contains_key(key)
390    }
391    fn load(&self, key: &str, ids: Vec<String>) -> ExtensionFuture<Vec<Label>> {
392        let result = self
393            .dictionaries
394            .get(key)
395            .map(|dict| {
396                ids.into_iter()
397                    .filter_map(|id| {
398                        dict.get(&id).map(|label| Label {
399                            id: id.clone(),
400                            label: label.clone(),
401                        })
402                    })
403                    .collect()
404            })
405            .ok_or_else(|| ExtensionError::rejected("label", "missing dictionary"));
406        Box::pin(async move { result })
407    }
408}
409fn row_key(value: &RowValue) -> Option<String> {
410    match value {
411        RowValue::Text(value) => Some(value.clone()),
412        RowValue::I32(value) => Some(value.to_string()),
413        RowValue::I64(value) => Some(value.to_string()),
414        #[cfg(feature = "uuid")]
415        RowValue::Uuid(value) => Some(value.to_string()),
416        _ => None,
417    }
418}
sift_query_sqlx/enrichment/label.rs

sift_query_sqlx/enrichment/
label.rs
