patched insert custom type registration

viant · Jul 8, 2023 · 71067d1 · 71067d1
1 parent aad8f81
commit 71067d1
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 20 deletions.
diff --git a/internal/codegen/entity.go b/internal/codegen/entity.go
@@ -58,22 +58,17 @@ func (t *Template) GenerateEntity(ctx context.Context, pkg string, info *plugin.
 }
 
 func (t *Template) generateRegisterType() string {
- var initElements []string
+ registry := &customTypeRegistry{}
  for _, param := range t.State {
  if param.In.Kind != view.KindDataView {
  continue
  }
- registerFragment := t.RegisterFragment(param.Schema.DataType)
- initElements = append(initElements, registerFragment)
+ registry.register(param.Schema.DataType)
  }
-
- initCode := strings.Join(initElements, "\n")
- return initCode
-}
-
-func (t *Template) RegisterFragment(dataType string) string {
- registerFragment := fmt.Sprintf(registerTypeTemplate, dataType, dataType)
- return registerFragment
+ for _, param := range t.State.FilterByKind(view.KindRequestBody) {
+ registry.register(param.Schema.DataType)
+ }
+ return registry.stringify()
 }
 
 func (t *Template) generateMapTypeBody() string {

diff --git a/internal/codegen/handler.go b/internal/codegen/handler.go
@@ -37,7 +37,10 @@ func (t *Template) GenerateHandler(opts *options.Gen, info *plugin.Info) (string
 
  handlerContent := strings.Replace(handlerTemplate, "$Package", opts.Package, 1)
  handlerContent = strings.Replace(handlerContent, "$LocalVariable", localVariableDeclaration, 1)
- registerTypes := t.RegisterFragment("Handler")
+
+ registry := &customTypeRegistry{}
+ registry.register("Handler")
+ registerTypes := registry.stringify()
  handlerContent = strings.Replace(handlerContent, "$RegisterTypes", registerTypes, 1)
  imports := NewImports()
  imports.AddPackage(info.ChecksumPkg())

diff --git a/internal/codegen/registry.go b/internal/codegen/registry.go
@@ -0,0 +1,30 @@
+package codegen
+
+import (
+ "fmt"
+ "strings"
+)
+
+type customTypeRegistry struct {
+ Elements []string
+ index map[string]bool
+}
+
+func (c *customTypeRegistry) register(dataType string) {
+ c.ensureIndex()
+ if _, ok := c.index[dataType]; ok {
+ return
+ }
+ c.index[dataType] = true
+ c.Elements = append(c.Elements, fmt.Sprintf(registerTypeTemplate, dataType, dataType))
+}
+
+func (c *customTypeRegistry) stringify() string {
+ return strings.Join(c.Elements, "\n")
+}
+
+func (c *customTypeRegistry) ensureIndex() {
+ if len(c.index) == 0 {
+ c.index = map[string]bool{}
+ }
+}
diff --git a/internal/codegen/state.go b/internal/codegen/state.go
@@ -218,18 +218,18 @@ func (t *Template) GenerateState(pkg string, info *plugin.Info) string {
  fields = append(fields, input.FieldDeclaration())
  }
  output = strings.Replace(output, "$Fields", strings.Join(fields, "\n\n"), 1)
-
- registerTypes := t.RegisterFragment("State")
+ registerTypes := ""
  importFragment := ""
- imports := NewImports()
- imports.AddPackage(info.ChecksumPkg())
- imports.AddPackage("reflect")
- imports.AddPackage(info.TypeCorePkg())
  switch info.IntegrationMode {
  case plugin.ModeExtension, plugin.ModeCustomTypeModule:
+ imports := NewImports()
+ imports.AddPackage(info.ChecksumPkg())
+ imports.AddPackage("reflect")
+ imports.AddPackage(info.TypeCorePkg())
  importFragment = imports.PackageImports()
- default:
- registerTypes = ""
+ registry := &customTypeRegistry{}
+ registry.register("State")
+ registerTypes = registry.stringify()
  }
  output = strings.Replace(output, "$Imports", importFragment, 1)
  output = strings.Replace(output, "$RegisterTypes", registerTypes, 1)