2021-05-03 14:35:24 +00:00
|
|
|
package postgres
|
2021-04-25 15:02:58 +00:00
|
|
|
|
|
|
|
import (
|
|
|
|
"fmt"
|
2021-05-03 14:33:12 +00:00
|
|
|
"github.com/Kichiyaki/go-pg-logrus-query-logger/v10"
|
|
|
|
"github.com/Kichiyaki/goutil/envutil"
|
2021-04-25 15:02:58 +00:00
|
|
|
"github.com/go-pg/pg/v10"
|
|
|
|
"github.com/go-pg/pg/v10/orm"
|
|
|
|
"github.com/pkg/errors"
|
|
|
|
"github.com/sirupsen/logrus"
|
2021-05-03 14:33:12 +00:00
|
|
|
"github.com/tribalwarshelp/shared/tw/twmodel"
|
2021-04-25 15:02:58 +00:00
|
|
|
)
|
|
|
|
|
2021-07-11 05:43:44 +00:00
|
|
|
var log = logrus.WithField("package", "pkg/postgres")
|
2021-04-25 15:02:58 +00:00
|
|
|
|
|
|
|
type Config struct {
|
2021-07-11 06:33:39 +00:00
|
|
|
SkipDBInitialization bool
|
2021-04-25 15:02:58 +00:00
|
|
|
}
|
|
|
|
|
2021-05-03 14:35:24 +00:00
|
|
|
func Connect(cfg *Config) (*pg.DB, error) {
|
2021-04-25 15:02:58 +00:00
|
|
|
db := pg.Connect(prepareOptions())
|
|
|
|
|
2021-07-11 08:56:44 +00:00
|
|
|
if envutil.GetenvBool("LOG_DB_QUERIES") {
|
2021-05-03 14:33:12 +00:00
|
|
|
db.AddQueryHook(querylogger.Logger{
|
2021-04-28 18:13:14 +00:00
|
|
|
Log: log,
|
2021-05-12 14:42:29 +00:00
|
|
|
MaxQueryLength: 2000,
|
2021-04-25 15:02:58 +00:00
|
|
|
})
|
|
|
|
}
|
|
|
|
|
2021-07-11 06:33:39 +00:00
|
|
|
if cfg == nil || !cfg.SkipDBInitialization {
|
|
|
|
if err := prepareDB(db); err != nil {
|
|
|
|
return nil, err
|
|
|
|
}
|
2021-04-25 15:02:58 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
return db, nil
|
|
|
|
}
|
|
|
|
|
|
|
|
func prepareOptions() *pg.Options {
|
|
|
|
return &pg.Options{
|
2021-05-03 14:33:12 +00:00
|
|
|
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"),
|
2021-04-25 15:02:58 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func prepareDB(db *pg.DB) error {
|
|
|
|
tx, err := db.Begin()
|
|
|
|
if err != nil {
|
2021-05-14 12:57:05 +00:00
|
|
|
return errors.Wrap(err, "couldn't start a transaction")
|
2021-04-25 15:02:58 +00:00
|
|
|
}
|
2021-04-27 18:06:38 +00:00
|
|
|
defer func() {
|
|
|
|
if err := tx.Close(); err != nil {
|
2021-05-14 12:57:05 +00:00
|
|
|
log.Warn(errors.Wrap(err, "prepareDB: couldn't rollback the transaction"))
|
2021-04-27 18:06:38 +00:00
|
|
|
}
|
|
|
|
}()
|
2021-04-25 15:02:58 +00:00
|
|
|
|
|
|
|
dbModels := []interface{}{
|
2021-05-03 14:33:12 +00:00
|
|
|
(*twmodel.SpecialServer)(nil),
|
|
|
|
(*twmodel.Server)(nil),
|
|
|
|
(*twmodel.Version)(nil),
|
|
|
|
(*twmodel.PlayerToServer)(nil),
|
|
|
|
(*twmodel.PlayerNameChange)(nil),
|
2021-04-25 15:02:58 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
for _, model := range dbModels {
|
|
|
|
err := tx.Model(model).CreateTable(&orm.CreateTableOptions{
|
|
|
|
IfNotExists: true,
|
|
|
|
})
|
|
|
|
if err != nil {
|
2021-05-14 12:57:05 +00:00
|
|
|
return errors.Wrap(err, "couldn't create the table")
|
2021-04-25 15:02:58 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
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 {
|
2021-05-14 12:57:05 +00:00
|
|
|
return errors.Wrap(err, "couldn't prepare the db")
|
2021-04-25 15:02:58 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
if err := tx.Commit(); err != nil {
|
2021-05-14 12:57:05 +00:00
|
|
|
return errors.Wrap(err, "couldn't commit changes")
|
2021-04-25 15:02:58 +00:00
|
|
|
}
|
|
|
|
|
2021-05-03 14:33:12 +00:00
|
|
|
var servers []*twmodel.Server
|
2021-04-25 15:02:58 +00:00
|
|
|
if err := db.Model(&servers).Select(); err != nil {
|
2021-05-14 12:57:05 +00:00
|
|
|
return errors.Wrap(err, "couldn't load servers")
|
2021-04-25 15:02:58 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
for _, server := range servers {
|
2021-07-11 06:27:31 +00:00
|
|
|
if err := createServerSchema(db, server, true); err != nil {
|
2021-04-25 15:02:58 +00:00
|
|
|
return err
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
2021-07-11 06:27:31 +00:00
|
|
|
func CreateServerSchema(db *pg.DB, server *twmodel.Server) error {
|
|
|
|
return createServerSchema(db, server, false)
|
2021-04-25 15:02:58 +00:00
|
|
|
}
|
|
|
|
|
2021-07-11 06:27:31 +00:00
|
|
|
func SchemaExists(db pg.DBI, schemaName string) bool {
|
2021-04-25 15:02:58 +00:00
|
|
|
exists, err := db.
|
|
|
|
Model().
|
|
|
|
Table("information_schema.schemata").
|
|
|
|
Where("schema_name = ?", schemaName).
|
|
|
|
Exists()
|
|
|
|
if err != nil {
|
|
|
|
return false
|
|
|
|
}
|
|
|
|
return exists
|
|
|
|
}
|
|
|
|
|
2021-07-11 06:27:31 +00:00
|
|
|
func createServerSchema(db *pg.DB, server *twmodel.Server, init bool) error {
|
2021-04-25 15:02:58 +00:00
|
|
|
if !init && SchemaExists(db, server.Key) {
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
|
|
|
tx, err := db.WithParam("SERVER", pg.Safe(server.Key)).Begin()
|
|
|
|
if err != nil {
|
2021-05-14 12:57:05 +00:00
|
|
|
return errors.Wrap(err, "couldn't start a transaction")
|
2021-04-25 15:02:58 +00:00
|
|
|
}
|
2021-04-27 18:06:38 +00:00
|
|
|
defer func() {
|
|
|
|
if err := tx.Close(); err != nil {
|
2021-07-11 06:27:31 +00:00
|
|
|
log.Warn(errors.Wrap(err, "createServerSchema: Couldn't rollback the transaction"))
|
2021-04-27 18:06:38 +00:00
|
|
|
}
|
|
|
|
}()
|
2021-04-25 15:02:58 +00:00
|
|
|
|
|
|
|
if _, err := tx.Exec(fmt.Sprintf("CREATE SCHEMA IF NOT EXISTS %s", server.Key)); err != nil {
|
2021-05-14 12:57:05 +00:00
|
|
|
return errors.Wrap(err, "couldn't create for the server '"+server.Key+"'")
|
2021-04-25 15:02:58 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
dbModels := []interface{}{
|
2021-05-03 14:33:12 +00:00
|
|
|
(*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),
|
2021-04-25 15:02:58 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
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 {
|
2021-05-14 12:57:05 +00:00
|
|
|
return errors.Wrap(err, "couldn't initialize the schema")
|
2021-04-25 15:02:58 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
if err := tx.Commit(); err != nil {
|
2021-05-14 12:57:05 +00:00
|
|
|
return errors.Wrap(err, "couldn't commit changes")
|
2021-04-25 15:02:58 +00:00
|
|
|
}
|
|
|
|
return nil
|
|
|
|
}
|