Golang Script to Migrate Postgres Last Sequence

main.go

package main

import (
	"database/sql"
	"os"

	"fmt"
	"log"

	_ "github.com/lib/pq"
)

const (
	SourceHost = "source-host"
	SourcePort = "source-port"

	TargetHost = "target-host"
	TargetPort = "target-port"

	MigrateList = "migrate_list.csv"
	SourceList  = "source_list.csv"
)

func main() {

  /*
    Manually list source schema
    Ensure all sequence in source exist in target
  */
  
  
	listSchema := []string{
		"schema_name"
	}

	for _, schema := range listSchema {
		res := GetSchemaSequence(schema, SourceHost, SourcePort)
		MigrateSchemaSequence(schema, SourceHost, SourcePort, res)
	}
}

const getAllSequenceInSchema = `-- getAllSequenceInSchema
SELECT
	schemaname AS schema, 
	sequencename AS sequence, 
	last_value
FROM pg_sequences
`

type Sequence struct {
	Role         string
	SequenceName string
	LastValue    sql.NullInt64
}

func GetSchemaSequence(schema, host, port string) (result []Sequence) {
	connStr := fmt.Sprintf("user=source_user dbname=%s host=%s port=%s password=source_password", schema, host, port)
	db, err := sql.Open("postgres", connStr)
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()

	rows, err := db.Query(getAllSequenceInSchema)
	if err != nil {
		log.Fatal(err)
	}

	for rows.Next() {
		var role, sequence string
		var lastValue sql.NullInt64

		err := rows.Scan(&role, &sequence, &lastValue)
		if err != nil {
			log.Fatal(err, role, sequence, lastValue)
		}

		if lastValue.Valid {
			WriteToFile(SourceList, fmt.Sprintf("VALID,%s,%s,%d", schema, sequence, lastValue.Int64))
			log.Println("VALID", schema, sequence, lastValue.Int64)
		} else {
			WriteToFile(SourceList, fmt.Sprintf("INVALID,%s,%s,%s", schema, sequence, "NULL"))
			log.Println("INVALID", schema, sequence, "NULL")
		}

		result = append(result, Sequence{
			Role:         role,
			SequenceName: sequence,
			LastValue:    lastValue,
		})
	}

	return result
}

const setLastSequence = `-- setLastSequence
SELECT setval($1, $2, true); 
`

func MigrateSchemaSequence(schema, host, port string, sequences []Sequence) {
	connStr := fmt.Sprintf("user=target_user dbname=%s host=%s port=%s password=target_password", schema, host, port)
	db, err := sql.Open("postgres", connStr)
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()

	for _, sequence := range sequences {
		if sequence.LastValue.Valid {
			WriteToFile(MigrateList, fmt.Sprintf("UPDATE,%s,%s,%d", schema, sequence.SequenceName, sequence.LastValue.Int64))
			log.Println("UPDATE", schema, sequence.SequenceName, sequence.LastValue.Int64)
			_, err = db.Exec(setLastSequence, sequence.SequenceName, sequence.LastValue.Int64)
			if err != nil {
				log.Fatal(err)
			}
		} else {
			WriteToFile(MigrateList, fmt.Sprintf("SKIP,%s,%s,%s", schema, sequence.SequenceName, "NULL"))
			log.Println("SKIP", schema, sequence.SequenceName, "NULL")
		}
	}
}

func WriteToFile(filename, data string) {
	file, err := os.OpenFile(filename, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		fmt.Println("Could not open", filename)
		return
	}

	defer file.Close()

	_, errWrite := file.WriteString(data + "\n")
	if errWrite != nil {
		fmt.Println("Could not write text to", filename)
	}
}