go-libs/pkg/storage/bun/connect/flags.go

package connect

import (
	"context"
	"database/sql/driver"
	"errors"
	"fmt"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/jackc/pgx/v5"
	"github.com/spf13/pflag"

	"github.com/hanzo-fi/go-libs/v5/pkg/cloud/aws/iam"
	logging "github.com/hanzo-fi/go-libs/v5/pkg/observe/log"
)

const (
	PostgresURIFlag             = "postgres-uri"
	PostgresAWSEnableIAMFlag    = "postgres-aws-enable-iam"
	PostgresMaxIdleConnsFlag    = "postgres-max-idle-conns"
	PostgresMaxOpenConnsFlag    = "postgres-max-open-conns"
	PostgresConnMaxIdleTimeFlag = "postgres-conn-max-idle-time"
	PostgresConnMaxLifetimeFlag = "postgres-conn-max-lifetime"
)

func AddFlags(flags *pflag.FlagSet) {
	flags.String(PostgresURIFlag, "", "Postgres URI")
	flags.Bool(PostgresAWSEnableIAMFlag, false, "Enable AWS IAM authentication")
	flags.Int(PostgresMaxIdleConnsFlag, 0, "Max Idle connections")
	flags.Duration(PostgresConnMaxIdleTimeFlag, time.Minute, "Max Idle time for connections")
	flags.Duration(PostgresConnMaxLifetimeFlag, 0, "Max lifetime for connections")
	flags.Int(PostgresMaxOpenConnsFlag, 20, "Max opened connections")
}

type Option func(options *pgx.ConnConfig)

func WithRuntimeParams(params map[string]string) Option {
	return func(options *pgx.ConnConfig) {
		for k, v := range params {
			options.RuntimeParams[k] = v
		}
	}
}

func GetAWSIAMAuthConnector(flags *pflag.FlagSet, ctx context.Context, opts ...Option) (func(s string) (driver.Connector, error), error) {
	var cfg aws.Config
	var err error

	if flags != nil {
		cfg, err = config.LoadDefaultConfig(ctx, iam.LoadOptionFromFlags(flags))
	} else {
		cfg, err = config.LoadDefaultConfig(ctx)
	}

	if err != nil {
		return nil, err
	}

	connector := func(s string) (driver.Connector, error) {
		return &iamConnector{
			dsn: s,
			driver: &iamDriver{
				awsConfig: cfg,
			},
			options: opts,
			logger:  logging.FromContext(ctx),
		}, nil
	}

	return connector, nil
}

func ConnectionOptionsFromFlags(flags *pflag.FlagSet, ctx context.Context, opts ...Option) (*ConnectionOptions, error) {
	var err error
	var connector func(string) (driver.Connector, error)

	awsEnable, _ := flags.GetBool(PostgresAWSEnableIAMFlag)
	if awsEnable {
		connector, err = GetAWSIAMAuthConnector(flags, ctx, opts...)
		if err != nil {
			return nil, err
		}
	} else {
		connector = func(dsn string) (driver.Connector, error) {
			parseConfig, err := pgx.ParseConfig(dsn)
			if err != nil {
				return nil, fmt.Errorf("failed to parse dsn: %w", err)
			}

			for _, opt := range opts {
				opt(parseConfig)
			}

			return buildPGXConnector(parseConfig), nil
		}
	}

	postgresUri, _ := flags.GetString(PostgresURIFlag)
	if postgresUri == "" {
		return nil, errors.New("missing postgres uri")
	}
	maxIdleConns, _ := flags.GetInt(PostgresMaxIdleConnsFlag)
	connMaxIdleConns, _ := flags.GetDuration(PostgresConnMaxIdleTimeFlag)
	connMaxLifetime, _ := flags.GetDuration(PostgresConnMaxLifetimeFlag)
	maxOpenConns, _ := flags.GetInt(PostgresMaxOpenConnsFlag)

	return &ConnectionOptions{
		DatabaseSourceName: postgresUri,
		MaxIdleConns:       maxIdleConns,
		ConnMaxIdleTime:    connMaxIdleConns,
		ConnMaxLifetime:    connMaxLifetime,
		MaxOpenConns:       maxOpenConns,
		Connector:          connector,
	}, nil
}