From 6132c306c00d8a3b020db1494d60f51a99c08f58 Mon Sep 17 00:00:00 2001
From: Avinash Sajjanshetty <640792+avinassh@users.noreply.github.com>
Date: Tue, 10 Oct 2023 14:27:46 +0530
Subject: [PATCH] verify after db restore is done (#751)

---
 bottomless-cli/src/main.rs | 29 +++++++++++++++++++++++------
 1 file changed, 23 insertions(+), 6 deletions(-)

diff --git a/bottomless-cli/src/main.rs b/bottomless-cli/src/main.rs
index 963c1da5..2e9dc38b 100644
--- a/bottomless-cli/src/main.rs
+++ b/bottomless-cli/src/main.rs
@@ -2,6 +2,7 @@ use anyhow::Result;
 use aws_sdk_s3::Client;
 use chrono::NaiveDateTime;
 use clap::{Parser, Subcommand};
+use std::path::PathBuf;
 
 mod replicator_extras;
 use crate::replicator_extras::detect_db;
@@ -168,6 +169,12 @@ async fn run() -> Result<()> {
         } => {
             tokio::fs::create_dir_all(&database_dir).await?;
             client.restore(generation, utc_time).await?;
+            let db_path = PathBuf::from(&database);
+            if let Err(e) = verify_db(&db_path) {
+                println!("Verification failed: {e}");
+                std::process::exit(1)
+            }
+            println!("Verification: ok");
         }
         Commands::Verify {
             generation,
@@ -179,15 +186,13 @@ async fn run() -> Result<()> {
             client.restore(generation, utc_time).await?;
             let size = tokio::fs::metadata(&temp).await?.len();
             println!("Snapshot size: {size}");
-            let conn = rusqlite::Connection::open(&temp)?;
-            let mut stmt = conn.prepare("PRAGMA integrity_check")?;
-            let mut rows = stmt.query(())?;
-            let result: String = rows.next()?.unwrap().get(0)?;
-            println!("Verification: {result}");
+            let result = verify_db(&temp);
             let _ = tokio::fs::remove_file(&temp).await;
-            if result != "ok" {
+            if let Err(e) = result {
+                println!("Verification failed: {e}");
                 std::process::exit(1)
             }
+            println!("Verification: ok");
         }
         Commands::Rm {
             generation,
@@ -205,6 +210,18 @@ async fn run() -> Result<()> {
     Ok(())
 }
 
+fn verify_db(path: &PathBuf) -> Result<()> {
+    let conn = rusqlite::Connection::open(path)?;
+    let mut stmt = conn.prepare("PRAGMA integrity_check")?;
+    let mut rows = stmt.query(())?;
+    let result: String = rows.next()?.unwrap().get(0)?;
+    if result == "ok" {
+        Ok(())
+    } else {
+        Err(anyhow::anyhow!(result.to_string()))
+    }
+}
+
 #[tokio::main]
 async fn main() {
     if let Err(e) = run().await {