From 81a56a10695b2f8dcbdcc2ce054df4c79a4e4174 Mon Sep 17 00:00:00 2001
From: Norberto Lopes <nlopes.ml@gmail.com>
Date: Tue, 29 Aug 2023 23:23:43 +0100
Subject: [PATCH] feat: add support for postgres (sqlx and diesel)

---
 Cargo.toml          |  2 +-
 src/lib.rs          | 13 ++++-----
 src/postgres/mod.rs | 64 +++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 72 insertions(+), 7 deletions(-)
 create mode 100644 src/postgres/mod.rs

diff --git a/Cargo.toml b/Cargo.toml
index 3ef7492..d81f9e8 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -26,7 +26,7 @@ runtime-async-std = ["sqlx?/runtime-async-std", "dep:async-std"]
 runtime-tokio = ["sqlx?/runtime-tokio", "dep:tokio"]
 sqlite = ["sqlx?/sqlite", "diesel?/sqlite", "diesel_migrations?/sqlite"]
 mysql = ["sqlx?/mysql", "diesel?/mysql", "diesel_migrations?/mysql", "url", "percent-encoding"]
-postgres = ["sqlx?/postgres", "url", "percent-encoding"]
+postgres = ["sqlx?/postgres", "diesel?/postgres", "url", "percent-encoding"]
 
 [dependencies]
 async-std = { version = "1", optional = true } # this has to include default due to task::block_on usage
diff --git a/src/lib.rs b/src/lib.rs
index 7f6cd76..d893297 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -29,8 +29,9 @@ mod sqlite;
 #[cfg_attr(docsrs, doc(cfg(feature = "mysql")))]
 mod mysql;
 
-//#[cfg(feature = "postgres")]
-//mod postgres;
+#[cfg(feature = "postgres")]
+#[cfg_attr(docsrs, doc(cfg(feature = "mysql")))]
+mod postgres;
 
 #[cfg(all(
     feature = "macros",
@@ -72,8 +73,8 @@ impl Default for ConnectionUrl {
         let conn = ConnectionUrl(String::from(sqlite::DEFAULT_CONNECTION_URL));
         #[cfg(feature = "mysql")]
         let conn = ConnectionUrl(String::from(mysql::DEFAULT_CONNECTION_URL));
-        //#[cfg(feature = "postgres")]
-        //let conn = ConnectionUrl(String::from(postgres::DEFAULT_CONNECTION_URL));
+        #[cfg(feature = "postgres")]
+        let conn = ConnectionUrl(String::from(postgres::DEFAULT_CONNECTION_URL));
 
         conn
     }
@@ -158,10 +159,10 @@ impl DatabaseSchema {
     pub async fn dump(&self) -> Result<(), Error> {
         #[cfg(all(feature = "mysql", any(feature = "sqlx", feature = "diesel")))]
         use crate::mysql::write_structure_sql;
+        #[cfg(all(feature = "postgres", any(feature = "sqlx", feature = "diesel")))]
+        use crate::postgres::write_structure_sql;
         #[cfg(all(feature = "sqlite", any(feature = "sqlx", feature = "diesel")))]
         use crate::sqlite::write_structure_sql;
-        //#[cfg(all(feature = "postgres", any(feature = "sqlx", feature = "diesel")))]
-        //use crate::postgres::write_structure_sql;
 
         write_structure_sql(
             &self.0.connection_url.0,
diff --git a/src/postgres/mod.rs b/src/postgres/mod.rs
new file mode 100644
index 0000000..182f52a
--- /dev/null
+++ b/src/postgres/mod.rs
@@ -0,0 +1,64 @@
+//! Implementation of the `postgres` feature
+
+pub(crate) const DEFAULT_CONNECTION_URL: &str = "postgresql://root:@localhost:5432/postgres";
+
+use crate::error::Error;
+
+#[allow(unused_results)]
+#[cfg(any(feature = "sqlx", feature = "diesel"))]
+pub(crate) async fn write_structure_sql<P: AsRef<std::path::Path>, Q: AsRef<std::path::Path>>(
+    connection_url: &str,
+    migrations_path: P,
+    destination_path: Q,
+) -> Result<(), Error> {
+    migrate(connection_url, migrations_path).await?;
+
+    let mut cmd = std::process::Command::new("pg_dump");
+    cmd.arg("--schema-only")
+        .arg("--no-owner")
+        .arg("--no-privileges")
+        .arg("--file")
+        .arg(destination_path.as_ref())
+        .arg(connection_url);
+    crate::process::run(&mut cmd).await?;
+    Ok(())
+}
+
+#[cfg(feature = "sqlx")]
+async fn migrate<P: AsRef<std::path::Path>>(
+    connection_url: &str,
+    migrations_path: P,
+) -> Result<(), sqlx::Error> {
+    use sqlx::{
+        migrate::{Migrate, Migrator},
+        postgres::PgConnectOptions,
+        ConnectOptions,
+    };
+    use std::str::FromStr;
+
+    let mut conn = PgConnectOptions::from_str(connection_url)?
+        .connect()
+        .await?;
+
+    // Ensure the migrations table exists before we run the migrations
+    conn.ensure_migrations_table().await?;
+
+    let migrator = Migrator::new(migrations_path.as_ref()).await?;
+    migrator.run_direct(&mut conn).await?;
+    Ok(())
+}
+
+#[cfg(feature = "diesel")]
+async fn migrate<P: AsRef<std::path::Path>>(
+    connection_url: &str,
+    migrations_path: P,
+) -> Result<(), Error> {
+    use diesel::Connection;
+    use diesel_migrations::{FileBasedMigrations, HarnessWithOutput, MigrationHarness};
+    let mut conn = diesel::PgConnection::establish(connection_url)?;
+    let migrations = FileBasedMigrations::from_path(migrations_path)?;
+    let _ = HarnessWithOutput::write_to_stdout(&mut conn)
+        .run_pending_migrations(migrations)
+        .map(|_| ());
+    Ok(())
+}