dataupdater/postgres/postgres.go

package postgres

import (
	"fmt"
	"github.com/Kichiyaki/go-pg-logrus-query-logger/v10"
	"github.com/Kichiyaki/goutil/envutil"
	"github.com/go-pg/pg/v10"
	"github.com/go-pg/pg/v10/orm"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"github.com/tribalwarshelp/shared/tw/twmodel"
)

var log = logrus.WithField("package", "pkg/postgres")

type Config struct {
	SkipDBInitialization bool
}

func Connect(cfg *Config) (*pg.DB, error) {
	db := pg.Connect(prepareOptions())

	if envutil.GetenvBool("LOG_DB_QUERIES") {
		db.AddQueryHook(querylogger.Logger{
			Log:            log,
			MaxQueryLength: 2000,
		})
	}

	if cfg == nil || !cfg.SkipDBInitialization {
		if err := prepareDB(db); err != nil {
			return nil, err
		}
	}

	return db, nil
}

func prepareOptions() *pg.Options {
	return &pg.Options{
		User:     envutil.GetenvString("DB_USER"),
		Password: envutil.GetenvString("DB_PASSWORD"),
		Database: envutil.GetenvString("DB_NAME"),
		Addr:     envutil.GetenvString("DB_HOST") + ":" + envutil.GetenvString("DB_PORT"),
		PoolSize: envutil.GetenvInt("DB_POOL_SIZE"),
	}
}

func prepareDB(db *pg.DB) error {
	tx, err := db.Begin()
	if err != nil {
		return errors.Wrap(err, "couldn't start a transaction")
	}
	defer func() {
		if err := tx.Close(); err != nil {
			log.Warn(errors.Wrap(err, "prepareDB: couldn't rollback the transaction"))
		}
	}()

	dbModels := []interface{}{
		(*twmodel.SpecialServer)(nil),
		(*twmodel.Server)(nil),
		(*twmodel.Version)(nil),
		(*twmodel.PlayerToServer)(nil),
		(*twmodel.PlayerNameChange)(nil),
	}

	for _, model := range dbModels {
		err := tx.Model(model).CreateTable(&orm.CreateTableOptions{
			IfNotExists: true,
		})
		if err != nil {
			return errors.Wrap(err, "couldn't create the table")
		}
	}

	type statementWithParams struct {
		statement string
		params    []interface{}
	}

	for _, s := range []statementWithParams{
		{
			statement: pgDefaultValues,
		},
		{
			statement: allVersionsPGInsertStatements,
		},
		{
			statement: allSpecialServersPGInsertStatements,
		},
		{
			statement: pgDropSchemaFunctions,
			params:    []interface{}{pg.Safe("public")},
		},
		{
			statement: pgFunctions,
		},
	} {
		if _, err := tx.Exec(s.statement, s.params...); err != nil {
			return errors.Wrap(err, "couldn't prepare the db")
		}
	}

	if err := tx.Commit(); err != nil {
		return errors.Wrap(err, "couldn't commit changes")
	}

	var servers []*twmodel.Server
	if err := db.Model(&servers).Select(); err != nil {
		return errors.Wrap(err, "couldn't load servers")
	}

	for _, server := range servers {
		if err := createServerSchema(db, server, true); err != nil {
			return err
		}
	}

	return nil
}

func CreateServerSchema(db *pg.DB, server *twmodel.Server) error {
	return createServerSchema(db, server, false)
}

func SchemaExists(db pg.DBI, schemaName string) bool {
	exists, err := db.
		Model().
		Table("information_schema.schemata").
		Where("schema_name = ?", schemaName).
		Exists()
	if err != nil {
		return false
	}
	return exists
}

func createServerSchema(db *pg.DB, server *twmodel.Server, init bool) error {
	if !init && SchemaExists(db, server.Key) {
		return nil
	}

	tx, err := db.WithParam("SERVER", pg.Safe(server.Key)).Begin()
	if err != nil {
		return errors.Wrap(err, "couldn't start a transaction")
	}
	defer func() {
		if err := tx.Close(); err != nil {
			log.Warn(errors.Wrap(err, "createServerSchema: Couldn't rollback the transaction"))
		}
	}()

	if _, err := tx.Exec(fmt.Sprintf("CREATE SCHEMA IF NOT EXISTS %s", server.Key)); err != nil {
		return errors.Wrap(err, "couldn't create for the server '"+server.Key+"'")
	}

	dbModels := []interface{}{
		(*twmodel.Tribe)(nil),
		(*twmodel.Player)(nil),
		(*twmodel.Village)(nil),
		(*twmodel.Ennoblement)(nil),
		(*twmodel.ServerStats)(nil),
		(*twmodel.TribeHistory)(nil),
		(*twmodel.PlayerHistory)(nil),
		(*twmodel.TribeChange)(nil),
		(*twmodel.DailyPlayerStats)(nil),
		(*twmodel.DailyTribeStats)(nil),
	}

	for _, model := range dbModels {
		err := tx.Model(model).CreateTable(&orm.CreateTableOptions{
			IfNotExists: true,
		})
		if err != nil {
			return err
		}
	}

	statements := []string{
		serverPGFunctions,
		serverPGTriggers,
		serverPGDefaultValues,
	}
	if init {
		statements = append([]string{pgDropSchemaFunctions}, statements...)
	}
	for _, statement := range statements {
		if _, err := tx.Exec(statement, pg.Safe(server.Key), server.VersionCode); err != nil {
			return errors.Wrap(err, "couldn't initialize the schema")
		}
	}

	if err := tx.Commit(); err != nil {
		return errors.Wrap(err, "couldn't commit changes")
	}
	return nil
}