Store image dimensions in DB

1 files changed, 47 insertions, 4 deletions

diff --git a/main.go b/main.go
index 8c154e8..4cad26f 100644
--- a/main.go
+++ b/main.go
@@ -275,6 +275,12 @@ func handleAPIBrowse(w http.ResponseWriter, r *http.Request) {
 
 // - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
 
+func getImageDimensions(sha1 string) (w int64, h int64, err error) {
+	err = db.QueryRow(`SELECT width, height FROM image WHERE sha1 = ?`,
+		sha1).Scan(&w, &h)
+	return
+}
+
 func getImagePaths(sha1 string) (paths []string, err error) {
 	rows, err := db.Query(`WITH RECURSIVE paths(parent, path) AS (
 		SELECT parent, name AS path FROM entry WHERE sha1 = ?
@@ -338,12 +344,19 @@ func handleAPIInfo(w http.ResponseWriter, r *http.Request) {
 	}
 
 	var result struct {
-		Paths []string                      `json:"paths"`
-		Tags  map[string]map[string]float32 `json:"tags"`
+		Width  int64                         `json:"width"`
+		Height int64                         `json:"height"`
+		Paths  []string                      `json:"paths"`
+		Tags   map[string]map[string]float32 `json:"tags"`
 		// TODO: Maybe add perceptual hash collisions.
 	}
 
 	var err error
+	result.Width, result.Height, err = getImageDimensions(params.SHA1)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
 	result.Paths, err = getImagePaths(params.SHA1)
 	if err != nil {
 		http.Error(w, err.Error(), http.StatusInternalServerError)
@@ -489,6 +502,31 @@ func isImage(path string) (bool, error) {
 	return bytes.HasPrefix(out, []byte("image/")), nil
 }
 
+func pingImage(path string) (int, int, error) {
+	cmd := exec.Command("identify", "-limit", "thread", "1", "-ping",
+		"-format", "%w %h", path+"[0]")
+
+	// XXX: Early returns may leak resources.
+	stdout, err := cmd.StdoutPipe()
+	if err != nil {
+		return 0, 0, err
+	}
+	if err := cmd.Start(); err != nil {
+		return 0, 0, err
+	}
+	out, err := io.ReadAll(stdout)
+	if err != nil {
+		return 0, 0, err
+	}
+	if err := cmd.Wait(); err != nil {
+		return 0, 0, err
+	}
+
+	var w, h int
+	_, err = fmt.Fscanf(bytes.NewReader(out), "%d %d", &w, &h)
+	return w, h, err
+}
+
 type importer struct {
 	dm      directoryManager
 	dmMutex sync.Mutex
@@ -512,6 +550,11 @@ func (i *importer) Import(path string) error {
 		return nil
 	}
 
+	width, height, err := pingImage(path)
+	if err != nil {
+		return err
+	}
+
 	f, err := os.Open(path)
 	if err != nil {
 		return err
@@ -551,8 +594,8 @@ func (i *importer) Import(path string) error {
 	}
 	defer tx.Rollback()
 
-	if _, err = tx.Exec(`INSERT INTO image(sha1) VALUES (?)
-		ON CONFLICT(sha1) DO NOTHING`, hexSHA1); err != nil {
+	if _, err = tx.Exec(`INSERT INTO image(sha1, width, height) VALUES (?, ?, ?)
+		ON CONFLICT(sha1) DO NOTHING`, hexSHA1, width, height); err != nil {
 		return err
 	}