sift_query_sqlx/
observe.rs

1//! PostgreSQL query observation and SDK-independent completion callbacks.
2//!
3//! The PostgreSQL observer adds execution-specific span fields. Portable
4//! backends use `sift-query-sqlx-common::Observers`, which records the same
5//! low-sensitivity outcome contract without PostgreSQL execution metadata.
6use std::future::Future;
7use std::sync::Arc;
8use std::task::Poll;
9use std::time::{Duration, Instant};
10
11use crate::{PgSearchError, PgSiftQuery, SearchStream};
12use sift_query_core::SearchEntity;
13use tracing::Instrument;
14
15#[derive(Debug, Clone, Copy)]
16pub struct QueryObservation {
17    pub entity: &'static str,
18    pub operation: &'static str,
19    pub duration: Duration,
20    pub rows_returned: usize,
21    /// `success`, `cancelled`, or a low-cardinality error category.
22    pub outcome: &'static str,
23    pub slow: bool,
24}
25
26/// Metrics/event bridge. Called once per completed or cancelled operation;
27/// panics are isolated. The callback must be nonblocking and cannot see SQL,
28/// binds, raw input, tenant identifiers, or cursor tokens.
29pub trait QueryObserver: Send + Sync {
30    fn completed(&self, observation: &QueryObservation);
31}
32
33#[derive(Clone, Default)]
34pub(crate) struct Observers {
35    observer: Option<Arc<dyn QueryObserver>>,
36    slow_threshold: Option<Duration>,
37    entity_thresholds: std::collections::HashMap<u64, Duration>,
38}
39impl std::fmt::Debug for Observers {
40    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
41        f.debug_struct("Observers")
42            .field("configured", &self.observer.is_some())
43            .field("slow_threshold", &self.slow_threshold)
44            .finish_non_exhaustive()
45    }
46}
47
48impl PgSiftQuery {
49    #[must_use]
50    pub fn with_query_observer<O: QueryObserver + 'static>(mut self, observer: O) -> Self {
51        self.observers.observer = Some(Arc::new(observer));
52        self
53    }
54    #[must_use]
55    pub fn with_slow_query_threshold(mut self, threshold: Duration) -> Self {
56        self.observers.slow_threshold = Some(threshold);
57        self
58    }
59    #[must_use]
60    pub fn with_entity_slow_query_threshold<E: SearchEntity>(
61        mut self,
62        threshold: Duration,
63    ) -> Self {
64        self.observers
65            .entity_thresholds
66            .insert(E::META.schema_fingerprint, threshold);
67        self
68    }
69
70    pub(crate) fn observation<E: SearchEntity>(&self, operation: &'static str) -> ObservationGuard {
71        let span = tracing::info_span!("sift_query", entity = E::META.entity_name, backend = "postgres", operation,
72            execution_strategy = ?self.execution.strategy, count_strategy = ?self.execution.count_strategy,
73            rows_returned = tracing::field::Empty, duration_us = tracing::field::Empty,
74            outcome = tracing::field::Empty, query_shape_id = tracing::field::Empty, rows_written = tracing::field::Empty);
75        ObservationGuard {
76            span,
77            started: Instant::now(),
78            entity: E::META.entity_name,
79            operation,
80            rows: 0,
81            outcome: "cancelled",
82            observer: self.observers.observer.clone(),
83            threshold: self
84                .observers
85                .entity_thresholds
86                .get(&E::META.schema_fingerprint)
87                .copied()
88                .or(self.observers.slow_threshold),
89        }
90    }
91
92    pub(crate) async fn observe<E: SearchEntity, T>(
93        &self,
94        operation: &'static str,
95        future: impl Future<Output = Result<T, PgSearchError>>,
96        rows: impl FnOnce(&T) -> usize,
97    ) -> Result<T, PgSearchError> {
98        let mut guard = self.observation::<E>(operation);
99        let result = future.instrument(guard.span.clone()).await;
100        match &result {
101            Ok(value) => {
102                guard.rows = rows(value);
103                guard.outcome = "success";
104            }
105            Err(error) => guard.outcome = error_class(error),
106        }
107        result
108    }
109
110    pub(crate) fn observe_stream<E: SearchEntity, T: Send + 'static>(
111        &self,
112        operation: &'static str,
113        build: impl FnOnce() -> Result<SearchStream<T>, PgSearchError>,
114    ) -> Result<SearchStream<T>, PgSearchError> {
115        let mut guard = self.observation::<E>(operation);
116        let result = guard.span.in_scope(build);
117        let mut stream = match result {
118            Ok(stream) => stream,
119            Err(error) => {
120                guard.outcome = error_class(&error);
121                return Err(error);
122            }
123        };
124        Ok(Box::pin(futures_util::stream::poll_fn(move |cx| {
125            let span = guard.span.clone();
126            let _entered = span.enter();
127            match stream.as_mut().poll_next(cx) {
128                Poll::Ready(Some(Ok(row))) => {
129                    guard.rows += 1;
130                    Poll::Ready(Some(Ok(row)))
131                }
132                Poll::Ready(Some(Err(error))) => {
133                    guard.outcome = error_class(&error);
134                    Poll::Ready(Some(Err(error)))
135                }
136                Poll::Ready(None) => {
137                    if guard.outcome == "cancelled" {
138                        guard.outcome = "success";
139                    }
140                    Poll::Ready(None)
141                }
142                Poll::Pending => Poll::Pending,
143            }
144        })))
145    }
146}
147
148pub(crate) struct ObservationGuard {
149    pub(crate) span: tracing::Span,
150    started: Instant,
151    entity: &'static str,
152    operation: &'static str,
153    rows: usize,
154    outcome: &'static str,
155    observer: Option<Arc<dyn QueryObserver>>,
156    threshold: Option<Duration>,
157}
158impl Drop for ObservationGuard {
159    fn drop(&mut self) {
160        let duration = self.started.elapsed();
161        let slow = self
162            .threshold
163            .is_some_and(|threshold| duration >= threshold);
164        self.span.record("rows_returned", self.rows);
165        self.span.record(
166            "duration_us",
167            u64::try_from(duration.as_micros()).unwrap_or(u64::MAX),
168        );
169        self.span.record("outcome", self.outcome);
170        tracing::info!(parent: &self.span, outcome = self.outcome, rows_returned = self.rows, slow, "SiftQuery operation completed");
171        if slow {
172            tracing::warn!(parent: &self.span, "slow SiftQuery operation");
173        }
174        if let Some(observer) = &self.observer {
175            let observation = QueryObservation {
176                entity: self.entity,
177                operation: self.operation,
178                duration,
179                rows_returned: self.rows,
180                outcome: self.outcome,
181                slow,
182            };
183            if std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
184                observer.completed(&observation);
185            }))
186            .is_err()
187            {
188                tracing::warn!(parent: &self.span, "query observer panicked");
189            }
190        }
191    }
192}
193
194pub(crate) fn error_class(error: &PgSearchError) -> &'static str {
195    match error {
196        PgSearchError::Param(_) | PgSearchError::Plan(_) => "parameter",
197        PgSearchError::Timeout(_) => "timeout",
198        PgSearchError::Sql(_) => "database",
199        PgSearchError::Decode(_) => "decode",
200        PgSearchError::Extension(_) => "extension",
201        PgSearchError::ScanLimitExceeded { .. } | PgSearchError::ScanQueueFull => "budget",
202        PgSearchError::ExportIo(_) => "output",
203        _ => "configuration",
204    }
205}
sift_query_sqlx/observe.rs

sift_query_sqlx/
observe.rs
