api/cmd/rpdata-restore/main.go

package main

import (
	"archive/zip"
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"git.aiterp.net/rpdata/api/database"
	"git.aiterp.net/rpdata/api/database/mongodb"
	"git.aiterp.net/rpdata/api/internal/config"
	"git.aiterp.net/rpdata/api/models"
	"log"
	"strings"
	"time"
)

var flagDriver = flag.String("driver", "postgres", "The database driver to use.")
var flagHost = flag.String("host", "127.0.0.1", "The host to connect to.")
var flagPort = flag.Int("port", 5432, "The port to connect on.")
var flagDb = flag.String("db", "rpdata", "The database name")
var flagUsername = flag.String("username", "", "")
var flagPassword = flag.String("password", "", "")
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()

	cfg := config.Database{
		Driver:     *flagDriver,
		Host:       *flagHost,
		Port:       *flagPort,
		Db:         *flagDb,
		Username:   *flagUsername,
		Password:   *flagPassword,
		Mechanism:  *flagMechanism,
		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)
	if err != nil {
		log.Fatalln("Failed to open database:", err)
	}

	cfg2 := cfg
	cfg2.RestoreIDs = false

	zipReader, err := zip.OpenReader(*flagInputFile)
	if err != nil {
		log.Fatalln("Failed to open input file:", err)
	}
	defer func() {
		err = zipReader.Close()
		if err != nil {
			log.Fatalln("Failed to close input file:", err)
		}
	}()

	ctx, cancel := context.WithTimeout(context.Background(), time.Minute*30)
	defer cancel()

	postMap := make(map[string][]*models.Post)

	for _, file := range zipReader.File {
		if strings.HasSuffix(file.Name, "/") {
			continue
		}

		parts := strings.Split(file.Name, "/")
		if len(parts) < 3 || parts[0] != "rpdata_dump_v1" {
			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)
		}

		hideList := make(map[string]bool)

		log.Println("Loaded", parts[1], parts[2], "from archive.")

		switch parts[1] {
		case "character":
			{
				character := models.Character{}
				err := json.NewDecoder(reader).Decode(&character)
				if err != nil {
					log.Fatalln("Could not parse character:", parts[2], err)
				}

				if *flagReplace {
					_ = db.Characters().Delete(ctx, character)
				}

				_, err = db.Characters().Insert(ctx, character)
				if err != nil {
					log.Fatalln("Could not insert character:", parts[2], err)
				}

				log.Println("Character", character.Name, "inserted.")
			}
		case "channel":
			{
				channel := models.Channel{}
				err := json.NewDecoder(reader).Decode(&channel)
				if err != nil {
					log.Fatalln("Could not parse channel:", parts[2], err)
				}

				if *flagReplace {
					_ = db.Channels().Remove(ctx, channel)
				}

				_, err = db.Channels().Insert(ctx, channel)
				if err != nil {
					log.Fatalln("Could not insert channel:", parts[2], err)
				}

				log.Println("Channel", channel.Name, "inserted.")
			}
		case "change":
			{
				change := models.Change{}
				err := json.NewDecoder(reader).Decode(&change)
				if err != nil {
					log.Fatalln("Could not parse character:", parts[2], err)
				}

				if *flagReplace {
					_ = db.Changes().Remove(ctx, change)
				}

				_, err = db.Changes().Insert(ctx, change)
				if err != nil {
					log.Fatalln("Could not insert character:", parts[2], err)
				}

				if change.Listed {
					log.Println("Change", change.ID, "inserted.")
				} else {
					log.Println("Unlisted change inserted.")
				}
			}
		case "story":
			{
				story := models.Story{}
				err := json.NewDecoder(reader).Decode(&story)
				if err != nil {
					log.Fatalln("Could not parse story:", parts[2], err)
				}

				if *flagReplace {
					_ = db.Stories().Delete(ctx, story)
				}

				_, err = db.Stories().Insert(ctx, story)
				if err != nil {
					log.Fatalln("Could not insert story:", parts[2], err)
				}

				if story.Listed {
					log.Println("Story", story.Name, "inserted.")
				} else {
					log.Println("Unlisted story inserted.")
					hideList[story.ID] = true
				}
			}
		case "chapter":
			{
				chapter := models.Chapter{}
				err := json.NewDecoder(reader).Decode(&chapter)
				if err != nil {
					log.Fatalln("Could not parse story:", parts[2], err)
				}

				if *flagReplace {
					_ = db.Chapters().Delete(ctx, chapter)
				}

				_, err = db.Chapters().Insert(ctx, chapter)
				if err != nil {
					log.Fatalln("Could not insert story:", parts[2], err)
				}

				if !hideList[chapter.StoryID] {
					log.Println("Chapter", fmt.Sprintf("%s (id: %s)", chapter.Title, chapter.ID), "inserted.")
				} else {
					log.Println("Unlisted chapter inserted.")
					hideList[chapter.ID] = true
				}
			}
		case "comment":
			{
				comment := models.Comment{}
				err := json.NewDecoder(reader).Decode(&comment)
				if err != nil {
					log.Fatalln("Could not parse story:", parts[2], err)
				}

				if *flagReplace {
					_ = db.Comments().Delete(ctx, comment)
				}

				_, err = db.Comments().Insert(ctx, comment)
				if err != nil {
					log.Fatalln("Could not insert story:", parts[2], err)
				}

				if !hideList[comment.ChapterID] {
					log.Println("Comment", comment.Subject, "inserted.")
				} else {
					log.Println("Unlisted comment inserted.")
				}
			}

		case "log":
			{
				logg := models.Log{}
				err := json.NewDecoder(reader).Decode(&logg)
				if err != nil {
					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)
				}

				log.Println("Log", logg.Date.Format(time.RFC3339)[:16], logg.ChannelName, "inserted.")
			}
		case "post":
			{
				post := models.Post{}
				err := json.NewDecoder(reader).Decode(&post)
				if err != nil {
					log.Fatalln("Could not parse post:", parts[2], err)
				}

				postMap[post.LogID] = append(postMap[post.LogID], &post)
			}

		case "user":
			{
				user := models.User{}
				err := json.NewDecoder(reader).Decode(&user)
				if err != nil {
					log.Fatalln("Could not parse post:", parts[2], err)
				}

				_, err = db.Users().Insert(ctx, user)
				if err != nil {
					log.Fatalln("Could not insert user:", parts[2], err)
				}

				log.Println("User", user.ID, "inserted.")
			}
		case "key":
			{
				if !*flagIncludeKeys {
					break
				}

				key := models.Key{}
				err := json.NewDecoder(reader).Decode(&key)
				if err != nil {
					log.Fatalln("Could not parse key:", parts[2], err)
				}

				_, err = db.Keys().Insert(ctx, key)
				if err != nil {
					log.Fatalln("Could not insert key:", parts[2], err)
				}

				log.Println("Key", key.ID, "inserted.")
			}
		}

		_ = reader.Close()
	}

	if only == nil || only["post"] {
		for _, posts := range postMap {
			log.Printf("Inserting %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!")
		}
	}

}