Add helper function for adding endpoint middleware

kujtimiihoxha · Oct 14, 2017 · c941ae5 · c941ae5
1 parent cf75337
commit c941ae5
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 1 deletion.
diff --git a/cmd/g_service.go b/cmd/g_service.go
@@ -47,7 +47,7 @@ var initserviceCmd = &cobra.Command{
 func init() {
 	generateCmd.AddCommand(initserviceCmd)
 	initserviceCmd.Flags().StringP("transport", "t", "http", "The transport you want your service to be initiated with")
-	initserviceCmd.Flags().Bool("dmw", false, "Generate default middleware for service and endpoint")
+	initserviceCmd.Flags().BoolP("dmw","w", false, "Generate default middleware for service and endpoint")
 	initserviceCmd.Flags().Bool("gorilla", false, "Generate http using gorilla mux")
 	initserviceCmd.Flags().StringArrayVarP(&methods, "methods", "m", []string{}, "Specify methods to be generated")
 	initserviceCmd.Flags().Bool("svc-mdw", false, "If set a default Logging and Instrumental middleware will be created and attached to the service")

diff --git a/generator/generate_service.go b/generator/generate_service.go
@@ -1445,6 +1445,25 @@ func (g *generateCmdBase) Generate() (err error) {
 		)
 		g.code.NewLine()
 	}
+	mth := []jen.Code{}
+	for _, v := range g.serviceInterface.Methods {
+		mth = append(mth, jen.Lit(v.Name))
+	}
+	g.code.appendFunction(
+		"addEndpointMiddlewareToAllMethods",
+		nil,
+		[]jen.Code{
+			jen.Id("mw").Map(jen.String()).Index().Qual("github.com/go-kit/kit/endpoint", "Middleware"),
+			jen.Id("m").Qual("github.com/go-kit/kit/endpoint", "Middleware"),
+		},
+		[]jen.Code{},
+		"",
+		jen.Id("methods").Op(":=").Index().String().Values(mth...),
+		jen.For(jen.List(jen.Id("_"), jen.Id("v")).Op(":=").Range().Id("methods")).Block(
+			jen.Id("mw").Index(jen.Id("v")).Op("=").Append(jen.Id("mw").Index(jen.Id("v")), jen.Id("m")),
+		),
+	)
+	g.code.NewLine()
 	return g.fs.WriteFile(g.filePath, g.srcFile.GoString(), true)
 }