package cmd import ( "ariga.io/atlas/sql/sqltool" "context" "entgo.io/ent/dialect" "entgo.io/ent/dialect/sql/schema" "fmt" entmigrate "framework_v2/internal/ent/migrate" _ "github.com/jackc/pgx/v5/stdlib" "github.com/pressly/goose/v3" "github.com/spf13/cobra" "log" "os" ) var dsnCommand = &cobra.Command{ Use: "dsn", Short: "生成 DSN", Long: "生成 DSN", Run: func(cmd *cobra.Command, args []string) { fmt.Print(config.DB.Driver + "://" + config.DB.DSN) }, } var migrateCommand = &cobra.Command{ Use: "migrate [command]", Short: "goose 迁移,用法 ", Long: "适用于生产环境的数据库迁移", Run: func(cmd *cobra.Command, args []string) { if len(args) == 0 { _ = cmd.Help() return } RunMigrate(args) }, } var createMigrateCommand = &cobra.Command{ Use: "create-migrate", Short: "新建迁移", Long: "从 internal/ent 中新建迁移。在这之前,需要运行 go generate ./internal/ent", Run: func(cmd *cobra.Command, args []string) { generateMigration() }, } // RunMigrate 为数据库函数 func RunMigrate(args []string) { db, err := goose.OpenDBWithDriver("postgres", config.DB.DSN2) if err != nil { log.Fatalf("goose: failed to open DB: %v\n", err) } defer func() { if err := db.Close(); err != nil { log.Fatalf("goose: failed to close DB: %v\n", err) } }() command := args[0] var arguments []string if len(args) > 3 { arguments = append(arguments, args[3:]...) } if err := goose.RunContext(context.Background(), command, db, "internal/migrations", arguments...); err != nil { log.Fatalf("goose %v: %v", command, err) } } func generateMigration() { ctx := context.Background() dir, err := sqltool.NewGooseDir("internal/migrations") if err != nil { log.Fatalf("failed creating atlas migration directory: %v", err) } // Migrate diff options. opts := []schema.MigrateOption{ schema.WithDir(dir), // provide migration directory schema.WithMigrationMode(schema.ModeInspect), // provide migration mode schema.WithDialect(dialect.Postgres), // Ent dialect to use } if len(os.Args) != 3 { log.Fatalln("migration name is required. Use: 'go run -mod=mod internal/ent/migrate/main.go '") } err = entmigrate.NamedDiff(ctx, config.DB.Driver+"://"+config.DB.DSN, os.Args[2], opts...) if err != nil { log.Fatalf("failed generating migration file: %v", err) } }