9a8a93061d
- Add PostgreSQL database connection setup with GORM - Include database initialization and migration functionality - Add connection pooling and error handling - Support for environment-based database configuration
121 lines
2.9 KiB
Go
121 lines
2.9 KiB
Go
package db
|
|
|
|
import (
|
|
"fmt"
|
|
"os"
|
|
"time"
|
|
|
|
"git.secnex.io/secnex/idp-api/models"
|
|
"gorm.io/driver/postgres"
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/logger"
|
|
)
|
|
|
|
// Global database instance
|
|
var DB *gorm.DB
|
|
|
|
type Database struct {
|
|
Host string
|
|
Port string
|
|
User string
|
|
Password string
|
|
DBName string
|
|
SSLMode string
|
|
Conn *gorm.DB
|
|
}
|
|
|
|
type DatabaseConnection struct {
|
|
DB *gorm.DB
|
|
}
|
|
|
|
func NewDatabase(host, port, user, password, dbName, sslMode string) *Database {
|
|
return &Database{Host: host, Port: port, User: user, Password: password, DBName: dbName, SSLMode: sslMode}
|
|
}
|
|
|
|
// InitializeDB initializes the global database connection with connection pooling
|
|
func InitializeDB() error {
|
|
// Get database configuration from environment variables
|
|
host := getEnv("DB_HOST", "10.2.1.2")
|
|
port := getEnv("DB_PORT", "5432")
|
|
user := getEnv("DB_USER", "postgres")
|
|
password := getEnv("DB_PASSWORD", "postgres")
|
|
dbName := getEnv("DB_NAME", "idp")
|
|
sslMode := getEnv("DB_SSLMODE", "disable")
|
|
|
|
dsn := fmt.Sprintf("host=%s port=%s user=%s password=%s dbname=%s sslmode=%s",
|
|
host, port, user, password, dbName, sslMode)
|
|
|
|
// Configure GORM with connection pooling
|
|
config := &gorm.Config{
|
|
Logger: logger.Default.LogMode(logger.Info),
|
|
}
|
|
|
|
db, err := gorm.Open(postgres.Open(dsn), config)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to connect to database: %w", err)
|
|
}
|
|
|
|
// Configure connection pool
|
|
sqlDB, err := db.DB()
|
|
if err != nil {
|
|
return fmt.Errorf("failed to get underlying sql.DB: %w", err)
|
|
}
|
|
|
|
// Set connection pool settings
|
|
sqlDB.SetMaxIdleConns(10) // Maximum number of idle connections
|
|
sqlDB.SetMaxOpenConns(100) // Maximum number of open connections
|
|
sqlDB.SetConnMaxLifetime(time.Hour) // Maximum lifetime of a connection
|
|
|
|
// Auto migrate models
|
|
if err := db.AutoMigrate(&models.User{}, &models.Session{}, &models.ApiKey{}); err != nil {
|
|
return fmt.Errorf("failed to auto migrate: %w", err)
|
|
}
|
|
|
|
// Set global database instance
|
|
DB = db
|
|
|
|
return nil
|
|
}
|
|
|
|
// GetDB returns the global database instance
|
|
func GetDB() *gorm.DB {
|
|
return DB
|
|
}
|
|
|
|
// CloseDB closes the database connection
|
|
func CloseDB() error {
|
|
if DB != nil {
|
|
sqlDB, err := DB.DB()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return sqlDB.Close()
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Helper function to get environment variables with fallback
|
|
func getEnv(key, fallback string) string {
|
|
if value := os.Getenv(key); value != "" {
|
|
return value
|
|
}
|
|
return fallback
|
|
}
|
|
|
|
// Legacy methods for backward compatibility
|
|
func (d *Database) Connect() (*gorm.DB, error) {
|
|
dsn := fmt.Sprintf("host=%s port=%s user=%s password=%s dbname=%s sslmode=%s",
|
|
d.Host, d.Port, d.User, d.Password, d.DBName, d.SSLMode)
|
|
|
|
db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if err := db.AutoMigrate(&models.User{}, &models.Session{}, &models.ApiKey{}); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return db, nil
|
|
}
|