Try avoiding a full clean in cargo sqlx prepare --merged (#1802)

* refactor(sqlx-cli): Try avoiding a full clean with `--merged`

* docs(sqlx-cli): Sprinkle some comments on the metadata changes

* refactor(sqlx-cli): Make the new recompiltion setup unit-testable

* fix(sqlx-cli): Only pass in `$RUSTFLAGS` when set when using `--merged`

* refactor(sqlx-cli): `cargo clean -p` works by name so rip out pkgid code

* chore(sqlx-cli): Remove unused imports

Cargo.lock

@@ -346,6 +346,37 @@ version = "1.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c1db59621ec70f09c5e9b597b220c7a2b43611f4710dc03ceb8748637775692c"
 
+[[package]]
+name = "camino"
+version = "1.0.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6f3132262930b0522068049f5870a856ab8affc80c70d08b6ecb785771a6fc23"
+dependencies = [
+ "serde",
+]
+
+[[package]]
+name = "cargo-platform"
+version = "0.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cbdb825da8a5df079a43676dbe042702f1707b1109f713a01420fbb4cc71fa27"
+dependencies = [
+ "serde",
+]
+
+[[package]]
+name = "cargo_metadata"
+version = "0.14.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4acbb09d9ee8e23699b9634375c72795d095bf268439da88562cf9b501f181fa"
+dependencies = [
+ "camino",
+ "cargo-platform",
+ "semver",
+ "serde",
+ "serde_json",
+]
+
 [[package]]
 name = "cast"
 version = "0.2.7"
@@ -842,6 +873,18 @@ dependencies = [
  "windows-sys 0.30.0",
 ]
 
+[[package]]
+name = "filetime"
+version = "0.2.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "975ccf83d8d9d0d84682850a38c8169027be83368805971cc4f238c2b245bc98"
+dependencies = [
+ "cfg-if",
+ "libc",
+ "redox_syscall",
+ "winapi",
+]
+
 [[package]]
 name = "float-cmp"
 version = "0.9.0"
@@ -2249,6 +2292,9 @@ name = "semver"
 version = "1.0.10"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a41d061efea015927ac527063765e73601444cdc344ba855bc7bd44578b25e1c"
+dependencies = [
+ "serde",
+]
 
 [[package]]
 name = "serde"
@@ -2434,10 +2480,12 @@ dependencies = [
  "anyhow",
  "async-trait",
  "backoff",
+ "cargo_metadata",
  "chrono",
  "clap 3.2.5",
  "console",
  "dotenv",
+ "filetime",
  "futures",
  "glob",
  "openssl",

sqlx-cli/Cargo.toml

@@ -46,6 +46,8 @@ glob = "0.3.0"
 openssl = { version = "0.10.38", optional = true }
 # workaround for https://github.com/rust-lang/rust/issues/29497
 remove_dir_all = "0.7.0"
+cargo_metadata = "0.14"
+filetime = "0.2"
 
 backoff = { version = "0.4.0", features = ["futures", "tokio"] }
 

sqlx-cli/src/lib.rs

@@ -7,6 +7,7 @@ use std::time::Duration;
 use crate::opt::{Command, ConnectOpts, DatabaseCommand, MigrateCommand};
 
 mod database;
+mod metadata;
 // mod migration;
 // mod migrator;
 mod migrate;

sqlx-cli/src/metadata.rs

@@ -0,0 +1,141 @@
+use anyhow::{Context, Result};
+use cargo_metadata::{
+    Metadata as CargoMetadata, Package as MetadataPackage, PackageId as MetadataId,
+};
+
+use std::{
+    collections::{btree_map, BTreeMap, BTreeSet},
+    path::{Path, PathBuf},
+    str::FromStr,
+};
+
+/// The minimal amount of package information we care about
+///
+/// The package's `name` is used to `cargo clean -p` specific crates while the `src_paths` are
+/// are used to trigger recompiles of packages within the workspace
+#[derive(Debug)]
+pub struct Package {
+    name: String,
+    src_paths: Vec<PathBuf>,
+}
+
+impl Package {
+    pub fn name(&self) -> &str {
+        &self.name
+    }
+
+    pub fn src_paths(&self) -> &[PathBuf] {
+        &self.src_paths
+    }
+}
+
+impl From<&MetadataPackage> for Package {
+    fn from(package: &MetadataPackage) -> Self {
+        let name = package.name.clone();
+        let src_paths = package
+            .targets
+            .iter()
+            .map(|target| target.src_path.clone().into_std_path_buf())
+            .collect();
+
+        Self { name, src_paths }
+    }
+}
+
+/// Contains metadata for the current project
+pub struct Metadata {
+    /// Maps packages metadata id to the package
+    ///
+    /// Currently `MetadataId` is used over `PkgId` because pkgid is not a UUID
+    packages: BTreeMap<MetadataId, Package>,
+    /// All of the crates in the current workspace
+    workspace_members: Vec<MetadataId>,
+    /// Maps each dependency to its set of dependents
+    reverse_deps: BTreeMap<MetadataId, BTreeSet<MetadataId>>,
+    /// The target directory of the project
+    ///
+    /// Typically `target` at the workspace root, but can be overridden
+    target_directory: PathBuf,
+}
+
+impl Metadata {
+    pub fn package(&self, id: &MetadataId) -> Option<&Package> {
+        self.packages.get(id)
+    }
+
+    pub fn entries<'this>(&'this self) -> btree_map::Iter<'this, MetadataId, Package> {
+        self.packages.iter()
+    }
+
+    pub fn workspace_members(&self) -> &[MetadataId] {
+        &self.workspace_members
+    }
+
+    pub fn target_directory(&self) -> &Path {
+        &self.target_directory
+    }
+
+    /// Gets all dependents (direct and transitive) of `id`
+    pub fn all_dependents_of(&self, id: &MetadataId) -> BTreeSet<&MetadataId> {
+        let mut dependents = BTreeSet::new();
+        self.all_dependents_of_helper(id, &mut dependents);
+        dependents
+    }
+
+    fn all_dependents_of_helper<'this>(
+        &'this self,
+        id: &MetadataId,
+        dependents: &mut BTreeSet<&'this MetadataId>,
+    ) {
+        if let Some(immediate_dependents) = self.reverse_deps.get(id) {
+            for immediate_dependent in immediate_dependents {
+                if dependents.insert(immediate_dependent) {
+                    self.all_dependents_of_helper(&immediate_dependent, dependents);
+                }
+            }
+        }
+    }
+}
+
+impl FromStr for Metadata {
+    type Err = anyhow::Error;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        let CargoMetadata {
+            packages: metadata_packages,
+            workspace_members,
+            resolve,
+            target_directory,
+            ..
+        } = serde_json::from_str(s)?;
+
+        let mut packages = BTreeMap::new();
+        for metadata_package in metadata_packages {
+            let package = Package::from(&metadata_package);
+            packages.insert(metadata_package.id, package);
+        }
+
+        let mut reverse_deps: BTreeMap<_, BTreeSet<_>> = BTreeMap::new();
+        let resolve =
+            resolve.context("Resolving the dependency graph failed (old version of cargo)")?;
+        for node in resolve.nodes {
+            for dep in node.deps {
+                let dependent = node.id.clone();
+                let dependency = dep.pkg;
+                reverse_deps
+                    .entry(dependency)
+                    .or_default()
+                    .insert(dependent);
+            }
+        }
+
+        let target_directory = target_directory.into_std_path_buf();
+
+        Ok(Self {
+            packages,
+            workspace_members,
+            reverse_deps,
+            target_directory,
+        })
+    }
+}

