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 }