2025-03-17 22:19:17 +08:00

62 lines
1006 B
Go

package database
import (
"fmt"
"legalgo-BE-go/config"
"gorm.io/driver/postgres"
"gorm.io/gorm"
)
type DB struct {
*gorm.DB
}
func NewDB(cfg *config.Config) (*DB, error) {
dsn := fmt.Sprintf(
"host=%s user=%s password=%s dbname=%s port=%v sslmode=disable",
cfg.Database.Host,
cfg.Database.Username,
cfg.Database.Password,
cfg.Database.DB,
cfg.Database.Port,
)
if dsn == "" {
return nil, fmt.Errorf("DATABASE_URL environment is not set")
}
db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{})
if err != nil {
return nil, fmt.Errorf("failed to connect to database: %w", err)
}
return &DB{db}, nil
}
func (db *DB) DropTables() error {
// Auto Migrate the User model
return db.Migrator().DropTable(
&Tag{},
&Category{},
&News{},
)
}
func (db *DB) Migrate() error {
// Auto Migrate the User model
return db.AutoMigrate(
&Staff{},
&SubscribePlan{},
&Subscribe{},
&User{},
&News{},
&Tag{},
&Category{},
&Ads{},
&LogAds{},
&LogNews{},
)
}