From e9e375ebe48996e7a48ac4b7ebb4e3db6472b20a Mon Sep 17 00:00:00 2001
From: cg33 <chg80333@163.com>
Date: Sun, 30 Oct 2022 08:45:38 +0800
Subject: [PATCH] feat(db): primary key identify

---
 plugins/admin/modules/tools/generator.go | 38 +++++++++++++++++++-----
 plugins/admin/modules/tools/template.go  | 12 ++++++--
 2 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/plugins/admin/modules/tools/generator.go b/plugins/admin/modules/tools/generator.go
index 54c4e2537..219666e88 100644
--- a/plugins/admin/modules/tools/generator.go
+++ b/plugins/admin/modules/tools/generator.go
@@ -45,7 +45,7 @@ type Param struct {
 	HideResetButton          bool `json:"hide_reset_button"`
 	HideBackButton           bool `json:"hide_back_button"`
 
-	TablePageTitle    string `json:"table_title"`
+	TablePageTitle    string `json:"table_page_title"`
 	TableDescription  string `json:"table_description"`
 	FormTitle         string `json:"form_title"`
 	FormDescription   string `json:"form_description"`
@@ -59,6 +59,9 @@ type Param struct {
 	FormFields   Fields `json:"form_fields"`
 	DetailFields Fields `json:"detail_fields"`
 
+	PrimaryKey     string `json:"primary_key"`
+	PrimaryKeyType string `json:"primary_key_type"`
+
 	DetailDisplay uint8 `json:"detail_display"`
 
 	Output string `json:"output"`
@@ -123,6 +126,8 @@ func NewParam(cfg Config) *Param {
 	fields := getFieldsFromConn(cfg.Conn, dbTable, cfg.Driver)
 	tt := strings.Title(ta)
 
+	pkey, ptype := fields.GetPrimaryKey()
+
 	return &Param{
 		Connection:               cfg.Connection,
 		Driver:                   cfg.Driver,
@@ -156,6 +161,9 @@ func NewParam(cfg Config) *Param {
 		TableDescription:         utils.SetDefault(cfg.TableDescription, "", tt),
 		FormTitle:                utils.SetDefault(cfg.FormTitle, "", tt),
 		FormDescription:          utils.SetDefault(cfg.FormDescription, "", tt),
+
+		PrimaryKey:     pkey,
+		PrimaryKeyType: ptype,
 	}
 }
 
@@ -218,6 +226,15 @@ func NewParamWithFields(cfg Config, fields ...Fields) *Param {
 
 type Fields []Field
 
+func (fs Fields) GetPrimaryKey() (string, string) {
+	for _, field := range fs {
+		if field.IsPrimaryKey {
+			return field.Name, field.DBType
+		}
+	}
+	return "", ""
+}
+
 type Field struct {
 	Head         string `json:"head"`
 	Name         string `json:"name"`
@@ -234,6 +251,7 @@ type Field struct {
 	Default      string `json:"default"`
 	CanAdd       bool   `json:"can_add"`
 	ExtraFun     string `json:"extra_fun"`
+	IsPrimaryKey bool   `json:"is_primary_key"`
 }
 
 func Generate(param *Param) error {
@@ -442,13 +460,19 @@ func getFieldsFromConn(conn db.Connection, table, driver string) Fields {
 
 	for i, model := range columnsModel {
 		typeName := getType(model[typeField].(string))
+		isPrimaryKey := false
+		if columnKey, ok := model["Key"].(string); ok {
+			isPrimaryKey = columnKey == "PRI"
+		}
+
 		fields[i] = Field{
-			Head:     strings.Title(model[fieldField].(string)),
-			Name:     model[fieldField].(string),
-			DBType:   typeName,
-			CanAdd:   true,
-			Editable: true,
-			FormType: form.GetFormTypeFromFieldType(db.DT(strings.ToUpper(typeName)), model[fieldField].(string)),
+			Head:         strings.Title(model[fieldField].(string)),
+			Name:         model[fieldField].(string),
+			DBType:       typeName,
+			CanAdd:       true,
+			Editable:     true,
+			IsPrimaryKey: isPrimaryKey,
+			FormType:     form.GetFormTypeFromFieldType(db.DT(strings.ToUpper(typeName)), model[fieldField].(string)),
 		}
 		if model[fieldField].(string) == "id" {
 			fields[i].Filterable = true
diff --git a/plugins/admin/modules/tools/template.go b/plugins/admin/modules/tools/template.go
index c550b7c24..bf6428a1a 100644
--- a/plugins/admin/modules/tools/template.go
+++ b/plugins/admin/modules/tools/template.go
@@ -14,9 +14,17 @@ import (
 func Get{{.TableTitle}}Table(ctx *context.Context) table.Table {
 	
 	{{if eq .Connection "default"}}
-	{{.Table}} := table.NewDefaultTable(table.DefaultConfigWithDriver("{{.Driver}}"))
+		{{if eq .PrimaryKey ""}}
+			{{.Table}} := table.NewDefaultTable(table.DefaultConfigWithDriver("{{.Driver}}"))
+		{{else}}
+			{{.Table}} := table.NewDefaultTable(table.DefaultConfigWithDriver("{{.Driver}}").SetPrimaryKey("{{.PrimaryKey}}", db.{{.PrimaryKeyType}}))
+		{{end}}
 	{{else}}
-	{{.Table}} := table.NewDefaultTable(table.DefaultConfigWithDriverAndConnection("{{.Driver}}", "{{.Connection}}"))
+		{{if eq .PrimaryKey ""}}
+			{{.Table}} := table.NewDefaultTable(table.DefaultConfigWithDriverAndConnection("{{.Driver}}", "{{.Connection}}"))
+		{{else}}
+			{{.Table}} := table.NewDefaultTable(table.DefaultConfigWithDriverAndConnection("{{.Driver}}", "{{.Connection}}").SetPrimaryKey("{{.PrimaryKey}}", db.{{.PrimaryKeyType}}))
+		{{end}}
 	{{end}}
 
 	info := {{.Table}}.GetInfo(){{if .HideFilterArea}}.HideFilterArea(){{end}}