1//! Writes a Cargo workspace for the first-party crates from the buck
2//! graph, so `cargo pgrx install` works without buck while buck stays the
3//! source of truth.
4//!
5//! Every `first_party_*` target records its kind, crate root and deps as a
6//! `<name>-cargo` JSON file (`build/defs.bzl`). From those, the version,
7//! edition and supported majors in `build/defs.bzl`, and the version
8//! requirements in `third-party/Cargo.toml`, this tool writes the root
9//! `Cargo.toml`, one manifest per package, and the extension's
10//! `pgrx_embed` binary. Integration tests (a package with no library or
11//! binary, such as `tests/`) stay buck-only: they need buck's `$(location)`
12//! of the built extension.
13//!
14//! Usage:
15//!   cargo-gen generate <out-dir> <version> <edition> <majors> <third-party/Cargo.toml> <meta.json>...
16//!   cargo-gen write <generated-dir>    (copies into the checkout holding the cwd)
17
18use std::collections::BTreeMap;
19use std::error::Error;
20use std::fmt::Write as _;
21use std::path::{Path, PathBuf};
22
23use serde::Deserialize;
24
25type Res<T> = Result<T, Box<dyn Error>>;
26
27const LICENSE: &str = "MIT OR Apache-2.0";
28const AUTHORS: &str = "The postjevsql Authors";
29const THIRD_PARTY: &str = "//third-party:";
30
31const HEADER: &str = "# @generated by tools/cargo-gen from the BUCK files, build/defs.bzl and
32# third-party/Cargo.toml. Do not edit: change the BUCK file and run
33#   buck2 run //tools/cargo-gen:update
34# The drift test (//tools/cargo-gen:drift) fails until you do.
35";
36
37#[derive(Deserialize)]
38struct Meta {
39    package: String,
40    name: String,
41    kind: Kind,
42    crate_root: String,
43    deps: Vec<String>,
44    pgrx_extension: bool,
45}
46
47#[derive(Deserialize, PartialEq, Clone, Copy)]
48#[serde(rename_all = "lowercase")]
49enum Kind {
50    Library,
51    Binary,
52    Test,
53}
54
55/// One Cargo package: the library or binary target of a buck package,
56/// with the deps of the unit tests that share its crate root.
57struct Package<'a> {
58    main: &'a Meta,
59    /// Binaries beside a library, as `[[bin]]`s that link it.
60    bins: Vec<&'a Meta>,
61    dev_deps: Vec<&'a str>,
62}
63
64fn main() {
65    let args: Vec<String> = std::env::args().skip(1).collect();
66    let result = match args.first().map(String::as_str) {
67        Some("generate") if args.len() >= 6 => generate(&args[1..]),
68        Some("write") if args.len() == 2 => write(Path::new(&args[1])),
69        _ => Err("usage: cargo-gen generate <out> <version> <edition> <majors> <third-party/Cargo.toml> <meta.json>... | cargo-gen write <generated-dir>".into()),
70    };
71    if let Err(e) = result {
72        eprintln!("cargo-gen: {e}");
73        std::process::exit(1);
74    }
75}
76
77fn generate(args: &[String]) -> Res<()> {
78    let out = PathBuf::from(&args[0]);
79    let version = &args[1];
80    let edition = &args[2];
81    let majors: Vec<&str> = args[3].split(',').filter(|m| !m.is_empty()).collect();
82    let Some(default_major) = majors.last() else {
83        return Err("no supported PostgreSQL majors".into());
84    };
85    let third_party: toml::Table = toml::from_str(&std::fs::read_to_string(&args[4])?)?;
86    let third_party = third_party
87        .get("dependencies")
88        .and_then(toml::Value::as_table)
89        .ok_or("third-party/Cargo.toml has no [dependencies]")?;
90
91    let metas = args[5..]
92        .iter()
93        .map(|p| Ok(serde_json::from_str::<Meta>(&std::fs::read_to_string(p)?)?))
94        .collect::<Res<Vec<Meta>>>()?;
95    let packages = packages(&metas)?;
96
97    // `//pkg:name` of each package's main target -> its Cargo name.
98    let first_party: BTreeMap<String, &str> = packages
99        .values()
100        .map(|p| (format!("//{}:{}", p.main.package, p.main.name), p.main.name.as_str()))
101        .collect();
102    let crate_dep = |label: &str| -> Res<Dep> {
103        if let Some(name) = label.strip_prefix(THIRD_PARTY) {
104            if !third_party.contains_key(name) {
105                return Err(format!("{label} is not in third-party/Cargo.toml").into());
106            }
107            Ok(Dep::Third(name.to_owned()))
108        } else if let Some(name) = first_party.get(label) {
109            Ok(Dep::First((*name).to_owned()))
110        } else {
111            Err(format!("{label} is neither third-party nor a first-party library").into())
112        }
113    };
114
115    // A package needs a pgNN feature when it reaches pgrx at all: a
116    // fixpoint over the (acyclic) first-party graph.
117    let mut uses_pgrx: BTreeMap<&str, bool> =
118        packages.values().map(|p| (p.main.name.as_str(), false)).collect();
119    loop {
120        let mut changed = false;
121        for p in packages.values() {
122            let reaches = p.main.deps.iter().any(|d| {
123                d == "//third-party:pgrx" || first_party.get(d.as_str()).is_some_and(|n| uses_pgrx[n])
124            });
125            if reaches && !uses_pgrx[p.main.name.as_str()] {
126                uses_pgrx.insert(p.main.name.as_str(), true);
127                changed = true;
128            }
129        }
130        if !changed {
131            break;
132        }
133    }
134
135    let mut used_third_party = std::collections::BTreeSet::new();
136    for (dir, pkg) in &packages {
137        let mut s = String::from(HEADER);
138        let name = &pkg.main.name;
139        writeln!(s, "\n[package]\nname = {}", q(name))?;
140        for key in ["version", "edition", "license", "authors", "publish"] {
141            writeln!(s, "{key}.workspace = true")?;
142        }
143        let root = &pkg.main.crate_root;
144        match pkg.main.kind {
145            Kind::Library if pkg.main.pgrx_extension => {
146                writeln!(s, "\n[lib]\npath = {}\ncrate-type = [\"cdylib\", \"lib\"]", q(root))?;
147                writeln!(
148                    s,
149                    "\n[[bin]]\nname = {}\npath = \"src/bin/pgrx_embed.rs\"",
150                    q(&format!("pgrx_embed_{}", name.replace('-', "_")))
151                )?;
152                let embed = out.join(dir).join("src/bin/pgrx_embed.rs");
153                std::fs::create_dir_all(embed.parent().unwrap())?;
154                std::fs::write(
155                    &embed,
156                    "// @generated by tools/cargo-gen: the schema binary `cargo pgrx` runs.\n::pgrx::pgrx_embed!();\n",
157                )?;
158            }
159            Kind::Library => writeln!(s, "\n[lib]\npath = {}", q(root))?,
160            Kind::Binary => writeln!(s, "\n[[bin]]\nname = {}\npath = {}", q(name), q(root))?,
161            Kind::Test => unreachable!("packages() never picks a test"),
162        }
163        for bin in &pkg.bins {
164            // One binary takes the package's name, as src/main.rs would.
165            let bin_name = if pkg.bins.len() == 1 { name } else { &bin.name };
166            writeln!(s, "\n[[bin]]\nname = {}\npath = {}", q(bin_name), q(&bin.crate_root))?;
167        }
168
169        // A binary's deps join the package's, less the library it links.
170        let own = format!("//{}:{}", pkg.main.package, pkg.main.name);
171        let mut labels: Vec<&str> = pkg.main.deps.iter().map(String::as_str).collect();
172        for d in pkg.bins.iter().flat_map(|b| &b.deps) {
173            if *d != own && !labels.contains(&d.as_str()) {
174                labels.push(d);
175            }
176        }
177        let deps = labels.iter().map(|d| crate_dep(d)).collect::<Res<Vec<Dep>>>()?;
178        if uses_pgrx[name.as_str()] {
179            writeln!(s, "\n[features]")?;
180            if pkg.main.pgrx_extension {
181                writeln!(s, "default = [{}]\npg_test = []", q(&format!("pg{default_major}")))?;
182            }
183            for major in &majors {
184                let fwd: Vec<String> = deps
185                    .iter()
186                    .filter(|d| match d {
187                        Dep::Third(n) => n == "pgrx",
188                        Dep::First(n) => uses_pgrx[n.as_str()],
189                    })
190                    .map(|d| q(&format!("{}/pg{major}", d.name())))
191                    .collect();
192                writeln!(s, "pg{major} = [{}]", fwd.join(", "))?;
193            }
194        }
195        write_deps(&mut s, "dependencies", &deps, &mut used_third_party)?;
196        let dev = pkg
197            .dev_deps
198            .iter()
199            .filter(|d| !labels.contains(d))
200            .map(|d| crate_dep(d))
201            .collect::<Res<Vec<Dep>>>()?;
202        write_deps(&mut s, "dev-dependencies", &dev, &mut used_third_party)?;
203        write_file(&out.join(dir).join("Cargo.toml"), &s)?;
204    }
205
206    // The workspace root.
207    let mut s = String::from(HEADER);
208    writeln!(s, "\n[workspace]\nresolver = \"3\"\nmembers = [")?;
209    for dir in packages.keys() {
210        writeln!(s, "    {},", q(dir))?;
211    }
212    writeln!(s, "]\n# buck's crates.io resolution, not a member.\nexclude = [\"third-party\"]")?;
213    writeln!(
214        s,
215        "\n[workspace.package]\nversion = {}\nedition = {}\nlicense = {}\nauthors = [{}]\npublish = false",
216        q(version),
217        q(edition),
218        q(LICENSE),
219        q(AUTHORS)
220    )?;
221    writeln!(s, "\n[workspace.dependencies]")?;
222    for (dir, pkg) in &packages {
223        writeln!(s, "{} = {{ path = {}, version = {} }}", pkg.main.name, q(dir), q(&format!("={version}")))?;
224    }
225    for name in &used_third_party {
226        let mut value = third_party[name.as_str()].clone();
227        if name == "pgrx" {
228            pin_pgrx(&mut value)?;
229        }
230        // A path in third-party/Cargo.toml (the jevcrates submodule) is
231        // relative to third-party/; the root manifest is one level up.
232        if let Some(path) = value.get_mut("path")
233            && let toml::Value::String(p) = path
234        {
235            *p = format!("third-party/{p}");
236        }
237        writeln!(s, "{name} = {value}")?;
238    }
239    // pgrx's template: a Postgres ERROR unwinds through Rust frames.
240    writeln!(
241        s,
242        "\n[profile.dev]\npanic = \"unwind\"\n\n[profile.release]\npanic = \"unwind\"\nopt-level = 3\nlto = \"fat\"\ncodegen-units = 1"
243    )?;
244    write_file(&out.join("Cargo.toml"), &s)
245}
246
247/// The buck packages with a library or binary, keyed by directory.
248fn packages(metas: &[Meta]) -> Res<BTreeMap<String, Package<'_>>> {
249    let mut out: BTreeMap<String, Package> = BTreeMap::new();
250    // A library is the package's main target, and binaries beside it
251    // become its `[[bin]]`s; otherwise a package has one target.
252    let mut metas_sorted: Vec<&Meta> = metas.iter().filter(|m| m.kind != Kind::Test).collect();
253    metas_sorted.sort_by_key(|m| m.kind != Kind::Library);
254    for m in metas_sorted {
255        match out.get_mut(&m.package) {
256            None => {
257                out.insert(m.package.clone(), Package { main: m, bins: vec![], dev_deps: vec![] });
258            }
259            Some(p) if p.main.kind == Kind::Library && m.kind == Kind::Binary && !p.main.pgrx_extension => {
260                p.bins.push(m);
261            }
262            Some(p) => {
263                return Err(format!(
264                    "{} has two Cargo targets, {} and {}; a Cargo package has one library and binaries",
265                    m.package, p.main.name, m.name
266                )
267                .into());
268            }
269        }
270    }
271    for t in metas.iter().filter(|m| m.kind == Kind::Test) {
272        if let Some(p) = out.get_mut(&t.package)
273            && p.main.crate_root == t.crate_root
274        {
275            p.dev_deps.extend(t.deps.iter().map(String::as_str));
276        }
277    }
278    Ok(out)
279}
280
281enum Dep {
282    First(String),
283    Third(String),
284}
285
286impl Dep {
287    fn name(&self) -> &str {
288        match self {
289            Dep::First(n) | Dep::Third(n) => n,
290        }
291    }
292}
293
294fn write_deps(
295    s: &mut String,
296    section: &str,
297    deps: &[Dep],
298    used: &mut std::collections::BTreeSet<String>,
299) -> Res<()> {
300    if deps.is_empty() {
301        return Ok(());
302    }
303    let mut names: Vec<&str> = deps.iter().map(Dep::name).collect();
304    names.sort_unstable();
305    names.dedup();
306    writeln!(s, "\n[{section}]")?;
307    for n in names {
308        writeln!(s, "{n}.workspace = true")?;
309    }
310    for d in deps {
311        if let Dep::Third(n) = d {
312            used.insert(n.clone());
313        }
314    }
315    Ok(())
316}
317
318/// pgrx must be an exact pin (CLAUDE.md §3), and its major feature is the
319/// member's to choose, so the workspace entry drops third-party's.
320fn pin_pgrx(value: &mut toml::Value) -> Res<()> {
321    let t = value.as_table_mut().ok_or("third-party pgrx must be a table")?;
322    let pinned = t.get("version").and_then(toml::Value::as_str).is_some_and(|v| v.starts_with('='));
323    if !pinned {
324        return Err("third-party pgrx must be pinned with `=`".into());
325    }
326    if let Some(toml::Value::Array(features)) = t.get_mut("features") {
327        features.retain(|f| !f.as_str().is_some_and(|f| f.starts_with("pg") && f[2..].bytes().all(|b| b.is_ascii_digit())));
328        if features.is_empty() {
329            t.remove("features");
330        }
331    }
332    Ok(())
333}
334
335fn q(s: &str) -> String {
336    toml::Value::String(s.to_owned()).to_string()
337}
338
339fn write_file(path: &Path, contents: &str) -> Res<()> {
340    std::fs::create_dir_all(path.parent().ok_or("no parent")?)?;
341    Ok(std::fs::write(path, contents)?)
342}
343
344/// Copies a generated tree into the checkout holding the cwd, found by its
345/// `.buckroot`.
346fn write(generated: &Path) -> Res<()> {
347    let cwd = std::env::current_dir()?;
348    let root = cwd
349        .ancestors()
350        .find(|d| d.join(".buckroot").exists())
351        .ok_or("not inside a buck checkout (no .buckroot above the cwd)")?;
352    copy_tree(generated, generated, root)
353}
354
355fn copy_tree(base: &Path, dir: &Path, root: &Path) -> Res<()> {
356    for entry in std::fs::read_dir(dir)? {
357        let path = entry?.path();
358        if path.is_dir() {
359            copy_tree(base, &path, root)?;
360        } else {
361            let dest = root.join(path.strip_prefix(base)?);
362            std::fs::create_dir_all(dest.parent().ok_or("no parent")?)?;
363            std::fs::copy(&path, &dest)?;
364            println!("wrote {}", dest.display());
365        }
366    }
367    Ok(())
368}