Implement Add Feature

cli/src/commands/add.rs

@@ -0,0 +1,167 @@
+use std::{collections::HashSet, path::PathBuf};
+
+use sqlx::{types::chrono::Utc, PgConnection};
+use tokio::fs;
+
+use crate::{
+    models::{Payload, UpdatePath},
+    util::{extension_versions, update_paths},
+};
+
+pub async fn add(
+    payload: &Payload,
+    output_path: &PathBuf,
+    mut conn: PgConnection,
+) -> anyhow::Result<()> {
+    let existing_versions = extension_versions(&mut conn, &payload.metadata.extension_name).await?;
+    let mut installed_extension_once = !existing_versions.is_empty();
+    let mut versions_installed_now = HashSet::new();
+
+    let timestamp = Utc::now().format("%Y%m%d%H%M%S");
+    let mut migration_content = String::new();
+
+    // Header with metadata
+    migration_content.push_str("-- Migration generated by dbdev add at:");
+    migration_content.push_str(&Utc::now().format("%Y-%m-%d %H:%M:%S").to_string());
+    migration_content.push('\n');
+
+    migration_content.push_str("-- Extension: ");
+    migration_content.push_str(&payload.metadata.extension_name);
+    migration_content.push('\n');
+
+    migration_content.push_str("-- Default version:");
+    migration_content.push_str(&payload.metadata.default_version);
+    migration_content.push('\n');
+
+    if let Some(comment) = &payload.metadata.comment {
+        migration_content.push_str("-- Comment:");
+        migration_content.push_str(&comment);
+        migration_content.push('\n');
+    }
+
+    // Add prerequisites first
+    if !&payload.metadata.requires.is_empty() {
+        migration_content.push_str("-- Install prerequisites\n");
+        for req in &payload.metadata.requires {
+            migration_content.push_str("CREATE EXTENSION IF NOT EXISTS ");
+            migration_content.push_str(req);
+            migration_content.push_str(";\n");
+        }
+        migration_content.push('\n');
+    }
+
+    for install_file in &payload.install_files {
+        if !existing_versions.contains(&install_file.version) {
+            if installed_extension_once {
+                // For subsequent versions
+                migration_content.push_str("-- Installing version ");
+                migration_content.push_str(&install_file.version);
+                migration_content.push('\n');
+
+                migration_content.push_str("SELECT pgtle.install_extension_version_sql('");
+                migration_content.push_str(&payload.metadata.extension_name);
+                migration_content.push_str("', '");
+                migration_content.push_str(&install_file.version);
+                migration_content.push_str("', $SQL$");
+                migration_content.push_str(&install_file.body);
+                migration_content.push_str("$SQL$);\n\n");
+
+                versions_installed_now.insert(install_file.version.clone());
+            } else {
+                // For initial installation
+                migration_content.push_str("-- Initial installation of version:");
+                migration_content.push_str(&install_file.version);
+                migration_content.push('\n');
+
+                let requirements = payload
+                    .metadata
+                    .requires
+                    .iter()
+                    .map(|r| format!("'{}'", r))
+                    .collect::<Vec<_>>()
+                    .join(",");
+
+                migration_content.push_str("SELECT pgtle.install_extension('");
+                migration_content.push_str(&payload.metadata.extension_name);
+                migration_content.push_str("', '");
+                migration_content.push_str(&install_file.version);
+                migration_content.push_str("', $COMMENT$");
+                migration_content.push_str(&payload.metadata.comment.as_deref().unwrap_or(""));
+                migration_content.push_str("$COMMENT$, $SQL$");
+                migration_content.push_str(&install_file.body);
+                migration_content.push_str("$SQL$, ARRAY[");
+                migration_content.push_str(&requirements);
+                migration_content.push_str("]::text[] );\n\n");
+
+                versions_installed_now.insert(install_file.version.clone());
+                installed_extension_once = true;
+            }
+        }
+    }
+
+    let existing_update_paths =
+        match update_paths(&mut conn, &payload.metadata.extension_name).await {
+            Ok(paths) => paths,
+            Err(_) => HashSet::new(),
+        };
+
+    for upgrade_file in &payload.upgrade_files {
+        let update_path = UpdatePath {
+            source: upgrade_file.from_version.clone(),
+            target: upgrade_file.to_version.clone(),
+        };
+        if !existing_update_paths.contains(&update_path) {
+            migration_content.push_str("SELECT pgtle.install_update_path('");
+            migration_content.push_str(&payload.metadata.extension_name);
+            migration_content.push_str("', '");
+            migration_content.push_str(&upgrade_file.from_version);
+            migration_content.push_str("', '");
+            migration_content.push_str(&upgrade_file.to_version);
+            migration_content.push_str("', $SQL$");
+            migration_content.push_str(&upgrade_file.body);
+            migration_content.push_str("$SQL$);\n\n");
+        }
+    }
+
+    // Create the extension
+    migration_content.push_str("-- Create the extension\n");
+    match &payload.metadata.schema {
+        Some(schema) => {
+            migration_content.push_str("CREATE EXTENSION IF NOT EXISTS ");
+            migration_content.push_str(&payload.metadata.extension_name);
+            migration_content.push_str(" SCHEMA ");
+            migration_content.push_str(schema);
+            migration_content.push_str(";\n");
+        }
+        None => {
+            migration_content.push_str("CREATE EXTENSION IF NOT EXISTS ");
+            migration_content.push_str(&payload.metadata.extension_name);
+            migration_content.push_str(";\n");
+        }
+    }
+
+    // Set default version
+    migration_content.push_str("-- Setting default version to:");
+    migration_content.push_str(&payload.metadata.default_version);
+    migration_content.push('\n');
+
+    migration_content.push_str("SELECT pgtle.set_default_version('");
+    migration_content.push_str(&payload.metadata.extension_name);
+    migration_content.push_str("', '");
+    migration_content.push_str(&payload.metadata.default_version);
+    migration_content.push_str("');\n");
+
+    // Write to file
+    let mut filename = String::new();
+    filename.push_str(&timestamp.to_string());
+    filename.push('_');
+    filename.push_str(&payload.metadata.extension_name);
+    filename.push_str("_install.sql");
+
+    let file_path = output_path.join(filename);
+
+    fs::write(&file_path, migration_content).await?;
+
+    println!("Generated migration file at: {}", file_path.display());
+    Ok(())
+}

