mirror of
https://github.com/Portabase/agent.git
synced 2026-10-14 08:39:52 +00:00
refactor(mssql): move restore drop into the connection module
This commit is contained in:
@@ -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(())
|
||||
}
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user