Files
xingrin/go-backend/internal/database/database.go
yyhuni 5b0416972a feat(go-backend): complete all Go models
- Add scan-related models: ScanLog, ScanInputTarget, ScheduledScan, SubfinderProviderSettings
- Add engine models: Wordlist, NucleiTemplateRepo
- Add notification models: Notification, NotificationSettings
- Add config model: BlacklistRule
- Add statistics models: AssetStatistics, StatisticsHistory
- Add auth models: User (auth_user), Session (django_session)
- Add shopspring/decimal dependency for Vulnerability model
- Update model_test.go with all 33 model table name tests
- All tests passing
2026-01-11 20:29:11 +08:00

85 lines
1.8 KiB
Go

package database
import (
"fmt"
"time"
"github.com/xingrin/go-backend/internal/config"
"gorm.io/driver/postgres"
"gorm.io/gorm"
"gorm.io/gorm/schema"
)
// DB is the global database instance
var DB *gorm.DB
// NewDatabase creates a new database connection
func NewDatabase(cfg *config.DatabaseConfig) (*gorm.DB, error) {
dsn := cfg.DSN()
db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{
NamingStrategy: schema.NamingStrategy{
SingularTable: true, // Use singular table names (e.g., "target" not "targets")
},
// Disable auto-migration to prevent creating new tables
DisableForeignKeyConstraintWhenMigrating: true,
})
if err != nil {
return nil, fmt.Errorf("failed to connect to database: %w", err)
}
// Get underlying sql.DB to configure connection pool
sqlDB, err := db.DB()
if err != nil {
return nil, fmt.Errorf("failed to get underlying sql.DB: %w", err)
}
// Configure connection pool
sqlDB.SetMaxOpenConns(cfg.MaxOpenConns)
sqlDB.SetMaxIdleConns(cfg.MaxIdleConns)
sqlDB.SetConnMaxLifetime(time.Duration(cfg.ConnMaxLifetime) * time.Second)
// Verify connection
if err := sqlDB.Ping(); err != nil {
return nil, fmt.Errorf("failed to ping database: %w", err)
}
return db, nil
}
// InitDatabase initializes the global database instance
func InitDatabase(cfg *config.DatabaseConfig) error {
var err error
DB, err = NewDatabase(cfg)
return err
}
// Close closes the database connection
func Close() error {
if DB != nil {
sqlDB, err := DB.DB()
if err != nil {
return err
}
return sqlDB.Close()
}
return nil
}
// Ping checks if the database connection is alive
func Ping() error {
if DB == nil {
return fmt.Errorf("database not initialized")
}
sqlDB, err := DB.DB()
if err != nil {
return err
}
return sqlDB.Ping()
}
// GetDB returns the global database instance
func GetDB() *gorm.DB {
return DB
}