github.com/safedep/dry@v0.0.0-20241016050132-a15651f0548b/db/postgres_adapter.go (about)

     1  package db
     2  
     3  import (
     4  	"fmt"
     5  	"time"
     6  
     7  	"github.com/safedep/dry/log"
     8  	"github.com/safedep/dry/retry"
     9  	"gorm.io/driver/postgres"
    10  	"gorm.io/gorm"
    11  	"gorm.io/plugin/opentelemetry/tracing"
    12  	"gorm.io/plugin/prometheus"
    13  )
    14  
    15  // GORM internally uses pgx as the default driver for PostgreSQL. The DSN
    16  // format can be anything supported by pgx. For example:
    17  // postgres://${user}:${password}@${host}/${db}?sslmode=verify-none
    18  type PostgreSqlAdapterConfig struct {
    19  	DSN string
    20  
    21  	// A name to report as part of the tracer plugin
    22  	TracingDBName string
    23  
    24  	EnableTracing bool
    25  	EnableMetrics bool
    26  
    27  	// This is an optional pointer to the SqlAdapterConfig struct.
    28  	// If not supplied, we will use the defaultSqlAdapterConfig.
    29  	SqlAdapterConfig *SqlAdapterConfig
    30  }
    31  
    32  type PostgreSqlAdapter struct {
    33  	*baseSqlAdapter
    34  
    35  	db     *gorm.DB
    36  	config PostgreSqlAdapterConfig
    37  }
    38  
    39  func NewPostgreSqlAdapter(config PostgreSqlAdapterConfig) (SqlDataAdapter, error) {
    40  	dsn := config.DSN
    41  	if dsn == "" {
    42  		dsnFromEnv, err := DatabaseURL()
    43  		if err != nil {
    44  			return nil, fmt.Errorf("failed to get database URL from env: %w", err)
    45  		}
    46  
    47  		log.Debugf("Connecting to PostgreSQL database with DSN from env")
    48  		dsn = dsnFromEnv
    49  	} else {
    50  		log.Debugf("Connecting to PostgreSQL database with DSN from config")
    51  	}
    52  
    53  	var db *gorm.DB
    54  	var err error
    55  
    56  	retry.InvokeWithRetry(retry.RetryConfig{
    57  		Count: 30,
    58  		Sleep: 1 * time.Second,
    59  	}, func(arg retry.RetryFuncArg) error {
    60  		// https://gorm.io/docs/connecting_to_the_database.html#PostgreSQL
    61  		db, err = gorm.Open(postgres.Open(dsn), &gorm.Config{})
    62  		if err != nil {
    63  			log.Debugf("[%d/%d] Failed to connect to PostgreSQL server: %v",
    64  				arg.Current, arg.Total, err)
    65  		}
    66  
    67  		return err
    68  	})
    69  
    70  	if err != nil {
    71  		return nil, err
    72  	}
    73  
    74  	log.Debugf("PostgreSQL database connection established")
    75  
    76  	if config.EnableTracing {
    77  		if err := db.Use(tracing.NewPlugin(tracing.WithDBName(config.TracingDBName))); err != nil {
    78  			return nil, err
    79  		}
    80  	}
    81  
    82  	if config.EnableMetrics {
    83  		metricsMiddleware := prometheus.New(prometheus.Config{
    84  			DBName:          config.TracingDBName,
    85  			RefreshInterval: 15,
    86  		})
    87  
    88  		if err := db.Use(metricsMiddleware); err != nil {
    89  			return nil, err
    90  		}
    91  	}
    92  
    93  	baseSqlAdapter := &baseSqlAdapter{
    94  		db:     db,
    95  		config: config.SqlAdapterConfig,
    96  	}
    97  
    98  	err = baseSqlAdapter.SetupConnectionPool()
    99  	if err != nil {
   100  		return nil, err
   101  	}
   102  
   103  	postgreSqlAdapter := &PostgreSqlAdapter{db: db, config: config, baseSqlAdapter: baseSqlAdapter}
   104  
   105  	err = postgreSqlAdapter.Ping()
   106  	return postgreSqlAdapter, nil
   107  }