From 587a5f5f833b1d7a8f6a3fa3e1288f180997cd2f Mon Sep 17 00:00:00 2001
From: Mayur Kale <mayurkale@google.com>
Date: Thu, 19 Dec 2019 22:31:55 -0800
Subject: [PATCH] fix(grpc): patch original client methods (#631)

* fix(grpc): patch original client methods

* fix: review comments and build

* fix: add test
---
 .../opentelemetry-plugin-grpc/src/grpc.ts     | 22 ++++++++++++---
 .../test/grpc.test.ts                         | 27 +++++++++++++++++++
 2 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/packages/opentelemetry-plugin-grpc/src/grpc.ts b/packages/opentelemetry-plugin-grpc/src/grpc.ts
index b6e5fd5b77..2fba28743e 100644
--- a/packages/opentelemetry-plugin-grpc/src/grpc.ts
+++ b/packages/opentelemetry-plugin-grpc/src/grpc.ts
@@ -320,9 +320,9 @@ export class GrpcPlugin extends BasePlugin<grpc> {
     const plugin = this;
     return (original: typeof grpcTypes.makeGenericClientConstructor): never => {
       plugin._logger.debug('patching client');
-      return function makeClientConstructor<ImplementationType>(
+      return function makeClientConstructor(
         this: typeof grpcTypes.Client,
-        methods: grpcTypes.ServiceDefinition<ImplementationType>,
+        methods: { [key: string]: { originalName?: string } },
         serviceName: string,
         options: grpcTypes.GenericClientOptions
       ) {
@@ -330,7 +330,7 @@ export class GrpcPlugin extends BasePlugin<grpc> {
         const client = original.apply(this, arguments as any);
         shimmer.massWrap(
           client.prototype as never,
-          Object.keys(methods) as never[],
+          plugin._getMethodsToWrap(client, methods) as never[],
           // tslint:disable-next-line:no-any
           plugin._getPatchedClientMethods() as any
         );
@@ -339,6 +339,22 @@ export class GrpcPlugin extends BasePlugin<grpc> {
     };
   }
 
+  private _getMethodsToWrap(
+    client: typeof grpcTypes.Client,
+    methods: { [key: string]: { originalName?: string } }
+  ): string[] {
+    const methodsToWrap = [
+      ...Object.keys(methods),
+      ...(Object.keys(methods)
+        .map(methodName => methods[methodName].originalName)
+        .filter(
+          originalName =>
+            !!originalName && client.prototype.hasOwnProperty(originalName)
+        ) as string[]),
+    ];
+    return methodsToWrap;
+  }
+
   private _getPatchedClientMethods() {
     const plugin = this;
     return (original: GrpcClientFunc) => {
diff --git a/packages/opentelemetry-plugin-grpc/test/grpc.test.ts b/packages/opentelemetry-plugin-grpc/test/grpc.test.ts
index 4c7d5cb73f..a6d4f35450 100644
--- a/packages/opentelemetry-plugin-grpc/test/grpc.test.ts
+++ b/packages/opentelemetry-plugin-grpc/test/grpc.test.ts
@@ -45,6 +45,8 @@ type TestGrpcClient = grpc.Client & {
   // tslint:disable-next-line:no-any
   unaryMethod: any;
   // tslint:disable-next-line:no-any
+  UnaryMethod: any;
+  // tslint:disable-next-line:no-any
   clientStreamMethod: any;
   // tslint:disable-next-line:no-any
   serverStreamMethod: any;
@@ -93,6 +95,24 @@ const grpcClient = {
     });
   },
 
+  UnaryMethod: (
+    client: TestGrpcClient,
+    request: TestRequestResponse
+  ): Promise<TestRequestResponse> => {
+    return new Promise((resolve, reject) => {
+      return client.UnaryMethod(
+        request,
+        (err: grpc.ServiceError, response: TestRequestResponse) => {
+          if (err) {
+            reject(err);
+          } else {
+            resolve(response);
+          }
+        }
+      );
+    });
+  },
+
   clientStreamMethod: (
     client: TestGrpcClient,
     request: TestRequestResponse[]
@@ -318,6 +338,13 @@ describe('GrpcPlugin', () => {
       request: requestList[0],
       result: requestList[0],
     },
+    {
+      description: 'Unary call',
+      methodName: 'UnaryMethod',
+      method: grpcClient.UnaryMethod,
+      request: requestList[0],
+      result: requestList[0],
+    },
     {
       description: 'clientStream call',
       methodName: 'ClientStreamMethod',