From bfbf9777e9a9a63787546dc5c3903d3e5e147c22 Mon Sep 17 00:00:00 2001
From: Alireza Ahmadi <alireza7@gmail.com>
Date: Thu, 6 Jun 2024 22:08:13 +0200
Subject: [PATCH] db migration

---
 backend/cmd/cmd.go       |  4 ++
 backend/cmd/migration.go | 85 ++++++++++++++++++++++++++++++++++++++++
 backend/database/db.go   |  7 +++-
 3 files changed, 95 insertions(+), 1 deletion(-)
 create mode 100644 backend/cmd/migration.go

diff --git a/backend/cmd/cmd.go b/backend/cmd/cmd.go
index dd00166..361b349 100644
--- a/backend/cmd/cmd.go
+++ b/backend/cmd/cmd.go
@@ -40,6 +40,7 @@ func ParseCmd() {
 		fmt.Println()
 		fmt.Println("Commands:")
 		fmt.Println("    admin          set/reset/show first admin credentials")
+		fmt.Println("    migrate        migrate form older version")
 		fmt.Println("    setting        set/reset/show settings")
 		fmt.Println()
 		adminCmd.Usage()
@@ -70,6 +71,9 @@ func ParseCmd() {
 			showAdmin()
 		}
 
+	case "migrate":
+		migrateDb()
+
 	case "setting":
 		err := settingCmd.Parse(os.Args[2:])
 		if err != nil {
diff --git a/backend/cmd/migration.go b/backend/cmd/migration.go
new file mode 100644
index 0000000..760856f
--- /dev/null
+++ b/backend/cmd/migration.go
@@ -0,0 +1,85 @@
+package cmd
+
+import (
+	"encoding/json"
+	"fmt"
+	"log"
+	"s-ui/config"
+	"s-ui/database"
+	"s-ui/database/model"
+	"strings"
+
+	"gorm.io/gorm"
+)
+
+func migrateDb() {
+	err := database.OpenDB(config.GetDBPath())
+	if err != nil {
+		log.Fatal(err)
+	}
+	db := database.GetDB()
+	tx := db.Begin()
+	defer func() {
+		if err == nil {
+			tx.Commit()
+		} else {
+			tx.Rollback()
+		}
+	}()
+	fmt.Println("Start migrating database...")
+	err = migrateClientSchema(tx)
+	fmt.Println("Migration done!")
+}
+
+func migrateClientSchema(db *gorm.DB) error {
+	rows, err := db.Raw("PRAGMA table_info(clients)").Rows()
+	if err != nil {
+		fmt.Println(err)
+		return err
+	}
+	defer rows.Close()
+
+	for rows.Next() {
+		var (
+			cid       int
+			cname     string
+			ctype     string
+			notnull   int
+			dfltValue interface{}
+			pk        int
+		)
+
+		rows.Scan(&cid, &cname, &ctype, &notnull, &dfltValue, &pk)
+		if cname == "config" || cname == "inbounds" || cname == "links" {
+			if ctype == "text" {
+				fmt.Printf("Column %s has type TEXT\n", cname)
+				oldData := make([]struct {
+					Id   uint
+					Data string
+				}, 0)
+				db.Model(model.Client{}).Select("id", cname+" as data").Scan(&oldData)
+				for _, data := range oldData {
+					var newData []byte
+					switch cname {
+					case "inbounds":
+						inbounds := strings.Split(data.Data, ",")
+						newData, _ = json.MarshalIndent(inbounds, " ", "  ")
+					case "config":
+						jsonData := map[string]interface{}{}
+						json.Unmarshal([]byte(data.Data), &jsonData)
+						newData, _ = json.MarshalIndent(jsonData, " ", "  ")
+					case "links":
+						jsonData := make([]interface{}, 0)
+						json.Unmarshal([]byte(data.Data), &jsonData)
+						newData, _ = json.MarshalIndent(jsonData, " ", "  ")
+					}
+					err = db.Model(model.Client{}).Where("id = ?", data.Id).UpdateColumn(cname, newData).Error
+					if err != nil {
+						return err
+					}
+				}
+			}
+		}
+	}
+	return nil
+}
diff --git a/backend/database/db.go b/backend/database/db.go
index efbe99b..a6b055d 100644
--- a/backend/database/db.go
+++ b/backend/database/db.go
@@ -29,7 +29,7 @@ func initUser() error {
 	return nil
 }
 
-func InitDB(dbPath string) error {
+func OpenDB(dbPath string) error {
 	dir := path.Dir(dbPath)
 	err := os.MkdirAll(dir, 01740)
 	if err != nil {
@@ -48,6 +48,11 @@ func InitDB(dbPath string) error {
 		Logger: gormLogger,
 	}
 	db, err = gorm.Open(sqlite.Open(dbPath), c)
+	return err
+}
+
+func InitDB(dbPath string) error {
+	err := OpenDB(dbPath)
 	if err != nil {
 		return err
 	}