api/database/mongodb/stories.go

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
	restoreIDs bool
}

func newStoryRepository(db *mgo.Database, restoreIDs bool) (repositories.StoryRepository, error) {
	collection := db.C("story.stories")

	err := collection.EnsureIndexKey("tags")
	if err != nil {
		return nil, err
	}
	err = collection.EnsureIndexKey("author")
	if err != nil {
		return nil, err
	}
	err = collection.EnsureIndexKey("updatedDate")
	if err != nil {
		return nil, err
	}
	err = collection.EnsureIndexKey("fictionalDate")
	if err != nil {
		return nil, err
	}
	err = collection.EnsureIndexKey("listed")
	if err != nil {
		return nil, err
	}

	return &storyRepository{
		stories:    collection,
		restoreIDs: restoreIDs,
		chapters:   db.C("story.chapters"),
		comments:   db.C("story.comments"),
	}, nil
}

func (r *storyRepository) Find(ctx context.Context, id string) (*models.Story, error) {
	story := new(models.Story)
	err := r.stories.FindId(id).One(story)
	if err != nil {
		return nil, err
	}

	return story, nil
}

func (r *storyRepository) List(ctx context.Context, filter models.StoryFilter) ([]*models.Story, error) {
	query := bson.M{}
	if filter.Author != nil {
		query["author"] = *filter.Author
	}
	if filter.Category != nil {
		query["category"] = *filter.Category
	}
	if filter.Open != nil {
		query["open"] = *filter.Open
	}
	if len(filter.Tags) > 0 {
		query["tags"] = bson.M{"$all": filter.Tags}
	}
	if filter.Unlisted != nil {
		query["listed"] = !*filter.Unlisted
	} else {
		query["listed"] = true
	}
	if !filter.EarliestFictionalDate.IsZero() && !filter.LatestFictionalDate.IsZero() {
		query["fictionalDate"] = bson.M{
			"$gte": filter.EarliestFictionalDate,
			"$lte": filter.LatestFictionalDate,
		}
	} else if !filter.EarliestFictionalDate.IsZero() {
		query["fictionalDate"] = bson.M{
			"$gte": filter.EarliestFictionalDate,
		}
	} else if !filter.LatestFictionalDate.IsZero() {
		query["fictionalDate"] = bson.M{
			"$lte": filter.LatestFictionalDate,
		}
	}

	stories := make([]*models.Story, 0, 32)
	err := r.stories.Find(query).Sort("-updatedDate").Limit(filter.Limit).All(&stories)
	if err != nil {
		if err == mgo.ErrNotFound {
			return stories, nil
		}

		return nil, err
	}

	return stories, nil
}

func (r *storyRepository) Insert(ctx context.Context, story models.Story) (*models.Story, error) {
	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 {
		return nil, err
	}

	return &story, nil
}

func (r *storyRepository) Update(ctx context.Context, story models.Story, update models.StoryUpdate) (*models.Story, error) {
	updateBson := bson.M{}
	if update.Name != nil {
		updateBson["name"] = *update.Name
		story.Name = *update.Name
	}
	if update.Open != nil {
		updateBson["open"] = *update.Open
		story.Open = *update.Open
	}
	if update.Category != nil {
		updateBson["category"] = *update.Category
		story.Category = *update.Category
	}
	if update.Author != nil {
		updateBson["author"] = *update.Author
		story.Author = *update.Author
	}
	if update.FictionalDate != nil {
		updateBson["fictionalDate"] = *update.FictionalDate
		story.FictionalDate = *update.FictionalDate
	}
	if update.UpdatedDate != nil {
		updateBson["updatedDate"] = *update.UpdatedDate
		story.UpdatedDate = *update.UpdatedDate
	}
	if update.Listed != nil {
		updateBson["listed"] = *update.Listed
		story.Listed = *update.Listed
	}
	if update.SortByFictionalDate != nil {
		updateBson["sortByFictionalDate"] = *update.SortByFictionalDate
		story.SortByFictionalDate = *update.SortByFictionalDate
	}

	err := r.stories.UpdateId(story.ID, bson.M{"$set": updateBson})
	if err != nil {
		return nil, err
	}

	return &story, nil
}

func (r *storyRepository) AddTag(ctx context.Context, story models.Story, tag models.Tag) error {
	ci, err := r.stories.UpdateAll(bson.M{"_id": story.ID}, bson.M{"$addToSet": bson.M{"tags": tag}})
	if err != nil {
		return err
	}
	if ci.Updated == 0 {
		return repositories.ErrTagExists
	}

	return nil
}

func (r *storyRepository) RemoveTag(ctx context.Context, story models.Story, tag models.Tag) error {
	ci, err := r.stories.UpdateAll(bson.M{"_id": story.ID}, bson.M{"$pull": bson.M{"tags": tag}})
	if err != nil {
		return err
	}
	if ci.Updated == 0 {
		return repositories.ErrTagDoesNotExist
	}

	return nil
}

func (r *storyRepository) Delete(ctx context.Context, story models.Story) error {
	err := r.stories.RemoveId(story.ID)
	if err != nil {
		return err
	}

	chapterIds := make([]string, 0, 8)
	err = r.chapters.Find(bson.M{"storyId": story.ID}).Distinct("_id", &chapterIds)
	if err != nil {
		log.Println("Failed to find chapterIds:", err)
		return nil
	}
	if len(chapterIds) > 0 {
		c1, err := r.chapters.RemoveAll(bson.M{"storyId": story.ID})
		if err != nil {
			log.Println("Failed to remove chapters:", err)
			return nil
		}

		c2, err := r.comments.RemoveAll(bson.M{"chapterId": bson.M{"$in": chapterIds}})
		if err != nil {
			log.Println("Failed to remove comments:", err)
			return nil
		}

		log.Printf("Removed story %s (%d chapters, %d comments)", story.ID, c1.Removed, c2.Removed)
	}

	return nil
}