magnetico-go-migrator/magnetico-migrator.go

package main

import (
	"database/sql"
	"encoding/hex"
	"gitlab.com/skobkin/magnetico-go-migrator/old"

	"github.com/cheggaaa/pb/v3"
	_ "github.com/mattn/go-sqlite3"
	"github.com/skobkin/magnetico/pkg/persistence"
	"go.uber.org/zap"
	"gopkg.in/alecthomas/kingpin.v2"
	"net/url"
)

var (
	goDatabaseUrl = kingpin.Flag("go-database", "magneticod Go version database URL.").Short('g').String()
	pyDatabaseUrl = kingpin.Flag("py-database", "Python version database URL.").Short('p').Required().String()
	chunkSize     = kingpin.Flag("chunk-size", "How many torrents to retrieve from old DB at once").Short('c').Default("200").Uint64()

	progress pb.ProgressBar
)

func main() {
	logger, _ := zap.NewDevelopment()
	defer logger.Sync()
	zap.ReplaceGlobals(logger)

	kingpin.Parse()

	goDatabase, err := persistence.MakeDatabase(*goDatabaseUrl, logger)
	if err != nil {
		logger.Panic("Could not open Go database.", zap.String("url", *goDatabaseUrl), zap.Error(err))
	}
	defer goDatabase.Close()

	oldSqliteUrl, err := url.Parse(*pyDatabaseUrl)
	if err != nil {
		logger.Panic("Can't parse Python database URL.", zap.String("url", *pyDatabaseUrl), zap.Error(err))
	}

	pyDatabase, err := old.OpenSqliteDb(*oldSqliteUrl)
	if err != nil {
		logger.Panic("Couldn't open Python database.", zap.String("url", *pyDatabaseUrl), zap.Error(err))
	}
	defer pyDatabase.Close()

	torrentsTotal, err := old.GetNumberOfTorrents(pyDatabase)
	if err != nil {
		zap.L().Panic("Couldn't count torrents", zap.Error(err))
	}

	progress = *pb.Full.New(int(torrentsTotal))
	progress.Start()

	err = mergeDatabases(pyDatabase, goDatabase)
	if err != nil {
		logger.Error("Error while processing data", zap.Error(err))
	}

	progress.Finish()
}

func mergeDatabases(pyDb *sql.DB, goDb persistence.Database) error {
	maxId, err := old.GetLastTorrentId(pyDb)
	if err != nil {
		zap.L().Fatal("Error while getting last torrent ID from old database")
	}

	var lastId uint64 = 0

	for lastId < maxId {
		chunk, err := old.GetTorrentsChunk(pyDb, lastId, *chunkSize)
		if err != nil {
			zap.L().Error("Error while getting torrents chunk", zap.Uint64("last-id", lastId))

			return err
		}

		if len(chunk) == 0 {
			if (maxId - lastId) <= *chunkSize {
				zap.L().Warn("Probably finished processing all data (earlier). Finishing.", zap.Uint64("last", lastId), zap.Uint64("max", maxId))

				return nil
			} else {
				zap.L().Error("Prematurely received empty chunk, aborting", zap.Uint64("last", lastId), zap.Uint64("max", maxId))
			}
		}

		for _, piece := range chunk {
			lastId = piece.Torrent.ID
			progress.Increment()

			err = importTorrent(goDb, piece.Torrent, piece.Files)
			if err != nil {
				zap.L().Error("Error when importing torrent to Go database.", zap.String("hash", hex.EncodeToString(piece.Torrent.InfoHash)), zap.Error(err))
			}
		}
	}

	return nil
}

func importTorrent(goDb persistence.Database, torrent persistence.TorrentMetadata, files []persistence.File) error {
	exists, err := goDb.DoesTorrentExist(torrent.InfoHash)
	if err != nil {
		return err
	}

	if !exists {
		err = goDb.AddNewTorrent(torrent.InfoHash, torrent.Name, files)
		if err != nil {
			return err
		}
	}

	return nil
}