Add operation flag to ensure truncate table occurs

GSA-TTS · May 30, 2024 · 5b10a88 · 5b10a88
1 parent 465d266
commit 5b10a88
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 4 deletions.
diff --git a/cmd/db_to_db.go b/cmd/db_to_db.go
@@ -21,6 +21,7 @@ import (
 var (
 	source_db string
 	dest_db   string
+	operation string
 )
 
 func get_table_and_schema_names_db(source_creds vcap.Credentials) map[string]string {
@@ -57,14 +58,24 @@ func get_table_and_schema_names_db(source_creds vcap.Credentials) map[string]str
 func LocalDatabaseSync(
 	source_db_creds vcap.Credentials,
 	dest_db_creds vcap.Credentials,
-	table_names []string) {
+	table_names []string,
+	operation string) {
 	var BACKUP_ALL = len(table_names) == 0
 
 	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})
+			}
+
 			psql_write := pipes.Psql(
 				pipes.PG_Dump_Table(source_db_creds, schema, table),
 				dest_db_creds,
@@ -82,14 +93,24 @@ func LocalDatabaseSync(
 func CgovDatabaseSync(
 	source_db_creds vcap.Credentials,
 	dest_db_creds vcap.Credentials,
-	table_names []string) {
+	table_names []string,
+	operation string) {
 	var BACKUP_ALL = len(table_names) == 0
 
 	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})
+			}
+
 			psql_write := pipes.Psql(
 				pipes.PG_Dump_Table(source_db_creds, schema, table),
 				dest_db_creds,
@@ -117,10 +138,10 @@ var DbToDb = &cobra.Command{
 
 		ch := structs.Choice{
 			Local: func() {
-				LocalDatabaseSync(source_db_creds, dest_db_creds, table_names)
+				LocalDatabaseSync(source_db_creds, dest_db_creds, table_names, operation)
 			},
 			Remote: func() {
-				CgovDatabaseSync(source_db_creds, dest_db_creds, table_names)
+				CgovDatabaseSync(source_db_creds, dest_db_creds, table_names, operation)
 			}}
 		runLocalOrRemote(ch)
 

diff --git a/cmd/root.go b/cmd/root.go
@@ -49,8 +49,10 @@ func parseFlags(cmd_name string, cmd *cobra.Command) {
 		fmt.Println("RUNNING DB_TO_DB FLAGS")
 		cmd.Flags().StringVarP(&source_db, "src_db", "", "", "source database name")
 		cmd.Flags().StringVarP(&dest_db, "dest_db", "", "", "destination database name")
+		cmd.Flags().StringVarP(&operation, "operation", "", "", "operation (initial/backup/restore)")
 		cmd.MarkFlagRequired("src_db")
 		cmd.MarkFlagRequired("dest_db")
+		cmd.MarkFlagRequired("operation")
 	case "truncate":
 		fmt.Println("RUNNING TRUNCATE FLAGS")
 		cmd.Flags().StringVarP(&truncate_db, "db", "", "", "target database name")