fix(grpc): patch original client methods (open-telemetry#631)

* fix(grpc): patch original client methods * fix: review comments and build * fix: add test
cthulhu-bot · Dec 20, 2019 · 587a5f5 · 587a5f5
1 parent bc583b8
commit 587a5f5
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 3 deletions.
diff --git a/packages/opentelemetry-plugin-grpc/src/grpc.ts b/packages/opentelemetry-plugin-grpc/src/grpc.ts
@@ -320,17 +320,17 @@ 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
       ) {
         // tslint:disable-next-line:no-any
         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
@@ -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',