dataupdater/db/db.go

package db

import (
	"fmt"
	gopglogrusquerylogger "github.com/Kichiyaki/go-pg-logrus-query-logger/v10"
	"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/models"

	envutils "github.com/tribalwarshelp/cron/utils/env"
)

var log = logrus.WithField("package", "db")

type Config struct {
	LogQueries bool
}

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

	if cfg != nil && cfg.LogQueries {
		db.AddQueryHook(gopglogrusquerylogger.QueryLogger{
			Entry:          log,
			MaxQueryLength: 5000,
		})
	}

	if err := prepareDB(db); err != nil {
		return nil, errors.Wrap(err, "New")
	}

	return db, nil
}

func prepareOptions() *pg.Options {
	return &pg.Options{
		User:     envutils.GetenvString("DB_USER"),
		Password: envutils.GetenvString("DB_PASSWORD"),
		Database: envutils.GetenvString("DB_NAME"),
		Addr:     envutils.GetenvString("DB_HOST") + ":" + envutils.GetenvString("DB_PORT"),
		PoolSize: envutils.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 tx.Close()

	dbModels := []interface{}{
		(*models.SpecialServer)(nil),
		(*models.Server)(nil),
		(*models.Version)(nil),
		(*models.PlayerToServer)(nil),
		(*models.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 initialize the db")
		}
	}

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

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

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

	return nil
}

func CreateSchema(db *pg.DB, server *models.Server) error {
	return createSchema(db, server, false)
}

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

func createSchema(db *pg.DB, server *models.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, "CreateSchema: couldn't start a transaction")
	}
	defer tx.Close()

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

	dbModels := []interface{}{
		(*models.Tribe)(nil),
		(*models.Player)(nil),
		(*models.Village)(nil),
		(*models.Ennoblement)(nil),
		(*models.ServerStats)(nil),
		(*models.TribeHistory)(nil),
		(*models.PlayerHistory)(nil),
		(*models.TribeChange)(nil),
		(*models.DailyPlayerStats)(nil),
		(*models.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, "CreateSchema: couldn't initialize the schema")
		}
	}

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