Michael McGuinness
111 lines
2.1 KiB
Go
111 lines
2.1 KiB
Go
package database
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"log"
|
|
"net/url"
|
|
"queryorchestration/internal/server/env"
|
|
|
|
"github.com/docker/go-connections/nat"
|
|
"github.com/jackc/pgx/v5"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
)
|
|
|
|
type config struct {
|
|
user string
|
|
password string
|
|
host string
|
|
port nat.Port
|
|
name string
|
|
disableSSL bool
|
|
}
|
|
|
|
func mustGetConfig() *config {
|
|
user := env.GetPanic("DB_USER")
|
|
pass := env.GetPanic("DB_PASS")
|
|
host := env.GetPanic("DB_HOST")
|
|
port := env.GetPanic("DB_PORT")
|
|
natPort, err := nat.NewPort("tcp", port)
|
|
if err != nil {
|
|
log.Panicf("Failed to create port: %v", err)
|
|
}
|
|
name := env.GetPanic("DB_NAME")
|
|
|
|
disableSSL, _ := env.Get("DB_NOSSL")
|
|
|
|
return &config{
|
|
user: user,
|
|
password: pass,
|
|
host: host,
|
|
port: natPort,
|
|
name: name,
|
|
disableSSL: disableSSL == "1",
|
|
}
|
|
}
|
|
|
|
func mustGetURI() *url.URL {
|
|
driver := "postgres"
|
|
conf := mustGetConfig()
|
|
opts := ""
|
|
|
|
if conf.disableSSL {
|
|
opts += "sslmode=disable"
|
|
}
|
|
|
|
connStr := fmt.Sprintf("%s://%s:%s@%s:%d/%s?%s", driver, conf.user, conf.password, conf.host, conf.port.Int(), conf.name, opts)
|
|
|
|
uri, err := url.Parse(connStr)
|
|
if err != nil {
|
|
log.Panicf("Unable to parse URI: %s", err)
|
|
}
|
|
|
|
return uri
|
|
}
|
|
|
|
func getPoolConfig() (*pgxpool.Config, error) {
|
|
connStr := mustGetURI()
|
|
|
|
config, err := pgxpool.ParseConfig(connStr.String())
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
enumTypes := []string{
|
|
"querytype",
|
|
}
|
|
|
|
config.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error {
|
|
for _, enumName := range enumTypes {
|
|
dt, err := conn.LoadType(ctx, enumName)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to load suitable enum type: %w", err)
|
|
}
|
|
conn.TypeMap().RegisterType(dt)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
return config, nil
|
|
}
|
|
|
|
func mustGetPoolConfig() *pgxpool.Config {
|
|
config, err := getPoolConfig()
|
|
if err != nil {
|
|
log.Panicf("Unable to create database config: %v\n", err)
|
|
}
|
|
|
|
return config
|
|
}
|
|
|
|
func GetDBPool(ctx context.Context) *pgxpool.Pool {
|
|
config := mustGetPoolConfig()
|
|
|
|
pool, err := pgxpool.NewWithConfig(ctx, config)
|
|
if err != nil {
|
|
log.Panicf("Unable to create database pool: %v\n", err)
|
|
}
|
|
|
|
return pool
|
|
}
|