guregu/dynamo migrate script

LICENSE

           DO WHAT THE FUCK YOU WANT TO PUBLIC LICENSE
                   Version 2, December 2004
 
Copyright (C) 2020 mohemohe <[email protected]>

Everyone is permitted to copy and distribute verbatim or modified
copies of this license document, and changing it is allowed as long
as the name is changed.
 
           DO WHAT THE FUCK YOU WANT TO PUBLIC LICENSE
  TERMS AND CONDITIONS FOR COPYING, DISTRIBUTION AND MODIFICATION
 
 0. You just DO WHAT THE FUCK YOU WANT TO.

main.go

package main

import (
	"bufio"
	"fmt"
	"github.com/guregu/dynamo"
	"github.com/joho/godotenv"
	"github.com/mohemohe/hoge-agent/core/models"
	"log"
	"os"
	"reflect"
	"strings"
	"time"
)

func main() {
	exitOnErr(godotenv.Load("../../.env"))
	if ok := prompt(); !ok {
		os.Exit(1)
	}

	migrate(models.TenantTableName, models.Tenant{})
	migrate(models.UserTableName, models.User{})
	migrate(models.GuestTableName, models.Guest{})
	migrate(models.LineTableName, models.Line{})
	migrate(models.RoomTableName, models.Room{})
	migrate(models.MessageTableName, models.Message{})

	log.Println("done!")
}

func prompt() bool {
	accessKeyToDebug := os.Getenv("AWS_SECRET_ACCESS_KEY")
	if len(accessKeyToDebug) > 5 {
		accessKeyToDebug = accessKeyToDebug[0:5] + "******************"
	} else {
		accessKeyToDebug = "***********************"
	}
	fmt.Println("migrate to:")
	fmt.Println("  AWS_REGION            :", os.Getenv("AWS_REGION"))
	fmt.Println("  AWS_ACCESS_KEY_ID     :", os.Getenv("AWS_ACCESS_KEY_ID"))
	fmt.Println("  AWS_SECRET_ACCESS_KEY :", accessKeyToDebug)
	fmt.Println("  DYNAMODB_ENDPOINT     :", os.Getenv("DYNAMODB_ENDPOINT"))

	if strings.ToLower(os.Getenv("CI")) != "true" {
		reader := bufio.NewReader(os.Stdin)
		for {
			fmt.Printf("continue? [y/N]: ")
			result, err := reader.ReadString('\n')
			if err != nil {
				log.Fatal(err)
			}
			result = strings.TrimSpace(strings.ToLower(result))
			if result == "y" || result == "yes" {
				return true
			} else if result == "n" || result == "no" || result == "" {
				return false
			}
		}
	}
	panic("invalid state")
}

func migrate(name string, value interface{}) {
	if hasTable(name) {
		index := structToIndex(value)
		setGSI(name, index)
		return
	}
	createTable(name, value)
}

func hasTable(name string) bool {
	db := models.NewDB()
	tableName := models.TableName(name)
	_, err := db.Table(tableName).Describe().Run()
	return err == nil
}

func createTable(name string, value interface{}) {
	db := models.NewDB()
	tableName := models.TableName(name)
	if err := db.CreateTable(tableName, value).OnDemand(true).Run(); err != nil {
		log.Println("error:", name, err.Error())
	} else {
		log.Println("create:", name)
	}
}

func deleteTable(name string) {
	db := models.NewDB()
	tableName := models.TableName(name)
	db.Table(tableName).DeleteTable().Run()
}

