任务三

internal/data/data.go

@@ -1,43 +1,44 @@
 package data
 
 import (
-	"context"
-	"log"
-	"os"
-
 	"kra/internal/conf"
-	"kra/internal/data/ent"
-	"kra/internal/data/ent/migrate"
+	"kra/internal/data/query"
+
 	_ "github.com/go-sql-driver/mysql"
 	"github.com/google/wire"
+	"gorm.io/driver/mysql"
+	"gorm.io/gorm"
+	"gorm.io/gorm/logger"
 )
 
 // ProviderSet is data providers.
-var ProviderSet = wire.NewSet(NewData, NewAdminRepo)
+var ProviderSet = wire.NewSet(NewDB, NewData)
 
-// Data is a struct that contains the database client.
+// Data 数据层
 type Data struct {
-	db *ent.Client
+	DB *gorm.DB
 }
 
-// NewData creates a new Data instance.
-func NewData(c *conf.Data) (*Data, func(), error) {
-	db, err := ent.Open(c.Database.Driver, c.Database.Source)
+// NewDB 创建数据库连接
+func NewDB(c *conf.Data) (*gorm.DB, error) {
+	db, err := gorm.Open(mysql.Open(c.Database.Source), &gorm.Config{
+		Logger: logger.Default.LogMode(logger.Info),
+	})
 	if err != nil {
-		log.Fatalf("failed opening connection to database: %v", err)
+		return nil, err
 	}
-	if os.Getenv("DEPLOY_ENV") == "dev" {
-		// Enable debug mode for detailed logging.
-		db = db.Debug()
-		// Run the auto migration tool.
-		if err = db.Schema.Create(context.Background(), migrate.WithDropIndex(true)); err != nil {
-			return nil, nil, err
+	// 初始化 gen query
+	query.SetDefault(db)
+	return db, nil
+}
+
+// NewData 创建 Data 实例
+func NewData(db *gorm.DB) (*Data, func(), error) {
+	cleanup := func() {
+		sqlDB, _ := db.DB()
+		if sqlDB != nil {
+			sqlDB.Close()
 		}
 	}
-	cleanup := func() {
-		db.Close()
-	}
-	return &Data{
-		db: db,
-	}, cleanup, nil
+	return &Data{DB: db}, cleanup, nil
 }