cli/src/commands/install.rs

@@ -1,6 +1,6 @@
 use std::collections::HashSet;
 
-use crate::models::Payload;
+use crate::{models::Payload, util::extension_versions};
 use anyhow::Context;
 use futures::TryStreamExt;
 use sqlx::postgres::PgConnection;
@@ -91,29 +91,6 @@ pub async fn install(payload: &Payload, mut conn: PgConnection) -> anyhow::Resul
     Ok(())
 }
 
-#[derive(sqlx::FromRow)]
-struct ExtensionVersion {
-    version: String,
-}
-
-async fn extension_versions(
-    conn: &mut PgConnection,
-    extension_name: &str,
-) -> anyhow::Result<HashSet<String>> {
-    let mut rows = sqlx::query_as::<_, ExtensionVersion>(
-        "select version from pgtle.available_extension_versions() where name = $1",
-    )
-    .bind(extension_name)
-    .fetch(conn);
-
-    let mut versions = HashSet::new();
-    while let Some(installed_version) = rows.try_next().await? {
-        versions.insert(installed_version.version);
-    }
-
-    Ok(versions)
-}
-
 #[derive(sqlx::FromRow, PartialEq, Eq, Hash)]
 pub(crate) struct UpdatePath {
     pub(crate) source: String,

cli/src/commands/mod.rs

@@ -1,5 +1,6 @@
+pub mod add;
 pub mod install;
+pub mod list;
 pub mod login;
 pub mod publish;
 pub mod uninstall;
-pub mod list;

cli/src/main.rs

@@ -34,6 +34,20 @@ enum Commands {
         install_args: InstallArgs,
     },
 
+    /// Generate a migration file that installs a package to a database
+    Add {
+        /// PostgreSQL connection string
+        #[arg(short, long)]
+        connection: String,
+
+        /// Location to create the migration SQL file
+        #[arg(short, long, value_parser)]
+        output_path: PathBuf,
+
+        #[clap(flatten)]
+        install_args: InstallArgs,
+    },
+
     /// Uninstall a package from a database
     Uninstall {
         /// PostgreSQL connection string
@@ -136,6 +150,27 @@ async fn main() -> anyhow::Result<()> {
             Ok(())
         }
 
+        Commands::Add {
+            connection,
+            output_path,
+            install_args: InstallArgs { path, package },
+        } => {
+            if let Some(_package) = package {
+                return Err(anyhow::anyhow!(
+                    "Generating migrations from packages is not yet supported"
+                ));
+            }
+
+            let current_dir = env::current_dir()?;
+            let extension_dir = path.as_ref().unwrap_or(&current_dir);
+            let payload = models::Payload::from_path(extension_dir)?;
+            let conn = util::get_connection(connection).await?;
+
+            commands::add::add(&payload, &output_path, conn).await?;
+
+            Ok(())
+        }
+
         Commands::Login { registry_name } => {
             let config = Config::read_from_default_file()?;
             let registry_name = registry_name

cli/src/models.rs

@@ -10,10 +10,12 @@ pub struct ControlFileRef {
     pub contents: String,
 }
 
+#[derive(Debug)]
 pub struct Metadata {
     pub extension_name: String,
     pub default_version: String,
     pub comment: Option<String>,
+    pub schema: Option<String>,
     pub relocatable: bool,
     pub requires: Vec<String>,
 }
@@ -26,16 +28,19 @@ impl Metadata {
             comment: control_file_ref.comment()?.clone(),
             relocatable: control_file_ref.relocatable()?,
             requires: control_file_ref.requires()?.clone(),
+            schema: control_file_ref.schema()?.clone(),
         })
     }
 }
 
