fix(macros): smarter .env loading, caching, and invalidation

Author: abonander

Cargo.lock

@@ -190,6 +190,7 @@ uuid = "1.1.2"
 # Common utility crates
 cfg-if = "1.0.0"
 dotenvy = { version = "0.15.0", default-features = false }
+thiserror = { version = "2.0.17", default-features = false, features = ["std"] }
 
 # Runtimes
 [workspace.dependencies.async-global-executor]

Cargo.toml

@@ -62,22 +62,22 @@ rustls-native-certs = { version = "0.8.0", optional = true }
 # Type Integrations
 bit-vec = { workspace = true, optional = true }
 bigdecimal = { workspace = true, optional = true }
+chrono = { workspace = true, optional = true }
 rust_decimal = { workspace = true, optional = true }
 time = { workspace = true, optional = true }
 ipnet = { workspace = true, optional = true }
 ipnetwork = { workspace = true, optional = true }
 mac_address = { workspace = true, optional = true }
 uuid = { workspace = true, optional = true }
 
+# work around bug in async-fs 2.0.0, which references futures-lite dependency wrongly, see https://github.com/launchbadge/sqlx/pull/3791#issuecomment-3043363281
+async-fs = { version = "2.1", optional = true }
 async-io = { version = "2.4.1", optional = true }
 async-task = { version = "4.7.1", optional = true }
 
-# work around bug in async-fs 2.0.0, which references futures-lite dependency wrongly, see https://github.com/launchbadge/sqlx/pull/3791#issuecomment-3043363281
-async-fs = { version = "2.1", optional = true }
 base64 = { version = "0.22.0", default-features = false, features = ["std"] }
 bytes = "1.1.0"
 cfg-if = { workspace = true }
-chrono = { version = "0.4.34", default-features = false, features = ["clock"], optional = true }
 crc = { version = "3", optional = true }
 crossbeam-queue = "0.3.2"
 either = "1.6.1"
@@ -93,7 +93,6 @@ serde_json = { version = "1.0.73", features = ["raw_value"], optional = true }
 toml = { version = "0.8.16", optional = true }
 sha2 = { version = "0.10.0", default-features = false, optional = true }
 #sqlformat = "0.2.0"
-thiserror = "2.0.0"
 tokio-stream = { version = "0.1.8", features = ["fs"], optional = true }
 tracing = { version = "0.1.37", features = ["log"] }
 smallvec = "1.7.0"