func structToIndex(value interface{}) []dynamo.Index {
	indexes := map[string]*dynamo.Index{}

	t := reflect.TypeOf(value)
	for i := 0; i < t.NumField(); i++ {
		field := t.Field(i)

		tagName := field.Tag.Get("dynamo")
		if tagName == "" {
			tagName = field.Name
		}
		tagKind := field.Type.Kind()
		var tagType dynamo.KeyType
		switch tagKind {
		case reflect.String:
			tagType = dynamo.StringType
			break
		case reflect.Int:
		case reflect.Float64:
			tagType = dynamo.NumberType
		default:
			tagType = dynamo.BinaryType
		}

		index := field.Tag.Get("index")
		if index == "" {
			continue
		}
		v := strings.Split(index, ",")
		if len(v) != 2 {
			continue
		}
		indexName := v[0]
		indexType := v[1]

		if indexes[indexName] == nil {
			indexes[indexName] = &dynamo.Index{
				Name:           indexName,
				Local:          false,
				ProjectionType: "ALL",
			}
		}

		if indexType == "hash" {
			indexes[indexName].HashKey = tagName
			indexes[indexName].HashKeyType = tagType
		}
		if indexType == "range" {
			indexes[indexName].RangeKey = tagName
			indexes[indexName].RangeKeyType = tagType
		}
	}

	result := make([]dynamo.Index, 0)
	for _, v := range indexes {
		if v.HashKey != "" {
			result = append(result, *v)
		}
	}
	return result
}

func setGSI(name string, indexes []dynamo.Index) {
	db := models.NewDB()
	tableName := models.TableName(name)
	current, err := db.Table(tableName).Describe().Run()
	exitOnErr(err)

	currentIndex := map[string]dynamo.Index{}
	for _, gsi := range current.GSI {
		currentIndex[gsi.Name] = gsi
	}

	nextIndex := map[string]dynamo.Index{}
	for _, gsi := range indexes {
		nextIndex[gsi.Name] = gsi
	}

	deleteIndex := make([]dynamo.Index, 0)
	for k, v := range currentIndex {
		if _, ok := nextIndex[k]; !ok {
			deleteIndex = append(deleteIndex, v)
			continue
		}
		if currentIndex[k].HashKey != nextIndex[k].HashKey {
			deleteIndex = append(deleteIndex, v)
			continue
		}
		if currentIndex[k].HashKeyType != nextIndex[k].HashKeyType {
			deleteIndex = append(deleteIndex, v)
			continue
		}
		if currentIndex[k].RangeKey != nextIndex[k].RangeKey {
			deleteIndex = append(deleteIndex, v)
			continue
		}
		if currentIndex[k].RangeKeyType != nextIndex[k].RangeKeyType {
			deleteIndex = append(deleteIndex, v)
			continue
		}
	}

	createIndex := make([]dynamo.Index, 0)
	for k, v := range nextIndex {
		if _, ok := currentIndex[k]; !ok {
			createIndex = append(createIndex, v)
			continue
		}
		if currentIndex[k].HashKey != nextIndex[k].HashKey {
			createIndex = append(createIndex, v)
			continue
		}
		if currentIndex[k].HashKeyType != nextIndex[k].HashKeyType {
			createIndex = append(createIndex, v)
			continue
		}
		if currentIndex[k].RangeKey != nextIndex[k].RangeKey {
			createIndex = append(createIndex, v)
			continue
		}
		if currentIndex[k].RangeKeyType != nextIndex[k].RangeKeyType {
			createIndex = append(createIndex, v)
			continue
		}
	}

	for _, gsi := range deleteIndex {
		for {
			time.Sleep(time.Second)
			current, err := db.Table(tableName).Describe().Run()
			exitOnErr(err)
			if current.Status == dynamo.ActiveStatus {
				break
			}
		}
		if _, err := db.Table(tableName).UpdateTable().DeleteIndex(gsi.Name).Run(); err != nil {
			log.Println("error:", name+"."+gsi.Name, err)
		} else {
			log.Println("delete:", name+"."+gsi.Name)
		}
	}

	for _, gsi := range createIndex {
		for {
			time.Sleep(time.Second)
			current, err := db.Table(tableName).Describe().Run()
			exitOnErr(err)
			if current.Status == dynamo.ActiveStatus {
				break
			}
		}
		if _, err := db.Table(tableName).UpdateTable().CreateIndex(gsi).Run(); err != nil {
			log.Println("error:", name+"."+gsi.Name, err)
		} else {
			log.Println("create:", name+"."+gsi.Name)
		}
	}
}

func exitOnErr(err error) {
	if err != nil {
		panic(err)
	}
}