database.go 759 B

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849
  1. package database
  2. import (
  3. "log"
  4. "os"
  5. "git.familybaby.top/flight/csdn/models"
  6. "gorm.io/driver/sqlite"
  7. "gorm.io/gorm"
  8. )
  9. var DB *gorm.DB
  10. func InitDB() error {
  11. dbPath := "device.db"
  12. if os.Getenv("DB_PATH") != "" {
  13. dbPath = os.Getenv("DB_PATH")
  14. }
  15. var err error
  16. DB, err = gorm.Open(sqlite.Open(dbPath), &gorm.Config{})
  17. if err != nil {
  18. return err
  19. }
  20. // 自动迁移表结构
  21. err = DB.AutoMigrate(
  22. &models.User{},
  23. &models.Parse{},
  24. )
  25. if err != nil {
  26. return err
  27. }
  28. log.Println("Database initialized successfully")
  29. return nil
  30. }
  31. func CloseDB() {
  32. db, err := DB.DB()
  33. if err != nil {
  34. log.Println("Error getting database instance:", err)
  35. return
  36. }
  37. err = db.Close()
  38. if err != nil {
  39. log.Println("Error closing database:", err)
  40. }
  41. }