+#[derive(Debug)]
 pub struct InstallFile {
     pub filename: String,
     pub version: String,
     pub body: String,
 }
 
+#[derive(Debug)]
 pub struct UpgradeFile {
     pub filename: String,
     pub from_version: String,
@@ -59,6 +64,7 @@ impl HasFilename for UpgradeFile {
     }
 }
 
+#[derive(Debug)]
 pub struct ReadmeFile {
     pub body: String,
 }
@@ -79,6 +85,18 @@ impl ReadmeFile {
     }
 }
 
+#[derive(sqlx::FromRow)]
+pub(crate) struct ExtensionVersion {
+    pub(crate) version: String,
+}
+
+#[derive(sqlx::FromRow, PartialEq, Eq, Hash)]
+pub(crate) struct UpdatePath {
+    pub(crate) source: String,
+    pub(crate) target: String,
+}
+
+#[derive(Debug)]
 pub struct Payload {
     /// Absolute path to extension directory
     pub abs_path: Option<PathBuf>,
@@ -275,6 +293,17 @@ impl ControlFileRef {
         Ok(vec![])
     }
 
+    // The schema the extension wants to be installed in, if any
+    fn schema(&self) -> anyhow::Result<Option<String>> {
+        for line in self.contents.lines() {
+            if line.starts_with("schema") {
+                let value = self.read_control_line_value(line)?;
+                return Ok(Some(value.trim().to_string()));
+            }
+        }
+        Ok(None)
+    }
+
     fn relocatable(&self) -> anyhow::Result<bool> {
         for line in self.contents.lines() {
             if line.starts_with("relocatable") {

cli/src/util.rs

@@ -1,10 +1,14 @@
 use anyhow::Context;
+use futures::TryStreamExt;
 use regex::Regex;
 use sqlx::postgres::PgConnection;
 use sqlx::Connection;
+use std::collections::HashSet;
 use std::fs::{File, OpenOptions};
 use std::path::Path;
 
+use crate::models::{ExtensionVersion, UpdatePath};
+
 pub async fn get_connection(connection_str: &str) -> anyhow::Result<PgConnection> {
     PgConnection::connect(connection_str)
         .await
@@ -36,6 +40,42 @@ pub fn is_valid_version(version: &str) -> bool {
     true
 }
 
+pub async fn extension_versions(
+    conn: &mut PgConnection,
+    extension_name: &str,
+) -> anyhow::Result<HashSet<String>> {
+    let mut rows = sqlx::query_as::<_, ExtensionVersion>(
+        "select version from pgtle.available_extension_versions() where name = $1",
+    )
+    .bind(extension_name)
+    .fetch(conn);
+
+    let mut versions = HashSet::new();
+    while let Some(installed_version) = rows.try_next().await? {
+        versions.insert(installed_version.version);
+    }
+
+    Ok(versions)
+}
+
+pub(crate) async fn update_paths(
+    conn: &mut PgConnection,
+    extension_name: &str,
+) -> anyhow::Result<HashSet<UpdatePath>> {
+    let mut rows = sqlx::query_as::<_, UpdatePath>(
+        "select source, target from pgtle.extension_update_paths($1) where path is not null;",
+    )
+    .bind(extension_name)
+    .fetch(conn);
+
+    let mut paths = HashSet::new();
+    while let Some(update_path) = rows.try_next().await? {
+        paths.insert(update_path);
+    }
+
+    Ok(paths)
+}
+
 #[cfg(target_family = "unix")]
 pub(crate) fn create_file(path: &Path) -> Result<File, std::io::Error> {
     use std::os::unix::fs::OpenOptionsExt;