pgrx: add automatic upgrade script generation via schema snapshots

cargo-pgrx/Cargo.toml

@@ -53,6 +53,7 @@ tracing-subscriber = { version = "0.3.19", features = ["env-filter"] }
 bzip2 = "0.5.2"
 env_proxy = "0.4.1"
 serde.workspace = true
+serde_json.workspace = true
 serde-xml-rs = "0.6.0"
 tar = "0.4.44"
 ureq = { version = "3.0.10", default-features = false, features = ["gzip", "platform-verifier", "rustls"] }

cargo-pgrx/src/command/install.rs

@@ -356,11 +356,18 @@ fn copy_sql_files(
     output_tracking: &mut Vec<PathBuf>,
 ) -> eyre::Result<()> {
     let (_, extname) = find_control_file(package_manifest_path)?;
+    let version = get_version(package_manifest_path)?;
+
+    // Generate main SQL file and snapshot
     {
-        let version = get_version(package_manifest_path)?;
         let filename = format!("{extname}--{version}.sql");
         let dest = extdir.join(filename);
 
+        // Also generate a snapshot for this version
+        let project_dir = package_manifest_path.parent().unwrap();
+        let snapshot_dir = project_dir.join("sql/snapshots");
+        let snapshot_path = snapshot_dir.join(format!("{extname}--{version}.json"));
+
         crate::command::schema::generate_schema(
             user_manifest_path,
             user_package,
@@ -371,6 +378,7 @@ fn copy_sql_files(
             target,
             Some(&dest),
             None,
+            Some(&snapshot_path),
             None,
             skip_build,
             output_tracking,

cargo-pgrx/src/command/schema.rs

@@ -54,6 +54,9 @@ pub(crate) struct Schema {
     /// A path to output a produced GraphViz DOT file
     #[clap(long, short, value_parser)]
     dot: Option<PathBuf>,
+    /// A path to output a schema snapshot JSON file (for upgrade script generation)
+    #[clap(long, value_parser)]
+    snapshot: Option<PathBuf>,
     #[clap(long)]
     target: Option<String>,
     #[clap(from_global, action = ArgAction::Count)]
@@ -107,6 +110,7 @@ impl CommandExecute for Schema {
             self.target.as_deref(),
             self.out.as_deref(),
             self.dot.as_deref(),
+            self.snapshot.as_deref(),
             log_level,
             self.skip_build,
             &mut vec![],
@@ -119,6 +123,7 @@ impl CommandExecute for Schema {
     test = is_test,
     path = path.map(|path| tracing::field::display(path.display())),
     dot,
+    snapshot,
     features = ?features.features,
 ))]
 pub(crate) fn generate_schema_for_cli(
@@ -131,6 +136,7 @@ pub(crate) fn generate_schema_for_cli(
     target: Option<&str>,
     path: Option<&Path>,
     dot: Option<&Path>,
+    snapshot: Option<&Path>,
     log_level: Option<String>,
     skip_build: bool,
     output_tracking: &mut Vec<PathBuf>,
@@ -163,6 +169,7 @@ pub(crate) fn generate_schema_for_cli(
         target,
         path,
         dot,
+        snapshot,
         output_tracking,
         manifest,
     )
@@ -177,6 +184,7 @@ pub(crate) fn generate_schema_implicit(
     target: Option<&str>,
     path: Option<&Path>,
     dot: Option<&Path>,
+    snapshot: Option<&Path>,
     output_tracking: &mut Vec<PathBuf>,
     manifest: cargo_toml::Manifest,
 ) -> eyre::Result<()> {
@@ -186,8 +194,15 @@ pub(crate) fn generate_schema_implicit(
 
     let symbols = find_and_compute_symbols(profile, &lib_filename, target)?;
 
-    let codegen =
-        compute_codegen(&control_file, package_manifest_path, &symbols, &lib_name, path, dot)?;
+    let codegen = compute_codegen(
+        &control_file,
+        package_manifest_path,
+        &symbols,
+        &lib_name,
+        path,
+        dot,
+        snapshot,
+    )?;
 
     let embed = {
         let mut embed = tempfile::NamedTempFile::new()?;
@@ -358,6 +373,7 @@ fn compute_codegen(
     lib_name: &str,
     path: Option<&Path>,
     dot: Option<&Path>,
+    snapshot: Option<&Path>,
 ) -> eyre::Result<String> {
     use proc_macro2::{Ident, Span, TokenStream};
     let lib_name_ident = Ident::new(lib_name, Span::call_site());
@@ -434,6 +450,104 @@ fn compute_codegen(
                     .expect("Could not write Graphviz DOT");
             });
         }
+        if let Some(snapshot) = snapshot {
+            let snapshot = str_from_path("snapshot", snapshot)?;
+            let writing = "     Writing".bold().green().to_string();
+            let generating = "   Generating".bold().green().to_string();
+            let wrote = "        Wrote".bold().green().to_string();
+            out.extend(quote::quote! {
+                eprintln!("{} schema snapshot to {}", #writing, #snapshot);
+                let current_snapshot = ::pgrx::pgrx_sql_entity_graph::SchemaSnapshot::from_pgrx_sql(&pgrx_sql);
+                current_snapshot
+                    .save(std::path::Path::new(#snapshot))
+                    .expect(&format!("Could not write snapshot to {}", #snapshot));
+
+                // Generate upgrade script from the previous version only
+                let snapshot_path = std::path::Path::new(#snapshot);
+                let snapshot_dir = snapshot_path.parent().unwrap();
+                let sql_dir = snapshot_dir.parent().unwrap();
+
+                if snapshot_dir.exists() {
+                    let extname = &pgrx_sql.extension_name;
+                    let current_version = &current_snapshot.version;
+
+                    // Collect all version snapshots and sort them
+                    let mut versions: Vec<(String, std::path::PathBuf)> = Vec::new();
+
+                    if let Ok(entries) = std::fs::read_dir(snapshot_dir) {
+                        for entry in entries.flatten() {
+                            let path = entry.path();
+                            if !path.is_file() || path.extension().and_then(|s| s.to_str()) != Some("json") {
+                                continue;
+                            }
+
+                            let filename = path.file_name().unwrap().to_str().unwrap();
+
+                            // Extract version from filename (format: extname--version.json)
+                            if let Some(version) = filename
+                                .strip_prefix(&format!("{extname}--"))
+                                .and_then(|s| s.strip_suffix(".json"))
+                            {
+                                versions.push((version.to_string(), path));
+                            }
+                        }
+                    }
+
+                    // Sort versions (simple string comparison)
+                    // Note: This is a simple lexicographic sort. For proper semantic
+                    // versioning, ensure version strings use consistent formatting (e.g., "0.1.0", "0.10.0")
+                    versions.sort_by(|a, b| a.0.cmp(&b.0));
+
+                    // Find the immediate previous version
+                    let current_idx = versions.iter().position(|(v, _)| v == current_version);
+                    if let Some(idx) = current_idx {
+                        if idx > 0 {
+                            // There is a previous version
+                            let (prev_version, prev_path) = &versions[idx - 1];
+
+                            // Load previous snapshot
+                            let prev_json_str = std::fs::read_to_string(prev_path)
+                                .expect(&format!("Failed to load snapshot for version {}", prev_version));
+
+                            // Leak the JSON string to get 'static lifetime for deserialization
+                            let prev_json_str_static: &'static str = Box::leak(prev_json_str.into_boxed_str());
+
+                            // Deserialize the snapshot
+                            let prev_snapshot_result: Result<::pgrx::pgrx_sql_entity_graph::SchemaSnapshot, _> =
+                                ::pgrx::pgrx_sql_entity_graph::deserialize_snapshot(prev_json_str_static);
+                            let prev_snapshot = prev_snapshot_result
+                                .expect(&format!("Failed to deserialize snapshot for version {}", prev_version));
+
+                            // Generate upgrade script
+                            eprintln!(
+                                "{} upgrade script: {} -> {}",
+                                #generating,
+                                prev_version,
+                                current_version
+                            );
+
+                            let diff = ::pgrx::pgrx_sql_entity_graph::SchemaDiff::compare(&prev_snapshot, &current_snapshot);
+
+                            let upgrade_script = ::pgrx::pgrx_sql_entity_graph::generate_upgrade_script_with_sql(
+                                prev_version,
+                                current_version,
+                                &diff,
+                                &pgrx_sql,
+                            )
+                            .expect(&format!("Failed to generate upgrade script from {} to {}", prev_version, current_version));
+
+                            // Write upgrade script
+                            let upgrade_filename = format!("{extname}--{prev_version}--{current_version}.sql");
+                            let upgrade_path = sql_dir.join(&upgrade_filename);
+                            std::fs::write(&upgrade_path, upgrade_script)
+                                .expect(&format!("Failed to write upgrade script {}", upgrade_filename));
+
+                            eprintln!("{} {}", #wrote, upgrade_filename);
+                        }
+                    }
+                }
+            });
+        }
         out
     };
     Ok(quote::quote! {

pgrx-examples/versioned_so/src/lib.rs

@@ -16,6 +16,14 @@ fn hello_versioned_so() -> &'static str {
     "Hello, versioned_so"
 }
 
+
+#[pg_extern]
+fn peos_pne(e: i32) -> String {
+    let p = e + 0;
+    format!("Hello, versioned_so {}", p)
+}
+
+
 #[cfg(any(test, feature = "pg_test"))]
 #[pg_schema]
 mod tests {

pgrx-sql-entity-graph/Cargo.toml

@@ -31,6 +31,8 @@ proc-macro2.workspace = true
 quote.workspace = true
 syn.workspace = true
 thiserror.workspace = true
+serde.workspace = true
+serde_json.workspace = true
 
 convert_case = "0.8.0"
 petgraph = "0.8.1"

pgrx-sql-entity-graph/src/aggregate/entity.rs

@@ -27,20 +27,32 @@ use crate::{SqlGraphEntity, SqlGraphIdentifier, UsedTypeEntity};
 use core::any::TypeId;
 use eyre::{WrapErr, eyre};
 
-#[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
+fn default_type_id() -> TypeId {
+    TypeId::of::<crate::__PgrxInternalTypeIdPlaceholder>()
+}
+
+#[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd, Ord, serde::Serialize, serde::Deserialize)]
+#[serde(bound(deserialize = ""))]
 pub struct AggregateTypeEntity {
     pub used_ty: UsedTypeEntity,
+    #[serde(deserialize_with = "crate::serde_helpers::deserialize_option_static_str")]
     pub name: Option<&'static str>,
 }
 
-#[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
+#[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd, Ord, serde::Serialize, serde::Deserialize)]
+#[serde(bound(deserialize = ""))]
 pub struct PgAggregateEntity {
+    #[serde(deserialize_with = "crate::serde_helpers::deserialize_static_str")]
     pub full_path: &'static str,
+    #[serde(deserialize_with = "crate::serde_helpers::deserialize_static_str")]
     pub module_path: &'static str,
+    #[serde(deserialize_with = "crate::serde_helpers::deserialize_static_str")]
     pub file: &'static str,
     pub line: u32,
+    #[serde(skip, default = "default_type_id")]
     pub ty_id: TypeId,
 
+    #[serde(deserialize_with = "crate::serde_helpers::deserialize_static_str")]
     pub name: &'static str,
 
     /// If the aggregate is an ordered set aggregate.
@@ -66,11 +78,13 @@ pub struct PgAggregateEntity {
     /// The `SFUNC` parameter for [`CREATE AGGREGATE`](https://www.postgresql.org/docs/current/sql-createaggregate.html)
     ///
     /// Corresponds to `state` in `pgrx::aggregate::Aggregate`.
+    #[serde(deserialize_with = "crate::serde_helpers::deserialize_static_str")]
     pub sfunc: &'static str,
 
     /// The `FINALFUNC` parameter for [`CREATE AGGREGATE`](https://www.postgresql.org/docs/current/sql-createaggregate.html)
     ///
     /// Corresponds to `finalize` in `pgrx::aggregate::Aggregate`.
+    #[serde(deserialize_with = "crate::serde_helpers::deserialize_option_static_str")]
     pub finalfunc: Option<&'static str>,
 
     /// The `FINALFUNC_MODIFY` parameter for [`CREATE AGGREGATE`](https://www.postgresql.org/docs/current/sql-createaggregate.html)
@@ -81,31 +95,37 @@ pub struct PgAggregateEntity {
     /// The `COMBINEFUNC` parameter for [`CREATE AGGREGATE`](https://www.postgresql.org/docs/current/sql-createaggregate.html)
     ///
     /// Corresponds to `combine` in `pgrx::aggregate::Aggregate`.
+    #[serde(deserialize_with = "crate::serde_helpers::deserialize_option_static_str")]
     pub combinefunc: Option<&'static str>,
 
     /// The `SERIALFUNC` parameter for [`CREATE AGGREGATE`](https://www.postgresql.org/docs/current/sql-createaggregate.html)
     ///
     /// Corresponds to `serial` in `pgrx::aggregate::Aggregate`.
+    #[serde(deserialize_with = "crate::serde_helpers::deserialize_option_static_str")]
     pub serialfunc: Option<&'static str>,
 
     /// The `DESERIALFUNC` parameter for [`CREATE AGGREGATE`](https://www.postgresql.org/docs/current/sql-createaggregate.html)
     ///
     /// Corresponds to `deserial` in `pgrx::aggregate::Aggregate`.
+    #[serde(deserialize_with = "crate::serde_helpers::deserialize_option_static_str")]
     pub deserialfunc: Option<&'static str>,
 
     /// The `INITCOND` parameter for [`CREATE AGGREGATE`](https://www.postgresql.org/docs/current/sql-createaggregate.html)
     ///
     /// Corresponds to `INITIAL_CONDITION` in `pgrx::aggregate::Aggregate`.
+    #[serde(deserialize_with = "crate::serde_helpers::deserialize_option_static_str")]
     pub initcond: Option<&'static str>,
 
     /// The `MSFUNC` parameter for [`CREATE AGGREGATE`](https://www.postgresql.org/docs/current/sql-createaggregate.html)
     ///
     /// Corresponds to `moving_state` in `pgrx::aggregate::Aggregate`.
+    #[serde(deserialize_with = "crate::serde_helpers::deserialize_option_static_str")]
     pub msfunc: Option<&'static str>,
 
     /// The `MINVFUNC` parameter for [`CREATE AGGREGATE`](https://www.postgresql.org/docs/current/sql-createaggregate.html)
     ///
     /// Corresponds to `moving_state_inverse` in `pgrx::aggregate::Aggregate`.
+    #[serde(deserialize_with = "crate::serde_helpers::deserialize_option_static_str")]
     pub minvfunc: Option<&'static str>,
 
     /// The `MSTYPE` parameter for [`CREATE AGGREGATE`](https://www.postgresql.org/docs/current/sql-createaggregate.html)
@@ -120,6 +140,7 @@ pub struct PgAggregateEntity {
     /// The `MFINALFUNC` parameter for [`CREATE AGGREGATE`](https://www.postgresql.org/docs/current/sql-createaggregate.html)
     ///
     /// Corresponds to `moving_state_finalize` in `pgrx::aggregate::Aggregate`.
+    #[serde(deserialize_with = "crate::serde_helpers::deserialize_option_static_str")]
     pub mfinalfunc: Option<&'static str>,
 
     /// The `MFINALFUNC_MODIFY` parameter for [`CREATE AGGREGATE`](https://www.postgresql.org/docs/current/sql-createaggregate.html)
@@ -130,11 +151,13 @@ pub struct PgAggregateEntity {
     /// The `MINITCOND` parameter for [`CREATE AGGREGATE`](https://www.postgresql.org/docs/current/sql-createaggregate.html)
     ///
     /// Corresponds to `MOVING_INITIAL_CONDITION` in `pgrx::aggregate::Aggregate`.
+    #[serde(deserialize_with = "crate::serde_helpers::deserialize_option_static_str")]
     pub minitcond: Option<&'static str>,
 
     /// The `SORTOP` parameter for [`CREATE AGGREGATE`](https://www.postgresql.org/docs/current/sql-createaggregate.html)
     ///
     /// Corresponds to `SORT_OPERATOR` in `pgrx::aggregate::Aggregate`.
+    #[serde(deserialize_with = "crate::serde_helpers::deserialize_option_static_str")]
     pub sortop: Option<&'static str>,
 
     /// The `PARALLEL` parameter for [`CREATE AGGREGATE`](https://www.postgresql.org/docs/current/sql-createaggregate.html)