Add protected db and drop functionality to restore

GSA-TTS · Jun 17, 2024 · 04474cf · 04474cf
1 parent b3625bb
commit 04474cf
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 39 deletions.
diff --git a/cmd/db_to_db.go b/cmd/db_to_db.go
@@ -106,33 +106,41 @@ func CgovDatabaseSync(
 	table_names []string,
 	operation string) {
 	var BACKUP_ALL = len(table_names) == 0
+	var PROTECTED_DB = "fac-db"
 
 	logging.Logger.Println("DBTODB " + source_db_creds.Get("name").String() + " to " + dest_db_creds.Get("name").String() + " starting")
 	table_to_schema := get_table_and_schema_names_db(source_db_creds)
 	//pg_dump -t table_to_copy source_db | psql target_db
 	for table, schema := range table_to_schema {
 		if slices.Contains(table_names, table) || BACKUP_ALL {
-			switch operation {
-			case "initial":
-				logging.Logger.Printf("Initial db2db operation, truncate not required")
-			case "backup":
-				fallthrough
-			case "restore":
-				truncate_tables(dest_db_creds, []string{table})
-			default:
-				logging.Logger.Printf("Correct operation not supplied. Please supply initial, backup, or restore")
-				os.Exit(-1)
-			}
 
-			psql_write := pipes.Psql(
-				pipes.PG_Dump_Table(source_db_creds, schema, table, "--format plain"),
-				dest_db_creds,
-			)
-			psql_write.Wait()
-			stdout, _ := psql_write.String()
-			if strings.Contains(stdout, "ERR") {
-				logging.Logger.Println("DBTODB " + source_db_creds.Get("name").String() + " to " + dest_db_creds.Get("name").String() + " pipe failed")
-				os.Exit(logging.PIPE_FAILURE)
+			if dest_db == PROTECTED_DB {
+				logging.Logger.Printf("Protected Database '%s' found to be target database. Aborting...", PROTECTED_DB)
+				os.Exit(logging.PROTECTED_DATABASE)
+			} else {
+				switch operation {
+				case "initial":
+					logging.Logger.Printf("Initial db2db operation, truncate not required")
+				case "backup":
+					fallthrough
+				case "restore":
+					//truncate_tables(dest_db_creds, []string{table})
+					drop_tables(dest_db_creds, []string{table})
+				default:
+					logging.Logger.Printf("Correct operation not supplied. Please supply initial, backup, or restore")
+					os.Exit(-1)
+				}
+
+				psql_write := pipes.Psql(
+					pipes.PG_Dump_Table(source_db_creds, schema, table, "--format plain"),
+					dest_db_creds,
+				)
+				psql_write.Wait()
+				stdout, _ := psql_write.String()
+				if strings.Contains(stdout, "ERR") {
+					logging.Logger.Println("DBTODB " + source_db_creds.Get("name").String() + " to " + dest_db_creds.Get("name").String() + " pipe failed")
+					os.Exit(logging.PIPE_FAILURE)
+				}
 			}
 		}
 	}

diff --git a/cmd/s3_to_db.go b/cmd/s3_to_db.go
@@ -22,6 +22,7 @@ func bucket_to_local_tables(
 	bucket_creds vcap.Credentials,
 	s3path *structs.S3Path,
 ) {
+	var PROTECTED_DB = "fac-db"
 	table_to_schema := get_table_and_schema_names(db_creds)
 	//fmt.Sprintf("%s%s/%s-%s.dump", s3path.Bucket, s3path.Key, schema, table)
 	for table, schema := range table_to_schema {
@@ -36,17 +37,23 @@ func bucket_to_local_tables(
 			exit_code = logging.PIPE_FAILURE
 		}
 
-		truncate_tables(db_creds, []string{table})
+		if s3_to_db_db == PROTECTED_DB {
+			logging.Logger.Printf("Protected Database '%s' found to be target database. Aborting...", PROTECTED_DB)
+			os.Exit(logging.PROTECTED_DATABASE)
+		} else {
+			//truncate_tables(db_creds, []string{table})
+			drop_tables(db_creds, []string{table})
 
-		pg_restore := pipes.PG_Restore(db_creds, schema, table)
-		restoreOut, restoreError := pg_restore.String()
-		util.ErrorCheck(restoreOut, restoreError)
+			pg_restore := pipes.PG_Restore(db_creds, schema, table)
+			restoreOut, restoreError := pg_restore.String()
+			util.ErrorCheck(restoreOut, restoreError)
 
-		os.Remove(fmt.Sprintf("./pg_dump_tables/%s", dump_file_name))
-		logging.Logger.Printf("REMOVING FILE: %s", dump_file_name)
+			os.Remove(fmt.Sprintf("./pg_dump_tables/%s", dump_file_name))
+			logging.Logger.Printf("REMOVING FILE: %s", dump_file_name)
 
-		if exit_code != 0 {
-			os.Exit(exit_code)
+			if exit_code != 0 {
+				os.Exit(exit_code)
+			}
 		}
 	}
 
@@ -58,6 +65,7 @@ func bucket_to_cgov_tables(
 	db_creds vcap.Credentials,
 	s3path *structs.S3Path,
 ) {
+	var PROTECTED_DB = "fac-db"
 	table_to_schema := get_table_and_schema_names(db_creds)
 	//fmt.Sprintf("%s%s/%s-%s.dump", s3path.Bucket, s3path.Key, schema, table)
 	for table, schema := range table_to_schema {
@@ -72,17 +80,23 @@ func bucket_to_cgov_tables(
 			exit_code = logging.PIPE_FAILURE
 		}
 
-		truncate_tables(db_creds, []string{table})
-
-		pg_restore := pipes.PG_Restore(db_creds, schema, table)
-		restoreOut, restoreError := pg_restore.String()
-		util.ErrorCheck(restoreOut, restoreError)
-		logging.Logger.Printf("RESTORE of table %s complete.", table)
-
-		os.Remove(fmt.Sprintf("./pg_dump_tables/%s", dump_file_name))
-		logging.Logger.Printf("REMOVING FILE: %s", dump_file_name)
-		if exit_code != 0 {
-			os.Exit(exit_code)
+		if s3_to_db_db == PROTECTED_DB {
+			logging.Logger.Printf("Protected Database '%s' found to be target database. Aborting...", PROTECTED_DB)
+			os.Exit(logging.PROTECTED_DATABASE)
+		} else {
+			//truncate_tables(db_creds, []string{table})
+			drop_tables(db_creds, []string{table})
+
+			pg_restore := pipes.PG_Restore(db_creds, schema, table)
+			restoreOut, restoreError := pg_restore.String()
+			util.ErrorCheck(restoreOut, restoreError)
+			logging.Logger.Printf("RESTORE of table %s complete.", table)
+
+			os.Remove(fmt.Sprintf("./pg_dump_tables/%s", dump_file_name))
+			logging.Logger.Printf("REMOVING FILE: %s", dump_file_name)
+			if exit_code != 0 {
+				os.Exit(exit_code)
+			}
 		}
 	}
 }