refactor(mssql): move restore drop into the connection module

This commit is contained in:
Charles Gauthereau
2026-10-09 17:54:08 +02:00
parent 7e9b276685
commit a1a84bb76d
2 changed files with 39 additions and 42 deletions
+35 -1
View File
@@ -1,9 +1,32 @@
use crate::services::config::DatabaseConfig;
use anyhow::Result;
use anyhow::{Context, Result};
use tiberius::{AuthMethod, Client, Config};
use tokio::net::TcpStream;
use tokio_util::compat::{Compat, TokioAsyncWriteCompatExt};
const DROP_DATABASE_SQL: &str = "DECLARE @q nvarchar(258) = QUOTENAME(@P1);
IF DB_ID(@P1) IS NOT NULL
BEGIN
IF DB_ID(@P1) <= 4
THROW 50000, N'Refusing to restore over a system database; nothing was dropped', 1;
IF ISNULL(HAS_PERMS_BY_NAME(NULL, 'DATABASE', 'CREATE DATABASE'), 0) = 0
THROW 50000, N'Restore needs CREATE DATABASE permission (dbcreator) to recreate the target database; nothing was dropped', 1;
DECLARE @single bit = 0;
IF CAST(SERVERPROPERTY('EngineEdition') AS int) <> 5
BEGIN
EXEC(N'ALTER DATABASE ' + @q + N' SET SINGLE_USER WITH ROLLBACK IMMEDIATE');
SET @single = 1;
END
BEGIN TRY
EXEC(N'DROP DATABASE ' + @q);
END TRY
BEGIN CATCH
IF @single = 1 AND DB_ID(@P1) IS NOT NULL
EXEC(N'ALTER DATABASE ' + @q + N' SET MULTI_USER');
THROW;
END CATCH
END";
pub async fn build_client(cfg: &DatabaseConfig, database: Option<&str>) -> Result<Client<Compat<TcpStream>>> {
let mut config = Config::new();
config.host(&cfg.host);
@@ -19,3 +42,14 @@ pub async fn build_client(cfg: &DatabaseConfig, database: Option<&str>) -> Resul
let client = Client::connect(config, tcp.compat_write()).await?;
Ok(client)
}
pub async fn drop_database(cfg: &DatabaseConfig) -> Result<()> {
let mut client = build_client(cfg, Some("master"))
.await
.with_context(|| format!("Failed to connect to master for {}", cfg.name))?;
client
.execute(DROP_DATABASE_SQL, &[&cfg.database.as_str()])
.await
.with_context(|| format!("Failed to drop MSSQL database {}", cfg.database))?;
Ok(())
}
+4 -41
View File
@@ -7,29 +7,6 @@ use std::process::Command;
use std::sync::Arc;
use std::time::Instant;
const DROP_DATABASE_SQL: &str = "DECLARE @q nvarchar(258) = QUOTENAME(@P1);
IF DB_ID(@P1) IS NOT NULL
BEGIN
IF DB_ID(@P1) <= 4
THROW 50000, N'Refusing to restore over a system database; nothing was dropped', 1;
IF ISNULL(HAS_PERMS_BY_NAME(NULL, 'DATABASE', 'CREATE DATABASE'), 0) = 0
THROW 50000, N'Restore needs CREATE DATABASE permission (dbcreator) to recreate the target database; nothing was dropped', 1;
DECLARE @single bit = 0;
IF CAST(SERVERPROPERTY('EngineEdition') AS int) <> 5
BEGIN
EXEC(N'ALTER DATABASE ' + @q + N' SET SINGLE_USER WITH ROLLBACK IMMEDIATE');
SET @single = 1;
END
BEGIN TRY
EXEC(N'DROP DATABASE ' + @q);
END TRY
BEGIN CATCH
IF @single = 1 AND DB_ID(@P1) IS NOT NULL
EXEC(N'ALTER DATABASE ' + @q + N' SET MULTI_USER');
THROW;
END CATCH
END";
pub async fn run(cfg: DatabaseConfig, restore_file: PathBuf, logger: Arc<JobLogger>) -> Result<()> {
logger.log("debug", format!("Starting MSSQL restore for database {}", cfg.name));
@@ -42,7 +19,10 @@ pub async fn run(cfg: DatabaseConfig, restore_file: PathBuf, logger: Arc<JobLogg
anyhow::bail!(message);
};
drop_database(&cfg, &logger).await?;
logger.log("info", format!("Replacing MSSQL database {} (drop if it exists, then recreate from the backup)", cfg.database));
let start = Instant::now();
super::connection::drop_database(&cfg).await?;
logger.log_command(format!("DROP DATABASE IF EXISTS [{}]", cfg.database), None, Some(0), Some(start.elapsed().as_millis() as f64));
tokio::task::spawn_blocking(move || -> Result<()> {
let connection_string = format!(
@@ -91,20 +71,3 @@ pub async fn run(cfg: DatabaseConfig, restore_file: PathBuf, logger: Arc<JobLogg
})
.await?
}
async fn drop_database(cfg: &DatabaseConfig, logger: &JobLogger) -> Result<()> {
logger.log("info", format!("Replacing MSSQL database {} (drop if it exists, then recreate from the backup)", cfg.database));
let start = Instant::now();
let mut client = super::connection::build_client(cfg, Some("master"))
.await
.with_context(|| format!("Failed to connect to master for {}", cfg.name))?;
client
.execute(DROP_DATABASE_SQL, &[&cfg.database.as_str()])
.await
.with_context(|| format!("Failed to drop MSSQL database {}", cfg.database))?;
let duration_ms = start.elapsed().as_millis() as f64;
logger.log_command(format!("DROP DATABASE IF EXISTS [{}]", cfg.database), None, Some(0), Some(duration_ms));
Ok(())
}