sqlx-cli/src/prepare.rs

@@ -2,10 +2,9 @@ use crate::opt::ConnectOpts;
 use anyhow::{bail, Context};
 use console::style;
 use remove_dir_all::remove_dir_all;
-use serde::Deserialize;
 use sqlx::any::{AnyConnectOptions, AnyKind};
 use sqlx::Connection;
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, BTreeSet};
 use std::fs::File;
 use std::io::{BufReader, BufWriter};
 use std::path::{Path, PathBuf};
@@ -14,6 +13,8 @@ use std::str::FromStr;
 use std::time::SystemTime;
 use std::{env, fs};
 
+use crate::metadata::Metadata;
+
 type QueryData = BTreeMap<String, serde_json::Value>;
 type JsonObject = serde_json::Map<String, serde_json::Value>;
 
@@ -116,38 +117,46 @@ hint: This command only works in the manifest directory of a Cargo package."#
         .output()
         .context("Could not fetch metadata")?;
 
-    #[derive(Deserialize)]
-    struct Metadata {
-        target_directory: PathBuf,
-    }
-
-    let metadata: Metadata =
-        serde_json::from_slice(&output.stdout).context("Invalid `cargo metadata` output")?;
+    let output_str =
+        std::str::from_utf8(&output.stdout).context("Invalid `cargo metadata` output")?;
+    let metadata: Metadata = output_str.parse()?;
 
     // try removing the target/sqlx directory before running, as stale files
     // have repeatedly caused issues in the past.
