kjol/go/cmd/migrate/main.go

package main

import (
	"errors"
	"flag"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"kjol/config"
	"kjol/dbutil"

	_ "time/tzdata"

	_ "github.com/lib/pq"
)

func invalidInput() {
	fmt.Println("Usage: [up, down, drop, goto {V}, new {migration name}]")
	os.Exit(1)
}

const (
	defaultTimeFormat = "20060102150405"
	defaultTimezone   = "UTC"
)

var (
	errInvalidSequenceWidth     = errors.New("Digits must be positive")
	errIncompatibleSeqAndFormat = errors.New("The seq and format options are mutually exclusive")
	errInvalidTimeFormat        = errors.New("Time format may not be empty")
)

func createFile(filename string) error {
	// create exclusive (fails if file already exists)
	// os.Create() specifies 0666 as the FileMode, so we're doing the same
	f, err := os.OpenFile(filename, os.O_RDWR|os.O_CREATE|os.O_EXCL, 0666)

	if err != nil {
		return err
	}

	return f.Close()
}

func nextSeqVersion(matches []string, seqDigits int) (string, error) {
	if seqDigits <= 0 {
		return "", errInvalidSequenceWidth
	}

	nextSeq := uint64(1)

	if len(matches) > 0 {
		filename := matches[len(matches)-1]
		matchSeqStr := filepath.Base(filename)
		idx := strings.Index(matchSeqStr, "_")

		if idx < 1 { // Using 1 instead of 0 since there should be at least 1 digit
			return "", fmt.Errorf("Malformed migration filename: %s", filename)
		}

		var err error
		matchSeqStr = matchSeqStr[0:idx]
		nextSeq, err = strconv.ParseUint(matchSeqStr, 10, 64)

		if err != nil {
			return "", err
		}

		nextSeq++
	}

	version := fmt.Sprintf("%0[2]*[1]d", nextSeq, seqDigits)

	if len(version) > seqDigits {
		return "", fmt.Errorf("Next sequence number %s too large. At most %d digits are allowed", version, seqDigits)
	}

	return version, nil
}

func timeVersion(startTime time.Time, format string) (version string, err error) {
	switch format {
	case "":
		err = errInvalidTimeFormat
	case "unix":
		version = strconv.FormatInt(startTime.Unix(), 10)
	case "unixNano":
		version = strconv.FormatInt(startTime.UnixNano(), 10)
	default:
		version = startTime.Format(format)
	}

	return
}

func newCmd(dir string, startTime time.Time, format string, name string, ext string, seq bool, seqDigits int, print bool) error {
	if seq && format != defaultTimeFormat {
		return errIncompatibleSeqAndFormat
	}

	var version string
	var err error

	dir = filepath.Clean(dir)
	ext = "." + strings.TrimPrefix(ext, ".")

	if seq {
		matches, err := filepath.Glob(filepath.Join(dir, "*"+ext))

		if err != nil {
			return err
		}

		version, err = nextSeqVersion(matches, seqDigits)

		if err != nil {
			return err
		}
	} else {
		version, err = timeVersion(startTime, format)

		if err != nil {
			return err
		}
	}

	versionGlob := filepath.Join(dir, version+"_*"+ext)
	matches, err := filepath.Glob(versionGlob)

	if err != nil {
		return err
	}

	if len(matches) > 0 {
		return fmt.Errorf("duplicate migration version: %s", version)
	}

	if err = os.MkdirAll(dir, os.ModePerm); err != nil {
		return err
	}

	for _, direction := range []string{"up", "down"} {
		basename := fmt.Sprintf("%s_%s.%s%s", version, name, direction, ext)
		filename := filepath.Join(dir, basename)

		if err = createFile(filename); err != nil {
			return err
		}

		if print {
			absPath, _ := filepath.Abs(filename)
			log.Println(absPath)
		}
	}

	return nil
}

func main() {
	envfile := flag.String("env-file", "", "Load environment variables from this file instead of the implicit ./.env. Real environment variables always take precedence.")
	dir := flag.String("dir", "./migrations", "Directory containing migration files.")

	flag.Parse()

	var dbc struct {
		Username string `env:"DATABASE_USERNAME"`
		Password string `env:"DATABASE_PASSWORD"`
		Host     string `env:"DATABASE_HOST"`
		Port     int    `env:"DATABASE_PORT"`
		Name     string `env:"DATABASE_NAME"`
		Schema   string `env:"DATABASE_SCHEMA"`
		SSLMode  string `env:"DATABASE_SSL_MODE"`
	}
	if err := config.Load(*envfile, &dbc); err != nil {
		fmt.Println(err.Error())
		os.Exit(1)
	}

	connectionString := dbutil.BuildConnectionString(dbutil.ConnConfig{
		Username: dbc.Username,
		Password: dbc.Password,
		Host:     dbc.Host,
		Port:     dbc.Port,
		Name:     dbc.Name,
		Schema:   dbc.Schema,
		SSLMode:  dbc.SSLMode,
	})

	// parse CLI args and do actions
	args := flag.Args()

	if len(args) < 1 {
		invalidInput()
	}

	// handle "new" command early — it doesn't need a database connection
	if args[0] == "new" {
		if len(args) < 2 {
			fmt.Println("Please provide a name for the new migration.")
			os.Exit(1)
		}
		newCmd(*dir, time.Now(), defaultTimeFormat, args[1], "sql", true, 7, true)
		return
	}

	eng, err := newEngine(*dir, connectionString, dbc.Schema)
	if err != nil {
		fmt.Println(err.Error())
		os.Exit(1)
	}
	defer eng.close()

	migrateNum := 0

	if len(args) >= 2 {
		var parseErr error
		migrateNum, parseErr = strconv.Atoi(args[1])
		if parseErr != nil {
			fmt.Println("Please provide a valid migration number.")
			os.Exit(1)
		}
	}

	switch args[0] {
	case "up":
		if err := eng.up(); err != nil {
			if errors.Is(err, errNoChange) {
				fmt.Println("No change")
				os.Exit(0)
			}
			fmt.Println(err.Error())
			os.Exit(1)
		}
		fmt.Println("Database successfully migrated to latest version")

	case "down":
		if err := eng.down(); err != nil {
			if errors.Is(err, errNoChange) {
				fmt.Println("No change")
				os.Exit(0)
			}
			fmt.Println(err.Error())
			os.Exit(1)
		}
		fmt.Println("Database successfully migrated to previous version")

	case "goto":
		if err := eng.goTo(uint(migrateNum)); err != nil {
			if errors.Is(err, errNoChange) {
				fmt.Println("No change")
				os.Exit(0)
			}
			fmt.Println(err.Error())
			os.Exit(1)
		}
		fmt.Printf("Database successfully migrated to version `%d`\n", migrateNum)

	case "drop":
		if err := eng.dropAll(); err != nil {
			fmt.Println(err.Error())
			os.Exit(1)
		}
		fmt.Println("Database tables, views, and types successfully dropped")

	default:
		invalidInput()
	}
}