add rpdata-restore command.

.gitignore

@@ -16,3 +16,5 @@ rpdata-*
 /graph2/graphcore/*_gen.go
 generated.gql
 generated.go
+/dump.zip
+/rpdata-dump-v1/

cmd/rpdata-dump/main.go

@@ -24,6 +24,8 @@ 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()
@@ -41,6 +43,8 @@ func main() {
 		Mechanism: *flagMechanism,
 	}
 
+	blacklist := make(map[string]bool)
+
 	mongodb.DisableFixes = true
 
 	db, err := database.Init(cfg)
@@ -105,6 +109,10 @@ func main() {
 		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)
@@ -127,6 +135,11 @@ func main() {
 		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)
@@ -143,6 +156,11 @@ func main() {
 		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)
@@ -159,6 +177,10 @@ func main() {
 		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)
@@ -191,26 +213,28 @@ func main() {
 		log.Println("No logs to dump.")
 	}
 	for i, logEntry := range logs {
-		err := writeJsonFile(zipWriter, "post", i, logEntry.ID, logEntry.Date, logEntry)
+		err := writeJsonFile(zipWriter, "log", i, logEntry.ID, logEntry.Date, logEntry)
 		if err != nil {
 			log.Println("Failed to write post", logEntry.ID, ":", err)
 			continue
 		}
 	}
 
-	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 *flagIncludeUsers {
+		log.Println("Dumping users...")
+		users, err := db.Users().List(ctx)
 		if err != nil {
-			log.Println("Failed to write user", userEntry.ID, ":", err)
-			continue
+			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
+			}
 		}
 	}
 

cmd/rpdata-restore/main.go

@@ -0,0 +1,270 @@
+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", "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 flagInputFile = flag.String("infile", "dump.zip", "The file to read from.")
+var flagIncludeKeys = flag.Bool("include-keys", false, "Whether to include the keys.")
+
+func main() {
+	flag.Parse()
+
+	cfg := config.Database{
+		Driver:     *flagDriver,
+		Host:       *flagHost,
+		Port:       *flagPort,
+		Db:         *flagDb,
+		Username:   *flagUsername,
+		Password:   *flagPassword,
+		Mechanism:  *flagMechanism,
+		RestoreIDs: true,
+	}
+
+	mongodb.DisableFixes = true
+
+	db, err := database.Init(cfg)
+	if err != nil {
+		log.Fatalln("Failed to open database:", err)
+	}
+
+	cfg2 := cfg
+	cfg2.RestoreIDs = false
+
+	db.Tags()
+
+	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)
+		}
+
+		reader, err := file.Open()
+		if err != nil {
+			log.Fatalln("Unrecognized file:", file.Name, err)
+		}
+
+		hideList := make(map[string]bool)
+
+		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)
+				}
+
+				_, 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)
+				}
+
+				_, 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)
+				}
+
+				_, 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)
+				}
+
+				_, 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)
+				}
+
+				_, 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)
+				}
+
+				_, 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)
+				}
+
+				_, 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()
+	}
+
+	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)
+		}
+
+		log.Printf("Inserted %d posts for log %s.", len(posts), posts[0].LogID)
+	}
+}

database/mongodb/changes.go

@@ -2,17 +2,20 @@ package mongodb
 
 import (
 	"context"
+	"errors"
 	"git.aiterp.net/rpdata/api/models"
 	"git.aiterp.net/rpdata/api/repositories"
 	"github.com/globalsign/mgo"
 	"github.com/globalsign/mgo/bson"
 	"strconv"
+	"strings"
 	"time"
 )
 
 type changeRepository struct {
-	changes   *mgo.Collection
-	idCounter *counter
+	restoreIDs bool
+	changes    *mgo.Collection
+	idCounter  *counter
 }
 
 func (r *changeRepository) Find(ctx context.Context, id string) (*models.Change, error) {
@@ -33,7 +36,7 @@ func (r *changeRepository) List(ctx context.Context, filter models.ChangeFilter)
 	}
 	if len(filter.Keys) > 0 {
 		query["keys"] = bson.M{"$in": filter.Keys}
-	} else {
+	} else if !r.restoreIDs {
 		query["listed"] = true
 	}
 	if filter.Author != nil && *filter.Author != "" {
@@ -58,14 +61,28 @@ func (r *changeRepository) List(ctx context.Context, filter models.ChangeFilter)
 }
 
 func (r *changeRepository) Insert(ctx context.Context, change models.Change) (*models.Change, error) {
-	next, err := r.idCounter.Increment(1)
-	if err != nil {
-		return nil, err
-	}
+	if !r.restoreIDs || change.ID == "" {
+		next, err := r.idCounter.Increment(1)
+		if err != nil {
+			return nil, err
+		}
 
-	change.ID = "Change_" + strconv.Itoa(next)
+		change.ID = "Change_" + strconv.Itoa(next)
+	} else {
+		tokens := strings.Split(change.ID, "_")
+		if len(tokens) != 2 || tokens[0] != "Change" {
+			return nil, errors.New("Invalid change ID")
+		}
+
+		n, err := strconv.Atoi(tokens[1])
+		if err != nil {
+			return nil, err
+		}
+
+		_ = r.idCounter.Bump(n)
+	}
 
-	err = r.changes.Insert(&change)
+	err := r.changes.Insert(&change)
 	if err != nil {
 		return nil, err
 	}
@@ -77,7 +94,7 @@ func (r *changeRepository) Remove(ctx context.Context, change models.Change) err
 	return r.changes.RemoveId(change.ID)
 }
 
-func newChangeRepository(db *mgo.Database) (repositories.ChangeRepository, error) {
+func newChangeRepository(db *mgo.Database, restoreIDs bool) (repositories.ChangeRepository, error) {
 	collection := db.C("common.changes")
 
 	// Delete the old index if it exists.
@@ -103,7 +120,8 @@ func newChangeRepository(db *mgo.Database) (repositories.ChangeRepository, error
 	}
 
 	return &changeRepository{
-		changes:   collection,
-		idCounter: newCounter(db, "auto_increment", "Change"),
+		restoreIDs: restoreIDs,
+		changes:    collection,
+		idCounter:  newCounter(db, "auto_increment", "Change"),
 	}, nil
 }

database/mongodb/chapters.go

@@ -2,20 +2,23 @@ package mongodb
 
 import (
 	"context"
+	"errors"
 	"git.aiterp.net/rpdata/api/internal/generate"
 	"git.aiterp.net/rpdata/api/models"
 	"git.aiterp.net/rpdata/api/repositories"
 	"github.com/globalsign/mgo"
 	"github.com/globalsign/mgo/bson"
 	"log"
+	"strings"
 )
 
 type chapterRepository struct {
-	chapters *mgo.Collection
-	comments *mgo.Collection
+	restoreIDs bool
+	chapters   *mgo.Collection
+	comments   *mgo.Collection
 }
 
-func newChapterRepository(db *mgo.Database) (repositories.ChapterRepository, error) {
+func newChapterRepository(db *mgo.Database, restoreIDs bool) (repositories.ChapterRepository, error) {
 	collection := db.C("story.chapters")
 
 	err := collection.EnsureIndexKey("storyId")
@@ -32,8 +35,9 @@ func newChapterRepository(db *mgo.Database) (repositories.ChapterRepository, err
 	}
 
 	return &chapterRepository{
-		chapters: collection,
-		comments: db.C("story.comments"),
+		restoreIDs: restoreIDs,
+		chapters:   collection,
+		comments:   db.C("story.comments"),
 	}, nil
 }
 
@@ -67,7 +71,13 @@ func (r *chapterRepository) List(ctx context.Context, filter models.ChapterFilte
 }
 
 func (r *chapterRepository) Insert(ctx context.Context, chapter models.Chapter) (*models.Chapter, error) {
-	chapter.ID = generate.StoryID()
+	if !r.restoreIDs {
+		chapter.ID = generate.ChapterID()
+	} else {
+		if len(chapter.ID) != len(generate.ChapterID()) && strings.HasPrefix(chapter.ID, "S") {
+			return nil, errors.New("invalid story id")
+		}
+	}
 
 	err := r.chapters.Insert(chapter)
 	if err != nil {

database/mongodb/characters.go

@@ -15,9 +15,10 @@ import (
 type characterRepository struct {
 	characters *mgo.Collection
 	cidCounter *counter
+	restoreIDs bool
 }
 
-func newCharacterRepository(db *mgo.Database) (repositories.CharacterRepository, error) {
+func newCharacterRepository(db *mgo.Database, restoreIDs bool) (repositories.CharacterRepository, error) {
 	collection := db.C("common.characters")
 
 	err := collection.EnsureIndexKey("name")
@@ -52,6 +53,7 @@ func newCharacterRepository(db *mgo.Database) (repositories.CharacterRepository,
 	}
 
 	return &characterRepository{
+		restoreIDs: restoreIDs,
 		characters: collection,
 		cidCounter: newCounter(db, "auto_increment", "Character"),
 	}, nil
@@ -150,13 +152,25 @@ func (r *characterRepository) List(ctx context.Context, filter models.CharacterF
 }
 
 func (r *characterRepository) Insert(ctx context.Context, character models.Character) (*models.Character, error) {
-	nextId, err := r.cidCounter.Increment(1)
-	if err != nil {
-		return nil, err
+	if !r.restoreIDs {
+		nextId, err := r.cidCounter.Increment(1)
+		if err != nil {
+			return nil, err
+		}
+		character.ID = "C" + strconv.Itoa(nextId)
+	} else {
+		n, err := strconv.Atoi(character.ID[1:])
+		if err != nil {
+			return nil, err
+		}
+
+		err = r.cidCounter.Bump(n)
+		if err != nil {
+			return nil, err
+		}
 	}
-	character.ID = "C" + strconv.Itoa(nextId)
 
-	err = r.characters.Insert(&character)
+	err := r.characters.Insert(&character)
 	if err != nil {
 		return nil, err
 	}

database/mongodb/comments.go

@@ -2,19 +2,22 @@ package mongodb
 
 import (
 	"context"
+	"errors"
 	"git.aiterp.net/rpdata/api/internal/generate"
 	"git.aiterp.net/rpdata/api/models"
 	"git.aiterp.net/rpdata/api/repositories"
 	"github.com/globalsign/mgo"
 	"github.com/globalsign/mgo/bson"
 	"log"
+	"strings"
 )
 
 type commentRepository struct {
-	comments *mgo.Collection
+	restoreIDs bool
+	comments   *mgo.Collection
 }
 
-func newCommentRepository(db *mgo.Database) (repositories.CommentRepository, error) {
+func newCommentRepository(db *mgo.Database, restoreIDs bool) (repositories.CommentRepository, error) {
 	collection := db.C("story.comments")
 
 	err := collection.EnsureIndexKey("chapterId")
@@ -31,7 +34,8 @@ func newCommentRepository(db *mgo.Database) (repositories.CommentRepository, err
 	}
 
 	r := &commentRepository{
-		comments: collection,
+		restoreIDs: restoreIDs,
+		comments:   collection,
 	}
 
 	go r.fixFieldTypo()
@@ -69,7 +73,13 @@ func (r *commentRepository) List(ctx context.Context, filter models.CommentFilte
 }
 
 func (r *commentRepository) Insert(ctx context.Context, comment models.Comment) (*models.Comment, error) {
-	comment.ID = generate.CommentID()
+	if !r.restoreIDs {
+		comment.ID = generate.CommentID()
+	} else {
+		if len(comment.ID) != len(generate.CommentID()) && strings.HasPrefix(comment.ID, "SSC") {
+			return nil, errors.New("invalid story id")
+		}
+	}
 
 	err := r.comments.Insert(comment)
 	if err != nil {

database/mongodb/db.go

@@ -105,7 +105,7 @@ func Init(cfg config.Database) (*MongoDB, error) {
 
 	db := session.DB(cfg.Db)
 
-	characters, err := newCharacterRepository(db)
+	characters, err := newCharacterRepository(db, cfg.RestoreIDs)
 	if err != nil {
 		session.Close()
 		return nil, err
@@ -117,19 +117,19 @@ func Init(cfg config.Database) (*MongoDB, error) {
 		return nil, err
 	}
 
-	changes, err := newChangeRepository(db)
+	changes, err := newChangeRepository(db, cfg.RestoreIDs)
 	if err != nil {
 		session.Close()
 		return nil, err
 	}
 
-	logs, err := newLogRepository(db)
+	logs, err := newLogRepository(db, cfg.RestoreIDs)
 	if err != nil {
 		session.Close()
 		return nil, err
 	}
 
-	posts, err := newPostRepository(db)
+	posts, err := newPostRepository(db, cfg.RestoreIDs)
 	if err != nil {
 		session.Close()
 		return nil, err
@@ -141,25 +141,25 @@ func Init(cfg config.Database) (*MongoDB, error) {
 		return nil, err
 	}
 
-	stories, err := newStoryRepository(db)
+	stories, err := newStoryRepository(db, cfg.RestoreIDs)
 	if err != nil {
 		session.Close()
 		return nil, err
 	}
 
-	chapters, err := newChapterRepository(db)
+	chapters, err := newChapterRepository(db, cfg.RestoreIDs)
 	if err != nil {
 		session.Close()
 		return nil, err
 	}
 
-	comments, err := newCommentRepository(db)
+	comments, err := newCommentRepository(db, cfg.RestoreIDs)
 	if err != nil {
 		session.Close()
 		return nil, err
 	}
 
-	keys, err := newKeyRepository(db)
+	keys, err := newKeyRepository(db, cfg.RestoreIDs)
 	if err != nil {
 		session.Close()
 		return nil, err
@@ -231,6 +231,29 @@ func (c *counter) With(category, name string) *counter {
 	}
 }
 
+func (c *counter) Bump(amount int) error {
+	id := c.category + "." + c.name
+
+	err := c.coll.Update(bson.M{
+		"_id":   id,
+		"value": bson.M{"$lt": amount},
+	}, bson.M{
+		"value": amount,
+	})
+	if err != nil {
+		if err == mgo.ErrNotFound {
+			return c.coll.Insert(bson.M{
+				"_id":   id,
+				"value": amount,
+			})
+		}
+
+		return err
+	}
+
+	return nil
+}
+
 func (c *counter) Increment(amount int) (int, error) {
 	type counterDoc struct {
 		ID    string `bson:"_id"`

database/mongodb/keys.go

@@ -10,10 +10,11 @@ import (
 )
 
 type keyRepository struct {
-	keys *mgo.Collection
+	restoreIDs bool
+	keys       *mgo.Collection
 }
 
-func newKeyRepository(db *mgo.Database) (repositories.KeyRepository, error) {
+func newKeyRepository(db *mgo.Database, restoreIDs bool) (repositories.KeyRepository, error) {
 	collection := db.C("auth.keys")
 
 	err := collection.EnsureIndexKey("user")
@@ -21,7 +22,10 @@ func newKeyRepository(db *mgo.Database) (repositories.KeyRepository, error) {
 		return nil, err
 	}
 
-	return &keyRepository{keys: collection}, nil
+	return &keyRepository{
+		keys:       collection,
+		restoreIDs: restoreIDs,
+	}, nil
 }
 
 func (r *keyRepository) Find(ctx context.Context, id string) (*models.Key, error) {
@@ -54,7 +58,9 @@ func (r *keyRepository) List(ctx context.Context, filter models.KeyFilter) ([]*m
 }
 
 func (r *keyRepository) Insert(ctx context.Context, key models.Key) (*models.Key, error) {
-	key.ID = generate.KeyID()
+	if !r.restoreIDs {
+		key.ID = generate.KeyID()
+	}
 
 	err := r.keys.Insert(&key)
 	if err != nil {

database/mongodb/logs.go

@@ -12,14 +12,15 @@ import (
 )
 
 type logRepository struct {
-	openMutex sync.Mutex
+	openMutex  sync.Mutex
+	restoreIds bool
 
 	logs           *mgo.Collection
 	posts          *mgo.Collection
 	shortIdCounter *counter
 }
 
-func newLogRepository(db *mgo.Database) (*logRepository, error) {
+func newLogRepository(db *mgo.Database, restoreIds bool) (*logRepository, error) {
 	logs := db.C("logbot3.logs")
 	posts := db.C("logbot3.posts")
 
@@ -54,6 +55,7 @@ func newLogRepository(db *mgo.Database) (*logRepository, error) {
 	}
 
 	return &logRepository{
+		restoreIds:     restoreIds,
 		logs:           logs,
 		posts:          posts,
 		shortIdCounter: newCounter(db, "auto_increment", "Log"),
@@ -123,26 +125,35 @@ func (r *logRepository) List(ctx context.Context, filter models.LogFilter) ([]*m
 }
 
 func (r *logRepository) Insert(ctx context.Context, log models.Log) (*models.Log, error) {
-	nextShortId, err := r.shortIdCounter.Increment(1)
-	if err != nil {
-		return nil, err
-	}
+	if !r.restoreIds || log.ID == "" || log.ShortID == "" {
+		nextShortId, err := r.shortIdCounter.Increment(1)
+		if err != nil {
+			return nil, err
+		}
 
-	log.ID = generate.LogID(log)
-	log.ShortID = "L" + strconv.Itoa(nextShortId)
+		log.ID = generate.LogID(log)
+		log.ShortID = "L" + strconv.Itoa(nextShortId)
+	} else {
+		n, err := strconv.Atoi(log.ShortID[1:])
+		if err != nil {
+			return nil, err
+		}
+
+		_ = r.shortIdCounter.Bump(n)
+	}
 
 	if log.Open {
 		// There can be only one open log in the same channel.
 		r.openMutex.Lock()
 		defer r.openMutex.Unlock()
 
-		_, err = r.logs.UpdateAll(bson.M{"channel": log.ChannelName, "open": true}, bson.M{"$set": bson.M{"open": false}})
+		_, err := r.logs.UpdateAll(bson.M{"channel": log.ChannelName, "open": true}, bson.M{"$set": bson.M{"open": false}})
 		if err != nil {
 			return nil, errors.New("Cannot close other logs: " + err.Error())
 		}
 	}
 
-	err = r.logs.Insert(&log)
+	err := r.logs.Insert(&log)
 	if err != nil {
 		return nil, err
 	}

database/mongodb/posts.go

@@ -2,6 +2,7 @@ package mongodb
 
 import (
 	"context"
+	"errors"
 	"git.aiterp.net/rpdata/api/internal/generate"
 	"git.aiterp.net/rpdata/api/models"
 	"git.aiterp.net/rpdata/api/repositories"
@@ -14,13 +15,14 @@ import (
 )
 
 type postRepository struct {
-	logs  *mgo.Collection
-	posts *mgo.Collection
+	restoreIDs bool
+	logs       *mgo.Collection
+	posts      *mgo.Collection
 
 	orderMutex sync.Mutex
 }
 
-func newPostRepository(db *mgo.Database) (*postRepository, error) {
+func newPostRepository(db *mgo.Database, restoreIDs bool) (*postRepository, error) {
 	posts := db.C("logbot3.posts")
 
 	err := posts.EnsureIndexKey("logId")
@@ -48,8 +50,9 @@ func newPostRepository(db *mgo.Database) (*postRepository, error) {
 	}
 
 	return &postRepository{
-		posts: posts,
-		logs:  db.C("logbot3.logs"),
+		restoreIDs: restoreIDs,
+		posts:      posts,
+		logs:       db.C("logbot3.logs"),
 	}, nil
 }
 
@@ -118,7 +121,14 @@ func (r *postRepository) Insert(ctx context.Context, post models.Post) (*models.
 		return nil, err
 	}
 
-	post.ID = generate.PostID()
+	if !r.restoreIDs {
+		post.ID = generate.PostID()
+	} else {
+		if len(post.ID) != len(generate.PostID()) && strings.HasPrefix(post.ID, "P") {
+			return nil, errors.New("invalid story id")
+		}
+	}
+
 	post.Position = lastPost.Position + 1 // Position 1 is first position, so this is safe.
 
 	err = r.posts.Insert(post)
@@ -140,7 +150,9 @@ func (r *postRepository) InsertMany(ctx context.Context, posts ...*models.Post)
 			return nil, repositories.ErrParentMismatch
 		}
 
-		post.ID = generate.PostID()
+		if !r.restoreIDs || post.ID == "" {
+			post.ID = generate.PostID()
+		}
 	}
 
 	r.orderMutex.Lock()

database/mongodb/stories.go

@@ -2,21 +2,24 @@ package mongodb
 
 import (
 	"context"
+	"errors"
 	"git.aiterp.net/rpdata/api/internal/generate"
 	"git.aiterp.net/rpdata/api/models"
 	"git.aiterp.net/rpdata/api/repositories"
 	"github.com/globalsign/mgo"
 	"github.com/globalsign/mgo/bson"
 	"log"
+	"strings"
 )
 
 type storyRepository struct {
-	stories  *mgo.Collection
-	chapters *mgo.Collection
-	comments *mgo.Collection
+	stories    *mgo.Collection
+	chapters   *mgo.Collection
+	comments   *mgo.Collection
+	restoreIDs bool
 }
 
-func newStoryRepository(db *mgo.Database) (repositories.StoryRepository, error) {
+func newStoryRepository(db *mgo.Database, restoreIDs bool) (repositories.StoryRepository, error) {
 	collection := db.C("story.stories")
 
 	err := collection.EnsureIndexKey("tags")
@@ -41,9 +44,10 @@ func newStoryRepository(db *mgo.Database) (repositories.StoryRepository, error)
 	}
 
 	return &storyRepository{
-		stories:  collection,
-		chapters: db.C("story.chapters"),
-		comments: db.C("story.comments"),
+		stories:    collection,
+		restoreIDs: restoreIDs,
+		chapters:   db.C("story.chapters"),
+		comments:   db.C("story.comments"),
 	}, nil
 }
 
@@ -105,7 +109,13 @@ func (r *storyRepository) List(ctx context.Context, filter models.StoryFilter) (
 }
 
 func (r *storyRepository) Insert(ctx context.Context, story models.Story) (*models.Story, error) {
-	story.ID = generate.StoryID()
+	if !r.restoreIDs {
+		story.ID = generate.StoryID()
+	} else {
+		if len(story.ID) != len(generate.StoryID()) && strings.HasPrefix(story.ID, "S") {
+			return nil, errors.New("invalid story id")
+		}
+	}
 
 	err := r.stories.Insert(story)
 	if err != nil {

internal/config/config.go

@@ -35,13 +35,14 @@ type Space struct {
 
 // Database is configuration for spaces.
 type Database struct {
-	Driver    string `json:"driver" yaml:"driver"`
-	Host      string `json:"host" yaml:"host"`
-	Port      int    `json:"port" yaml:"port"`
-	Db        string `json:"db" yaml:"db"`
-	Username  string `json:"username" yaml:"username"`
-	Password  string `json:"password" yaml:"password"`
-	Mechanism string `json:"mechanism" yaml:"mechanism"`
+	Driver     string `json:"driver" yaml:"driver"`
+	Host       string `json:"host" yaml:"host"`
+	Port       int    `json:"port" yaml:"port"`
+	Db         string `json:"db" yaml:"db"`
+	Username   string `json:"username" yaml:"username"`
+	Password   string `json:"password" yaml:"password"`
+	Mechanism  string `json:"mechanism" yaml:"mechanism"`
+	RestoreIDs bool   `json:"restoreIDs" yaml:"restoreIDs"`
 }
 
 // Wiki is the wiki stuff.