stash/pkg/movie/import.go

package movie

import (
	"context"
	"database/sql"
	"fmt"

	"github.com/stashapp/stash/pkg/hash/md5"
	"github.com/stashapp/stash/pkg/models"
	"github.com/stashapp/stash/pkg/models/jsonschema"
	"github.com/stashapp/stash/pkg/studio"
	"github.com/stashapp/stash/pkg/utils"
)

type NameFinderCreatorUpdater interface {
	NameFinderCreator
	UpdateFull(ctx context.Context, updatedMovie models.Movie) (*models.Movie, error)
	UpdateImages(ctx context.Context, movieID int, frontImage []byte, backImage []byte) error
}

type Importer struct {
	ReaderWriter        NameFinderCreatorUpdater
	StudioWriter        studio.NameFinderCreator
	Input               jsonschema.Movie
	MissingRefBehaviour models.ImportMissingRefEnum

	movie          models.Movie
	frontImageData []byte
	backImageData  []byte
}

func (i *Importer) PreImport(ctx context.Context) error {
	i.movie = i.movieJSONToMovie(i.Input)

	if err := i.populateStudio(ctx); err != nil {
		return err
	}

	var err error
	if len(i.Input.FrontImage) > 0 {
		i.frontImageData, err = utils.ProcessBase64Image(i.Input.FrontImage)
		if err != nil {
			return fmt.Errorf("invalid front_image: %v", err)
		}
	}
	if len(i.Input.BackImage) > 0 {
		i.backImageData, err = utils.ProcessBase64Image(i.Input.BackImage)
		if err != nil {
			return fmt.Errorf("invalid back_image: %v", err)
		}
	}

	return nil
}

func (i *Importer) movieJSONToMovie(movieJSON jsonschema.Movie) models.Movie {
	checksum := md5.FromString(movieJSON.Name)

	newMovie := models.Movie{
		Checksum:  checksum,
		Name:      sql.NullString{String: movieJSON.Name, Valid: true},
		Aliases:   sql.NullString{String: movieJSON.Aliases, Valid: true},
		Date:      models.SQLiteDate{String: movieJSON.Date, Valid: true},
		Director:  sql.NullString{String: movieJSON.Director, Valid: true},
		Synopsis:  sql.NullString{String: movieJSON.Synopsis, Valid: true},
		URL:       sql.NullString{String: movieJSON.URL, Valid: true},
		CreatedAt: models.SQLiteTimestamp{Timestamp: movieJSON.CreatedAt.GetTime()},
		UpdatedAt: models.SQLiteTimestamp{Timestamp: movieJSON.UpdatedAt.GetTime()},
	}

	if movieJSON.Rating != 0 {
		newMovie.Rating = sql.NullInt64{Int64: int64(movieJSON.Rating), Valid: true}
	}

	if movieJSON.Duration != 0 {
		newMovie.Duration = sql.NullInt64{Int64: int64(movieJSON.Duration), Valid: true}
	}

	return newMovie
}

func (i *Importer) populateStudio(ctx context.Context) error {
	if i.Input.Studio != "" {
		studio, err := i.StudioWriter.FindByName(ctx, i.Input.Studio, false)
		if err != nil {
			return fmt.Errorf("error finding studio by name: %v", err)
		}

		if studio == nil {
			if i.MissingRefBehaviour == models.ImportMissingRefEnumFail {
				return fmt.Errorf("movie studio '%s' not found", i.Input.Studio)
			}

			if i.MissingRefBehaviour == models.ImportMissingRefEnumIgnore {
				return nil
			}

			if i.MissingRefBehaviour == models.ImportMissingRefEnumCreate {
				studioID, err := i.createStudio(ctx, i.Input.Studio)
				if err != nil {
					return err
				}
				i.movie.StudioID = sql.NullInt64{
					Int64: int64(studioID),
					Valid: true,
				}
			}
		} else {
			i.movie.StudioID = sql.NullInt64{Int64: int64(studio.ID), Valid: true}
		}
	}

	return nil
}

func (i *Importer) createStudio(ctx context.Context, name string) (int, error) {
	newStudio := *models.NewStudio(name)

	created, err := i.StudioWriter.Create(ctx, newStudio)
	if err != nil {
		return 0, err
	}

	return created.ID, nil
}

func (i *Importer) PostImport(ctx context.Context, id int) error {
	if len(i.frontImageData) > 0 {
		if err := i.ReaderWriter.UpdateImages(ctx, id, i.frontImageData, i.backImageData); err != nil {
			return fmt.Errorf("error setting movie images: %v", err)
		}
	}

	return nil
}

func (i *Importer) Name() string {
	return i.Input.Name
}

func (i *Importer) FindExistingID(ctx context.Context) (*int, error) {
	const nocase = false
	existing, err := i.ReaderWriter.FindByName(ctx, i.Name(), nocase)
	if err != nil {
		return nil, err
	}

	if existing != nil {
		id := existing.ID
		return &id, nil
	}

	return nil, nil
}

func (i *Importer) Create(ctx context.Context) (*int, error) {
	created, err := i.ReaderWriter.Create(ctx, i.movie)
	if err != nil {
		return nil, fmt.Errorf("error creating movie: %v", err)
	}

	id := created.ID
	return &id, nil
}

func (i *Importer) Update(ctx context.Context, id int) error {
	movie := i.movie
	movie.ID = id
	_, err := i.ReaderWriter.UpdateFull(ctx, movie)
	if err != nil {
		return fmt.Errorf("error updating existing movie: %v", err)
	}

	return nil
}