1 files changed, 55 insertions, 33 deletions

diff --git a/src/db/mod.rs b/src/db/mod.rs
index eb69d7b..cfa0599 100644
--- a/src/db/mod.rs
+++ b/src/db/mod.rs
@@ -1,12 +1,3 @@
-use std::{
-    convert::{
-        AsRef,
-        From,
-    },
-    env,
-    str::FromStr,
-};
-
 use anyhow::{
     Error,
     Result,
@@ -29,10 +20,12 @@ use diesel::{
     NotFound,
     PgTextExpressionMethods,
     QueryDsl,
+    migration::Migration,
     prelude::*,
 };
 use diesel_async::{
     AsyncConnection,
+    AsyncMigrationHarness,
     AsyncPgConnection,
     RunQueryDsl,
     pooled_connection::{
@@ -41,9 +34,21 @@ use diesel_async::{
     },
     scoped_futures::ScopedFutureExt,
 };
+use diesel_migrations::MigrationHarness;
 use grate::tracing;
 use rand::prelude::IndexedRandom;
-use serenity::FutureExt;
+use std::{
+    convert::{
+        AsRef,
+        From,
+    },
+    env,
+    str::FromStr,
+    sync::atomic::{
+        AtomicBool,
+        Ordering,
+    },
+};
 use tokio_postgres::types::FromSql;
 
 use self::schema::*;
@@ -55,12 +60,24 @@ mod schema;
 
 pub use self::models::*;
 
-static MIGRATIONS: diesel_async_migrations::EmbeddedMigrations =
-    diesel_async_migrations::embed_migrations!();
+static MIGRATIONS: diesel_migrations::EmbeddedMigrations = diesel_migrations::embed_migrations!();
 
-lazy_static::lazy_static! {
-    static ref MIGRATE: tokio::sync::OnceCell<()> = tokio::sync::OnceCell::new();
+struct EmbeddedMigrationWrapper<'a> {
+    inner: &'a diesel_migrations::EmbeddedMigrations,
+}
+
+impl<'a, D> diesel::migration::MigrationSource<D> for EmbeddedMigrationWrapper<'a>
+where
+    D: diesel::backend::Backend,
+{
+    fn migrations(&self) -> diesel::migration::Result<Vec<Box<dyn Migration<D>>>> {
+        self.inner.migrations()
+    }
+}
 
+static MIGRATED: AtomicBool = AtomicBool::new(false);
+
+lazy_static::lazy_static! {
     static ref DB_URL: String =
         env::var("DATABASE_URL").expect("no database url in environment");
 
@@ -82,35 +99,40 @@ lazy_static::lazy_static! {
     };
 }
 
+type PooledConn = diesel_async::pooled_connection::deadpool::Object<AsyncPgConnection>;
+
 #[inline]
-pub async fn connection()
--> Result<diesel_async::pooled_connection::deadpool::Object<AsyncPgConnection>> {
-    POOL.get()
-        .then(|mut conn| async move {
-            if let Ok(ref mut conn) = conn {
-                do_migrate(conn).await;
-            }
+pub async fn connection() -> Result<PooledConn> {
+    let conn = POOL.get().await?;
+    let conn = do_migrate(conn);
 
-            conn
-        })
-        .await
-        .map_err(Error::from)
+    Ok(conn)
 }
 
-async fn do_migrate(conn: &mut AsyncPgConnection) {
-    MIGRATE
-        .get_or_init(|| async move {
-            tracing::info!("running migrations");
-            MIGRATIONS.run_pending_migrations(conn).await.expect("failed running migrations");
-            tracing::info!("migrations complete");
+fn do_migrate(conn: PooledConn) -> PooledConn {
+    if MIGRATED.fetch_or(true, Ordering::SeqCst) {
+        return conn;
+    }
+
+    tracing::info!("running migrations");
+
+    let mut harness = AsyncMigrationHarness::new(conn);
+
+    let versions = harness
+        .run_pending_migrations(EmbeddedMigrationWrapper {
+            inner: &MIGRATIONS,
         })
-        .await;
+        .expect("failed running migrations");
+
+    tracing::info!(ran = ?versions, "migrations complete");
+
+    harness.into_inner()
 }
 
 #[inline]
 async fn raw_connection() -> Result<deadpool_postgres::Object> {
     // HACK
-    if !MIGRATE.initialized() {
+    if !AtomicBool::load(&MIGRATED, Ordering::SeqCst) {
         let _ = connection().await?;
     }