-    let _ = remove_dir_all(metadata.target_directory.join("sqlx"));
+    let _ = remove_dir_all(metadata.target_directory().join("sqlx"));
 
     let check_status = if merge {
-        let check_status = Command::new(&cargo).arg("clean").status()?;
+        // Try only triggering a recompile on crates that use `sqlx-macros` falling back to a full
+        // clean on error
+        match setup_minimal_project_recompile(&cargo, &metadata) {
+            Ok(()) => {}
+            Err(err) => {
+                println!(
+                    "Failed minimal recompile setup. Cleaning entire project. Err: {}",
+                    err
+                );
+                let clean_status = Command::new(&cargo).arg("clean").status()?;
+                if !clean_status.success() {
+                    bail!("`cargo clean` failed with status: {}", clean_status);
+                }
+            }
+        };
 
-        if !check_status.success() {
-            bail!("`cargo clean` failed with status: {}", check_status);
-        }
-
-        let mut rustflags = env::var("RUSTFLAGS").unwrap_or_default();
-        rustflags.push_str(&format!(
-            " --cfg __sqlx_recompile_trigger=\"{}\"",
-            SystemTime::UNIX_EPOCH.elapsed()?.as_millis()
-        ));
-
-        Command::new(&cargo)
+        let mut check_command = Command::new(&cargo);
+        check_command
             .arg("check")
             .args(cargo_args)
-            .env("RUSTFLAGS", rustflags)
             .env("SQLX_OFFLINE", "false")
-            .env("DATABASE_URL", url)
-            .status()?
+            .env("DATABASE_URL", url);
+
+        // `cargo check` recompiles on changed rust flags which can be set either via the env var
+        // or through the `rustflags` field in `$CARGO_HOME/config` when the env var isn't set.
+        // Because of this we only pass in `$RUSTFLAGS` when present
+        if let Ok(rustflags) = env::var("RUSTFLAGS") {
+            check_command.env("RUSTFLAGS", rustflags);
+        }
+
+        check_command.status()?
     } else {
         Command::new(&cargo)
             .arg("rustc")
@@ -170,7 +179,7 @@ hint: This command only works in the manifest directory of a Cargo package."#
         bail!("`cargo check` failed with status: {}", check_status);
     }
 
-    let pattern = metadata.target_directory.join("sqlx/query-*.json");
+    let pattern = metadata.target_directory().join("sqlx/query-*.json");
 
     let mut data = BTreeMap::new();
 
@@ -205,6 +214,95 @@ hint: This command only works in the manifest directory of a Cargo package."#
     Ok(data)
 }
 
