Merge pull request #169 from cipherstash/index-performance

feat(test): add benchmarking helpers and pg_stat_statements support

Author: tobyhede

.github/workflows/test-eql.yml

@@ -43,7 +43,7 @@ jobs:
 
       - uses: jdx/mise-action@v3
         with:
-          version: 2025.11.2 # [default: latest] mise version to install
+          version: 2026.4.0
           install: true # [default: true] run `mise install`
           cache: true # [default: true] cache mise using GitHub's cache
 

mise.toml

@@ -11,7 +11,7 @@
 "rust" = { version = "latest", components = "rustc,rust-std,cargo,rustfmt,rust-docs,clippy" }
 "cargo:cargo-binstall" = "latest"
 "cargo:sqlx-cli" = "latest"
-"python" = "latest"
+"python" = "3.13"
 
 [task_config]
 includes = ["tasks", "tasks/postgres.toml"]

tasks/test.sh

@@ -20,7 +20,7 @@ echo ""
 
 # Build first
 echo "Building EQL..."
-mise run --output prefix build --force
+mise run --output prefix --force build
 
 # Run lints on sqlx tests
 echo ""

tests/docker-compose.yml

@@ -2,7 +2,7 @@ services:
   postgres: &postgres
     container_name: postgres
     image: postgres:17
-    command: postgres -c track_functions=all
+    command: postgres -c track_functions=all -c shared_preload_libraries=pg_stat_statements -c pg_stat_statements.track=all
     ports:
       - 7432:7432
     environment:

tests/sqlx/src/helpers.rs

@@ -384,7 +384,9 @@ pub async fn analyze_table(pool: &PgPool, table: &str) -> Result<()> {
 ///
 /// # Arguments
 /// * `pool` - Database connection pool
-/// * `query` - SQL query to explain (without EXPLAIN prefix)
+/// * `query` - SQL query to explain (without EXPLAIN prefix).
+///   Must be a trusted/hardcoded string — not user-supplied input,
+///   as it is interpolated directly into the SQL statement.
 ///
 /// # Returns
 /// The EXPLAIN output as a newline-separated string
@@ -448,3 +450,319 @@ pub fn assert_uses_seq_scan(explain_output: &str) {
         explain_output
     );
 }
