api/cmd/rpdata-dump/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"
	"os"
	"time"
)

var flagDriver = flag.String("driver", "mongodb", "The database driver to use.")
var flagHost = flag.String("host", "127.0.0.1", "The host to connect to.")
var flagPort = flag.Int("port", 27017, "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 flagOutputFile = flag.String("outfile", "dump.zip", "The file to write to.")
var flagIncludeKeys = flag.Bool("include-keys", false, "Whether to include the keys.")
var flagIncludeUsers = flag.Bool("include-users", false, "Whether to include the keys.")
var flagIncludeUnlisted = flag.Bool("include-unlisted", false, "Whether to include unlisted stuff.")

func main() {
	flag.Parse()

	trueValue := true
	truePtr := &trueValue

	cfg := config.Database{
		Driver:    *flagDriver,
		Host:      *flagHost,
		Port:      *flagPort,
		Db:        *flagDb,
		Username:  *flagUsername,
		Password:  *flagPassword,
		Mechanism: *flagMechanism,
	}

	blacklist := make(map[string]bool)

	mongodb.DisableFixes = true

	db, err := database.Init(cfg)
	if err != nil {
		log.Fatalln("Failed to open database:", err)
	}

	file, err := os.OpenFile(*flagOutputFile, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0600)
	if err != nil {
		log.Fatalln("Failed to open output file:", err)
	}

	zipWriter := zip.NewWriter(file)
	defer func() {
		err = zipWriter.Close()
		if err != nil {
			log.Fatalln("Failed to close output file:", err)
		}
	}()

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

	log.Println("Dumping characters...")
	characters, err := db.Characters().List(ctx, models.CharacterFilter{})
	if err != nil {
		log.Println("Failed to get characters:", err)
	}
	if len(characters) == 0 {
		log.Println("No characters to dump.")
	}
	for i, character := range characters {
		err := writeJsonFile(zipWriter, "character", i, character.ID, time.Now(), character)
		if err != nil {
			log.Println("Failed to write character", character.ID, ":", err)
			continue
		}
	}

	log.Println("Dumping channels...")
	channels, err := db.Channels().List(ctx, models.ChannelFilter{})
	if err != nil {
		log.Println("Failed to get channels:", err)
	}
	if len(channels) == 0 {
		log.Println("No channels to dump.")
	}
	for i, channel := range channels {
		err := writeJsonFile(zipWriter, "channel", i, channel.Name, time.Now(), channel)
		if err != nil {
			log.Println("Failed to write channel", channel.Name, ":", err)
			continue
		}
	}

	log.Println("Dumping changes...")
	changes, err := db.Changes().List(ctx, models.ChangeFilter{})
	if err != nil {
		log.Println("Failed to get changes:", err)
	}
	if len(changes) == 0 {
		log.Println("No changes to dump.")
	}
	for i, change := range changes {
		if !*flagIncludeUnlisted && !change.Listed {
			continue
		}

		err := writeJsonFile(zipWriter, "change", i, change.ID, change.Date, change)
		if err != nil {
			log.Println("Failed to write change", change.ID, ":", err)
			continue
		}
	}

	log.Println("Dumping stories...")
	stories, err := db.Stories().List(ctx, models.StoryFilter{})
	if err != nil {
		log.Println("Failed to get stories:", err)
	}
	unlistedStories, err := db.Stories().List(ctx, models.StoryFilter{Unlisted: truePtr})
	if err != nil {
		log.Println("Failed to get unlisted stories:", err)
	} else {
		stories = append(stories, unlistedStories...)
	}
	if len(stories) == 0 {
		log.Println("No stories to dump.")
	}
	for i, story := range stories {
		if !*flagIncludeUnlisted && !story.Listed {
			blacklist[story.ID] = true
			continue
		}

		err := writeJsonFile(zipWriter, "story", i, story.ID, story.CreatedDate, story)
		if err != nil {
			log.Println("Failed to write story", story.ID, ":", err)
			continue
		}
	}

	log.Println("Dumping chapters...")
	chapters, err := db.Chapters().List(ctx, models.ChapterFilter{})
	if err != nil {
		log.Println("Failed to get chapters:", err)
	}
	if len(chapters) == 0 {
		log.Println("No chapters to dump.")
	}
	for i, chapter := range chapters {
		if !*flagIncludeUnlisted && blacklist[chapter.StoryID] {
			blacklist[chapter.ID] = true
			continue
		}

		err := writeJsonFile(zipWriter, "chapter", i, chapter.ID, chapter.CreatedDate, chapter)
		if err != nil {
			log.Println("Failed to write chapter", chapter.ID, ":", err)
			continue
		}
	}

	log.Println("Dumping comments...")
	comments, err := db.Comments().List(ctx, models.CommentFilter{})
	if err != nil {
		log.Println("Failed to get comments:", err)
	}
	if len(comments) == 0 {
		log.Println("No comments to dump.")
	}
	for i, comment := range comments {
		if !*flagIncludeUnlisted && blacklist[comment.ChapterID] {
			continue
		}

		err := writeJsonFile(zipWriter, "comment", i, comment.ID, comment.CreatedDate, comment)
		if err != nil {
			log.Println("Failed to write comment", comment.ID, ":", err)
			continue
		}
	}

	log.Println("Dumping posts...")
	posts, err := db.Posts().List(ctx, models.PostFilter{})
	if err != nil {
		log.Println("Failed to get posts:", err)
	}
	if len(posts) == 0 {
		log.Println("No posts to dump.")
	}
	for i, post := range posts {
		err := writeJsonFile(zipWriter, "post", i, post.ID, post.Time, post)
		if err != nil {
			log.Println("Failed to write post", post.ID, ":", err)
			continue
		}
	}

	log.Println("Dumping logs...")
	logs, err := db.Logs().List(ctx, models.LogFilter{})
	if err != nil {
		log.Println("Failed to get logs:", err)
	}
	if len(logs) == 0 {
		log.Println("No logs to dump.")
	}
	for i, logEntry := range logs {
		err := writeJsonFile(zipWriter, "log", i, logEntry.ID, logEntry.Date, logEntry)
		if err != nil {
			log.Println("Failed to write post", logEntry.ID, ":", err)
			continue
		}
	}

	if *flagIncludeUsers {
		log.Println("Dumping users...")
		users, err := db.Users().List(ctx)
		if err != nil {
			log.Println("Failed to get users:", err)
		}
		if len(users) == 0 {
			log.Println("No users to dump.")
		}
		for i, userEntry := range users {
			err := writeJsonFile(zipWriter, "user", i, userEntry.ID, time.Now(), userEntry)
			if err != nil {
				log.Println("Failed to write user", userEntry.ID, ":", err)
				continue
			}
		}
	}

	if *flagIncludeKeys {
		log.Println("Dumping keys...")
		keys, err := db.Keys().List(ctx, models.KeyFilter{})
		if err != nil {
			log.Println("Failed to get users:", err)
		}
		if len(keys) == 0 {
			log.Println("No users to dump.")
		}
		for i, keyEntry := range keys {
			err := writeJsonFile(zipWriter, "key", i, keyEntry.ID, time.Now(), keyEntry)
			if err != nil {
				log.Println("Failed to write key", keyEntry.ID, ":", err)
				continue
			}
		}
	}
}

func writeJsonFile(zw *zip.Writer, model string, idx int, id string, t time.Time, data interface{}) error {
	w, err := zw.CreateHeader(&zip.FileHeader{
		Name:     fmt.Sprintf("rpdata_dump_v1/%s/%06d_%s.json", model, idx, id),
		Modified: t,
	})
	if err != nil {
		return err
	}

	return json.NewEncoder(w).Encode(data)
}