query-orchestration/internal/database/connection.go

package database

import (
	"context"
	"fmt"
	"log"
	"net/url"
	"queryorchestration/internal/database/repository"
	"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
}

func ExecuteTransaction(ctx context.Context, db *Connection, executeQueries func(context.Context, *repository.Queries) error) error {
	tx, err := db.Pool.Begin(ctx)
	if err != nil {
		return err
	}
	defer func() {
		_ = tx.Rollback(ctx)
	}()

	qtx := db.Queries.WithTx(tx)

	err = executeQueries(ctx, qtx)
	if err != nil {
		return err
	}

	err = tx.Commit(ctx)
	if err != nil {
		return err
	}

	return nil
}