add log and post postgres repos.

cmd/rpdata-restore/main.go

@@ -25,6 +25,7 @@ var flagMechanism = flag.String("mechanism", "", "")
 var flagInputFile = flag.String("infile", "dump.zip", "The file to read from.")
 var flagIncludeKeys = flag.Bool("include-keys", false, "Whether to include the keys.")
 var flagReplace = flag.Bool("replace", false, "Replace existing content")
+var flagOnly = flag.String("only", "", "Only restore these models (e.g. \"log,post\")")
 
 func main() {
 	flag.Parse()
@@ -40,6 +41,15 @@ func main() {
 		RestoreIDs: true,
 	}
 
+	var only map[string]bool
+	if *flagOnly != "" {
+		only = make(map[string]bool)
+
+		for _, part := range strings.Split(*flagOnly, ",") {
+			only[strings.Trim(part, "  \t\n\r")] = true
+		}
+	}
+
 	mongodb.DisableFixes = true
 
 	db, err := database.Init(cfg)
@@ -76,6 +86,10 @@ func main() {
 			log.Fatalln("Unrecognized file path:", file.Name)
 		}
 
+		if only != nil && !only[parts[1]] {
+			continue
+		}
+
 		reader, err := file.Open()
 		if err != nil {
 			log.Fatalln("Unrecognized file:", file.Name, err)
@@ -227,6 +241,10 @@ func main() {
 					log.Fatalln("Could not parse log:", parts[2], err)
 				}
 
+				if *flagReplace {
+					_ = db.Logs().Delete(ctx, logg)
+				}
+
 				_, err = db.Logs().Insert(ctx, logg)
 				if err != nil {
 					log.Fatalln("Could not insert log:", parts[2], err)
@@ -284,12 +302,17 @@ func main() {
 		_ = reader.Close()
 	}
 
-	for _, posts := range postMap {
-		_, err = db.Posts().InsertMany(ctx, posts...)
-		if err != nil {
-			log.Fatalln("Could not insert post for logId:", posts[0].LogID, err)
-		}
+	if only == nil || only["post"] {
+		for _, posts := range postMap {
+			log.Printf("Inserting %d posts for log %s...", len(posts), posts[0].LogID)
 
-		log.Printf("Inserted %d posts for log %s.", len(posts), posts[0].LogID)
+			_, err = db.Posts().InsertMany(ctx, posts...)
+			if err != nil {
+				log.Fatalln("Could not insert post for logId:", posts[0].LogID, err)
+			}
+
+			log.Println("\tDone!")
+		}
 	}
+
 }