From 6608a9993a19169ad276247bc97063042f5aae09 Mon Sep 17 00:00:00 2001
From: Kavin <20838718+FireMasterK@users.noreply.github.com>
Date: Sun, 30 Oct 2022 16:06:27 +0000
Subject: [PATCH] Rollback on transaction failure, and run postgres as root.

---
 docker-compose.yml |  1 +
 src/main.rs        | 11 ++++++++---
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/docker-compose.yml b/docker-compose.yml
index 2e289d8..5eb0d76 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -15,6 +15,7 @@ services:
     image: postgres:15-alpine
     container_name: postgres-sb-mirror
     shm_size: 1g
+    user: 0:0
     # ports:
     #   - 5432:5432
     volumes:
diff --git a/src/main.rs b/src/main.rs
index 44ade8a..4ff7def 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -95,21 +95,26 @@ fn rocket() -> Rocket<Build> {
                                 let start = Instant::now();
                                 println!("Importing database...");
                                 // Execute a query of some kind
-                                db.run(move |c| {
+                                let res = db.run(move |c| {
                                     let result = c.batch_execute("BEGIN; DROP TABLE IF EXISTS \"sponsorTimesTemp\"; CREATE UNLOGGED TABLE \"sponsorTimesTemp\"(LIKE \"sponsorTimes\" INCLUDING defaults INCLUDING constraints INCLUDING indexes); COPY \"sponsorTimesTemp\" FROM '/mirror/sponsorTimes.csv' DELIMITER ',' CSV HEADER; DROP TABLE \"sponsorTimes\"; ALTER TABLE \"sponsorTimesTemp\" RENAME TO \"sponsorTimes\"; COMMIT;");
                                     if result.is_err() {
+                                        c.batch_execute("ROLLBACK;").unwrap();
                                         eprintln!("Failed to import database: {}", result.err().unwrap());
+                                        return false;
                                     }
                                     println!("Imported database in {}ms", start.elapsed().as_millis());
                                     // Vacuum the database
                                     let result = c.batch_execute("VACUUM \"sponsorTimes\";");
                                     if result.is_err() {
                                         eprintln!("Failed to vacuum database: {}", result.err().unwrap());
+                                        return false;
                                     }
+
+                                    true
                                 }).await;
 
-                                unsafe {
-                                    LAST_UPDATE = Some(last_modified);
+                                if res {
+                                    unsafe { LAST_UPDATE = Some(last_modified) };
                                 }
                             }