From a84cbec2ca9730487ef41859fff6d61280c54099 Mon Sep 17 00:00:00 2001
From: ucpr <contact@ucpr.dev>
Date: Wed, 10 Apr 2024 17:10:45 +0900
Subject: [PATCH] feat: support custom tag

---
 generator.go      | 3 ++-
 generator_test.go | 1 +
 main.go           | 4 +++-
 3 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/generator.go b/generator.go
index 4822093..3d6b71e 100644
--- a/generator.go
+++ b/generator.go
@@ -14,6 +14,7 @@ type input struct {
 	hclPath string
 	outPath string
 	pkg     string
+	tag     string
 }
 
 func generate(s schema.Schema, in input) ([]byte, error) {
@@ -29,7 +30,7 @@ func generate(s schema.Schema, in input) ([]byte, error) {
 		for j := range table.Columns {
 			column := table.Columns[j]
 			tp := toGoTypeString(column.Type)
-			fmt.Fprintf(buf, "\t%s\t%s\t`db:\"%s\"`\n", toCamelCase(column.Name), tp, column.Name)
+			fmt.Fprintf(buf, "\t%s\t%s\t`%s:\"%s\"`\n", toCamelCase(column.Name), tp, in.tag, column.Name)
 		}
 		fmt.Fprintf(buf, "}\n\n")
 	}
diff --git a/generator_test.go b/generator_test.go
index 6f65bdd..04bbe75 100644
--- a/generator_test.go
+++ b/generator_test.go
@@ -32,6 +32,7 @@ func Test_generate(t *testing.T) {
 				hclPath: "input.hcl",
 				outPath: "output.go",
 				pkg:     "main",
+				tag:     "db",
 			},
 			out: `// Code generated by github.com/ucpr/atlas-hcl-gen-go. DO NOT EDIT.
 // atlas-hcl-gen-go:
diff --git a/main.go b/main.go
index 2336211..b7f45b5 100644
--- a/main.go
+++ b/main.go
@@ -27,10 +27,11 @@ func main() {
 }
 
 func run() error {
-	var hclPath, outPath, target, pkg string
+	var hclPath, outPath, target, tag, pkg string
 	flag.StringVar(&hclPath, "i", "", "input file path")
 	flag.StringVar(&outPath, "o", "", "output file path")
 	flag.StringVar(&target, "t", "mysql", "target database")
+	flag.StringVar(&tag, "tag", "db", "tag name")
 	flag.StringVar(&pkg, "package", "main", "package name")
 	flag.Parse()
 
@@ -54,6 +55,7 @@ func run() error {
 		hclPath: hclPath,
 		outPath: outPath,
 		pkg:     pkg,
+		tag:     tag,
 	})
 	if err != nil {
 		return fmt.Errorf("failed to generate: %w", err)