WIP

go-reform · Dec 6, 2020 · 821dd24 · 821dd24
1 parent c76e205
commit 821dd24
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 40 deletions.
diff --git a/internal/test/models/extra_reform.go b/internal/test/models/extra_reform.go
diff --git a/internal/test/models/good_reform.go b/internal/test/models/good_reform.go
diff --git a/parse/base.go b/parse/base.go
@@ -121,6 +121,19 @@ func (s *StructInfo) IsTable() bool {
 	return s.PKFieldIndex >= 0
 }
 
+func (s *StructInfo) IsIntPK() bool {
+	switch s.PKField().Type {
+	case "int", "int8", "int16", "int32", "int64":
+		return true
+	case "uint", "uint8", "uint16", "uint32", "uint64":
+		return true
+	case "byte", "rune":
+		return true
+	default:
+		return false
+	}
+}
+
 // PKField returns a primary key field, panics for views.
 func (s *StructInfo) PKField() FieldInfo {
 	if !s.IsTable() {

diff --git a/reform/template.go b/reform/template.go
@@ -125,17 +125,26 @@ func (s *{{ .Type }}) HasPK() bool {
 	return s.{{ .PKField.Name }} != {{ .TableVar }}.z[{{ .TableVar }}.s.PKFieldIndex]
 }
 
+{{- if .IsIntPK }}
+
 // SetPK sets record primary key.
 //
 // Prefer direct field assignment where possible: s.{{ .PKField.Name }} = pk.
 func (s *{{ .Type }}) SetPK(pk interface{}) {
-	if i64, ok := pk.(int64); ok {
-		s.{{ .PKField.Name }} = {{ .PKField.Type }}(i64)
-	} else {
-		s.{{ .PKField.Name }} = pk.({{ .PKField.Type }})
-	}
+	s.{{ .PKField.Name }} = {{ .PKField.Type }}(pk.(int64))
 }
 
+{{ else }}
+
+// SetPK sets record primary key.
+//
+// Prefer direct field assignment where possible: s.{{ .PKField.Name }} = pk.
+func (s *{{ .Type }}) SetPK(pk interface{}) {
+	s.{{ .PKField.Name }} = pk.({{ .PKField.Type }})
+}
+
+{{- end }}
+
 {{- end }}
 
 // check interfaces