uploader.go

package main

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

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

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() {

	db, err := sql.Open("postgres", "postgres://postgres:postgres@localhost/postgres?sslmode=disable")

	if err != nil {
		panic(err)
	}

	defer db.Close()

	_, err = db.Query("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)
	}

	rows, err := db.Query("SELECT * from files where err is null")

	if err != nil {
		panic(err)
	}

	defer rows.Close()

	count, err := db.Query("SELECT COUNT(*) FROM files")

	if err != nil {
		panic(err)
	}

	var totalFiles int
	count.Next()
	err = count.Scan(&totalFiles)

	if err != nil {
		panic(err)
	}

	bar := progressbar.NewOptions(totalFiles, progressbar.OptionShowCount())
	uploadFilesChan := make(chan File, 10)

	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		for rows.Next() {
			var file File
			err = rows.Scan(&file.Path, &file.Type, &file.Err, &file.CreatedAt)

			if err != nil {
				panic(err)
			}

			uploadFilesChan <- file
		}
	}()

	for i := 0; i != 15; i++ {
		wg.Add(1)
		config := &oauth2.Config{}

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

		srv, _ := drive.NewService(context.Background(), option.WithHTTPClient(client))

		go func(srv *drive.Service) {
			defer wg.Done()

			file, more := <-uploadFilesChan

			if !more {
				return
			}

			for {
				bar.Add(1)

				fileData, err := os.Open(file.Path)

				if err != nil {
					db.Query("INSERT INTO files_uploaded (file_path, err) VALUES ($1, $2)", file.Path, err.Error())
					continue
				}
				stat, err := fileData.Stat()
				_, err = srv.Files.Insert(&drive.File{
					Title: fmt.Sprintf("%s.%s", stat.ModTime().Format("02-01-2006"), filepath.Ext(file.Path)),
					Parents: []*drive.ParentReference{
						{
							Id: GetFolderIDByMediaType(file.Type),
						},
					},
				}).Media(fileData).Do()

				if err != nil {
					db.Query("INSERT INTO files_uploaded (file_path, err) VALUES ($1, $2)", file.Path, err.Error())
					continue
				}

				db.Query("INSERT INTO files_uploaded (file_path) VALUES($1)", file.Path)
			}

		}(srv)
	}

	if err != nil {
		panic(err)
	}

	wg.Wait()
}