chore: regresql-stub

Author: radim

crates/qshape-cli/src/attribute.rs

@@ -2,31 +2,22 @@
 //! on each cluster's canonical and walking the plan tree
 
 use std::collections::HashMap;
-use std::fs;
-use std::io::{self, Read, Write};
+use std::io::{self, Write};
 use std::sync::OnceLock;
 
 use anyhow::{Context, Result, anyhow};
 use postgres::{Client, NoTls, SimpleQueryMessage};
 use qshape_core::{
-    CURRENT_SCHEMA_VERSION, ParamAttribution, cast_func_param_refs, load_clusters_doc,
-    max_param_number, normalize,
+    CURRENT_SCHEMA_VERSION, ParamAttribution, cast_func_param_refs, max_param_number, normalize,
 };
 use regex::Regex;
 use serde::Deserialize;
 
+use crate::loader::load_clusters;
 use crate::typecast::TypecastCache;
 
 pub fn run(in_path: Option<&str>, conn_str: &str, top: usize) -> Result<()> {
-    let bytes = match in_path {
-        Some(p) => fs::read(p).with_context(|| format!("read {p}"))?,
-        None => {
-            let mut buf = Vec::new();
-            io::stdin().read_to_end(&mut buf).context("read stdin")?;
-            buf
-        }
-    };
-    let mut doc = load_clusters_doc(&bytes).context("decode clusters.json")?;
+    let mut doc = load_clusters(in_path)?;
 
     let mut client = Client::connect(conn_str, NoTls).context("connect")?;
     let mut cache = TypecastCache::new(&mut client);

crates/qshape-cli/src/loader.rs

@@ -0,0 +1,20 @@
+//! Shared CLI helpers
+
+use std::fs;
+use std::io::{self, Read};
+
+use anyhow::{Context, Result};
+use qshape_core::{ClustersDoc, load_clusters_doc};
+
+/// Load and validate clusters.json from `--in <path>` or stdin
+pub fn load_clusters(in_path: Option<&str>) -> Result<ClustersDoc> {
+    let bytes = match in_path {
+        Some(p) => fs::read(p).with_context(|| format!("read {p}"))?,
+        None => {
+            let mut buf = Vec::new();
+            io::stdin().read_to_end(&mut buf).context("read stdin")?;
+            buf
+        }
+    };
+    load_clusters_doc(&bytes).context("decode clusters.json")
+}

crates/qshape-cli/src/main.rs

@@ -5,6 +5,8 @@ use clap::{Parser, Subcommand};
 
 mod attribute;
 mod capture;
+mod loader;
+mod regresql_stub;
 mod typecast;
 
 #[derive(Debug, Parser)]
@@ -53,6 +55,26 @@ enum Command {
         top: usize,
     },
 
+    /// Generate regresql sql/ skeletons from clusters.json
+    ///
+    /// For each of the top N clusters, writes <out>/sql/<slug>.sql with
+    /// the canonical SQL ($N rewritten to :paramN) and a header carrying
+    /// fingerprint + call count
+    RegresqlStub {
+        /// Input clusters.json (default: stdin)
+        #[arg(long)]
+        r#in: Option<String>,
+        /// Output directory
+        #[arg(long, default_value = "regresql-stubs")]
+        out: String,
+        /// Number of top clusters to emit
+        #[arg(long, default_value_t = 10)]
+        top: usize,
+        /// Skip clusters with total_calls <= this
+        #[arg(long, default_value_t = 0)]
+        min_calls: i64,
+    },
+
     /// Fetch pg_stat_statements (with timing) from a live PG and cluster it
     Capture {
         /// libpq connection string.
@@ -98,6 +120,9 @@ fn main() -> Result<()> {
         Command::Attribute { conn, r#in, top } => {
             attribute::run(r#in.as_deref(), &conn, top)?;
         }
+        Command::RegresqlStub { r#in, out, top, min_calls } => {
+            regresql_stub::run(r#in.as_deref(), &out, top, min_calls)?;
+        }
     }
 
     Ok(())

crates/qshape-cli/src/regresql_stub.rs

@@ -0,0 +1,64 @@
+//! Generate regresql `sql/` stubs for the top-N clusters
+
+use std::fs;
+use std::path::Path;
+use std::sync::OnceLock;
+
+use anyhow::{Context, Result};
+use qshape_core::Cluster;
+use regex::Regex;
+
+use crate::loader::load_clusters;
+
+pub fn run(in_path: Option<&str>, out_dir: &str, top: usize, min_calls: i64) -> Result<()> {
+    let doc = load_clusters(in_path)?;
+
+    let sql_dir = Path::new(out_dir).join("sql");
+    fs::create_dir_all(&sql_dir).with_context(|| format!("create {}", sql_dir.display()))?;
+
+    let mut emitted = 0usize;
+    for c in &doc.clusters {
+        if emitted >= top {
+            break;
+        }
+        if c.fingerprint.is_empty() || c.total_calls <= min_calls {
+            continue;
+        }
+        emitted += 1;
+        let slug = stub_slug(emitted, &c.fingerprint);
+        let sql = rewrite_params(&c.canonical);
+        let path = sql_dir.join(format!("{slug}.sql"));
+        write_sql_stub(&path, &slug, c, &sql)?;
+    }
+
+    eprintln!("wrote {emitted} stubs to {out_dir}");
+    Ok(())
+}
+
+fn stub_slug(rank: usize, fp: &str) -> String {
+    let body = fp.strip_prefix("sha1:").unwrap_or(fp);
+    let prefix: String = body.chars().take(8).collect();
+    format!("q{rank:02}-{prefix}")
+}
+
+// $N to:paramN
+fn rewrite_params(sql: &str) -> String {
+    static R: OnceLock<Regex> = OnceLock::new();
+    let re = R.get_or_init(|| Regex::new(r"\$(\d+)").expect("static regex"));
+    re.replace_all(sql, ":param$1").into_owned()
+}
+
+fn write_sql_stub(path: &Path, slug: &str, c: &Cluster, sql: &str) -> Result<()> {
+    let trailing = if sql.ends_with('\n') { "" } else { "\n" };
+    let content = format!(
+        "-- name: {slug}\n\
+         -- Generated from qshape cluster {fp}\n\
+         -- Total calls (prod): {tc} across {n} member variants\n\
+         -- TODO: rename this slug, review canonical SQL, replace :paramN with meaningful names\n\
+         {sql}{trailing}",
+        fp = c.fingerprint,
+        tc = c.total_calls,
+        n = c.members.len(),
+    );
+    fs::write(path, content).with_context(|| format!("write {}", path.display()))
+}