+
+// ============================================================================
+// Benchmarking / EXPLAIN Helpers
+// ============================================================================
+
+/// Statistics extracted from EXPLAIN ANALYZE JSON output
+///
+/// Contains timing and plan information for benchmarking queries.
+/// Used by `explain_analyze_avg` to return averaged statistics.
+#[derive(Debug, Clone)]
+pub struct ExplainStats {
+    /// Average execution time in milliseconds across runs
+    pub execution_time_ms: f64,
+    /// Average planning time in milliseconds across runs
+    pub planning_time_ms: f64,
+    /// Top-level node type from the query plan (e.g., "Index Scan", "Seq Scan")
+    pub node_type: String,
+}
+
+/// Run EXPLAIN with JSON format on a query and return the parsed plan
+///
+/// Executes `EXPLAIN (FORMAT JSON) {query}` and parses the result.
+/// PostgreSQL returns a single-element JSON array containing the plan tree.
+///
+/// This is distinct from `explain_query()` which returns plain text output.
+/// The JSON format provides structured access to plan nodes, costs, and types.
+///
+/// # Arguments
+/// * `pool` - Database connection pool
+/// * `query` - SQL query to explain (without EXPLAIN prefix).
+///   Must be a trusted/hardcoded string — not user-supplied input,
+///   as it is interpolated directly into the SQL statement.
+///
+/// # Returns
+/// The full EXPLAIN JSON output as a `serde_json::Value`
+///
+/// # Example
+/// ```ignore
+/// let plan = explain_json(&pool, "SELECT * FROM foo WHERE x = 1").await?;
+/// let node_type = plan[0]["Plan"]["Node Type"].as_str().unwrap();
+/// ```
+pub async fn explain_json(pool: &PgPool, query: &str) -> Result<serde_json::Value> {
+    let sql = format!("EXPLAIN (FORMAT JSON) {}", query);
+    let plan: serde_json::Value = sqlx::query_scalar(&sql)
+        .fetch_one(pool)
+        .await
+        .with_context(|| format!("running EXPLAIN (FORMAT JSON) on query: {}", query))?;
+
+    Ok(plan)
+}
+
+/// Run EXPLAIN ANALYZE multiple times and return averaged statistics
+///
+/// Executes `EXPLAIN (ANALYZE, FORMAT JSON) {query}` the specified number of times
+/// and returns the arithmetic mean of execution and planning times.
+///
+/// **Warning**: EXPLAIN ANALYZE actually executes the query. If the query has
+/// side effects (INSERT, UPDATE, DELETE), those effects will occur on every run.
+///
+/// **Note**: The first run may include cold-start overhead (buffer cache misses,
+/// plan cache population). No runs are discarded — callers should account for this
+/// when setting thresholds or increase the run count to dilute the effect.
+///
+/// # Arguments
+/// * `pool` - Database connection pool
+/// * `query` - SQL query to explain and execute (without EXPLAIN prefix).
+///   Must be a trusted/hardcoded string — not user-supplied input,
+///   as it is interpolated directly into the SQL statement.
+/// * `runs` - Number of times to execute (must be >= 1)
+///
+/// # Returns
+/// Averaged `ExplainStats` with mean execution_time_ms, mean planning_time_ms,
+/// and the node_type from the first run's top-level plan node
+///
+/// # Example
+/// ```ignore
+/// let stats = explain_analyze_avg(&pool, "SELECT * FROM foo WHERE x = 1", 5).await?;
+/// assert!(stats.execution_time_ms < 10.0, "Query too slow: {}ms", stats.execution_time_ms);
+/// assert_eq!(stats.node_type, "Index Scan");
+/// ```
+pub async fn explain_analyze_avg(pool: &PgPool, query: &str, runs: usize) -> Result<ExplainStats> {
+    anyhow::ensure!(runs >= 1, "runs must be >= 1, got {}", runs);
+
+    let sql = format!("EXPLAIN (ANALYZE, FORMAT JSON) {}", query);
+
+    let mut total_execution_ms = 0.0_f64;
+    let mut total_planning_ms = 0.0_f64;
+    let mut node_type = String::new();
+
+    for i in 0..runs {
+        let plan: serde_json::Value = sqlx::query_scalar(&sql)
+            .fetch_one(pool)
+            .await
+            .with_context(|| {
+                format!(
+                    "running EXPLAIN ANALYZE (run {}/{}) on query: {}",
+                    i + 1,
+                    runs,
+                    query
+                )
+            })?;
+
+        // EXPLAIN (ANALYZE, FORMAT JSON) returns:
+        // [{"Plan": {...}, "Planning Time": N, "Execution Time": N}]
+        let entry = &plan[0];
+
+        let exec_time = entry["Execution Time"]
+            .as_f64()
+            .with_context(|| format!("extracting Execution Time on run {}/{}", i + 1, runs))?;
+
+        let plan_time = entry["Planning Time"]
+            .as_f64()
+            .with_context(|| format!("extracting Planning Time on run {}/{}", i + 1, runs))?;
+
+        total_execution_ms += exec_time;
+        total_planning_ms += plan_time;
+
+        // Capture node type from first run only
+        if i == 0 {
+            node_type = entry["Plan"]["Node Type"]
+                .as_str()
+                .with_context(|| "extracting Node Type from first run")?
+                .to_string();
+        }
+    }
+
+    let n = runs as f64;
+    Ok(ExplainStats {
+        execution_time_ms: total_execution_ms / n,
+        planning_time_ms: total_planning_ms / n,
+        node_type,
+    })
+}
+
+/// Assert that a JSON EXPLAIN plan does not use any sequential scan
+///
+/// Recursively walks the JSON plan tree checking all "Node Type" fields.
+/// A plan can have nested nodes (e.g., Aggregate -> Seq Scan), so all levels
+/// are checked. Both "Seq Scan" and "Parallel Seq Scan" are rejected.
+///
+/// This is the structured (JSON) counterpart to `assert_uses_seq_scan()` which
+/// operates on plain text output.
+///
+/// # Arguments
+/// * `plan` - JSON EXPLAIN output from `explain_json()` or `EXPLAIN (FORMAT JSON)`
+///
+/// # Panics
+/// Panics if any node in the plan tree has a "Seq Scan" or "Parallel Seq Scan" node type
+///
+/// # Example
+/// ```ignore
+/// let plan = explain_json(&pool, "SELECT * FROM foo WHERE x = 1").await?;
+/// assert_no_seq_scan(&plan);
+/// ```
+pub fn assert_no_seq_scan(plan: &serde_json::Value) {
+    let mut seq_scan_nodes = Vec::new();
+    collect_seq_scan_nodes(plan, &mut seq_scan_nodes);
+
+    assert!(
+        seq_scan_nodes.is_empty(),
+        "Expected no sequential scans but found {} node(s): {:?}\nFull plan: {}",
+        seq_scan_nodes.len(),
+        seq_scan_nodes,
+        serde_json::to_string_pretty(plan).unwrap_or_else(|_| plan.to_string())
+    );
+}
+
+/// Recursively collect all sequential scan node types from a JSON EXPLAIN plan
+///
+/// Checks standard PostgreSQL node types only ("Seq Scan", "Parallel Seq Scan").
+/// Custom scan providers (e.g., from extensions) are not currently detected.
+fn collect_seq_scan_nodes(value: &serde_json::Value, found: &mut Vec<String>) {
+    match value {
+        serde_json::Value::Object(map) => {
+            if let Some(node_type) = map.get("Node Type").and_then(|v| v.as_str()) {
+                if node_type == "Seq Scan" || node_type == "Parallel Seq Scan" {
+                    let relation = map
+                        .get("Relation Name")
+                        .and_then(|v| v.as_str())
+                        .unwrap_or("unknown");
+                    found.push(format!("{} on {}", node_type, relation));
+                }
+            }
+            for v in map.values() {
+                collect_seq_scan_nodes(v, found);
+            }
+        }
+        serde_json::Value::Array(arr) => {
+            for item in arr {
+                collect_seq_scan_nodes(item, found);
+            }
+        }
+        _ => {}
+    }
+}
+
+// ============================================================================
+// pg_stat_statements Helpers (Tier 2)
+// ============================================================================
+
+/// Statistics from pg_stat_statements for a matched query
+///
+/// Contains key performance metrics from the pg_stat_statements view.
+/// See PostgreSQL documentation for pg_stat_statements column definitions.
+#[derive(Debug, Clone)]
+pub struct PgStatEntry {
+    /// Number of times the query was executed
+    pub calls: i64,
+    /// Mean execution time in milliseconds
+    pub mean_exec_time: f64,
+    /// Population standard deviation of execution time in milliseconds
+    pub stddev_exec_time: f64,
+    /// Total execution time in milliseconds across all calls
+    pub total_exec_time: f64,
+    /// The normalized query string from pg_stat_statements
+    pub query: String,
+}
+
+/// Ensure pg_stat_statements extension is available
+///
+/// Creates the extension if it doesn't exist. Should be called once
+/// at the start of benchmark tests that need pg_stat_statements.
+///
+/// Requires `shared_preload_libraries=pg_stat_statements` in the PostgreSQL
+/// server configuration (see docker-compose.yml).
+pub async fn ensure_pg_stat_statements(pool: &PgPool) -> Result<()> {
+    sqlx::query("CREATE EXTENSION IF NOT EXISTS pg_stat_statements")
+        .execute(pool)
+        .await
+        .with_context(|| "creating pg_stat_statements extension")?;
+    Ok(())
+}
+
+/// Reset all pg_stat_statements counters
+///
+/// Clears cumulative per-query statistics so the next sampling window starts
+/// from zero. Call this before the measurement phase of a benchmark case to
+/// ensure `read_pg_stat_statements` reflects only the queries executed after
+/// the reset — not leftovers from prior cases or setup work.
+///
+/// Requires the `pg_stat_statements` extension to be loaded
+/// (see `ensure_pg_stat_statements`).
+///
+/// # Example
+/// ```ignore
+/// ensure_pg_stat_statements(&pool).await?;
+/// reset_pg_stat_statements(&pool).await?;
+/// // ... run benchmark queries ...
+/// let stats = read_pg_stat_statements(&pool, "%FROM bench%").await?;
+/// ```
+pub async fn reset_pg_stat_statements(pool: &PgPool) -> Result<()> {
+    sqlx::query("SELECT pg_stat_statements_reset(NULL::oid, NULL::oid, (SELECT oid FROM pg_database WHERE datname = current_database()))")
+        .execute(pool)
+        .await
+        .with_context(|| "resetting pg_stat_statements counters for current database")?;
+    Ok(())
+}
+
+/// Read query statistics from pg_stat_statements
+///
+/// Looks up a query in the `pg_stat_statements` view using a SQL LIKE pattern.
+/// Requires the `pg_stat_statements` extension to be loaded
+/// (see `ensure_pg_stat_statements`).
+///
+/// # Arguments
+/// * `pool` - Database connection pool
+/// * `query_pattern` - SQL LIKE pattern to match against normalized query text
+///   (e.g., `"%FROM ore WHERE%"`).
+///   Note: `pg_stat_statements` normalizes queries by replacing literal values
+///   with `$N` placeholders. Patterns must match the normalized form
+///   (e.g., `"%FROM bench WHERE e = $1%"`, not `"%FROM bench WHERE e = 'abc'%"`).
+///
+/// # Returns
+/// `PgStatEntry` for the matched query. Returns error if no match or multiple matches.
+///
+/// # Example
+/// ```ignore
+/// ensure_pg_stat_statements(&pool).await?;
+/// let stats = read_pg_stat_statements(&pool, "%FROM ore WHERE%").await?;
+/// assert!(stats.mean_exec_time < 5.0, "Query regression: {}ms", stats.mean_exec_time);
+/// ```
+pub async fn read_pg_stat_statements(pool: &PgPool, query_pattern: &str) -> Result<PgStatEntry> {
+    let sql = "SELECT query, calls, mean_exec_time, stddev_exec_time, total_exec_time \
+               FROM pg_stat_statements \
+               WHERE query LIKE $1 \
+               AND dbid = (SELECT oid FROM pg_database WHERE datname = current_database())";
+
+    let rows: Vec<(String, i64, f64, f64, f64)> = sqlx::query_as(sql)
+        .bind(query_pattern)
+        .fetch_all(pool)
+        .await
+        .with_context(|| format!("reading pg_stat_statements for pattern: {}", query_pattern))?;
+
+    match rows.len() {
+        0 => Err(anyhow::anyhow!(
+            "No pg_stat_statements entry found matching pattern: {}",
+            query_pattern
+        )),
+        1 => {
+            let (query, calls, mean_exec_time, stddev_exec_time, total_exec_time) =
+                rows.into_iter().next().unwrap();
+            Ok(PgStatEntry {
+                calls,
+                mean_exec_time,
+                stddev_exec_time,
+                total_exec_time,
+                query,
+            })
+        }
+        n => Err(anyhow::anyhow!(
+            "Expected 1 pg_stat_statements entry but found {} matching pattern: {}",
+            n,
+            query_pattern
+        )),
+    }
+}

