sqlx/sqlx-cli/src/database.rs

use crate::migrate;
use crate::opt::ConnectOpts;
use console::{style, Term};
use dialoguer::Confirm;
use sqlx::any::Any;
use sqlx::migrate::MigrateDatabase;
use std::{io, mem};
use tokio::task;

pub async fn create(connect_opts: &ConnectOpts) -> anyhow::Result<()> {
    // NOTE: only retry the idempotent action.
    // We're assuming that if this succeeds, then any following operations should also succeed.
    let exists = crate::retry_connect_errors(connect_opts, Any::database_exists).await?;

    if !exists {
        #[cfg(feature = "_sqlite")]
        sqlx::sqlite::CREATE_DB_WAL.store(
            connect_opts.sqlite_create_db_wal,
            std::sync::atomic::Ordering::Release,
        );

        Any::create_database(connect_opts.required_db_url()?).await?;
    }

    Ok(())
}

pub async fn drop(connect_opts: &ConnectOpts, confirm: bool, force: bool) -> anyhow::Result<()> {
    if confirm && !ask_to_continue_drop(connect_opts.required_db_url()?.to_owned()).await {
        return Ok(());
    }

    // NOTE: only retry the idempotent action.
    // We're assuming that if this succeeds, then any following operations should also succeed.
    let exists = crate::retry_connect_errors(connect_opts, Any::database_exists).await?;

    if exists {
        if force {
            Any::force_drop_database(connect_opts.required_db_url()?).await?;
        } else {
            Any::drop_database(connect_opts.required_db_url()?).await?;
        }
    }

    Ok(())
}

pub async fn reset(
    migration_source: &str,
    connect_opts: &ConnectOpts,
    confirm: bool,
    force: bool,
) -> anyhow::Result<()> {
    drop(connect_opts, confirm, force).await?;
    setup(migration_source, connect_opts).await
}

pub async fn setup(migration_source: &str, connect_opts: &ConnectOpts) -> anyhow::Result<()> {
    create(connect_opts).await?;
    migrate::run(migration_source, connect_opts, false, false, None).await
}

async fn ask_to_continue_drop(db_url: String) -> bool {
    // If the setup operation is cancelled while we are waiting for the user to decide whether
    // or not to drop the database, this will restore the terminal's cursor to its normal state.
    struct RestoreCursorGuard {
        disarmed: bool,
    }

    impl Drop for RestoreCursorGuard {
        fn drop(&mut self) {
            if !self.disarmed {
                Term::stderr().show_cursor().unwrap()
            }
        }
    }

    let mut guard = RestoreCursorGuard { disarmed: false };

    let decision_result = task::spawn_blocking(move || {
        Confirm::new()
            .with_prompt(format!("Drop database at {}?", style(&db_url).cyan()))
            .wait_for_newline(true)
            .default(false)
            .show_default(true)
            .interact()
    })
    .await
    .expect("Confirm thread panicked");
    match decision_result {
        Ok(decision) => {
            guard.disarmed = true;
            decision
        }
        Err(dialoguer::Error::IO(err)) if err.kind() == io::ErrorKind::Interrupted => {
            // Sometimes CTRL + C causes this error to be returned
            mem::drop(guard);
            false
        }
        Err(err) => {
            mem::drop(guard);
            panic!("Confirm dialog failed with {err}")
        }
    }
}