From 18a76fbdbfcd3886abe6045adaef08663a9a9b26 Mon Sep 17 00:00:00 2001
From: fuzzbuck <54221024+fuzzbuck@users.noreply.github.com>
Date: Tue, 6 Sep 2022 22:42:56 +0200
Subject: [PATCH] customizable db locking during migration (#2063)

* customizable db locking during migration

* Update sqlx-core/src/migrate/migrator.rs

Co-authored-by: Austin Bonander <austin.bonander@gmail.com>

* [migrator.rs] cargo fmt

* fix Migrator 'locking' param doctest fail

Co-authored-by: Austin Bonander <austin.bonander@gmail.com>
---
 sqlx-core/src/migrate/migrator.rs | 31 +++++++++++++++++++++++++++----
 sqlx-macros/src/migrate.rs        |  1 +
 src/macros/test.md                |  1 +
 3 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/sqlx-core/src/migrate/migrator.rs b/sqlx-core/src/migrate/migrator.rs
index 5526dda7..f120e3e1 100644
--- a/sqlx-core/src/migrate/migrator.rs
+++ b/sqlx-core/src/migrate/migrator.rs
@@ -9,6 +9,7 @@ use std::slice;
 pub struct Migrator {
     pub migrations: Cow<'static, [Migration]>,
     pub ignore_missing: bool,
+    pub locking: bool,
 }
 
 fn validate_applied_migrations(
@@ -56,6 +57,7 @@ impl Migrator {
         Ok(Self {
             migrations: Cow::Owned(source.resolve().await.map_err(MigrateError::Source)?),
             ignore_missing: false,
+            locking: true,
         })
     }
 
@@ -65,6 +67,19 @@ impl Migrator {
         self
     }
 
+    /// Specify whether or not to lock database during migration. Defaults to `true`.
+    ///
+    /// ### Warning
+    /// Disabling locking can lead to errors or data loss if multiple clients attempt to apply migrations simultaneously
+    /// without some sort of mutual exclusion.
+    ///
+    /// This should only be used if the database does not support locking, e.g. CockroachDB which talks the Postgres
+    /// protocol but does not support advisory locks used by SQLx's migrations support for Postgres.
+    pub fn set_locking(&mut self, locking: bool) -> &Self {
+        self.locking = locking;
+        self
+    }
+
     /// Get an iterator over all known migrations.
     pub fn iter(&self) -> slice::Iter<'_, Migration> {
         self.migrations.iter()
@@ -103,7 +118,9 @@ impl Migrator {
         C: Migrate,
     {
         // lock the database for exclusive access by the migrator
-        conn.lock().await?;
+        if self.locking {
+            conn.lock().await?;
+        }
 
         // creates [_migrations] table only if needed
         // eventually this will likely migrate previous versions of the table
@@ -141,7 +158,9 @@ impl Migrator {
 
         // unlock the migrator to allow other migrators to run
         // but do nothing as we already migrated
-        conn.unlock().await?;
+        if self.locking {
+            conn.unlock().await?;
+        }
 
         Ok(())
     }
@@ -170,7 +189,9 @@ impl Migrator {
         let mut conn = migrator.acquire().await?;
 
         // lock the database for exclusive access by the migrator
-        conn.lock().await?;
+        if self.locking {
+            conn.lock().await?;
+        }
 
         // creates [_migrations] table only if needed
         // eventually this will likely migrate previous versions of the table
@@ -201,7 +222,9 @@ impl Migrator {
 
         // unlock the migrator to allow other migrators to run
         // but do nothing as we already migrated
-        conn.unlock().await?;
+        if self.locking {
+            conn.unlock().await?;
+        }
 
         Ok(())
     }
diff --git a/sqlx-macros/src/migrate.rs b/sqlx-macros/src/migrate.rs
index a463b9dc..58d3ef45 100644
--- a/sqlx-macros/src/migrate.rs
+++ b/sqlx-macros/src/migrate.rs
@@ -146,6 +146,7 @@ pub(crate) fn expand_migrator(path: &Path) -> crate::Result<TokenStream> {
                 #(#migrations),*
             ]),
             ignore_missing: false,
+            locking: true,
         }
     })
 }
diff --git a/src/macros/test.md b/src/macros/test.md
index dbcb7646..05de0ffa 100644
--- a/src/macros/test.md
+++ b/src/macros/test.md
@@ -132,6 +132,7 @@ use sqlx::PgPool;
 #   static MIGRATOR: sqlx::migrate::Migrator = sqlx::migrate::Migrator {
 #       migrations: Cow::Borrowed(&[]),
 #       ignore_missing: false,
+#       locking: true,
 #   };
 # }