Merge pull request #103 from uber/lu.ms

Allow multiple thrift services for tchannel client

codegen/gateway.go

@@ -99,7 +99,11 @@ type ClientSpec struct {
 	ClientName string
 	// ThriftServiceName, if the thrift file has multiple
 	// services then this is the service that describes the client
+	// TODO: this field needs to be deprecated for multi-service support
 	ThriftServiceName string
+	// ExposedMethods is a map of exposed method name to thrift "$service::$method"
+	// only the method values in this map are generated for the client
+	ExposedMethods map[string]string
 }
 
 // ModuleClassConfig represents the generic JSON config for
@@ -112,11 +116,10 @@ type ModuleClassConfig struct {
 
 // ClientClassConfig represents the specific config for
 // a client. This is a downcast of the moduleClassConfig.
-// Config here is a map[string]string but could be any type.
 type ClientClassConfig struct {
-	Name   string            `json:"name"`
-	Type   string            `json:"type"`
-	Config map[string]string `json:"config"`
+	Name   string                 `json:"name"`
+	Type   string                 `json:"type"`
+	Config map[string]interface{} `json:"config"`
 }
 
 // NewClientSpec creates a client spec from a json file.
@@ -161,7 +164,28 @@ func NewClientSpec(jsonFile string, h *PackageHelper) (*ClientSpec, error) {
 
 // NewTChannelClientSpec creates a client spec from a json file whose type is tchannel
 func NewTChannelClientSpec(jsonFile string, clientConfig *ClientClassConfig, h *PackageHelper) (*ClientSpec, error) {
-	return newClientSpec(jsonFile, clientConfig, false, h)
+	exposedMethods := clientConfig.Config["exposedMethods"].(map[string]interface{})
+	if len(exposedMethods) == 0 {
+		return nil, errors.Errorf("No methods are exposed in client config: %s", jsonFile)
+	}
+
+	cspec, err := newClientSpec(jsonFile, clientConfig, false, h)
+	if err != nil {
+		return nil, err
+	}
+
+	cspec.ExposedMethods = map[string]string{}
+	reversed := map[string]string{}
+	for key, val := range exposedMethods {
+		cspec.ExposedMethods[key] = val.(string)
+		reversed[val.(string)] = key
+	}
+
+	if len(cspec.ExposedMethods) != len(reversed) {
+		return nil, errors.Errorf("Keys or values of the exposedMethods of are not unique: %s", jsonFile)
+	}
+
+	return cspec, nil
 }
 
 // NewCustomClientSpec creates a client spec from a json file whose type is custom
@@ -177,11 +201,11 @@ func NewCustomClientSpec(jsonFile string, clientConfig *ClientClassConfig, h *Pa
 	clientSpec := &ClientSpec{
 		JSONFile:          jsonFile,
 		ClientType:        clientConfig.Type,
-		ClientID:          clientConfig.Config["clientId"],
-		ClientName:        clientConfig.Config["clientName"],
-		CustomImportPath:  clientConfig.Config["customImportPath"],
-		CustomClientType:  clientConfig.Config["customClientType"],
-		CustomPackageName: clientConfig.Config["customPackageName"],
+		ClientID:          clientConfig.Config["clientId"].(string),
+		ClientName:        clientConfig.Config["clientName"].(string),
+		CustomImportPath:  clientConfig.Config["customImportPath"].(string),
+		CustomClientType:  clientConfig.Config["customClientType"].(string),
+		CustomPackageName: clientConfig.Config["customPackageName"].(string),
 	}
 
 	return clientSpec, nil
@@ -206,7 +230,7 @@ func newClientSpec(jsonFile string, clientConfig *ClientClassConfig, wantAnnot b
 	}
 
 	thriftFile := filepath.Join(
-		h.ThriftIDLPath(), config["thriftFile"],
+		h.ThriftIDLPath(), config["thriftFile"].(string),
 	)
 
 	mspec, err := NewModuleSpec(thriftFile, wantAnnot, h)
@@ -247,9 +271,9 @@ func newClientSpec(jsonFile string, clientConfig *ClientClassConfig, wantAnnot b
 		GoPackageName:     goPackageName,
 		GoStructsFileName: goStructsFileName,
 		ThriftFile:        thriftFile,
-		ClientID:          config["clientId"],
-		ClientName:        config["clientName"],
-		ThriftServiceName: config["serviceName"],
+		ClientID:          config["clientId"].(string),
+		ClientName:        config["clientName"].(string),
+		ThriftServiceName: config["serviceName"].(string),
 	}, nil
 }
 
@@ -686,7 +710,7 @@ func (e *EndpointSpec) SetDownstream(
 
 	return e.ModuleSpec.SetDownstream(
 		e.ThriftServiceName, e.ThriftMethodName,
-		clientSpec, e.ClientName, e.ClientMethod, h,
+		clientSpec, e.ClientMethod, h,
 	)
 }
 

codegen/method.go

@@ -378,9 +378,7 @@ func findParamsAnnotation(
 	return "", false
 }
 
-func (ms *MethodSpec) setHTTPPath(
-	httpPath string, funcSpec *compile.FunctionSpec,
-) {
+func (ms *MethodSpec) setHTTPPath(httpPath string, funcSpec *compile.FunctionSpec) {
 	ms.HTTPPath = httpPath
 
 	segments := strings.Split(httpPath[1:], "/")
@@ -417,37 +415,37 @@ func (ms *MethodSpec) setHTTPPath(
 }
 
 func (ms *MethodSpec) setDownstream(
-	clientModule *ModuleSpec, clientService string, clientMethod string,
+	clientModule *ModuleSpec, clientThriftService, clientThriftMethod string,
 ) error {
 	var downstreamService *ServiceSpec
 	for _, service := range clientModule.Services {
-		if service.Name == clientService {
+		if service.Name == clientThriftService {
 			downstreamService = service
 			break
 		}
 	}
 	if downstreamService == nil {
 		return errors.Errorf(
 			"Downstream service '%s' is not found in '%s'",
-			clientService, clientModule.ThriftFile,
+			clientThriftService, clientModule.ThriftFile,
 		)
 	}
 	var downstreamMethod *MethodSpec
 	for _, method := range downstreamService.Methods {
-		if method.Name == clientMethod {
+		if method.Name == clientThriftMethod {
 			downstreamMethod = method
 			break
 		}
 	}
 	if downstreamMethod == nil {
 		return errors.Errorf(
 			"\n Downstream method '%s' is not found in '%s'",
-			clientMethod, clientModule.ThriftFile,
+			clientThriftMethod, clientModule.ThriftFile,
 		)
 	}
 	// Remove irrelevant services and methods.
 	ms.Downstream = clientModule
-	ms.DownstreamService = clientService
+	ms.DownstreamService = clientThriftService
 	ms.DownstreamMethod = downstreamMethod
 	return nil
 }

codegen/module_system.go

@@ -39,6 +39,7 @@ type EndpointMeta struct {
 	IncludedPackages   []GoPackageImport
 	Method             *MethodSpec
 	ClientName         string
+	ClientMethodName   string
 	WorkflowName       string
 	ReqHeaderMap       map[string]string
 	ReqHeaderMapKeys   []string
@@ -365,11 +366,18 @@ func (g *TChannelClientGenerator) Generate(
 
 	g.genSpecs[clientSpec.JSONFile] = clientSpec
 
+	// reverse index the exposed methods map
+	exposedMethods := map[string]string{}
+	for k, v := range clientSpec.ExposedMethods {
+		exposedMethods[v] = k
+	}
+
 	clientMeta := &ClientMeta{
 		PackageName:      clientSpec.ModuleSpec.PackageName,
 		Services:         clientSpec.ModuleSpec.Services,
 		IncludedPackages: clientSpec.ModuleSpec.IncludedPackages,
 		ClientID:         clientSpec.ClientID,
+		ExposedMethods:   exposedMethods,
 	}
 
 	client, err := g.templates.execTemplate(
@@ -689,6 +697,7 @@ func (g *EndpointGenerator) generateEndpointFile(
 			strings.Title(method.Name) + "Endpoint"
 	}
 
+	// TODO: http client needs to support multiple thrift services
 	meta := &EndpointMeta{
 		GatewayPackageName: g.packageHelper.GoGatewayPackageName(),
 		PackageName:        m.PackageName,
@@ -699,6 +708,7 @@ func (g *EndpointGenerator) generateEndpointFile(
 		ResHeaderMap:       e.ResHeaderMap,
 		ResHeaderMapKeys:   e.ResHeaderMapKeys,
 		ClientName:         e.ClientName,
+		ClientMethodName:   e.ClientMethod,
 		WorkflowName:       workflowName,
 	}
 

codegen/service.go

@@ -22,6 +22,7 @@ package codegen
 
 import (
 	"sort"
+	"strings"
 
 	"github.com/pkg/errors"
 	"go.uber.org/thriftrw/compile"
@@ -147,8 +148,10 @@ func NewServiceSpec(spec *compile.ServiceSpec, wantAnnot bool, packageHelper *Pa
 
 // SetDownstream ...
 func (ms *ModuleSpec) SetDownstream(
-	serviceName string, methodName string,
-	clientSpec *ClientSpec, clientService string, clientMethod string,
+	serviceName string,
+	methodName string,
+	clientSpec *ClientSpec,
+	clientMethod string,
 	h *PackageHelper,
 ) error {
 	var service *ServiceSpec
@@ -179,9 +182,20 @@ func (ms *ModuleSpec) SetDownstream(
 		)
 	}
 
-	err := method.setDownstream(
-		clientSpec.ModuleSpec, serviceName, clientMethod,
-	)
+	var err error
+	// TODO: we should do the same thing for http client as well
+	if clientSpec.ClientType == "tchannel" {
+		serviceMethod, ok := clientSpec.ExposedMethods[clientMethod]
+		if !ok {
+			return errors.Errorf("Client %q does not expose method %q", clientSpec.ClientName, clientMethod)
+		}
+		sm := strings.Split(serviceMethod, "::")
+
+		err = method.setDownstream(clientSpec.ModuleSpec, sm[0], sm[1])
+	} else {
+		err = method.setDownstream(clientSpec.ModuleSpec, clientSpec.ClientName, clientMethod)
+	}
+
 	if err != nil {
 		return err
 	}

codegen/template.go

@@ -130,6 +130,7 @@ type ClientMeta struct {
 	ClientID         string
 	IncludedPackages []GoPackageImport
 	Services         []*ServiceSpec
+	ExposedMethods   map[string]string
 }
 
 func findMethod(

codegen/templates/endpoint.tmpl

@@ -33,6 +33,7 @@ import (
 {{ $clientName := title .ClientName -}}
 {{ $handlerName := title .Method.Name | printf "%sHandler" }}
 {{ $responseType := .Method.ResponseType}}
+{{ $clientMethodName := title .ClientMethodName -}}
 {{with .Method -}}
 
 // {{$handlerName}} is the handler for "{{.HTTPPath}}"
@@ -122,7 +123,6 @@ func (handler *{{$handlerName}}) HandleRequest(
 {{- $clientMethod := .DownstreamMethod -}}
 {{- $clientReqType := fullTypeName ($clientMethod).RequestType ($clientPackage) -}}
 {{- $clientResType := fullTypeName  ($clientMethod).ResponseType ($clientPackage) -}}
-{{- $clientMethodName := title ($clientMethod).Name -}}
 {{- $clientExceptions := .DownstreamMethod.Exceptions -}}
 
 // {{$workflow}} calls thrift client {{$clientName}}.{{$clientMethodName}}