diff --git a/tools/importblobs/main.go b/tools/importblobs/main.go new file mode 100644 index 0000000..9f0dafd --- /dev/null +++ b/tools/importblobs/main.go @@ -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) +}