query-orchestration/internal/database/connection.go

package database

import (
	"context"
	"fmt"
	"log"
	"queryorchestration/internal/env"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

type config struct {
	user       string
	password   string
	host       string
	port       string
	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")
	name := env.GetPanic("DB_NAME")

	disableSSL, _ := env.Get("DB_NOSSL")

	return &config{
		user:       user,
		password:   pass,
		host:       host,
		port:       port,
		name:       name,
		disableSSL: disableSSL == "1",
	}
}

func mustGetConnectionString() string {
	driver := "postgres"
	conf := mustGetConfig()
	opts := ""

	if conf.disableSSL {
		opts += "sslmode=disable"
	}

	return fmt.Sprintf("%s://%s:%s@%s:%s/%s?%s", driver, conf.user, conf.password, conf.host, conf.port, conf.name, opts)
}

func getPoolConfig() (*pgxpool.Config, error) {
	connStr := mustGetConnectionString()

	config, err := pgxpool.ParseConfig(connStr)
	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 := mustGetConnectionString()

	conn, err := pgx.Connect(ctx, config)
	if err != nil {
		log.Panicf("Unable to connect to database: %v\n", err)
	}

	return conn
}