github.com/dhax/go-base@v0.0.0-20231004214136-8be7e5c1972b/database/postgres.go (about)

     1  // Package database implements postgres connection and queries.
     2  package database
     3  
     4  import (
     5  	"log"
     6  
     7  	"github.com/spf13/viper"
     8  
     9  	"github.com/go-pg/pg"
    10  )
    11  
    12  // DBConn returns a postgres connection pool.
    13  func DBConn() (*pg.DB, error) {
    14  	viper.SetDefault("db_network", "tcp")
    15  	viper.SetDefault("db_addr", "localhost:5432")
    16  	viper.SetDefault("db_user", "postgres")
    17  	viper.SetDefault("db_password", "postgres")
    18  	viper.SetDefault("db_database", "postgres")
    19  
    20  	db := pg.Connect(&pg.Options{
    21  		Network:  viper.GetString("db_network"),
    22  		Addr:     viper.GetString("db_addr"),
    23  		User:     viper.GetString("db_user"),
    24  		Password: viper.GetString("db_password"),
    25  		Database: viper.GetString("db_database"),
    26  	})
    27  
    28  	if err := checkConn(db); err != nil {
    29  		return nil, err
    30  	}
    31  
    32  	if viper.GetBool("db_debug") {
    33  		db.AddQueryHook(&logSQL{})
    34  	}
    35  
    36  	return db, nil
    37  }
    38  
    39  type logSQL struct{}
    40  
    41  func (l *logSQL) BeforeQuery(e *pg.QueryEvent) {}
    42  
    43  func (l *logSQL) AfterQuery(e *pg.QueryEvent) {
    44  	query, err := e.FormattedQuery()
    45  	if err != nil {
    46  		panic(err)
    47  	}
    48  	log.Println(query)
    49  }
    50  
    51  func checkConn(db *pg.DB) error {
    52  	var n int
    53  	_, err := db.QueryOne(pg.Scan(&n), "SELECT 1")
    54  	return err
    55  }