Merge branch 'main' of github.com:geoCML/geocml-base-deployment

geoCML · Sep 7, 2024 · 7976905 · 7976905
2 parents 136d309 + 1ae5708
commit 7976905
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 21 deletions.
diff --git a/build-resources/geocml-task-scheduler/geocml-task-scheduler/backup_geocml_db.py b/build-resources/geocml-task-scheduler/geocml-task-scheduler/backup_geocml_db.py
@@ -1,5 +1,6 @@
 import psycopg2
 import os
+import shutil
 import subprocess
 from time import time
 from task_logger import log
@@ -10,8 +11,8 @@
 def backup_geocml_db():
     try:
         conn = psycopg2.connect(dbname="geocml_db",
-                                user="geocml",
-                                password="geocml",
+                                user="postgres",
+                                password="admin",
                                 host="geocml-postgres",
                                 port=5432)
     except psycopg2.OperationalError:
@@ -25,14 +26,14 @@ def backup_geocml_db():
 
     # Write table schemata to .tabor file
     out = subprocess.run(["tabor", "write", "--db", "geocml_db",
-                             "--username", "postgres", "--password", "admin",
+                             "--username", "geocml", "--password", "geocml",
                              "--host", "geocml-postgres",
                              "--file", os.path.join(path_to_backup_dir, "geocml_db.tabor")],
                              capture_output=True)
 
     if out.stderr:
         log("Failed to generate .tabor file {}".format(out.stderr))
-        os.rmdir(path_to_backup_dir)
+        shutil.rmtree(path_to_backup_dir, ignore_errors=True)
         return
 
     cursor = conn.cursor()
@@ -56,7 +57,7 @@ def backup_geocml_db():
 
             data_file_path = os.path.join(path_to_backup_dir, "data:{}.{}.csv".format(schema[0], table[2]))
             data_file = open(data_file_path, "w")
-            cursor.copy_expert(f"""COPY {schema[0]}."{table[2]}" TO STDOUT WITH (FORMAT csv, DELIMITER ',', HEADER);""", data_file)
+            cursor.copy_expert(f"""COPY {schema[0]}."{table[2]}" TO STDOUT WITH (FORMAT csv, DELIMITER ',', HEADER, NULL 'NULL');""", data_file)
             data_file.close()
 
     if delete_backup_dir: # nothing to back up

diff --git a/...d-resources/geocml-task-scheduler/geocml-task-scheduler/restore_geocml_db_from_backups.py b/...d-resources/geocml-task-scheduler/geocml-task-scheduler/restore_geocml_db_from_backups.py
@@ -37,25 +37,15 @@ def restore_geocml_db_from_backups():
 
     log("Restoring geocml_db from {}".format(most_recent_backup))
 
-    cursor = conn.cursor()
-
     # Rebuild tables from .tabor file
 
-    out = subprocess.run(["tabor", "read", "--file", os.path.join(most_recent_backup, "geocml_db.tabor")],
-                         capture_output=True)
-
+    out = subprocess.run(["tabor", "load", "--file", os.path.join(most_recent_backup, "geocml_db.tabor"), "--db", "geocml_db", "--host", "geocml-postgres", "--username", "postgres", "--password", "admin"], capture_output=True)
     if out.stderr:
-        log("Failed to read .tabor file {}".format(out.stderr))
-
-    psql_data = ast.literal_eval(out.stdout.decode())
-
-    for table, psql_queries in psql_data.items():
-        log("Restoring table: {}".format(table))
-        for _, value in psql_queries.items():
-            cursor.execute(value)
+        log("Failed to load tables from .tabor file")
+        return 0
 
-    conn.commit() # commit schema changes to the database before loading data from the CSV
-    log("Tables restored!")
+    cursor = conn.cursor()
+    cursor.execute("SET session_replication_role = replica;")
 
     for csv_data_file in os.listdir(most_recent_backup): # load data from CSV backups
         file_name_split = csv_data_file.split(":")
@@ -65,10 +55,13 @@ def restore_geocml_db_from_backups():
             file_name_split = file_name_split[1].split(".")
             data_file = open(os.path.join(db_backups_dir, most_recent_backup, csv_data_file), "r").readlines()
             cursor.copy_from(StringIO("".join(data_file[1::])), f"{file_name_split[1]}", sep=",",
-                             columns=tuple(data_file[0].replace("\n", "").split(",")))
+                             columns=tuple(data_file[0].replace("\n", "").split(",")), null="NULL")
             log("Finished loading data!")
 
     conn.commit()
+
+    cursor.execute("SET session_replication_role = DEFAULT;")
+
     cursor.close()
     conn.close()
     return 0