@@ -102,7 +101,9 @@ bstr = { version = "1.0", default-features = false, features = ["std"], optional
 hashlink = "0.10.0"
 indexmap = "2.0"
 event-listener = "5.2.0"
-hashbrown = "0.15.0"
+hashbrown = "0.16.0"
+
+thiserror.workspace = true
 
 [dev-dependencies]
 sqlx = { workspace = true, features = ["postgres", "sqlite", "mysql", "migrate", "macros", "time", "uuid"] }

sqlx-core/Cargo.toml

@@ -158,7 +158,16 @@ impl Config {
     /// * If the file exists but could not be read or parsed.
     /// * If the file exists but the `sqlx-toml` feature is disabled.
     pub fn try_from_crate_or_default() -> Result<Self, ConfigError> {
-        Self::read_from(get_crate_path()?).or_else(|e| {
+        Self::try_from_path_or_default(get_crate_path()?)
+    }
+
+    /// Attempt to read `Config` from the path given, or return `Config::default()` if it does not exist.
+    ///
+    /// # Errors
+    /// * If the file exists but could not be read or parsed.
+    /// * If the file exists but the `sqlx-toml` feature is disabled.
+    pub fn try_from_path_or_default(path: PathBuf) -> Result<Self, ConfigError> {
+        Self::read_from(path).or_else(|e| {
             if let ConfigError::NotFound { .. } = e {
                 Ok(Config::default())
             } else {

sqlx-core/src/config/mod.rs

@@ -26,7 +26,7 @@ _sqlite = []
 
 # SQLx features
 derive = []
-macros = []
+macros = ["thiserror"]
 migrate = ["sqlx-core/migrate"]
 
 sqlx-toml = ["sqlx-core/sqlx-toml", "sqlx-sqlite?/sqlx-toml"]
@@ -66,6 +66,7 @@ tokio = { workspace = true, optional = true }
 
 cfg-if = { workspace = true}
 dotenvy = { workspace = true }
+thiserror = { workspace = true, optional = true }
 
 hex = { version = "0.4.3" }
 heck = { version = "0.5" }

sqlx-macros-core/Cargo.toml

@@ -0,0 +1,3 @@
+[[disallowed-methods]]
+path = "std::env::var"
+reason = "use `crate::env()` instead, which optionally calls `proc_macro::tracked_env::var()`"

sqlx-macros-core/clippy.toml

@@ -26,7 +26,7 @@ use crate::query::QueryDriver;
 
 pub type Error = Box<dyn std::error::Error>;
 
-pub type Result<T> = std::result::Result<T, Error>;
+pub type Result<T, E = Error> = std::result::Result<T, E>;
 
 mod common;
 pub mod database;
@@ -84,3 +84,24 @@ where
         }
     }
 }
+
+pub fn env(var: &str) -> Result<String> {
+    env_opt(var)?
+        .ok_or_else(|| format!("env var {var:?} must be set to use the query macros").into())
+}
+
+pub fn env_opt(var: &str) -> Result<Option<String>> {
+    use std::env::VarError;
+
+    #[cfg(any(sqlx_macros_unstable, procmacro2_semver_exempt))]
+    let res: Result<String, VarError> = proc_macro::tracked_env::var(var);
+
+    #[cfg(not(any(sqlx_macros_unstable, procmacro2_semver_exempt)))]
+    let res: Result<String, VarError> = std::env::var(var);
+
+    match res {
+        Ok(val) => Ok(Some(val)),
+        Err(VarError::NotPresent) => Ok(None),
+        Err(VarError::NotUnicode(_)) => Err(format!("env var {var:?} is not valid UTF-8").into()),
+    }
+}

sqlx-macros-core/src/lib.rs

@@ -0,0 +1,97 @@
+use std::path::{Path, PathBuf};
+use std::sync::Mutex;
+use std::time::SystemTime;
+
+/// A cached value derived from one or more files, which is automatically invalidated
+/// if the modified-time of any watched file changes.
+pub struct MtimeCache<T> {
+    inner: Mutex<Option<MtimeCacheInner<T>>>,
+}
+
+pub struct MtimeCacheBuilder {
+    file_mtimes: Vec<(PathBuf, Option<SystemTime>)>,
+}
+
+struct MtimeCacheInner<T> {
+    builder: MtimeCacheBuilder,
+    cached: T,
+}
+
+impl<T: Clone> MtimeCache<T> {
+    pub fn new() -> Self {
+        MtimeCache {
+            inner: Mutex::new(None),
+        }
+    }
+
+    /// Get the cached value, or (re)initialize it if it does not exist or a file's mtime has changed.
+    pub fn get_or_try_init<E>(
+        &self,
+        init: impl FnOnce(&mut MtimeCacheBuilder) -> Result<T, E>,
+    ) -> Result<T, E> {
+        let mut inner = self.inner.lock().unwrap_or_else(|e| {
+            // Reset the cache on-panic.
+            let mut locked = e.into_inner();
+            *locked = None;
+            locked
+        });
+
+        if let Some(inner) = &*inner {
+            if !inner.builder.any_modified() {
+                return Ok(inner.cached.clone());
+            }
+        }
+
+        let mut builder = MtimeCacheBuilder::new();
+
+        let value = init(&mut builder)?;
+
+        *inner = Some(MtimeCacheInner {
+            builder,
+            cached: value.clone(),
+        });
+
+        Ok(value)
+    }
+}
+
+impl MtimeCacheBuilder {
+    fn new() -> Self {
+        MtimeCacheBuilder {
+            file_mtimes: Vec::new(),
+        }
+    }
+
+    /// Add a file path to watch.
+    ///
+    /// The cached value will be automatically invalidated if the modified-time of the file changes,
+    /// or if the file does not exist but is created sometime after this call.
+    pub fn add_path(&mut self, path: PathBuf) {
+        let mtime = get_mtime(&path);
+
+        #[cfg(any(sqlx_macros_unstable, procmacro2_semver_exempt))]
+        {
+            proc_macro::tracked_path::path(&path);
+        }
+
+        self.file_mtimes.push((path, mtime));
+    }
+
+    fn any_modified(&self) -> bool {
+        for (path, expected_mtime) in &self.file_mtimes {
+            let actual_mtime = get_mtime(path);
+
+            if expected_mtime != &actual_mtime {
+                return true;
+            }
+        }
+
+        false
+    }
+}
+
+fn get_mtime(path: &Path) -> Option<SystemTime> {
+    std::fs::metadata(path)
+        .and_then(|metadata| metadata.modified())
+        .ok()
+}