From 517cabaa8a703cd5b0039543ceb75726357ea86b Mon Sep 17 00:00:00 2001
From: pillar <ymblender@gamil.com>
Date: Sun, 27 Jun 2021 21:54:39 +0800
Subject: [PATCH] feat: add custom tags

---
 internal/codegen/golang/custom_tags.go | 55 ++++++++++++++++++++++++++
 internal/codegen/golang/result.go      |  1 +
 2 files changed, 56 insertions(+)
 create mode 100644 internal/codegen/golang/custom_tags.go

diff --git a/internal/codegen/golang/custom_tags.go b/internal/codegen/golang/custom_tags.go
new file mode 100644
index 0000000000..d5ee5f6bcb
--- /dev/null
+++ b/internal/codegen/golang/custom_tags.go
@@ -0,0 +1,55 @@
+package golang
+
+import (
+	"fmt"
+	"regexp"
+	"strings"
+
+	"github.com/kyleconroy/sqlc/internal/sql/catalog"
+)
+
+const tagKeyValueSplit = "__@CUSTOM_TAG_VALUE__"
+
+var tagReg = regexp.MustCompile(`^@.*\:`)
+
+// Comment parse go struct tags
+//
+// Example:
+// COMMENT ON COLUMN account.id is '@gorm:primaryKey @validate:required,min=3,max=32';
+//
+// To:
+// type Account {
+//   id int64 `gorm:"primaryKey" validate:"required,min=3,max=32"`
+// }
+//
+func customTags(tags *map[string]string, column *catalog.Column) {
+	if column.Comment == "" {
+		return
+	}
+
+	comments := strings.Split(column.Comment, " ")
+	fliterComments := make([]string, len(comments))
+
+	for _, tag := range comments {
+		if tagReg.Match([]byte(tag)) {
+			tag = strings.Replace(tag, "@", "", 1)
+			tag := strings.Replace(tag, ":", tagKeyValueSplit, 1)
+			kv := strings.Split(tag, tagKeyValueSplit)
+			if len(kv) < 2 {
+				panic(fmt.Sprintf("comment tags JSON tags style error:  %s %s in %s", column.Type.Name, column.Name, column.Comment))
+			}
+			k := kv[0]
+			v := kv[1]
+			key := k + ":"
+			(*tags)[key] = v
+		} else {
+			fliterComments = append(fliterComments, tag)
+		}
+	}
+
+	// clear tags in Comment
+	if len(fliterComments) > 0 {
+		column.Comment = strings.Join(fliterComments, " ")
+		column.Comment = strings.Trim(column.Comment, " ")
+	}
+}
diff --git a/internal/codegen/golang/result.go b/internal/codegen/golang/result.go
index 0b91980185..cafd3bd290 100644
--- a/internal/codegen/golang/result.go
+++ b/internal/codegen/golang/result.go
@@ -80,6 +80,7 @@ func buildStructs(r *compiler.Result, settings config.CombinedSettings) []Struct
 			}
 			for _, column := range table.Columns {
 				tags := map[string]string{}
+				customTags(&tags, column)
 				if settings.Go.EmitDBTags {
 					tags["db:"] = column.Name
 				}