Adding upstream version 86.0.1.upstream/86.0.1upstream

Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>

1 files changed, 186 insertions, 0 deletions

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<QueryPlanStep>,
+}
+
+impl QueryPlan {
+    // TODO: support positional params (it's a pain...)
+    pub fn new(conn: &Connection, sql: &str, params: &[(&str, &dyn ToSql)]) -> SqlResult<Self> {
+        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::<Result<Vec<QueryPlanStep>, _>>()?;
+        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, "<no query plan>");
+        }
+        writeln!(f, "QUERY PLAN")?;
+        let children = self
+            .plan
+            .iter()
+            .filter(|e| e.parent_id == 0)
+            .collect::<Vec<_>>();
+        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::<Vec<_>>();
+        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<String, QueryPlan>,
+        out: Box<dyn Write + Send>,
+    }
+
+    impl PlanLogger {
+        fn new() -> Self {
+            let out_file = std::env::var("QUERY_PLAN_LOG").unwrap_or_default();
+            let output: Box<dyn Write + Send> = 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<PlanLogger> = 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);
+    }
+}