Generate thumbnails too

t2bot · Dec 24, 2023 · 45e7363 · 45e7363
1 parent 4f8f5fe
commit 45e7363
Show file tree

Hide file tree

Showing 2 changed files with 129 additions and 6 deletions.
diff --git a/cmd/homeserver_offline_exporters/import_to_synapse/main.go b/cmd/homeserver_offline_exporters/import_to_synapse/main.go
@@ -1,20 +1,56 @@
 package main
 
 import (
+	"fmt"
 	"io"
+	"math"
 	"os"
 	"path"
+	"sync"
 
+	"github.com/disintegration/imaging"
 	"github.com/turt2live/matrix-media-repo/archival/v2archive"
 	"github.com/turt2live/matrix-media-repo/cmd/homeserver_offline_exporters/_common"
+	"github.com/turt2live/matrix-media-repo/common/config"
 	"github.com/turt2live/matrix-media-repo/common/rcontext"
 	"github.com/turt2live/matrix-media-repo/homeserver_interop/synapse"
+	"github.com/turt2live/matrix-media-repo/thumbnailing"
 	"github.com/turt2live/matrix-media-repo/util"
 )
 
+type thumbnailSize struct {
+	width  int
+	height int
+	method string
+}
+
+var synapseDefaultSizes = []thumbnailSize{
+	{32, 32, "crop"},
+	{96, 96, "crop"},
+	{320, 240, "scale"},
+	{640, 480, "scale"},
+	{800, 600, "scale"},
+}
+
 func main() {
 	cfg := _common.InitExportPsqlFlatFile("Synapse", "media_store_path")
 	ctx := rcontext.InitialNoConfig()
+	ctx.Config.Thumbnails = config.ThumbnailsConfig{
+		Types: []string{
+			"image/jpeg",
+			"image/jpg",
+			"image/png",
+			"image/apng",
+			"image/gif",
+			"image/heif",
+			"image/heic",
+			"image/webp",
+			"image/bmp",
+			"image/tiff",
+		},
+		MaxPixels:      32000000, // 32M
+		MaxSourceBytes: 10485760, // 10MB
+	}
 
 	ctx.Log.Debug("Connecting to homeserver database...")
 	hsDb, err := synapse.OpenDatabase(cfg.ConnectionString)
@@ -26,6 +62,10 @@ func main() {
 		defer f.Close()
 		mxc := util.MxcUri(record.Origin, record.MediaId)
 
+		if record.SizeBytes > math.MaxInt32 {
+			ctx.Log.Warnf("%s is potentially too large for Synapse to handle. See https://github.com/matrix-org/synapse/issues/12023 for details.", mxc)
+		}
+
 		if ok, err := hsDb.HasMedia(record.MediaId); err != nil {
 			return err
 		} else if ok {
@@ -57,9 +97,80 @@ func main() {
 			return err
 		}
 
-		// TODO: Populate thumbnails
+		if !cfg.GenerateThumbnails {
+			return nil // we're done here
+		}
+
+		// Try to generate thumbnails concurrently
+		ctx.Log.Infof("Generating thumbnails for %s", mxc)
+		thumbWg := &sync.WaitGroup{}
+		dirLock := &sync.Mutex{}
+		thumbWg.Add(len(synapseDefaultSizes))
+		for _, size := range synapseDefaultSizes {
+			go func(s thumbnailSize) {
+				defer thumbWg.Done()
+
+				src, err := os.Open(filePath)
+				if err != nil {
+					ctx.Log.Warn("Error preparing thumbnail. ", s, err)
+					return
+				}
+				defer src.Close()
+
+				thumb, err := thumbnailing.GenerateThumbnail(src, record.ContentType, s.width, s.height, s.method, false, ctx)
+				if err != nil {
+					ctx.Log.Debug("Error generating thumbnail (you can probably ignore this). ", s, err)
+					return
+				}
+
+				img, err := imaging.Decode(thumb.Reader)
+				if err != nil {
+					ctx.Log.Warn("Error processing thumbnail. ", s, err)
+					return
+				}
+
+				// Same as media, but different directory
+				// $exportPath/local_thumbnails/AA/BB/CCDD/$thumbFile
+				dirLock.Lock()
+				defer dirLock.Unlock()
+				thumbDir := path.Join(cfg.ExportPath, "local_thumbnails", record.MediaId[0:2], record.MediaId[2:4], record.MediaId[4:])
+				err = os.MkdirAll(thumbDir, 0655)
+				if err != nil {
+					ctx.Log.Warn("Error creating thumbnail directories. ", s, err)
+					return
+				}
+				thumbFile, err := os.Create(path.Join(thumbDir, fmt.Sprintf("%d-%d-image-png-%s", img.Bounds().Max.X, img.Bounds().Max.Y, s.method)))
+				if err != nil {
+					ctx.Log.Warn("Error creating thumbnail. ", s, err)
+					return
+				}
+				defer thumbFile.Close()
+				err = imaging.Encode(thumbFile, img, imaging.PNG)
+				if err != nil {
+					ctx.Log.Warn("Error writing thumbnail. ", s, err)
+					return
+				}
+
+				pos, err := thumbFile.Seek(0, io.SeekCurrent)
+				if err != nil {
+					ctx.Log.Warn("Error seeking within thumbnail. ", s, err)
+					return
+				}
+
+				if err = hsDb.InsertThumbnail(record.MediaId, img.Bounds().Max.X, img.Bounds().Max.Y, "image/png", s.method, pos); err != nil {
+					ctx.Log.Warn("Error recording thumbnail. ", s, err)
+					return
+				}
+			}(size)
+		}
+		thumbWg.Wait()
+
+		// Done
+		return nil
 	})
 	if err != nil {
-		panic(err)
+		ctx.Log.Fatal(err)
 	}
+
+	ctx.Log.Info("Done! If there's no warnings above, you're probably fine.")
 }
