Add partial import functionality (#812)

2025-12-17 04:14:39 +03:00 · 2020-09-20 18:36:02 +10:00
parent 7a45943e8e
commit 8866670e53
56 changed files with 5030 additions and 624 deletions
--- a/pkg/movie/import.go
+++ b/pkg/movie/import.go
@@ -0,0 +1,166 @@
+package movie
+
+import (
+	"database/sql"
+	"fmt"
+
+	"github.com/stashapp/stash/pkg/manager/jsonschema"
+	"github.com/stashapp/stash/pkg/models"
+	"github.com/stashapp/stash/pkg/utils"
+)
+
+type Importer struct {
+	ReaderWriter        models.MovieReaderWriter
+	StudioWriter        models.StudioReaderWriter
+	Input               jsonschema.Movie
+	MissingRefBehaviour models.ImportMissingRefEnum
+
+	movie          models.Movie
+	frontImageData []byte
+	backImageData  []byte
+}
+
+func (i *Importer) PreImport() error {
+	i.movie = i.movieJSONToMovie(i.Input)
+
+	if err := i.populateStudio(); 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: %s", err.Error())
+		}
+	}
+	if len(i.Input.BackImage) > 0 {
+		_, i.backImageData, err = utils.ProcessBase64Image(i.Input.BackImage)
+		if err != nil {
+			return fmt.Errorf("invalid back_image: %s", err.Error())
+		}
+	}
+
+	return nil
+}
+
+func (i *Importer) movieJSONToMovie(movieJSON jsonschema.Movie) models.Movie {
+	checksum := utils.MD5FromString(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() error {
+	if i.Input.Studio != "" {
+		studio, err := i.StudioWriter.FindByName(i.Input.Studio, false)
+		if err != nil {
+			return fmt.Errorf("error finding studio by name: %s", err.Error())
+		}
+
+		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(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(name string) (int, error) {
+	newStudio := *models.NewStudio(name)
+
+	created, err := i.StudioWriter.Create(newStudio)
+	if err != nil {
+		return 0, err
+	}
+
+	return created.ID, nil
+}
+
+func (i *Importer) PostImport(id int) error {
+	if len(i.frontImageData) > 0 {
+		if err := i.ReaderWriter.UpdateMovieImages(id, i.frontImageData, i.backImageData); err != nil {
+			return fmt.Errorf("error setting movie images: %s", err.Error())
+		}
+	}
+
+	return nil
+}
+
+func (i *Importer) Name() string {
+	return i.Input.Name
+}
+
+func (i *Importer) FindExistingID() (*int, error) {
+	const nocase = false
+	existing, err := i.ReaderWriter.FindByName(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() (*int, error) {
+	created, err := i.ReaderWriter.Create(i.movie)
+	if err != nil {
+		return nil, fmt.Errorf("error creating movie: %s", err.Error())
+	}
+
+	id := created.ID
+	return &id, nil
+}
+
+func (i *Importer) Update(id int) error {
+	movie := i.movie
+	movie.ID = id
+	_, err := i.ReaderWriter.UpdateFull(movie)
+	if err != nil {
+		return fmt.Errorf("error updating existing movie: %s", err.Error())
+	}
+
+	return nil
+}
--- a/pkg/movie/import_test.go
+++ b/pkg/movie/import_test.go
@@ -0,0 +1,278 @@
+package movie
+
+import (
+	"errors"
+	"testing"
+
+	"github.com/stashapp/stash/pkg/manager/jsonschema"
+	"github.com/stashapp/stash/pkg/models"
+	"github.com/stashapp/stash/pkg/models/mocks"
+	"github.com/stashapp/stash/pkg/models/modelstest"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+)
+
+const invalidImage = "aW1hZ2VCeXRlcw&&"
+
+const (
+	movieNameErr      = "movieNameErr"
+	existingMovieName = "existingMovieName"
+
+	existingMovieID  = 100
+	existingStudioID = 101
+
+	existingStudioName = "existingStudioName"
+	existingStudioErr  = "existingStudioErr"
+	missingStudioName  = "existingStudioName"
+
+	errImageID = 3
+)
+
+func TestImporterName(t *testing.T) {
+	i := Importer{
+		Input: jsonschema.Movie{
+			Name: movieName,
+		},
+	}
+
+	assert.Equal(t, movieName, i.Name())
+}
+
+func TestImporterPreImport(t *testing.T) {
+	i := Importer{
+		Input: jsonschema.Movie{
+			Name:       movieName,
+			FrontImage: invalidImage,
+		},
+	}
+
+	err := i.PreImport()
+	assert.NotNil(t, err)
+
+	i.Input.FrontImage = frontImage
+	i.Input.BackImage = invalidImage
+
+	err = i.PreImport()
+	assert.NotNil(t, err)
+
+	i.Input.BackImage = ""
+
+	err = i.PreImport()
+	assert.Nil(t, err)
+
+	i.Input.BackImage = backImage
+
+	err = i.PreImport()
+	assert.Nil(t, err)
+}
+
+func TestImporterPreImportWithStudio(t *testing.T) {
+	studioReaderWriter := &mocks.StudioReaderWriter{}
+
+	i := Importer{
+		StudioWriter: studioReaderWriter,
+		Input: jsonschema.Movie{
+			Name:       movieName,
+			FrontImage: frontImage,
+			Studio:     existingStudioName,
+			Rating:     5,
+			Duration:   10,
+		},
+	}
+
+	studioReaderWriter.On("FindByName", existingStudioName, false).Return(&models.Studio{
+		ID: existingStudioID,
+	}, nil).Once()
+	studioReaderWriter.On("FindByName", existingStudioErr, false).Return(nil, errors.New("FindByName error")).Once()
+
+	err := i.PreImport()
+	assert.Nil(t, err)
+	assert.Equal(t, int64(existingStudioID), i.movie.StudioID.Int64)
+
+	i.Input.Studio = existingStudioErr
+	err = i.PreImport()
+	assert.NotNil(t, err)
+
+	studioReaderWriter.AssertExpectations(t)
+}
+
+func TestImporterPreImportWithMissingStudio(t *testing.T) {
+	studioReaderWriter := &mocks.StudioReaderWriter{}
+
+	i := Importer{
+		StudioWriter: studioReaderWriter,
+		Input: jsonschema.Movie{
+			Name:       movieName,
+			FrontImage: frontImage,
+			Studio:     missingStudioName,
+		},
+		MissingRefBehaviour: models.ImportMissingRefEnumFail,
+	}
+
+	studioReaderWriter.On("FindByName", missingStudioName, false).Return(nil, nil).Times(3)
+	studioReaderWriter.On("Create", mock.AnythingOfType("models.Studio")).Return(&models.Studio{
+		ID: existingStudioID,
+	}, nil)
+
+	err := i.PreImport()
+	assert.NotNil(t, err)
+
+	i.MissingRefBehaviour = models.ImportMissingRefEnumIgnore
+	err = i.PreImport()
+	assert.Nil(t, err)
+
+	i.MissingRefBehaviour = models.ImportMissingRefEnumCreate
+	err = i.PreImport()
+	assert.Nil(t, err)
+	assert.Equal(t, int64(existingStudioID), i.movie.StudioID.Int64)
+
+	studioReaderWriter.AssertExpectations(t)
+}
+
+func TestImporterPreImportWithMissingStudioCreateErr(t *testing.T) {
+	studioReaderWriter := &mocks.StudioReaderWriter{}
+
+	i := Importer{
+		StudioWriter: studioReaderWriter,
+		Input: jsonschema.Movie{
+			Name:       movieName,
+			FrontImage: frontImage,
+			Studio:     missingStudioName,
+		},
+		MissingRefBehaviour: models.ImportMissingRefEnumCreate,
+	}
+
+	studioReaderWriter.On("FindByName", missingStudioName, false).Return(nil, nil).Once()
+	studioReaderWriter.On("Create", mock.AnythingOfType("models.Studio")).Return(nil, errors.New("Create error"))
+
+	err := i.PreImport()
+	assert.NotNil(t, err)
+}
+
+func TestImporterPostImport(t *testing.T) {
+	readerWriter := &mocks.MovieReaderWriter{}
+
+	i := Importer{
+		ReaderWriter:   readerWriter,
+		frontImageData: frontImageBytes,
+		backImageData:  backImageBytes,
+	}
+
+	updateMovieImageErr := errors.New("UpdateMovieImage error")
+
+	readerWriter.On("UpdateMovieImages", movieID, frontImageBytes, backImageBytes).Return(nil).Once()
+	readerWriter.On("UpdateMovieImages", errImageID, frontImageBytes, backImageBytes).Return(updateMovieImageErr).Once()
+
+	err := i.PostImport(movieID)
+	assert.Nil(t, err)
+
+	err = i.PostImport(errImageID)
+	assert.NotNil(t, err)
+
+	readerWriter.AssertExpectations(t)
+}
+
+func TestImporterFindExistingID(t *testing.T) {
+	readerWriter := &mocks.MovieReaderWriter{}
+
+	i := Importer{
+		ReaderWriter: readerWriter,
+		Input: jsonschema.Movie{
+			Name: movieName,
+		},
+	}
+
+	errFindByName := errors.New("FindByName error")
+	readerWriter.On("FindByName", movieName, false).Return(nil, nil).Once()
+	readerWriter.On("FindByName", existingMovieName, false).Return(&models.Movie{
+		ID: existingMovieID,
+	}, nil).Once()
+	readerWriter.On("FindByName", movieNameErr, false).Return(nil, errFindByName).Once()
+
+	id, err := i.FindExistingID()
+	assert.Nil(t, id)
+	assert.Nil(t, err)
+
+	i.Input.Name = existingMovieName
+	id, err = i.FindExistingID()
+	assert.Equal(t, existingMovieID, *id)
+	assert.Nil(t, err)
+
+	i.Input.Name = movieNameErr
+	id, err = i.FindExistingID()
+	assert.Nil(t, id)
+	assert.NotNil(t, err)
+
+	readerWriter.AssertExpectations(t)
+}
+
+func TestCreate(t *testing.T) {
+	readerWriter := &mocks.MovieReaderWriter{}
+
+	movie := models.Movie{
+		Name: modelstest.NullString(movieName),
+	}
+
+	movieErr := models.Movie{
+		Name: modelstest.NullString(movieNameErr),
+	}
+
+	i := Importer{
+		ReaderWriter: readerWriter,
+		movie:        movie,
+	}
+
+	errCreate := errors.New("Create error")
+	readerWriter.On("Create", movie).Return(&models.Movie{
+		ID: movieID,
+	}, nil).Once()
+	readerWriter.On("Create", movieErr).Return(nil, errCreate).Once()
+
+	id, err := i.Create()
+	assert.Equal(t, movieID, *id)
+	assert.Nil(t, err)
+
+	i.movie = movieErr
+	id, err = i.Create()
+	assert.Nil(t, id)
+	assert.NotNil(t, err)
+
+	readerWriter.AssertExpectations(t)
+}
+
+func TestUpdate(t *testing.T) {
+	readerWriter := &mocks.MovieReaderWriter{}
+
+	movie := models.Movie{
+		Name: modelstest.NullString(movieName),
+	}
+
+	movieErr := models.Movie{
+		Name: modelstest.NullString(movieNameErr),
+	}
+
+	i := Importer{
+		ReaderWriter: readerWriter,
+		movie:        movie,
+	}
+
+	errUpdate := errors.New("Update error")
+
+	// id needs to be set for the mock input
+	movie.ID = movieID
+	readerWriter.On("UpdateFull", movie).Return(nil, nil).Once()
+
+	err := i.Update(movieID)
+	assert.Nil(t, err)
+
+	i.movie = movieErr
+
+	// need to set id separately
+	movieErr.ID = errImageID
+	readerWriter.On("UpdateFull", movieErr).Return(nil, errUpdate).Once()
+
+	err = i.Update(errImageID)
+	assert.NotNil(t, err)
+
+	readerWriter.AssertExpectations(t)
+}