Images section (#813)

* Add new configuration options
* Refactor scan/clean
* Schema changes
* Add details to galleries
* Remove redundant code
* Refine thumbnail generation
* Gallery overhaul
* Don't allow modifying zip gallery images
* Show gallery card overlays
* Hide zoom slider when not in grid mode

pkg/manager/task_scan.go

@@ -1,17 +1,24 @@
 package manager
 
 import (
+	"archive/zip"
 	"context"
 	"database/sql"
+	"os"
 	"path/filepath"
 	"strconv"
 	"strings"
 	"sync"
 	"time"
 
+	"github.com/facebookgo/symwalk"
+	"github.com/jmoiron/sqlx"
+
 	"github.com/stashapp/stash/pkg/database"
 	"github.com/stashapp/stash/pkg/ffmpeg"
+	"github.com/stashapp/stash/pkg/image"
 	"github.com/stashapp/stash/pkg/logger"
+	"github.com/stashapp/stash/pkg/manager/config"
 	"github.com/stashapp/stash/pkg/models"
 	"github.com/stashapp/stash/pkg/utils"
 )
@@ -21,13 +28,17 @@ type ScanTask struct {
 	UseFileMetadata     bool
 	calculateMD5        bool
 	fileNamingAlgorithm models.HashAlgorithm
+
+	zipGallery *models.Gallery
 }
 
 func (t *ScanTask) Start(wg *sync.WaitGroup) {
 	if isGallery(t.FilePath) {
 		t.scanGallery()
-	} else {
+	} else if isVideo(t.FilePath) {
 		t.scanScene()
+	} else if isImage(t.FilePath) {
+		t.scanImage()
 	}
 
 	wg.Done()
@@ -39,6 +50,20 @@ func (t *ScanTask) scanGallery() {
 
 	if gallery != nil {
 		// We already have this item in the database, keep going
+
+		// scan the zip files if the gallery has no images
+		iqb := models.NewImageQueryBuilder()
+		images, err := iqb.CountByGalleryID(gallery.ID)
+		if err != nil {
+			logger.Errorf("error getting images for zip gallery %s: %s", t.FilePath, err.Error())
+		}
+
+		if images == 0 {
+			t.scanZipImages(gallery)
+		} else {
+			// in case thumbnails have been deleted, regenerate them
+			t.regenerateZipImages(gallery)
+		}
 		return
 	}
 
@@ -57,27 +82,34 @@ func (t *ScanTask) scanGallery() {
 	tx := database.DB.MustBeginTx(ctx, nil)
 	gallery, _ = qb.FindByChecksum(checksum, tx)
 	if gallery != nil {
-		exists, _ := utils.FileExists(gallery.Path)
+		exists, _ := utils.FileExists(gallery.Path.String)
 		if exists {
-			logger.Infof("%s already exists.  Duplicate of %s ", t.FilePath, gallery.Path)
+			logger.Infof("%s already exists.  Duplicate of %s ", t.FilePath, gallery.Path.String)
 		} else {
 
 			logger.Infof("%s already exists.  Updating path...", t.FilePath)
-			gallery.Path = t.FilePath
-			_, err = qb.Update(*gallery, tx)
+			gallery.Path = sql.NullString{
+				String: t.FilePath,
+				Valid:  true,
+			}
+			gallery, err = qb.Update(*gallery, tx)
 		}
 	} else {
 		currentTime := time.Now()
 
 		newGallery := models.Gallery{
-			Checksum:  checksum,
-			Path:      t.FilePath,
+			Checksum: checksum,
+			Zip:      true,
+			Path: sql.NullString{
+				String: t.FilePath,
+				Valid:  true,
+			},
 			CreatedAt: models.SQLiteTimestamp{Timestamp: currentTime},
 			UpdatedAt: models.SQLiteTimestamp{Timestamp: currentTime},
 		}
 
 		// don't create gallery if it has no images
-		if newGallery.CountFiles() > 0 {
+		if countImagesInZip(t.FilePath) > 0 {
 			// only warn when creating the gallery
 			ok, err := utils.IsZipFileUncompressed(t.FilePath)
 			if err == nil && !ok {
@@ -85,15 +117,25 @@ func (t *ScanTask) scanGallery() {
 			}
 
 			logger.Infof("%s doesn't exist.  Creating new item...", t.FilePath)
-			_, err = qb.Create(newGallery, tx)
+			gallery, err = qb.Create(newGallery, tx)
 		}
 	}
 
 	if err != nil {
 		logger.Error(err.Error())
-		_ = tx.Rollback()
-	} else if err := tx.Commit(); err != nil {
+		tx.Rollback()
+		return
+	}
+
+	err = tx.Commit()
+	if err != nil {
 		logger.Error(err.Error())
+		return
+	}
+
+	// if the gallery has no associated images, then scan the zip for images
+	if gallery != nil {
+		t.scanZipImages(gallery)
 	}
 }
 
@@ -109,19 +151,24 @@ func (t *ScanTask) associateGallery(wg *sync.WaitGroup) {
 		return
 	}
 
-	if !gallery.SceneID.Valid { // gallery has no SceneID
+	// gallery has no SceneID
+	if !gallery.SceneID.Valid {
 		basename := strings.TrimSuffix(t.FilePath, filepath.Ext(t.FilePath))
 		var relatedFiles []string
-		for _, ext := range extensionsToScan { // make a list of media files that can be related to the gallery
+		vExt := config.GetVideoExtensions()
+		// make a list of media files that can be related to the gallery
+		for _, ext := range vExt {
 			related := basename + "." + ext
-			if !isGallery(related) { //exclude gallery extensions from the related files
+			// exclude gallery extensions from the related files
+			if !isGallery(related) {
 				relatedFiles = append(relatedFiles, related)
 			}
 		}
 		for _, scenePath := range relatedFiles {
 			qbScene := models.NewSceneQueryBuilder()
 			scene, _ := qbScene.FindByPath(scenePath)
-			if scene != nil { // found related Scene
+			// found related Scene
+			if scene != nil {
 				logger.Infof("associate: Gallery %s is related to scene: %d", t.FilePath, scene.ID)
 
 				gallery.SceneID.Int64 = int64(scene.ID)
@@ -138,12 +185,11 @@ func (t *ScanTask) associateGallery(wg *sync.WaitGroup) {
 					logger.Error(err.Error())
 				}
 
-				break // since a gallery can have only one related scene
+				// since a gallery can have only one related scene
 				// only first found is associated
+				break
 			}
-
 		}
-
 	}
 	wg.Done()
 }
@@ -371,6 +417,188 @@ func (t *ScanTask) makeScreenshots(probeResult *ffmpeg.VideoFile, checksum strin
 	}
 }
 
+func (t *ScanTask) scanZipImages(zipGallery *models.Gallery) {
+	err := walkGalleryZip(zipGallery.Path.String, func(file *zip.File) error {
+		// copy this task and change the filename
+		subTask := *t
+
+		// filepath is the zip file and the internal file name, separated by a null byte
+		subTask.FilePath = image.ZipFilename(zipGallery.Path.String, file.Name)
+		subTask.zipGallery = zipGallery
+
+		// run the subtask and wait for it to complete
+		var wg sync.WaitGroup
+		wg.Add(1)
+		subTask.Start(&wg)
+		return nil
+	})
+	if err != nil {
+		logger.Warnf("failed to scan zip file images for %s: %s", zipGallery.Path.String, err.Error())
+	}
+}
+
+func (t *ScanTask) regenerateZipImages(zipGallery *models.Gallery) {
+	iqb := models.NewImageQueryBuilder()
+
+	images, err := iqb.FindByGalleryID(zipGallery.ID)
+	if err != nil {
+		logger.Warnf("failed to find gallery images: %s", err.Error())
+		return
+	}
+
+	for _, img := range images {
+		t.generateThumbnail(img)
+	}
+}
+
+func (t *ScanTask) scanImage() {
+	qb := models.NewImageQueryBuilder()
+	i, _ := qb.FindByPath(t.FilePath)
+	if i != nil {
+		// We already have this item in the database
+		// check for thumbnails
+		t.generateThumbnail(i)
+
+		return
+	}
+
+	// Ignore directories.
+	if isDir, _ := utils.DirExists(t.FilePath); isDir {
+		return
+	}
+
+	var checksum string
+
+	logger.Infof("%s not found.  Calculating checksum...", t.FilePath)
+	checksum, err := t.calculateImageChecksum()
+	if err != nil {
+		logger.Errorf("error calculating checksum for %s: %s", t.FilePath, err.Error())
+		return
+	}
+
+	// check for scene by checksum and oshash - MD5 should be
+	// redundant, but check both
+	i, _ = qb.FindByChecksum(checksum)
+
+	ctx := context.TODO()
+	tx := database.DB.MustBeginTx(ctx, nil)
+	if i != nil {
+		exists := image.FileExists(i.Path)
+		if exists {
+			logger.Infof("%s already exists.  Duplicate of %s ", image.PathDisplayName(t.FilePath), image.PathDisplayName(i.Path))
+		} else {
+			logger.Infof("%s already exists.  Updating path...", image.PathDisplayName(t.FilePath))
+			imagePartial := models.ImagePartial{
+				ID:   i.ID,
+				Path: &t.FilePath,
+			}
+			_, err = qb.Update(imagePartial, tx)
+		}
+	} else {
+		logger.Infof("%s doesn't exist.  Creating new item...", image.PathDisplayName(t.FilePath))
+		currentTime := time.Now()
+		newImage := models.Image{
+			Checksum:  checksum,
+			Path:      t.FilePath,
+			CreatedAt: models.SQLiteTimestamp{Timestamp: currentTime},
+			UpdatedAt: models.SQLiteTimestamp{Timestamp: currentTime},
+		}
+		err = image.SetFileDetails(&newImage)
+		if err == nil {
+			i, err = qb.Create(newImage, tx)
+		}
+	}
+
+	if err == nil {
+		jqb := models.NewJoinsQueryBuilder()
+		if t.zipGallery != nil {
+			// associate with gallery
+			_, err = jqb.AddImageGallery(i.ID, t.zipGallery.ID, tx)
+		} else if config.GetCreateGalleriesFromFolders() {
+			// create gallery from folder or associate with existing gallery
+			logger.Infof("Associating image %s with folder gallery", i.Path)
+			err = t.associateImageWithFolderGallery(i.ID, tx)
+		}
+	}
+
+	if err != nil {
+		logger.Error(err.Error())
+		_ = tx.Rollback()
+		return
+	} else if err := tx.Commit(); err != nil {
+		logger.Error(err.Error())
+		return
+	}
+
+	t.generateThumbnail(i)
+}
+
+func (t *ScanTask) associateImageWithFolderGallery(imageID int, tx *sqlx.Tx) error {
+	// find a gallery with the path specified
+	path := filepath.Dir(t.FilePath)
+	gqb := models.NewGalleryQueryBuilder()
+	jqb := models.NewJoinsQueryBuilder()
+	g, err := gqb.FindByPath(path)
+	if err != nil {
+		return err
+	}
+
+	if g == nil {
+		checksum := utils.MD5FromString(path)
+
+		// create the gallery
+		currentTime := time.Now()
+
+		newGallery := models.Gallery{
+			Checksum: checksum,
+			Path: sql.NullString{
+				String: path,
+				Valid:  true,
+			},
+			CreatedAt: models.SQLiteTimestamp{Timestamp: currentTime},
+			UpdatedAt: models.SQLiteTimestamp{Timestamp: currentTime},
+		}
+
+		logger.Infof("Creating gallery for folder %s", path)
+		g, err = gqb.Create(newGallery, tx)
+		if err != nil {
+			return err
+		}
+	}
+
+	// associate image with gallery
+	_, err = jqb.AddImageGallery(imageID, g.ID, tx)
+	return err
+}
+
+func (t *ScanTask) generateThumbnail(i *models.Image) {
+	thumbPath := GetInstance().Paths.Generated.GetThumbnailPath(i.Checksum, models.DefaultGthumbWidth)
+	exists, _ := utils.FileExists(thumbPath)
+	if exists {
+		logger.Debug("Thumbnail already exists for this path... skipping")
+		return
+	}
+
+	srcImage, err := image.GetSourceImage(i)
+	if err != nil {
+		logger.Errorf("error reading image %s: %s", i.Path, err.Error())
+		return
+	}
+
+	if image.ThumbnailNeeded(srcImage, models.DefaultGthumbWidth) {
+		data, err := image.GetThumbnail(srcImage, models.DefaultGthumbWidth)
+		if err != nil {
+			logger.Errorf("error getting thumbnail for image %s: %s", i.Path, err.Error())
+			return
+		}
+
+		err = utils.WriteFile(thumbPath, data)
+		if err != nil {
+			logger.Errorf("error writing thumbnail for image %s: %s", i.Path, err)
+		}
+	}
+}
+
 func (t *ScanTask) calculateChecksum() (string, error) {
 	logger.Infof("Calculating checksum for %s...", t.FilePath)
 	checksum, err := utils.MD5FromFilePath(t.FilePath)
@@ -381,19 +609,67 @@ func (t *ScanTask) calculateChecksum() (string, error) {
 	return checksum, nil
 }
 
+func (t *ScanTask) calculateImageChecksum() (string, error) {
+	logger.Infof("Calculating checksum for %s...", image.PathDisplayName(t.FilePath))
+	// uses image.CalculateMD5 to read files in zips
+	checksum, err := image.CalculateMD5(t.FilePath)
+	if err != nil {
+		return "", err
+	}
+	logger.Debugf("Checksum calculated: %s", checksum)
+	return checksum, nil
+}
+
 func (t *ScanTask) doesPathExist() bool {
-	if filepath.Ext(t.FilePath) == ".zip" {
+	vidExt := config.GetVideoExtensions()
+	imgExt := config.GetImageExtensions()
+	gExt := config.GetGalleryExtensions()
+
+	if matchExtension(t.FilePath, gExt) {
 		qb := models.NewGalleryQueryBuilder()
 		gallery, _ := qb.FindByPath(t.FilePath)
 		if gallery != nil {
 			return true
 		}
-	} else {
+	} else if matchExtension(t.FilePath, vidExt) {
 		qb := models.NewSceneQueryBuilder()
 		scene, _ := qb.FindByPath(t.FilePath)
 		if scene != nil {
 			return true
 		}
+	} else if matchExtension(t.FilePath, imgExt) {
+		qb := models.NewImageQueryBuilder()
+		i, _ := qb.FindByPath(t.FilePath)
+		if i != nil {
+			return true
+		}
 	}
+
 	return false
 }
+
+func walkFilesToScan(s *models.StashConfig, f filepath.WalkFunc) error {
+	vidExt := config.GetVideoExtensions()
+	imgExt := config.GetImageExtensions()
+	gExt := config.GetGalleryExtensions()
+	excludeVid := config.GetExcludes()
+	excludeImg := config.GetImageExcludes()
+
+	return symwalk.Walk(s.Path, func(path string, info os.FileInfo, err error) error {
+		if info.IsDir() {
+			return nil
+		}
+
+		if !s.ExcludeVideo && matchExtension(path, vidExt) && !matchFile(path, excludeVid) {
+			return f(path, info, err)
+		}
+
+		if !s.ExcludeImage {
+			if (matchExtension(path, imgExt) || matchExtension(path, gExt)) && !matchFile(path, excludeImg) {
+				return f(path, info, err)
+			}
+		}
+
+		return nil
+	})
+}