Add a tag import command

1 files changed, 82 insertions, 2 deletions

diff --git a/main.go b/main.go
index 598c5c5..61f089a 100644
--- a/main.go
+++ b/main.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"bufio"
 	"bytes"
 	"context"
 	"crypto/sha1"
@@ -38,8 +39,8 @@ var (
 
 func openDB(directory string) error {
 	var err error
-	db, err = sql.Open("sqlite3",
-		"file:"+filepath.Join(directory, "gallery.db?_foreign_keys=1"))
+	db, err = sql.Open("sqlite3", "file:"+filepath.Join(directory,
+		"gallery.db?_foreign_keys=1&_busy_timeout=1000"))
 	galleryDirectory = directory
 	return err
 }
@@ -279,6 +280,7 @@ func cmdRun(args []string) error {
 	// but having an elementary level of security doesn't hurt either.
 	staticHandler = http.FileServer(http.Dir("public"))
 
+	// TODO: Make sure the database handle isn't used concurrently.
 	http.HandleFunc("/", handleRequest)
 	http.HandleFunc("/image/", handleImages)
 	http.HandleFunc("/thumb/", handleThumbs)
@@ -534,6 +536,83 @@ func cmdSync(args []string) error {
 	return nil
 }
 
+// --- Tagging -----------------------------------------------------------------
+
+// cmdTag mass imports tags from data passed on stdin as a TSV
+// of SHA1 TAG WEIGHT entries.
+func cmdTag(args []string) error {
+	if len(args) < 2 || len(args) > 3 {
+		return errors.New("usage: GD SPACE [DESCRIPTION]")
+	}
+	if err := openDB(args[0]); err != nil {
+		return err
+	}
+
+	space := args[1]
+
+	var description sql.NullString
+	if len(args) >= 3 {
+		description = sql.NullString{String: args[2], Valid: true}
+	}
+
+	// Note that starting as a write transaction prevents deadlocks.
+	// Imports are rare, and just bulk load data, so this scope is fine.
+	tx, err := db.Begin()
+	if err != nil {
+		return err
+	}
+	defer tx.Rollback()
+
+	if _, err := tx.Exec(`INSERT OR IGNORE INTO tag_space(name, description)
+		VALUES (?, ?)`, space, description); err != nil {
+		return err
+	}
+
+	var spaceID int64
+	if err := tx.QueryRow(`SELECT id FROM tag_space WHERE name = ?`,
+		space).Scan(&spaceID); err != nil {
+		return err
+	}
+
+	// TODO: Prepare statements for tag/assignment updates.
+
+	scanner := bufio.NewScanner(os.Stdin)
+	for scanner.Scan() {
+		fields := strings.Split(scanner.Text(), "\t")
+		if len(fields) != 3 {
+			return errors.New("invalid input format")
+		}
+
+		sha1, tag := fields[0], fields[1]
+		weight, err := strconv.ParseFloat(fields[2], 64)
+		if err != nil {
+			return err
+		}
+
+		if _, err := tx.Exec(`INSERT OR IGNORE INTO tag(space, name)
+			VALUES (?, ?);`, spaceID, tag); err != nil {
+			return nil
+		}
+
+		var tagID int64
+		if err := tx.QueryRow(`SELECT id FROM tag WHERE space = ? AND name = ?`,
+			spaceID, tag).Scan(&tagID); err != nil {
+			return err
+		}
+
+		if _, err := tx.Exec(`INSERT INTO tag_assignment(sha1, tag, weight)
+			VALUES (?, ?, ?) ON CONFLICT DO UPDATE SET weight = ?`,
+			sha1, tagID, weight, weight); err != nil {
+			return err
+		}
+	}
+	if err := scanner.Err(); err != nil {
+		return err
+	}
+
+	return tx.Commit()
+}
+
 // --- Check -------------------------------------------------------------------
 
 // cmdCheck checks if all files tracked in the DB are accessible.
@@ -735,6 +814,7 @@ var commands = map[string]struct {
 	"init":      {cmdInit},
 	"run":       {cmdRun},
 	"import":    {cmdImport},
+	"tag":       {cmdTag},
 	"sync":      {cmdSync},
 	"check":     {cmdCheck},
 	"thumbnail": {cmdThumbnail},