diff --git a/homeserver_interop/synapse/db.go b/homeserver_interop/synapse/db.go
@@ -11,6 +11,7 @@ import (
 const selectLocalMedia = "SELECT media_id, media_type, media_length, created_ts, upload_name, user_id, url_cache FROM local_media_repository;"
 const selectHasLocalMedia = "SELECT media_id FROM local_media_repository WHERE media_id = $1;"
 const insertLocalMedia = "INSERT INTO local_media_repository (media_id, media_type, media_length, created_ts, upload_name, user_id) VALUES ($1, $2, $3, $4, $5, $6);"
+const insertLocalThumbnail = "INSERT INTO local_media_repository_thumbnails (media_id, thumbnail_width, thumbnail_height, thumbnail_type, thumbnail_method, thumbnail_length) VALUES ($1, $2, $3, $4, $5, $6);"
 
 type LocalMedia struct {
 	homeserver_interop.ImportDbMedia
@@ -30,9 +31,10 @@ type SynDatabase struct {
 }
 
 type statements struct {
-	selectLocalMedia    *sql.Stmt
-	selectHasLocalMedia *sql.Stmt
-	insertLocalMedia    *sql.Stmt
+	selectLocalMedia     *sql.Stmt
+	selectHasLocalMedia  *sql.Stmt
+	insertLocalMedia     *sql.Stmt
+	insertLocalThumbnail *sql.Stmt
 }
 
 func OpenDatabase(connectionString string) (*SynDatabase, error) {
@@ -52,10 +54,18 @@ func OpenDatabase(connectionString string) (*SynDatabase, error) {
 	if d.statements.insertLocalMedia, err = d.db.Prepare(insertLocalMedia); err != nil {
 		return nil, err
 	}
+	if d.statements.insertLocalThumbnail, err = d.db.Prepare(insertLocalThumbnail); err != nil {
+		return nil, err
+	}
 
 	return &d, nil
 }
 
+func (d *SynDatabase) InsertThumbnail(mediaId string, width int, height int, contentType string, method string, sizeBytes int64) error {
+	_, err := d.statements.insertLocalThumbnail.Exec(mediaId, width, height, contentType, method, sizeBytes)
+	return err
+}
+
 func (d *SynDatabase) InsertMedia(mediaId string, contentType string, sizeBytes int64, createdTs int64, uploadName string, userId string) error {
 	_, err := d.statements.insertLocalMedia.Exec(mediaId, contentType, sizeBytes, createdTs, uploadName, userId)
 	return err
@@ -65,7 +75,9 @@ func (d *SynDatabase) HasMedia(mediaId string) (bool, error) {
 	row := d.statements.selectHasLocalMedia.QueryRow(mediaId)
 
 	val := ""
-	if err := row.Scan(&val); err != nil {
+	if err := row.Scan(&val); errors.Is(err, sql.ErrNoRows) {
+		return false, nil
+	} else if err != nil {
 		return false, err
 	}
 	return val == mediaId, nil