waf-platform/EdgeDNS/internal/dbs/db_migrate_sqlite.go

// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
//go:build plus

package dbs

import (
	"database/sql"
	"github.com/TeaOSLab/EdgeDNS/internal/models"
	"github.com/TeaOSLab/EdgeDNS/internal/remotelogs"
	_ "github.com/mattn/go-sqlite3"
	"os"
	"strings"
)

const sqliteMigrationMarkerKey = "meta:migrated_from_sqlite"

func (this *DB) migrateSQLiteIfNeeded() error {
	if len(this.path) == 0 || this.path == this.storePath {
		return nil
	}

	_, err := os.Stat(this.path)
	if err != nil {
		if os.IsNotExist(err) {
			return nil
		}
		return err
	}

	ok, err := this.hasMigrationMarker()
	if err != nil {
		return err
	}
	if ok {
		removeErr := removeSQLiteFiles(this.path)
		if removeErr != nil {
			remotelogs.Warn("DB", "remove sqlite files failed after migration: "+removeErr.Error())
		}
		return nil
	}

	remotelogs.Println("DB", "migrating sqlite database from '"+this.path+"' to '"+this.storePath+"' ...")

	err = this.truncateMigratedData()
	if err != nil {
		return err
	}

	err = this.importSQLiteData()
	if err != nil {
		return err
	}

	err = this.rawDB.Set([]byte(sqliteMigrationMarkerKey), []byte("1"), defaultWriteOptions)
	if err != nil {
		return err
	}

	err = flushRawDB(this.rawDB)
	if err != nil {
		return err
	}

	removeErr := removeSQLiteFiles(this.path)
	if removeErr != nil {
		remotelogs.Warn("DB", "remove sqlite files failed after migration: "+removeErr.Error())
	}

	remotelogs.Println("DB", "migrated sqlite database to pebble")
	return nil
}

func (this *DB) hasMigrationMarker() (bool, error) {
	_, closer, err := this.rawDB.Get([]byte(sqliteMigrationMarkerKey))
	if err != nil {
		if isNotFound(err) {
			return false, nil
		}
		return false, err
	}
	_ = closer.Close()
	return true, nil
}

func (this *DB) truncateMigratedData() error {
	for _, prefix := range []string{
		domainPrefix,
		domainClusterIndex,
		recordPrefix,
		routePrefix,
		keyPrefix,
		agentIPPrefix,
		"meta:",
	} {
		err := this.rawDB.DeleteRange([]byte(prefix), prefixUpperBound([]byte(prefix)), defaultWriteOptions)
		if err != nil {
			return err
		}
	}
	return nil
}

func (this *DB) importSQLiteData() error {
	sqliteDB, err := sql.Open("sqlite3", "file:"+this.path+"?mode=ro")
	if err != nil {
		return err
	}
	defer func() {
		_ = sqliteDB.Close()
	}()

	err = this.importSQLiteDomains(sqliteDB)
	if err != nil {
		return err
	}
	err = this.importSQLiteRecords(sqliteDB)
	if err != nil {
		return err
	}
	err = this.importSQLiteRoutes(sqliteDB)
	if err != nil {
		return err
	}
	err = this.importSQLiteKeys(sqliteDB)
	if err != nil {
		return err
	}
	return this.importSQLiteAgentIPs(sqliteDB)
}

func (this *DB) importSQLiteDomains(sqliteDB *sql.DB) error {
	rows, err := sqliteDB.Query(`SELECT "id", "clusterId", "userId", "name", "tsig", "version" FROM "domains_v2" ORDER BY "id" ASC`)
	if err != nil {
		return ignoreMissingTable(err)
	}
	defer func() {
		_ = rows.Close()
	}()

	for rows.Next() {
		value := &domainValue{}
		err = rows.Scan(&value.Id, &value.ClusterId, &value.UserId, &value.Name, &value.TSIGJSON, &value.Version)
		if err != nil {
			return err
		}

		err = this.saveDomain(value)
		if err != nil {
			return err
		}
	}

	return rows.Err()
}