+#[derive(Debug, PartialEq)]
+struct ProjectRecompileAction {
+    // The names of the packages
+    clean_packages: Vec<String>,
+    touch_paths: Vec<PathBuf>,
+}
+
+/// Sets up recompiling only crates that depend on `sqlx-macros`
+///
+/// This gets a listing of all crates that depend on `sqlx-macros` (direct and transitive). The
+/// crates within the current workspace have their source file's mtimes updated while crates
+/// outside the workspace are selectively `cargo clean -p`ed. In this way we can trigger a
+/// recompile of crates that may be using compile-time macros without forcing a full recompile
+fn setup_minimal_project_recompile(cargo: &str, metadata: &Metadata) -> anyhow::Result<()> {
+    let ProjectRecompileAction {
+        clean_packages,
+        touch_paths,
+    } = minimal_project_recompile_action(metadata)?;
+
+    for file in touch_paths {
+        let now = filetime::FileTime::now();
+        filetime::set_file_times(&file, now, now)
+            .with_context(|| format!("Failed to update mtime for {:?}", file))?;
+    }
+
+    for pkg_id in &clean_packages {
+        let clean_status = Command::new(cargo)
+            .args(&["clean", "-p", pkg_id])
+            .status()?;
+
+        if !clean_status.success() {
+            bail!("`cargo clean -p {}` failed", pkg_id);
+        }
+    }
+
+    Ok(())
+}
+
+fn minimal_project_recompile_action(metadata: &Metadata) -> anyhow::Result<ProjectRecompileAction> {
+    // Get all the packages that depend on `sqlx-macros`
+    let mut sqlx_macros_dependents = BTreeSet::new();
+    let sqlx_macros_ids: BTreeSet<_> = metadata
+        .entries()
+        // We match just by name instead of name and url because some people may have it installed
+        // through different means like vendoring
+        .filter(|(_, package)| package.name() == "sqlx-macros")
+        .map(|(id, _)| id)
+        .collect();
+    for sqlx_macros_id in sqlx_macros_ids {
+        sqlx_macros_dependents.extend(metadata.all_dependents_of(sqlx_macros_id));
+    }
+
+    // Figure out which `sqlx-macros` dependents are in the workspace vs out
+    let mut in_workspace_dependents = Vec::new();
+    let mut out_of_workspace_dependents = Vec::new();
+    for dependent in sqlx_macros_dependents {
+        if metadata.workspace_members().contains(&dependent) {
+            in_workspace_dependents.push(dependent);
+        } else {
+            out_of_workspace_dependents.push(dependent);
+        }
+    }
+
+    // In-workspace dependents have their source file's mtime updated. Out-of-workspace get
+    // `cargo clean -p <PKGID>`ed
+    let files_to_touch: Vec<_> = in_workspace_dependents
+        .iter()
+        .filter_map(|id| {
+            metadata
+                .package(id)
+                .map(|package| package.src_paths().to_owned())
+        })
+        .flatten()
+        .collect();
+    let packages_to_clean: Vec<_> = out_of_workspace_dependents
+        .iter()
+        .filter_map(|id| {
+            metadata
+                .package(id)
+                .map(|package| package.name().to_owned())
+        })
+        .collect();
+
+    Ok(ProjectRecompileAction {
+        clean_packages: packages_to_clean,
+        touch_paths: files_to_touch,
+    })
+}
+
 fn get_db_kind(url: &str) -> anyhow::Result<&'static str> {
     let options = AnyConnectOptions::from_str(&url)?;
 
@@ -277,4 +375,27 @@ mod tests {
         assert_eq!(data.get("a"), Some(&json!({"key1": "value1"})));
         assert_eq!(data.get("z"), Some(&json!({"key2": "value2"})));
     }
+
+    #[test]
+    fn minimal_project_recompile_action_works() -> anyhow::Result<()> {
+        let sample_metadata_path = Path::new("tests")
+            .join("assets")
+            .join("sample_metadata.json");
+        let sample_metadata = std::fs::read_to_string(sample_metadata_path)?;
+        let metadata: Metadata = sample_metadata.parse()?;
+
+        let action = minimal_project_recompile_action(&metadata)?;
+        assert_eq!(
+            action,
+            ProjectRecompileAction {
+                clean_packages: vec!["sqlx".into()],
+                touch_paths: vec![
+                    "/home/user/problematic/workspace/b_in_workspace_lib/src/lib.rs".into(),
+                    "/home/user/problematic/workspace/c_in_workspace_bin/src/main.rs".into(),
+                ]
+            }
+        );
+
+        Ok(())
+    }
 }