uploader.go

package main

import (
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"sync"

	"github.com/schollz/progressbar/v3"
	"golang.org/x/oauth2"
	"google.golang.org/api/drive/v2"
	"google.golang.org/api/option"

	_ "github.com/lib/pq"
)

type File struct {
	Path      string
	Type      string
	Err       sql.NullString
	CreatedAt string
	Uploaded  bool
}

func main() {
	UploadAllFromDatabase()
}

func tokenFromFile(file string) (*oauth2.Token, error) {
	f, err := os.Open(file)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	tok := &oauth2.Token{}
	err = json.NewDecoder(f).Decode(tok)
	return tok, err
}

func getClient(config *oauth2.Config) *http.Client {
	// The file token.json stores the user's access and refresh tokens, and is
	// created automatically when the authorization flow completes for the first
	// time.
	tokFile := "token.json"
	tok, err := tokenFromFile(tokFile)
	if err != nil {
		panic(err)
	}
	return config.Client(context.Background(), tok)
}

func GetFolderIDByMediaType(mediatype string) string {
	switch mediatype {
	case "image":
		return "a"
	case "video":
		return "b"
	case "audio":
		return "c"
	case "documents":
		return "d"
	default:
		return "e"
	}
}

func UploadAllFromDatabase() {
	client := (&oauth2.Config{}).Client(context.Background(), &oauth2.Token{
		AccessToken: "token",
		TokenType:   "bearer",
	})

	svc, err := drive.NewService(context.Background(), option.WithHTTPClient(client))
	if err != nil {
		panic(err)
	}

	db, err := sql.Open("postgres", "postgres://postgres:postgres@localhost/postgres?sslmode=disable")
	if err != nil {
		panic(err)
	}
	defer db.Close()

	_, err = db.Exec(`
	CREATE TABLE IF NOT EXISTS files_uploaded (
		id SERIAL, 
		file_path text NOT NULL, 
		err text, 
		FOREIGN KEY (file_path) REFERENCES files(path), 
		PRIMARY KEY(id)
	)`)
	if err != nil {
		panic(err)
	}

	var totalFiles int
	if err := db.QueryRow("SELECT COUNT(*) FROM files").Scan(&totalFiles); err != nil {
		panic(err)
	}

	rows, err := db.Query("SELECT * from files where err is null")
	if err != nil {
		panic(err)
	}
	defer rows.Close()

	uploadFilesChan := make(chan File)
	go func() {
		defer close(uploadFilesChan)
		for rows.Next() {
			var file File
			if err := rows.Scan(
				&file.Path,
				&file.Type,
				&file.Err,
				&file.CreatedAt,
			); err != nil {
				panic(err)
			}
			uploadFilesChan <- file
		}
	}()

	bar := progressbar.NewOptions(totalFiles, progressbar.OptionShowCount())

	const workers = 15
	var wg sync.WaitGroup
	wg.Add(workers)

	for i := 0; i < workers; i++ {
		go func() {
			defer wg.Done()
			for file := range uploadFilesChan {
				bar.Add(1)
				if err := uploadFile(svc, file.Path); err != nil {
					if _, err := db.Exec("INSERT INTO files_uploaded (file_path, err) VALUES ($1, $2)", file.Path, err.Error()); err != nil {
						panic(err)
					}
					continue
				}

				if _, err := db.Exec("INSERT INTO files_uploaded (file_path) VALUES($1)", file.Path); err != nil {
					panic(err)
				}
			}
		}()
	}

	wg.Wait()
}

func uploadFile(svc *drive.Service, path string) error {
	file, err := os.Open(path)
	if err != nil {
		return fmt.Errorf("unable to open file: %v", err)
	}

	stat, err := file.Stat()
	if err != nil {
		return fmt.Errorf("unable to get file stat: %v", err)
	}

	_, err = svc.Files.Insert(&drive.File{
		Title: fmt.Sprintf("%s.%s", stat.ModTime().Format("02-01-2006"), filepath.Ext(path)),
		Parents: []*drive.ParentReference{
			{
				Id: GetFolderIDByMediaType("image"),
			},
		},
	}).Media(file).Do()
	if err != nil {
		return fmt.Errorf("unable to upload file: %v", err)
	}

	return nil
}