From 8a29fc18751eb2db7e2f5e5cf3a982db4f2a3a84 Mon Sep 17 00:00:00 2001
From: Twilight <im@ivampiresp.com>
Date: Sun, 14 Jul 2024 19:48:07 +0800
Subject: [PATCH] =?UTF-8?q?=E6=94=B9=E8=BF=9B=20=E8=BF=81=E7=A7=BB?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/cmd/migrate.go                  | 102 ++++++++++++++++++++++-
 internal/cmd/root.go                     |  11 ++-
 internal/migrations/00002_rename_root.go |  21 +++++
 3 files changed, 127 insertions(+), 7 deletions(-)
 create mode 100644 internal/migrations/00002_rename_root.go

diff --git a/internal/cmd/migrate.go b/internal/cmd/migrate.go
index bd2be18..c2a8a96 100644
--- a/internal/cmd/migrate.go
+++ b/internal/cmd/migrate.go
@@ -7,11 +7,14 @@ import (
 	"entgo.io/ent/dialect/sql/schema"
 	"fmt"
 	entmigrate "framework_v2/internal/ent/migrate"
+	"framework_v2/internal/migrations"
 	_ "github.com/jackc/pgx/v5/stdlib"
 	"github.com/pressly/goose/v3"
 	"github.com/spf13/cobra"
 	"log"
 	"os"
+	"strings"
+	"time"
 )
 
 var dsnCommand = &cobra.Command{
@@ -36,15 +39,31 @@ var migrateCommand = &cobra.Command{
 	},
 }
 
-var createMigrateCommand = &cobra.Command{
-	Use:   "create-migrate",
-	Short: "新建迁移",
+var genMigrateCommand = &cobra.Command{
+	Use:   "gen-migrate [name]",
+	Short: "新建 ent 迁移",
 	Long:  "从 internal/ent 中新建迁移。在这之前，需要运行 go generate ./internal/ent",
 	Run: func(cmd *cobra.Command, args []string) {
 		generateMigration()
 	},
 }
 
+var createGoMigrateCommand = &cobra.Command{
+	Use:   "create-migrate",
+	Short: "新建 go 迁移",
+	Long:  "新建 goose 的 go 迁移。",
+	Run: func(cmd *cobra.Command, args []string) {
+		if len(args) == 0 {
+			_ = cmd.Help()
+			return
+		}
+
+		name := args[0]
+
+		createGooseMigration(name)
+	},
+}
+
 // RunMigrate 为数据库函数
 func RunMigrate(args []string) {
 	db, err := goose.OpenDBWithDriver("postgres", config.DB.DSN2)
@@ -65,7 +84,12 @@ func RunMigrate(args []string) {
 		arguments = append(arguments, args[3:]...)
 	}
 
-	if err := goose.RunContext(context.Background(), command, db, "internal/migrations", arguments...); err != nil {
+	goose.SetBaseFS(migrations.MigrationFS)
+	if err := goose.SetDialect("postgres"); err != nil {
+		panic(err)
+	}
+
+	if err := goose.RunContext(context.Background(), command, db, ".", arguments...); err != nil {
 		log.Fatalf("goose %v: %v", command, err)
 	}
 }
@@ -92,3 +116,73 @@ func generateMigration() {
 		log.Fatalf("failed generating migration file: %v", err)
 	}
 }
+
+func createGooseMigration(name string) {
+	// 在 internal/migrations 目录下新建一个迁移文件
+	// 文件名为 yyyy-mm-dd-hh-mm-ss-<name>.go
+	month := int(time.Now().Month())
+	monthString := fmt.Sprintf("%d", month)
+	if month < 10 {
+		// 转 string
+		monthString = "0" + monthString
+	}
+
+	day := time.Now().Day()
+	dayString := fmt.Sprintf("%d", day)
+	if day < 10 {
+		dayString = "0" + dayString
+	}
+
+	hour := time.Now().Hour()
+	hourString := fmt.Sprintf("%d", hour)
+	if hour < 10 {
+		hourString = "0" + hourString
+	}
+
+	minute := time.Now().Minute()
+	minuteString := fmt.Sprintf("%d", minute)
+	if minute < 10 {
+		minuteString = "0" + minuteString
+	}
+
+	// 秒
+	second := time.Now().Second()
+	secondString := fmt.Sprintf("%d", second)
+	if second < 10 {
+		secondString = "0" + secondString
+	}
+
+	funcName := fmt.Sprintf("%d%s%s%s%s%s", time.Now().Year(), monthString, dayString, hourString, minuteString, secondString)
+	fileName := fmt.Sprintf("%s_%s.go", funcName, name)
+
+	// 模板内容
+	var template = `package migrations
+
+import (
+	"context"
+	"database/sql"
+	"github.com/pressly/goose/v3"
+)
+
+func init() {
+	goose.AddMigrationContext(Up<FuncName>, Down<FuncName>)
+}
+
+func Up<FuncName>(ctx context.Context, tx *sql.Tx) error {
+	_, err := tx.ExecContext(ctx, "UPDATE users SET username='admin' WHERE username='root';")
+	return err
+}
+
+func Down<FuncName>(ctx context.Context, tx *sql.Tx) error {
+	_, err := tx.ExecContext(ctx, "UPDATE users SET username='root' WHERE username='admin';")
+	return err
+}
+`
+
+	template = strings.ReplaceAll(template, "<FuncName>", funcName+name)
+	err := os.WriteFile("internal/migrations/"+fileName, []byte(template), 0644)
+	if err != nil {
+		log.Fatalf("failed creating migration file: %v", err)
+	}
+
+}
diff --git a/internal/cmd/root.go b/internal/cmd/root.go
index 416d5a2..80aaa59 100644
--- a/internal/cmd/root.go
+++ b/internal/cmd/root.go
@@ -23,13 +23,18 @@ var rootCmd = &cobra.Command{
 
 func Execute() {
 	rootCmd.AddCommand(rpcCommand)
-	rootCmd.AddCommand(migrateCommand)
 	rootCmd.AddCommand(workerCommand)
 	rootCmd.AddCommand(scheduleCommand)
-	rootCmd.AddCommand(dsnCommand)
-	rootCmd.AddCommand(createMigrateCommand)
+
 	rootCmd.AddCommand(httpCommand)
 
+	if config.DebugMode.Enable {
+		rootCmd.AddCommand(migrateCommand)
+		rootCmd.AddCommand(dsnCommand)
+		rootCmd.AddCommand(createGoMigrateCommand)
+		rootCmd.AddCommand(genMigrateCommand)
+	}
+
 	if err := rootCmd.Execute(); err != nil {
 		fmt.Println(err)
 		os.Exit(1)
diff --git a/internal/migrations/00002_rename_root.go b/internal/migrations/00002_rename_root.go
new file mode 100644
index 0000000..409d7bd
--- /dev/null
+++ b/internal/migrations/00002_rename_root.go
@@ -0,0 +1,21 @@
+package migrations
+
+import (
+	"context"
+	"database/sql"
+	"github.com/pressly/goose/v3"
+)
+
+func init() {
+	goose.AddMigrationContext(Up00002, Down00002)
+}
+
+func Up00002(ctx context.Context, tx *sql.Tx) error {
+	_, err := tx.ExecContext(ctx, "UPDATE users SET username='admin' WHERE username='root';")
+	return err
+}
+
+func Down00002(ctx context.Context, tx *sql.Tx) error {
+	_, err := tx.ExecContext(ctx, "UPDATE users SET username='root' WHERE username='admin';")
+	return err
+}