mirror of
https://github.com/stashapp/stash.git
synced 2025-12-17 04:14:39 +03:00
Add partial import functionality (#812)
This commit is contained in:
WithoutPants
143
pkg/studio/import.go
Normal file
143
pkg/studio/import.go
Normal file
@@ -0,0 +1,143 @@
|
||||
package studio
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"errors"
|
||||
"fmt"
|
||||
|
||||
"github.com/stashapp/stash/pkg/manager/jsonschema"
|
||||
"github.com/stashapp/stash/pkg/models"
|
||||
"github.com/stashapp/stash/pkg/utils"
|
||||
)
|
||||
|
||||
var ErrParentStudioNotExist = errors.New("parent studio does not exist")
|
||||
|
||||
type Importer struct {
|
||||
ReaderWriter models.StudioReaderWriter
|
||||
Input jsonschema.Studio
|
||||
MissingRefBehaviour models.ImportMissingRefEnum
|
||||
|
||||
studio models.Studio
|
||||
imageData []byte
|
||||
}
|
||||
|
||||
func (i *Importer) PreImport() error {
|
||||
checksum := utils.MD5FromString(i.Input.Name)
|
||||
|
||||
i.studio = models.Studio{
|
||||
Checksum: checksum,
|
||||
Name: sql.NullString{String: i.Input.Name, Valid: true},
|
||||
URL: sql.NullString{String: i.Input.URL, Valid: true},
|
||||
CreatedAt: models.SQLiteTimestamp{Timestamp: i.Input.CreatedAt.GetTime()},
|
||||
UpdatedAt: models.SQLiteTimestamp{Timestamp: i.Input.UpdatedAt.GetTime()},
|
||||
}
|
||||
|
||||
if err := i.populateParentStudio(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
var err error
|
||||
if len(i.Input.Image) > 0 {
|
||||
_, i.imageData, err = utils.ProcessBase64Image(i.Input.Image)
|
||||
if err != nil {
|
||||
return fmt.Errorf("invalid image: %s", err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (i *Importer) populateParentStudio() error {
|
||||
if i.Input.ParentStudio != "" {
|
||||
studio, err := i.ReaderWriter.FindByName(i.Input.ParentStudio, false)
|
||||
if err != nil {
|
||||
return fmt.Errorf("error finding studio by name: %s", err.Error())
|
||||
}
|
||||
|
||||
if studio == nil {
|
||||
if i.MissingRefBehaviour == models.ImportMissingRefEnumFail {
|
||||
return ErrParentStudioNotExist
|
||||
}
|
||||
|
||||
if i.MissingRefBehaviour == models.ImportMissingRefEnumIgnore {
|
||||
return nil
|
||||
}
|
||||
|
||||
if i.MissingRefBehaviour == models.ImportMissingRefEnumCreate {
|
||||
parentID, err := i.createParentStudio(i.Input.ParentStudio)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
i.studio.ParentID = sql.NullInt64{
|
||||
Int64: int64(parentID),
|
||||
Valid: true,
|
||||
}
|
||||
}
|
||||
} else {
|
||||
i.studio.ParentID = sql.NullInt64{Int64: int64(studio.ID), Valid: true}
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (i *Importer) createParentStudio(name string) (int, error) {
|
||||
newStudio := *models.NewStudio(name)
|
||||
|
||||
created, err := i.ReaderWriter.Create(newStudio)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
return created.ID, nil
|
||||
}
|
||||
|
||||
func (i *Importer) PostImport(id int) error {
|
||||
if len(i.imageData) > 0 {
|
||||
if err := i.ReaderWriter.UpdateStudioImage(id, i.imageData); err != nil {
|
||||
return fmt.Errorf("error setting studio image: %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.studio)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error creating studio: %s", err.Error())
|
||||
}
|
||||
|
||||
id := created.ID
|
||||
return &id, nil
|
||||
}
|
||||
|
||||
func (i *Importer) Update(id int) error {
|
||||
studio := i.studio
|
||||
studio.ID = id
|
||||
_, err := i.ReaderWriter.UpdateFull(studio)
|
||||
if err != nil {
|
||||
return fmt.Errorf("error updating existing studio: %s", err.Error())
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
263
pkg/studio/import_test.go
Normal file
263
pkg/studio/import_test.go
Normal file
@@ -0,0 +1,263 @@
|
||||
package studio
|
||||
|
||||
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 (
|
||||
studioNameErr = "studioNameErr"
|
||||
existingStudioName = "existingTagName"
|
||||
|
||||
existingStudioID = 100
|
||||
|
||||
existingParentStudioName = "existingParentStudioName"
|
||||
existingParentStudioErr = "existingParentStudioErr"
|
||||
missingParentStudioName = "existingParentStudioName"
|
||||
)
|
||||
|
||||
func TestImporterName(t *testing.T) {
|
||||
i := Importer{
|
||||
Input: jsonschema.Studio{
|
||||
Name: studioName,
|
||||
},
|
||||
}
|
||||
|
||||
assert.Equal(t, studioName, i.Name())
|
||||
}
|
||||
|
||||
func TestImporterPreImport(t *testing.T) {
|
||||
i := Importer{
|
||||
Input: jsonschema.Studio{
|
||||
Name: studioName,
|
||||
Image: invalidImage,
|
||||
},
|
||||
}
|
||||
|
||||
err := i.PreImport()
|
||||
|
||||
assert.NotNil(t, err)
|
||||
|
||||
i.Input.Image = image
|
||||
|
||||
err = i.PreImport()
|
||||
|
||||
assert.Nil(t, err)
|
||||
}
|
||||
|
||||
func TestImporterPreImportWithParent(t *testing.T) {
|
||||
readerWriter := &mocks.StudioReaderWriter{}
|
||||
|
||||
i := Importer{
|
||||
ReaderWriter: readerWriter,
|
||||
Input: jsonschema.Studio{
|
||||
Name: studioName,
|
||||
Image: image,
|
||||
ParentStudio: existingParentStudioName,
|
||||
},
|
||||
}
|
||||
|
||||
readerWriter.On("FindByName", existingParentStudioName, false).Return(&models.Studio{
|
||||
ID: existingStudioID,
|
||||
}, nil).Once()
|
||||
readerWriter.On("FindByName", existingParentStudioErr, false).Return(nil, errors.New("FindByName error")).Once()
|
||||
|
||||
err := i.PreImport()
|
||||
assert.Nil(t, err)
|
||||
assert.Equal(t, int64(existingStudioID), i.studio.ParentID.Int64)
|
||||
|
||||
i.Input.ParentStudio = existingParentStudioErr
|
||||
err = i.PreImport()
|
||||
assert.NotNil(t, err)
|
||||
|
||||
readerWriter.AssertExpectations(t)
|
||||
}
|
||||
|
||||
func TestImporterPreImportWithMissingParent(t *testing.T) {
|
||||
readerWriter := &mocks.StudioReaderWriter{}
|
||||
|
||||
i := Importer{
|
||||
ReaderWriter: readerWriter,
|
||||
Input: jsonschema.Studio{
|
||||
Name: studioName,
|
||||
Image: image,
|
||||
ParentStudio: missingParentStudioName,
|
||||
},
|
||||
MissingRefBehaviour: models.ImportMissingRefEnumFail,
|
||||
}
|
||||
|
||||
readerWriter.On("FindByName", missingParentStudioName, false).Return(nil, nil).Times(3)
|
||||
readerWriter.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.studio.ParentID.Int64)
|
||||
|
||||
readerWriter.AssertExpectations(t)
|
||||
}
|
||||
|
||||
func TestImporterPreImportWithMissingParentCreateErr(t *testing.T) {
|
||||
readerWriter := &mocks.StudioReaderWriter{}
|
||||
|
||||
i := Importer{
|
||||
ReaderWriter: readerWriter,
|
||||
Input: jsonschema.Studio{
|
||||
Name: studioName,
|
||||
Image: image,
|
||||
ParentStudio: missingParentStudioName,
|
||||
},
|
||||
MissingRefBehaviour: models.ImportMissingRefEnumCreate,
|
||||
}
|
||||
|
||||
readerWriter.On("FindByName", missingParentStudioName, false).Return(nil, nil).Once()
|
||||
readerWriter.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.StudioReaderWriter{}
|
||||
|
||||
i := Importer{
|
||||
ReaderWriter: readerWriter,
|
||||
imageData: imageBytes,
|
||||
}
|
||||
|
||||
updateStudioImageErr := errors.New("UpdateStudioImage error")
|
||||
|
||||
readerWriter.On("UpdateStudioImage", studioID, imageBytes).Return(nil).Once()
|
||||
readerWriter.On("UpdateStudioImage", errImageID, imageBytes).Return(updateStudioImageErr).Once()
|
||||
|
||||
err := i.PostImport(studioID)
|
||||
assert.Nil(t, err)
|
||||
|
||||
err = i.PostImport(errImageID)
|
||||
assert.NotNil(t, err)
|
||||
|
||||
readerWriter.AssertExpectations(t)
|
||||
}
|
||||
|
||||
func TestImporterFindExistingID(t *testing.T) {
|
||||
readerWriter := &mocks.StudioReaderWriter{}
|
||||
|
||||
i := Importer{
|
||||
ReaderWriter: readerWriter,
|
||||
Input: jsonschema.Studio{
|
||||
Name: studioName,
|
||||
},
|
||||
}
|
||||
|
||||
errFindByName := errors.New("FindByName error")
|
||||
readerWriter.On("FindByName", studioName, false).Return(nil, nil).Once()
|
||||
readerWriter.On("FindByName", existingStudioName, false).Return(&models.Studio{
|
||||
ID: existingStudioID,
|
||||
}, nil).Once()
|
||||
readerWriter.On("FindByName", studioNameErr, false).Return(nil, errFindByName).Once()
|
||||
|
||||
id, err := i.FindExistingID()
|
||||
assert.Nil(t, id)
|
||||
assert.Nil(t, err)
|
||||
|
||||
i.Input.Name = existingStudioName
|
||||
id, err = i.FindExistingID()
|
||||
assert.Equal(t, existingStudioID, *id)
|
||||
assert.Nil(t, err)
|
||||
|
||||
i.Input.Name = studioNameErr
|
||||
id, err = i.FindExistingID()
|
||||
assert.Nil(t, id)
|
||||
assert.NotNil(t, err)
|
||||
|
||||
readerWriter.AssertExpectations(t)
|
||||
}
|
||||
|
||||
func TestCreate(t *testing.T) {
|
||||
readerWriter := &mocks.StudioReaderWriter{}
|
||||
|
||||
studio := models.Studio{
|
||||
Name: modelstest.NullString(studioName),
|
||||
}
|
||||
|
||||
studioErr := models.Studio{
|
||||
Name: modelstest.NullString(studioNameErr),
|
||||
}
|
||||
|
||||
i := Importer{
|
||||
ReaderWriter: readerWriter,
|
||||
studio: studio,
|
||||
}
|
||||
|
||||
errCreate := errors.New("Create error")
|
||||
readerWriter.On("Create", studio).Return(&models.Studio{
|
||||
ID: studioID,
|
||||
}, nil).Once()
|
||||
readerWriter.On("Create", studioErr).Return(nil, errCreate).Once()
|
||||
|
||||
id, err := i.Create()
|
||||
assert.Equal(t, studioID, *id)
|
||||
assert.Nil(t, err)
|
||||
|
||||
i.studio = studioErr
|
||||
id, err = i.Create()
|
||||
assert.Nil(t, id)
|
||||
assert.NotNil(t, err)
|
||||
|
||||
readerWriter.AssertExpectations(t)
|
||||
}
|
||||
|
||||
func TestUpdate(t *testing.T) {
|
||||
readerWriter := &mocks.StudioReaderWriter{}
|
||||
|
||||
studio := models.Studio{
|
||||
Name: modelstest.NullString(studioName),
|
||||
}
|
||||
|
||||
studioErr := models.Studio{
|
||||
Name: modelstest.NullString(studioNameErr),
|
||||
}
|
||||
|
||||
i := Importer{
|
||||
ReaderWriter: readerWriter,
|
||||
studio: studio,
|
||||
}
|
||||
|
||||
errUpdate := errors.New("Update error")
|
||||
|
||||
// id needs to be set for the mock input
|
||||
studio.ID = studioID
|
||||
readerWriter.On("UpdateFull", studio).Return(nil, nil).Once()
|
||||
|
||||
err := i.Update(studioID)
|
||||
assert.Nil(t, err)
|
||||
|
||||
i.studio = studioErr
|
||||
|
||||
// need to set id separately
|
||||
studioErr.ID = errImageID
|
||||
readerWriter.On("UpdateFull", studioErr).Return(nil, errUpdate).Once()
|
||||
|
||||
err = i.Update(errImageID)
|
||||
assert.NotNil(t, err)
|
||||
|
||||
readerWriter.AssertExpectations(t)
|
||||
}
|
||||
Reference in New Issue
Block a user