Codegen fixes (#3741)

* Fix importing user types

* Fix handling error for wrapped http stream

* Check for stream nil

* Fix tests

Author: nitinmohan87

codegen/generator/service.go

@@ -31,7 +31,9 @@ func Service(genpkg string, roots []eval.Root) ([]*codegen.File, error) {
 			}
 			for _, f := range files {
 				if len(f.SectionTemplates) > 0 {
-					service.AddServiceDataMetaTypeImports(f.SectionTemplates[0], s, services.Get(s.Name))
+					d := services.Get(s.Name)
+					service.AddServiceDataMetaTypeImports(f.SectionTemplates[0], s, d)
+					service.AddUserTypeImports(genpkg, f.SectionTemplates[0], d)
 				}
 			}
 			f, err := service.ConvertFile(r, s, services)

codegen/generator/transport.go

@@ -44,7 +44,9 @@ func Transport(genpkg string, roots []eval.Root) ([]*codegen.File, error) {
 		for _, f := range files {
 			if len(f.SectionTemplates) > 0 {
 				for _, s := range r.Services {
-					service.AddServiceDataMetaTypeImports(f.SectionTemplates[0], s, services.Get(s.Name))
+					d := services.Get(s.Name)
+					service.AddServiceDataMetaTypeImports(f.SectionTemplates[0], s, d)
+					service.AddUserTypeImports(genpkg, f.SectionTemplates[0], d)
 				}
 			}
 		}

codegen/service/client.go

@@ -26,7 +26,6 @@ func ClientFile(_ string, service *expr.ServiceExpr, services *ServicesData) *co
 			{Path: "io"},
 			codegen.GoaImport(""),
 		}
-		imports = append(imports, svc.UserTypeImports...)
 		header := codegen.Header(service.Name+" client", svc.PkgName, imports)
 		def := &codegen.SectionTemplate{
 			Name:   "client-struct",

codegen/service/endpoint.go

@@ -80,7 +80,6 @@ func EndpointFile(genpkg string, service *expr.ServiceExpr, services *ServicesDa
 			codegen.GoaImport("security"),
 			{Path: genpkg + "/" + svcName + "/" + "views", Name: svc.ViewsPkg},
 		}
-		imports = append(imports, svc.UserTypeImports...)
 		header := codegen.Header(service.Name+" endpoints", svc.PkgName, imports)
 		def := &codegen.SectionTemplate{
 			Name:   "endpoints-struct",

codegen/service/interceptors.go

@@ -78,7 +78,6 @@ func interceptorFile(svc *Data, server bool) *codegen.File {
 		},
 	}
 	if len(interceptors) > 0 {
-		codegen.AddImport(sections[0], svc.UserTypeImports...)
 		sections = append(sections, &codegen.SectionTemplate{
 			Name:   "interceptor-types",
 			Source: readTemplate("interceptors_types"),

codegen/service/service.go

@@ -15,7 +15,6 @@ import (
 // type definitions.
 func Files(genpkg string, service *expr.ServiceExpr, services *ServicesData, userTypePkgs map[string][]string) []*codegen.File {
 	svc := services.Get(service.Name)
-	svc.initUserTypeImports(genpkg)
 	svcName := svc.PathName
 	svcPath := filepath.Join(codegen.Gendir, svcName, "service.go")
 	seen := make(map[string]struct{})
@@ -161,7 +160,6 @@ func Files(genpkg string, service *expr.ServiceExpr, services *ServicesData, use
 		codegen.GoaImport("security"),
 		codegen.NewImport(svc.ViewsPkg, genpkg+"/"+svcName+"/views"),
 	}
-	imports = append(imports, svc.UserTypeImports...)
 	header := codegen.Header(service.Name+" service", svc.PkgName, imports)
 	def := &codegen.SectionTemplate{
 		Name:    "service",
@@ -249,6 +247,38 @@ func AddServiceDataMetaTypeImports(header *codegen.SectionTemplate, svcExpr *exp
 	}
 }
 
+// AddUserTypeImports sets the import paths for the user types defined in the
+// service.  User types may be declared in multiple packages when defined with
+// the Meta key "struct:pkg:path".
+func AddUserTypeImports(genpkg string, header *codegen.SectionTemplate, d *Data) {
+	importsByPath := make(map[string]*codegen.ImportSpec)
+
+	initLoc := func(loc *codegen.Location) {
+		if loc == nil {
+			return
+		}
+		importsByPath[loc.FilePath] = &codegen.ImportSpec{Name: loc.PackageName(), Path: genpkg + "/" + loc.RelImportPath}
+	}
+
+	for _, m := range d.Methods {
+		initLoc(m.PayloadLoc)
+		initLoc(m.ResultLoc)
+		for _, l := range m.ErrorLocs {
+			initLoc(l)
+		}
+		for _, ut := range d.userTypes {
+			initLoc(ut.Loc)
+		}
+		for _, et := range d.errorTypes {
+			initLoc(et.Loc)
+		}
+	}
+
+	for _, imp := range importsByPath { // Order does not matter, imports are sorted during formatting.
+		codegen.AddImport(header, imp)
+	}
+}
+
 func errorName(et *UserTypeData) string {
 	obj := expr.AsObject(et.Type)
 	if obj != nil {

codegen/service/service_data.go

@@ -73,9 +73,6 @@ type (
 		Scope *codegen.NameScope
 		// ViewScope initialized with all the viewed types.
 		ViewScope *codegen.NameScope
-		// UserTypeImports lists the import specifications for the user
-		// types used by the service.
-		UserTypeImports []*codegen.ImportSpec
 		// ProtoImports lists the import specifications for the custom
 		// proto types used by the service.
 		ProtoImports []*codegen.ImportSpec
@@ -599,42 +596,6 @@ func (d *Data) Method(name string) *MethodData {
 	return nil
 }
 
-// initUserTypeImports sets the import paths for the user types defined in the
-// service.  User types may be declared in multiple packages when defined with
-// the Meta key "struct:pkg:path".
-func (d *Data) initUserTypeImports(genpkg string) {
-	importsByPath := make(map[string]*codegen.ImportSpec)
-
-	initLoc := func(loc *codegen.Location) {
-		if loc == nil {
-			return
-		}
-		importsByPath[loc.FilePath] = &codegen.ImportSpec{Name: loc.PackageName(), Path: genpkg + "/" + loc.RelImportPath}
-	}
-
-	for _, m := range d.Methods {
-		initLoc(m.PayloadLoc)
-		initLoc(m.ResultLoc)
-		for _, l := range m.ErrorLocs {
-			initLoc(l)
-		}
-		for _, ut := range d.userTypes {
-			initLoc(ut.Loc)
-		}
-		for _, et := range d.errorTypes {
-			initLoc(et.Loc)
-		}
-	}
-
-	imports := make([]*codegen.ImportSpec, len(importsByPath))
-	i := 0
-	for _, imp := range importsByPath { // Order does not matter, imports are sorted during formatting.
-		imports[i] = imp
-		i++
-	}
-	d.UserTypeImports = imports
-}
-
 // Scheme returns the scheme data with the given scheme name.
 func (r RequirementsData) Scheme(name string) *SchemeData {
 	for _, req := range r {

codegen/service/templates/client_interceptor_stream_wrappers.go.tpl

@@ -2,6 +2,11 @@
 
 	{{- if ne .SendTypeRef "" }}
 
+{{ comment (print "Unwrap returns the underlying stream type.") }}
+func (w *wrapped{{ .Interface }}) Unwrap() interface{} {
+       return w.stream
+}
+
 {{ comment (printf "%s streams instances of \"%s\" after executing the applied interceptor." .SendName .Interface) }}
 func (w *wrapped{{ .Interface }}) {{ .SendName }}(v {{ .SendTypeRef }}) error {
 	return w.SendWithContext(w.ctx, v)

codegen/service/templates/server_interceptor_stream_wrapper_types.go.tpl

@@ -1,3 +1,9 @@
+{{- if .WrappedServerStreams }}
+type Unwrap interface {
+	Unwrap() interface{}
+}
+{{- end }}
+
 {{- range .WrappedServerStreams }}
 
 {{ comment (printf "wrapped%s is a server interceptor wrapper for the %s stream." .Interface .Interface) }}

codegen/service/templates/server_interceptor_stream_wrappers.go.tpl

@@ -1,5 +1,10 @@
 {{- range .WrappedServerStreams }}
 
+{{ comment (print "Unwrap returns the underlying stream type.") }}
+func (w *wrapped{{ .Interface }}) Unwrap() interface{} {
+	return w.stream
+}
+
 	{{- if ne .SendTypeRef "" }}
 
 {{ comment (printf "%s streams instances of \"%s\" after executing the applied interceptor." .SendName .Interface) }}