tests/sqlx/src/lib.rs

@@ -11,11 +11,13 @@ pub mod selectors;
 
 pub use assertions::QueryAssertion;
 pub use helpers::{
-    analyze_table, assert_sequential_ids, assert_uses_index, assert_uses_seq_scan,
-    create_jsonb_gin_index, explain_query, get_encrypted_term, get_ore_encrypted,
-    get_ore_encrypted_as_jsonb, get_ore_text_encrypted, get_ore_text_encrypted_as_jsonb,
-    get_ste_vec_encrypted, get_ste_vec_encrypted_pair, get_ste_vec_selector_term,
-    get_ste_vec_sv_element, get_ste_vec_term_by_id,
+    analyze_table, assert_no_seq_scan, assert_sequential_ids, assert_uses_index,
+    assert_uses_seq_scan, create_jsonb_gin_index, ensure_pg_stat_statements, explain_analyze_avg,
+    explain_json, explain_query, get_encrypted_term, get_ore_encrypted, get_ore_encrypted_as_jsonb,
+    get_ore_text_encrypted, get_ore_text_encrypted_as_jsonb, get_ste_vec_encrypted,
+    get_ste_vec_encrypted_pair, get_ste_vec_selector_term, get_ste_vec_sv_element,
+    get_ste_vec_term_by_id, read_pg_stat_statements, reset_pg_stat_statements, ExplainStats,
+    PgStatEntry,
 };
 pub use index_types as IndexTypes;
 pub use selectors::Selectors;