Files
query-orchestration/internal/database/connection.go
T
Michael McGuinness 7da57058d4 Merged in bugfix/flakydb (pull request #21)
Bugfix/flakydb

* init

* fixenvvar

* fixtest
2025-01-14 17:28:26 +00:00

107 lines
2.0 KiB
Go

package database
import (
"context"
"fmt"
"log"
"net/url"
"queryorchestration/internal/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
}
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
}
func GetDBConn(ctx context.Context) *pgx.Conn {
config := mustGetURI()
conn, err := pgx.Connect(ctx, config.String())
if err != nil {
log.Panicf("Unable to connect to database: %v\n", err)
}
return conn
}