From 7b349ba1cae4d522b310618afcfa771ab8ae2d36 Mon Sep 17 00:00:00 2001
From: Vasiliy Tolstov <v.tolstov@unistack.org>
Date: Sun, 28 Feb 2021 00:38:57 +0300
Subject: [PATCH] update micro-proto

Signed-off-by: Vasiliy Tolstov <v.tolstov@unistack.org>
---
 go.mod  |  4 +++-
 go.sum  |  4 ++--
 util.go | 47 ++++++++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 51 insertions(+), 4 deletions(-)

diff --git a/go.mod b/go.mod
index a329085..77a0681 100644
--- a/go.mod
+++ b/go.mod
@@ -4,7 +4,9 @@ go 1.16
 
 require (
 	github.com/google/go-cmp v0.5.4 // indirect
-	github.com/unistack-org/micro-proto v0.0.1
+	github.com/unistack-org/micro-proto v0.0.2-0.20210227213711-77c7563bd01e
 	golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 // indirect
 	google.golang.org/protobuf v1.25.0
 )
+
+//replace github.com/unistack-org/micro-proto => ../micro-proto
diff --git a/go.sum b/go.sum
index 93109e1..75fe1c9 100644
--- a/go.sum
+++ b/go.sum
@@ -24,8 +24,8 @@ github.com/google/go-cmp v0.5.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/
 github.com/google/go-cmp v0.5.4 h1:L8R9j+yAqZuZjsqh/z+F1NCffTKKLShY6zXTItVIZ8M=
 github.com/google/go-cmp v0.5.4/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
-github.com/unistack-org/micro-proto v0.0.1 h1:1d9O8d3/ZITQPwIQCvZ+2OCDjVP0iaX+gv5klbwkE/I=
-github.com/unistack-org/micro-proto v0.0.1/go.mod h1:GYO53DWmeldRIo90cAdQx8bLr/WJMxW62W4ja74p1Ac=
+github.com/unistack-org/micro-proto v0.0.2-0.20210227213711-77c7563bd01e h1:hQJ3V0QggeFdU5967wO5v6oWnaK42wUnG4UU4zWcyu4=
+github.com/unistack-org/micro-proto v0.0.2-0.20210227213711-77c7563bd01e/go.mod h1:GYO53DWmeldRIo90cAdQx8bLr/WJMxW62W4ja74p1Ac=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
 golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE=
diff --git a/util.go b/util.go
index 4abf865..73447e5 100644
--- a/util.go
+++ b/util.go
@@ -32,6 +32,9 @@ func unexport(s string) string {
 
 func generateServiceClient(gfile *protogen.GeneratedFile, service *protogen.Service) {
 	serviceName := service.GoName
+	if rule, ok := getMicroApiService(service); ok {
+		gfile.P("// client wrappers ", strings.Join(rule.ClientWrappers, ", "))
+	}
 	gfile.P("type ", unexport(serviceName), "Client struct {")
 	gfile.P("c ", microClientPackage.Ident("Client"))
 	gfile.P("name string")
@@ -177,7 +180,13 @@ func generateServiceServerMethods(gfile *protogen.GeneratedFile, service *protog
 	serviceName := service.GoName
 	for _, method := range service.Methods {
 		generateServerFuncSignature(gfile, serviceName, method, true)
-
+		if rule, ok := getMicroApiMethod(method); ok {
+			if rule.Timeout > 0 {
+				gfile.P("var cancel ", contextPackage.Ident("CancelFunc"))
+				gfile.P("ctx, cancel = ", contextPackage.Ident("WithTimeout"), "(ctx, time.Second*", rule.Timeout, ")")
+				gfile.P("defer cancel()")
+			}
+		}
 		if method.Desc.IsStreamingClient() || method.Desc.IsStreamingServer() {
 			if !method.Desc.IsStreamingClient() {
 				gfile.P("msg := &", gfile.QualifiedGoIdent(method.Input.GoIdent), "{}")
@@ -470,6 +479,42 @@ func generateEndpoints(method *protogen.Method) ([]*api_options.HttpRule, bool)
 	return rules, method.Desc.IsStreamingServer() || method.Desc.IsStreamingClient()
 }
 
+func getMicroApiMethod(method *protogen.Method) (*api_options.MicroMethod, bool) {
+	if method.Desc.Options() == nil {
+		return nil, false
+	}
+
+	if !proto.HasExtension(method.Desc.Options(), api_options.E_MicroMethod) {
+		return nil, false
+	}
+
+	r := proto.GetExtension(method.Desc.Options(), api_options.E_MicroMethod)
+	if r == nil {
+		return nil, false
+	}
+
+	rule := r.(*api_options.MicroMethod)
+	return rule, true
+}
+
+func getMicroApiService(service *protogen.Service) (*api_options.MicroService, bool) {
+	if service.Desc.Options() == nil {
+		return nil, false
+	}
+
+	if !proto.HasExtension(service.Desc.Options(), api_options.E_MicroService) {
+		return nil, false
+	}
+
+	r := proto.GetExtension(service.Desc.Options(), api_options.E_MicroService)
+	if r == nil {
+		return nil, false
+	}
+
+	rule := r.(*api_options.MicroService)
+	return rule, true
+}
+
 func getEndpoint(rule *api_options.HttpRule) (string, string, string) {
 	var meth string
 	var path string