func (this *DB) importSQLiteRecords(sqliteDB *sql.DB) error {
	rows, err := sqliteDB.Query(`SELECT "id", "domainId", "name", "type", "value", "mxPriority", "srvPriority", "srvWeight", "srvPort", "caaFlag", "caaTag", "ttl", "weight", "routeIds", "version" FROM "records_v2" ORDER BY "id" ASC`)
	if err != nil {
		return ignoreMissingTable(err)
	}
	defer func() {
		_ = rows.Close()
	}()

	for rows.Next() {
		value := &recordValue{}
		var routeIDs string
		err = rows.Scan(&value.Id, &value.DomainId, &value.Name, &value.Type, &value.Value, &value.MXPriority, &value.SRVPriority, &value.SRVWeight, &value.SRVPort, &value.CAAFlag, &value.CAATag, &value.TTL, &value.Weight, &routeIDs, &value.Version)
		if err != nil {
			return err
		}
		if len(routeIDs) > 0 {
			value.RouteIds = strings.Split(routeIDs, ",")
		}

		err = this.saveJSON(recordKey(value.Id), value)
		if err != nil {
			return err
		}
	}

	return rows.Err()
}

func (this *DB) importSQLiteRoutes(sqliteDB *sql.DB) error {
	rows, err := sqliteDB.Query(`SELECT "id", "userId", "ranges", "priority", "order", "version" FROM "routes_v2" ORDER BY "id" ASC`)
	if err != nil {
		return ignoreMissingTable(err)
	}
	defer func() {
		_ = rows.Close()
	}()

	for rows.Next() {
		value := &routeValue{}
		err = rows.Scan(&value.Id, &value.UserId, &value.RangesJSON, &value.Priority, &value.Order, &value.Version)
		if err != nil {
			return err
		}

		err = this.saveJSON(routeKey(value.Id), value)
		if err != nil {
			return err
		}
	}

	return rows.Err()
}

func (this *DB) importSQLiteKeys(sqliteDB *sql.DB) error {
	rows, err := sqliteDB.Query(`SELECT "id", "domainId", "zoneId", "algo", "secret", "secretType", "version" FROM "keys" ORDER BY "id" ASC`)
	if err != nil {
		return ignoreMissingTable(err)
	}
	defer func() {
		_ = rows.Close()
	}()

	for rows.Next() {
		value := &models.NSKey{}
		err = rows.Scan(&value.Id, &value.DomainId, &value.ZoneId, &value.Algo, &value.Secret, &value.SecretType, &value.Version)
		if err != nil {
			return err
		}

		err = this.saveJSON(keyKey(value.Id), value)
		if err != nil {
			return err
		}
	}

	return rows.Err()
}

func (this *DB) importSQLiteAgentIPs(sqliteDB *sql.DB) error {
	rows, err := sqliteDB.Query(`SELECT "id", "ip", "agentCode" FROM "agentIPs" ORDER BY "id" ASC`)
	if err != nil {
		return ignoreMissingTable(err)
	}
	defer func() {
		_ = rows.Close()
	}()

	for rows.Next() {
		value := &models.AgentIP{}
		err = rows.Scan(&value.Id, &value.IP, &value.AgentCode)
		if err != nil {
			return err
		}

		err = this.saveJSON(agentIPKey(value.Id), value)
		if err != nil {
			return err
		}
	}

	return rows.Err()
}

func ignoreMissingTable(err error) error {
	if err == nil {
		return nil
	}
	if strings.Contains(err.Error(), "no such table") {
		return nil
	}
	return err
}

func removeSQLiteFiles(path string) error {
	var lastErr error
	for _, filename := range []string{path, path + "-shm", path + "-wal"} {
		err := os.Remove(filename)
		if err != nil && !os.IsNotExist(err) {
			lastErr = err
		}
	}
	return lastErr
}