sqlite-find-replace.go

package main

import (
	"database/sql"
	"flag"
	"os"
	"regexp"
	"runtime"
	"strings"

	"git.tcp.direct/kayos/zwrap"
	_ "github.com/glebarez/go-sqlite"
	"github.com/rs/zerolog"
)

type config struct {
	sqliteDB            string
	tableRegex          string
	tableRegexCompiled  *regexp.Regexp
	tableSubStringMatch string
	findString          string
	replaceString       string
}

func flags() *config {
	cfg := &config{}
	flag.StringVar(&cfg.sqliteDB, "sqlite-db", "", "sqlite database file")
	flag.StringVar(&cfg.tableRegex, "table-regex", "", "regex to match table names")
	flag.StringVar(&cfg.tableSubStringMatch, "table-substring", "", "substring to match table names")
	flag.StringVar(&cfg.findString, "find-string", "", "string to find in table")
	flag.StringVar(&cfg.replaceString, "replace-string", "", "string to replace in table")
	flag.Parse()
	if cfg.tableRegex != "" {
		var err error
		if cfg.tableRegexCompiled, err = regexp.Compile(cfg.tableRegex); err != nil {
			println("\n" + colorRed + "invalid table-regex" + colorReset + "\n")
			println(err.Error())
			flag.Usage()
			os.Exit(1)
		}
	}
	return cfg
}

const (
	colorRed   = "\033[31m"
	colorReset = "\033[0m"
)

func (c *config) validate() {
	switch {
	case c.tableRegex != "" && c.tableSubStringMatch != "":
		println("\n" + colorRed + "cannot use both table-regex and table-substring-match" + colorReset + "\n")
	case c.findString == "":
		println("\n" + colorRed + "find-string is required" + colorReset + "\n")
	case c.replaceString == "":
		println("\n" + colorRed + "replace-string is required" + colorReset + "\n")
	case c.sqliteDB == "":
		println("\n" + colorRed + "sqlite-db is required" + colorReset + "\n")
	/*case c.tableRegex == "" && c.tableSubStringMatch == "":
	println("\n" + colorRed + "either table-regex or table-substring-match is required" + colorReset + "\n")*/
	default:
		return
	}

	flag.Usage()
	os.Exit(1)
}

func (c *config) matchTable(tableName string) bool {
	if c.tableRegex != "" {
		return c.tableRegexCompiled.MatchString(tableName)
	}
	return c.tableSubStringMatch == "" || strings.Contains(tableName, c.tableSubStringMatch)
}

func instantiateConsoleLogger() *zerolog.Logger {
	cl := zerolog.NewConsoleWriter()
	//goland:noinspection GoBoolExpressions
	cl.NoColor = runtime.GOOS == "windows"
	cl.FormatLevel = zwrap.LogLevelFmt(cl.NoColor)
	cl.Out = os.Stdout
	clogger := zerolog.New(cl).With().Timestamp().Logger()
	return &clogger
}

type columnInfo struct {
	CID     int    `db:"cid"`
	Name    string `db:"name"`
	Type    string `db:"type"`
	NotNull bool   `db:"notnull"`
	Default any    `db:"dflt_value"`
	PK      bool   `db:"pk"`
}

func main() {
	cfg := flags()
	cfg.validate()

	log := instantiateConsoleLogger()

	db, err := sql.Open("sqlite", cfg.sqliteDB)
	if err != nil {
		log.Fatal().Err(err).Msg("failed to open database")
	}
	defer func() {
		if err = db.Close(); err != nil {
			panic(err)
		}
	}()
	if err = db.Ping(); err != nil {
		log.Fatal().Err(err).Msg("failed to ping database")
	}

	rows, err := db.Query(`SELECT name FROM sqlite_master WHERE type='table';`)
	if err != nil {
		log.Fatal().Err(err).Msg("failed to query database tables")
	}

	var tables = make(map[string][]string)

	for rows.Next() {
		var table string
		if err = rows.Scan(&table); err != nil {
			log.Fatal().Err(err).Msg("failed to scan table name")
		}
		if !cfg.matchTable(table) {
			continue
		}
		log.Info().Str("table", table).Msg("found table")
		tables[table] = []string{}
	}

	for table := range tables {
		log.Debug().Str("table", table).Msg("processing table columns")

		var columns *sql.Rows

		if columns, err = db.Query(`PRAGMA table_info(` + table + `)`); err != nil {
			log.Fatal().Err(err).Msg("failed to query table columns")
		}

		for columns.Next() {
			var column columnInfo
			if err = columns.Scan(
				&column.CID, &column.Name, &column.Type, &column.NotNull, &column.Default, &column.PK,
			); err != nil {
				log.Fatal().Err(err).Msg("failed to scan column name")
			}

			if column.Name == "sqlite_master" || column.Name == "sqlite_sequence" {
				log.Debug().Str("table", table).Str("column", column.Name).Msg("skipping system column")
				continue
			}

			log.Debug().Str("table", table).Str("column", column.Name).Msg("found column")

			tables[table] = append(tables[table], column.Name)
		}
	}

	for table, columns := range tables {
		if len(columns) == 0 {
			log.Warn().Str("caller", table).Msg("no columns found, skipping table")
			continue
		}
		for _, column := range columns {
			var res sql.Result

			qBase := `UPDATE ` + table + ` SET ` + column + ` = replace(` + column + `, ?, ?) WHERE ` + column + ` LIKE '%` + cfg.findString + `%';`

			log.Trace().Str("caller", table).Str("column", column).
				Msg(qBase + ", " + cfg.findString + ", " + cfg.replaceString)

			if res, err = db.Exec(
				qBase, cfg.findString, cfg.replaceString,
			); err != nil {
				log.Fatal().Err(err).Msg("failed to update table")
			} else {
				rowsAffected, _ := res.RowsAffected()
				if rowsAffected > 0 {
					log.Info().Str("caller", table).Int64("rowsAffected", rowsAffected).Msg("updated table")
				}
			}
		}
	}
}