From 2aa4a82499d4becd2284cdb482213d541b8804dd Mon Sep 17 00:00:00 2001 From: Daniel Baumann Date: Sun, 28 Apr 2024 16:29:10 +0200 Subject: Adding upstream version 86.0.1. Signed-off-by: Daniel Baumann --- third_party/rust/sql-support/src/query_plan.rs | 186 +++++++++++++++++++++++++ 1 file changed, 186 insertions(+) create mode 100644 third_party/rust/sql-support/src/query_plan.rs (limited to 'third_party/rust/sql-support/src/query_plan.rs') diff --git a/third_party/rust/sql-support/src/query_plan.rs b/third_party/rust/sql-support/src/query_plan.rs new file mode 100644 index 0000000000..7bf6c705f5 --- /dev/null +++ b/third_party/rust/sql-support/src/query_plan.rs @@ -0,0 +1,186 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ + +use rusqlite::{types::ToSql, Connection, Result as SqlResult}; + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct QueryPlanStep { + pub node_id: i32, + pub parent_id: i32, + pub aux: i32, + pub detail: String, +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct QueryPlan { + pub query: String, + pub plan: Vec, +} + +impl QueryPlan { + // TODO: support positional params (it's a pain...) + pub fn new(conn: &Connection, sql: &str, params: &[(&str, &dyn ToSql)]) -> SqlResult { + let plan_sql = format!("EXPLAIN QUERY PLAN {}", sql); + let mut stmt = conn.prepare(&plan_sql)?; + let plan = stmt + .query_and_then_named(params, |row| -> SqlResult<_> { + Ok(QueryPlanStep { + node_id: row.get(0)?, + parent_id: row.get(1)?, + aux: row.get(2)?, + detail: row.get(3)?, + }) + })? + .collect::, _>>()?; + Ok(QueryPlan { + query: sql.into(), + plan, + }) + } + + pub fn print_pretty_tree(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if self.plan.is_empty() { + return writeln!(f, ""); + } + writeln!(f, "QUERY PLAN")?; + let children = self + .plan + .iter() + .filter(|e| e.parent_id == 0) + .collect::>(); + for (i, child) in children.iter().enumerate() { + let last = i == children.len() - 1; + self.print_tree(f, child, "", last)?; + } + Ok(()) + } + + fn print_tree( + &self, + f: &mut std::fmt::Formatter<'_>, + entry: &QueryPlanStep, + prefix: &str, + last_child: bool, + ) -> std::fmt::Result { + let children = self + .plan + .iter() + .filter(|e| e.parent_id == entry.node_id) + .collect::>(); + let next_prefix = if last_child { + writeln!(f, "{}`--{}", prefix, entry.detail)?; + format!("{} ", prefix) + } else { + writeln!(f, "{}|--{}", prefix, entry.detail)?; + format!("{}| ", prefix) + }; + for (i, child) in children.iter().enumerate() { + let last = i == children.len() - 1; + self.print_tree(f, child, &next_prefix, last)?; + } + Ok(()) + } +} + +impl std::fmt::Display for QueryPlan { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + writeln!(f, "### QUERY PLAN")?; + writeln!(f, "#### SQL:\n{}\n#### PLAN:", self.query)?; + self.print_pretty_tree(f)?; + writeln!(f, "### END QUERY PLAN") + } +} + +/// Log a query plan if the `log_query_plans` feature is enabled and it hasn't been logged yet. +#[inline] +pub fn maybe_log_plan(_conn: &Connection, _sql: &str, _params: &[(&str, &dyn ToSql)]) { + // Note: underscores ar needed becasue those go unused if the feature is not turned on. + #[cfg(feature = "log_query_plans")] + { + plan_log::log_plan(_conn, _sql, _params) + } +} + +#[cfg(feature = "log_query_plans")] +mod plan_log { + use super::*; + use std::collections::HashMap; + use std::io::Write; + use std::sync::Mutex; + + struct PlanLogger { + seen: HashMap, + out: Box, + } + + impl PlanLogger { + fn new() -> Self { + let out_file = std::env::var("QUERY_PLAN_LOG").unwrap_or_default(); + let output: Box = if out_file != "" { + let mut file = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(out_file) + .expect("QUERY_PLAN_LOG file does not exist!"); + writeln!( + file, + "\n\n# Query Plan Log starting at time: {:?}\n", + std::time::SystemTime::now() + ) + .expect("Failed to write to plan log file"); + Box::new(file) + } else { + println!("QUERY_PLAN_LOG was not set, logging to stdout"); + Box::new(std::io::stdout()) + }; + Self { + seen: Default::default(), + out: output, + } + } + + fn maybe_log(&mut self, plan: QueryPlan) { + use std::collections::hash_map::Entry; + match self.seen.entry(plan.query.clone()) { + Entry::Occupied(mut o) => { + if o.get() == &plan { + return; + } + // Ignore IO failures. + let _ = writeln!(self.out, "### QUERY PLAN CHANGED!\n{}", plan); + o.insert(plan); + } + Entry::Vacant(v) => { + let _ = writeln!(self.out, "{}", plan); + v.insert(plan); + } + } + let _ = self.out.flush(); + } + } + + lazy_static::lazy_static! { + static ref PLAN_LOGGER: Mutex = Mutex::new(PlanLogger::new()); + } + + pub fn log_plan(conn: &Connection, sql: &str, params: &[(&str, &dyn ToSql)]) { + if sql.starts_with("EXPLAIN") { + return; + } + let _ = conn.set_db_config( + rusqlite::config::DbConfig::SQLITE_DBCONFIG_TRIGGER_EQP, + true, + ); + let plan = match QueryPlan::new(conn, sql, params) { + Ok(plan) => plan, + Err(e) => { + // We're usually doing this during tests where logs often arent available + eprintln!("Failed to get query plan for {}: {}", sql, e); + return; + } + }; + let mut logger = PLAN_LOGGER.lock().unwrap(); + logger.maybe_log(plan); + } +} -- cgit v1.2.3