oh-my-url/app/db/base.go

package db

import (
	"fmt"
	"github.com/aykhans/oh-my-url/app/config"
	gormPostgres "gorm.io/driver/postgres"
	"gorm.io/gorm"
	"time"
)

type DB interface {
	Init()
	CreateURL(url string) (string, error)
	GetURL(key string) (string, error)
}

func GetDB() DB {
	db := config.GetDB()
	switch db {
	case "postgres":
		postgresConf := config.GetPostgresConfig()
		dsn := fmt.Sprintf(
			"host=%s user=%s password=%s dbname=%s port=%s sslmode=%s TimeZone=%s",
			postgresConf.HOST,
			postgresConf.USER,
			postgresConf.PASSWORD,
			postgresConf.DBNAME,
			postgresConf.PORT,
			"disable",
			"UTC",
		)

		var db *gorm.DB
		var err error
		for i := 0; i < 5; i++ {
			db, err = gorm.Open(gormPostgres.Open(dsn), &gorm.Config{})
			if err == nil {
				break
			}
			time.Sleep(3 * time.Second)
		}
		if err != nil {
			panic(err)
		}
		return &Postgres{gormDB: db}
	}
	return nil
}