sift_query_sqlx/
diagnostics.rs

1use std::sync::Arc;
2
3use sift_query_core::{
4    PredicateArena, PredicateNode, QueryInput, QueryTarget, SearchEntity, ServerScope,
5};
6use sqlx::{AssertSqlSafe, Row};
7use tokio::time::{Instant, timeout_at};
8
9use crate::cache::RenderRequest;
10use crate::execution::io::{configure_statement_timeout, finish_statement, prepare};
11use crate::{PgSearchError, PgSiftQuery, QueryShapeId, TimeoutStage};
12
13/// A conservative, server-side index hint. It is advisory only and never
14/// creates or changes database indexes.
15#[derive(Debug, Clone, PartialEq, Eq)]
16pub struct IndexAdvice {
17    pub field: &'static str,
18    pub reason: &'static str,
19}
20
21#[derive(Debug, Clone, PartialEq, Eq)]
22pub struct DiagnosticStage {
23    pub phase: &'static str,
24    pub shape_id: QueryShapeId,
25    pub sql: Arc<str>,
26    pub explain_json: String,
27}
28
29#[derive(Debug, Clone, PartialEq, Eq)]
30pub struct QueryDiagnosticReport {
31    pub entity: &'static str,
32    pub stages: Vec<DiagnosticStage>,
33    pub index_advice: Vec<IndexAdvice>,
34}
35
36impl PgSiftQuery {
37    /// Returns parameterized SQL, PostgreSQL JSON plans, and conservative
38    /// index hints for an already authorized query. Callers should expose this
39    /// admin-only operation behind their own authorization boundary.
40    pub async fn diagnose<E: SearchEntity>(
41        &self,
42        query: impl QueryInput<E>,
43        scope: &ServerScope<E>,
44    ) -> Result<QueryDiagnosticReport, PgSearchError> {
45        let plan = self.resolve_plan(query, scope, QueryTarget::MapPage, "diagnose")?;
46        if plan.force_empty {
47            return Ok(QueryDiagnosticReport {
48                entity: E::META.entity_name,
49                stages: Vec::new(),
50                index_advice: Vec::new(),
51            });
52        }
53        let list = self.render::<E>(&plan, "diagnose", "list", None, RenderRequest::List)?;
54        let cluster = self.render::<E>(
55            &plan,
56            "diagnose",
57            "cluster",
58            None,
59            RenderRequest::Cluster(None),
60        )?;
61        let deadline = self.deadline::<E>();
62        let list = explain_stage(self, "list", list, deadline).await?;
63        let cluster = explain_stage(self, "cluster", cluster, deadline).await?;
64        Ok(QueryDiagnosticReport {
65            entity: E::META.entity_name,
66            stages: vec![list, cluster],
67            index_advice: index_advice(E::META, &plan.row_predicate, &plan.order),
68        })
69    }
70}
71
72async fn explain_stage(
73    search: &PgSiftQuery,
74    phase: &'static str,
75    rendered: crate::RenderedPgQuery,
76    deadline: Instant,
77) -> Result<DiagnosticStage, PgSearchError> {
78    let shape_id = rendered.shape_id;
79    let sql = Arc::clone(&rendered.sql);
80    let prepared = prepare(rendered)?;
81    let explain_sql: Arc<str> = format!("EXPLAIN (FORMAT JSON) {}", prepared.sql).into();
82    let connection = search.acquire(deadline).await?;
83    let mut connection =
84        configure_statement_timeout(connection, deadline, TimeoutStage::SchemaValidation).await?;
85    let result = timeout_at(
86        deadline,
87        sqlx::query_with(AssertSqlSafe(explain_sql), prepared.arguments)
88            .fetch_one(&mut *connection),
89    )
90    .await;
91    let Ok(result) = result else {
92        crate::execution::io::close_hard(connection).await;
93        return Err(PgSearchError::Timeout(TimeoutStage::SchemaValidation));
94    };
95    let (mut connection, row) =
96        finish_statement(connection, result, TimeoutStage::SchemaValidation).await?;
97    let explain_json = row
98        .try_get::<sqlx::types::Json<serde_json::Value>, _>(0)
99        .map(|value| value.0.to_string())?;
100    let _ = &mut connection;
101    Ok(DiagnosticStage {
102        phase,
103        shape_id,
104        sql,
105        explain_json,
106    })
107}
108
109fn index_advice(
110    meta: &'static sift_query_core::EntityMeta,
111    predicate: &PredicateArena,
112    order: &[sift_query_core::OrderBy],
113) -> Vec<IndexAdvice> {
114    let mut advice = Vec::new();
115    for node in predicate.nodes() {
116        let PredicateNode::Condition(condition) = node else {
117            continue;
118        };
119        let Some(field) = meta.field(condition.field) else {
120            continue;
121        };
122        if matches!(
123            condition.operator,
124            sift_query_core::Operator::Eq | sift_query_core::Operator::InList
125        ) {
126            advice.push(IndexAdvice {
127                field: field.public_name,
128                reason: "equality or IN filter; consider an index on this field",
129            });
130        }
131    }
132    for item in order {
133        if let Some(field) = meta.field(item.field)
134            && !advice.iter().any(|hint| hint.field == field.public_name)
135        {
136            advice.push(IndexAdvice {
137                field: field.public_name,
138                reason: "ordering field; consider including it in a composite index",
139            });
140        }
141    }
142    advice
143}
sift_query_sqlx/diagnostics.rs

sift_query_sqlx/
diagnostics.rs
