magnetico-go-migrator/magnetico-migrator.go

package main

import (
	"database/sql"
	"encoding/hex"
	"fmt"

	"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"
	"time"
)

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

	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 := openSqliteDb(*oldSqliteUrl)
	if err != nil {
		logger.Panic("Couldn't open Python database.", zap.String("url", *pyDatabaseUrl), zap.Error(err))
	}
	defer pyDatabase.Close()

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

	progress = *pb.New(int(torrentsTotal))
	progress.SetRefreshRate(time.Second)
	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 {
	// Selecting torrents
	tRows, err := pyDb.Query("SELECT id, info_hash, name, total_size, discovered_on FROM torrents ORDER BY id ASC;")
	if err != nil {
		zap.L().Fatal("Error when querying torrents from old database.")
	}
	defer tRows.Close()

	for tRows.Next() {
		var torrent persistence.TorrentMetadata

		progress.Increment()

		err := tRows.Scan(&torrent.ID, &torrent.InfoHash, &torrent.Name, &torrent.Size, &torrent.DiscoveredOn)
		if err != nil {
			zap.L().Error("Error scanning torrent row.", zap.Error(err))

			continue
		}

		files, err := getFilesForTorrent(pyDb, torrent)
		if err != nil {
			zap.L().Error("Error getting files for torrent.", zap.String("hash", hex.EncodeToString(torrent.InfoHash)), zap.Error(err))

			continue
		}

		err = importTorrent(goDb, torrent, files)
		if err != nil {
			zap.L().Error("Error when importing torrent to Go database.", zap.String("hash", hex.EncodeToString(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
}

func getFilesForTorrent(pyDb sql.DB, torrent persistence.TorrentMetadata) ([]persistence.File, error) {
	files := make([]persistence.File, 0)

	// Selecting files for torrent
	fRows, err := pyDb.Query("SELECT f.path, f.size FROM files AS f WHERE f.torrent_id = ? ORDER BY f.id ASC;", torrent.ID)
	if err != nil {
		return files, nil
	}
	defer fRows.Close()

	for fRows.Next() {
		var file persistence.File
		err := fRows.Scan(&file.Path, &file.Size)
		if err != nil {
			return files, err
		}

		files = append(files, file)
	}

	return files, nil
}

func getNumberOfTorrents(db sql.DB) (uint, error) {
	rows, err := db.Query("SELECT COUNT(id) FROM torrents;")
	if err != nil {
		return 0, err
	}
	defer rows.Close()

	if rows.Next() != true {
		return 0, fmt.Errorf("No rows returned from `SELECT COUNT(id)`")
	}

	var n *uint
	if err = rows.Scan(&n); err != nil {
		return 0, err
	}

	// If the database is empty (i.e. 0 entries in 'torrents') then the query will return nil.
	if n == nil {
		return 0, nil
	} else {
		return *n, nil
	}
}

func openSqliteDb(url_ url.URL) (*sql.DB, error) {
	url_.Scheme = ""
	return sql.Open("sqlite3", url_.String())
}