Add drive blob import tool

2026-08-15 20:47:32 -07:00
parent 92628c4830
commit 00239d027c
1 changed files with 242 additions and 0 deletions
@@ -0,0 +1,242 @@
+// Command importblobs copies directory media into the drive folder under sane names.
+package main
+
+import (
+	"bytes"
+	"context"
+	"flag"
+	"fmt"
+	"io"
+	"log"
+	"net/http"
+	"path"
+	"strings"
+	"sync"
+	"time"
+
+	"heliosian/internal/data"
+	"heliosian/internal/directory"
+	"google.golang.org/api/drive/v3"
+	"google.golang.org/api/googleapi"
+	"google.golang.org/api/option"
+)
+
+const folderMime = "application/vnd.google-apps.folder"
+
+var extensions = map[string]string{
+	"image/jpeg": ".jpg",
+	"image/png":  ".png",
+	"image/webp": ".webp",
+	"image/gif":  ".gif",
+	"audio/mpeg": ".mp3",
+	"audio/mp3":  ".mp3",
+	"audio/mp4":  ".m4a",
+	"audio/wav":  ".wav",
+}
+
+type task struct {
+	folderName string
+	base       string
+	url        string
+}
+
+func localPart(email string) string {
+	name, _, _ := strings.Cut(email, "@")
+	return name
+}
+
+func findSharedFolder(svc *drive.Service) string {
+	list, err := svc.Files.List().
+		Q("mimeType = '" + folderMime + "' and sharedWithMe = true and trashed = false").
+		Fields("files(id, name)").Do()
+	if err != nil {
+		log.Fatalf("[ERROR] list shared folders: %v", err)
+	}
+	if len(list.Files) != 1 {
+		for _, f := range list.Files {
+			log.Printf("candidate folder: %s (%s)", f.Name, f.Id)
+		}
+		log.Fatalf("[ERROR] expected exactly one folder shared with the service account, found %d; pass -folder", len(list.Files))
+	}
+	log.Printf("using folder %q (%s)", list.Files[0].Name, list.Files[0].Id)
+	return list.Files[0].Id
+}
+
+func ensureFolder(svc *drive.Service, parent, name string) string {
+	list, err := svc.Files.List().
+		Q(fmt.Sprintf("name = '%s' and '%s' in parents and mimeType = '%s' and trashed = false", name, parent, folderMime)).
+		Fields("files(id)").Do()
+	if err != nil {
+		log.Fatalf("[ERROR] find folder %s: %v", name, err)
+	}
+	if len(list.Files) > 0 {
+		return list.Files[0].Id
+	}
+	created, err := svc.Files.Create(&drive.File{Name: name, MimeType: folderMime, Parents: []string{parent}}).Fields("id").Do()
+	if err != nil {
+		log.Fatalf("[ERROR] create folder %s: %v", name, err)
+	}
+	return created.Id
+}
+
+func listBases(svc *drive.Service, folderID string) map[string]bool {
+	bases := map[string]bool{}
+	token := ""
+	for {
+		call := svc.Files.List().
+			Q(fmt.Sprintf("'%s' in parents and trashed = false", folderID)).
+			Fields("nextPageToken, files(name)").PageSize(1000)
+		if token != "" {
+			call = call.PageToken(token)
+		}
+		list, err := call.Do()
+		if err != nil {
+			log.Fatalf("[ERROR] list folder contents: %v", err)
+		}
+		for _, f := range list.Files {
+			bases[strings.TrimSuffix(f.Name, path.Ext(f.Name))] = true
+		}
+		if list.NextPageToken == "" {
+			return bases
+		}
+		token = list.NextPageToken
+	}
+}
+
+func fetch(client *http.Client, url string) ([]byte, string, error) {
+	resp, err := client.Get(url)
+	if err != nil {
+		return nil, "", err
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusOK {
+		return nil, "", fmt.Errorf("status %s", resp.Status)
+	}
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, "", err
+	}
+	contentType, _, _ := strings.Cut(resp.Header.Get("Content-Type"), ";")
+	return body, strings.TrimSpace(contentType), nil
+}
+
+func extension(contentType, url string) string {
+	if ext, ok := extensions[contentType]; ok {
+		return ext
+	}
+	if ext := path.Ext(strings.SplitN(path.Base(url), "?", 2)[0]); ext != "" && len(ext) <= 5 {
+		return ext
+	}
+	return ".bin"
+}
+
+func main() {
+	sheetID := flag.String("sheet", "", "directory spreadsheet id")
+	folderID := flag.String("folder", "", "drive folder id (default: the folder shared with the service account)")
+	flag.Parse()
+	if *sheetID == "" {
+		log.Fatal("[ERROR] -sheet is required")
+	}
+
+	source, err := data.NewSheet(map[string]string{"directory": *sheetID})
+	if err != nil {
+		log.Fatalf("[ERROR] sheet source: %v", err)
+	}
+	model, err := directory.LoadModel(source)
+	if err != nil {
+		log.Fatalf("[ERROR] load model: %v", err)
+	}
+
+	svc, err := drive.NewService(context.Background(),
+		option.WithCredentialsFile(data.KeyFile),
+		option.WithScopes(drive.DriveScope))
+	if err != nil {
+		log.Fatalf("[ERROR] drive client: %v", err)
+	}
+
+	root := *folderID
+	if root == "" {
+		root = findSharedFolder(svc)
+	}
+	folders := map[string]string{
+		"people":   ensureFolder(svc, root, "people"),
+		"families": ensureFolder(svc, root, "families"),
+	}
+
+	tasks := []task{}
+	for _, p := range model.People {
+		if p.PhotoURL != "" {
+			tasks = append(tasks, task{"people", localPart(p.Email) + "-photo", p.PhotoURL})
+		}
+		if p.PronunciationURL != "" {
+			tasks = append(tasks, task{"people", localPart(p.Email) + "-pronunciation", p.PronunciationURL})
+		}
+	}
+	for key, f := range model.Families {
+		if f.PhotoURL != "" {
+			tasks = append(tasks, task{"families", localPart(key) + "-photo", f.PhotoURL})
+		}
+		if f.PronunciationURL != "" {
+			tasks = append(tasks, task{"families", localPart(key) + "-pronunciation", f.PronunciationURL})
+		}
+	}
+
+	pending := []task{}
+	skipped := 0
+	for folderName, id := range folders {
+		bases := listBases(svc, id)
+		for _, t := range tasks {
+			if t.folderName != folderName {
+				continue
+			}
+			if bases[t.base] {
+				skipped++
+			} else {
+				pending = append(pending, t)
+			}
+		}
+	}
+	log.Printf("%d media files in model, %d already imported, %d to fetch", len(tasks), skipped, len(pending))
+
+	client := &http.Client{Timeout: 60 * time.Second}
+	var mu sync.Mutex
+	uploaded, failed := 0, 0
+	work := make(chan task)
+	var wg sync.WaitGroup
+	for range 6 {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for t := range work {
+				body, contentType, err := fetch(client, t.url)
+				if err != nil {
+					mu.Lock()
+					failed++
+					log.Printf("[ERROR] fetch %s/%s: %v", t.folderName, t.base, err)
+					mu.Unlock()
+					continue
+				}
+				name := t.base + extension(contentType, t.url)
+				_, err = svc.Files.Create(&drive.File{Name: name, Parents: []string{folders[t.folderName]}}).
+					Media(bytes.NewReader(body), googleapi.ContentType(contentType)).Fields("id").Do()
+				mu.Lock()
+				if err != nil {
+					failed++
+					log.Printf("[ERROR] upload %s/%s: %v", t.folderName, name, err)
+				} else {
+					uploaded++
+					if uploaded%50 == 0 {
+						log.Printf("uploaded %d/%d", uploaded, len(pending))
+					}
+				}
+				mu.Unlock()
+			}
+		}()
+	}
+	for _, t := range pending {
+		work <- t
+	}
+	close(work)
+	wg.Wait()
+	log.Printf("done: %d uploaded, %d failed, %d skipped", uploaded, failed, skipped)
+}