fix(grpc): patch original client methods

open-telemetry · Dec 16, 2019 · d783b9e · d783b9e
1 parent 9458390
commit d783b9e
Showing 1 changed file with 15 additions and 9 deletions.
diff --git a/packages/opentelemetry-plugin-grpc/src/grpc.ts b/packages/opentelemetry-plugin-grpc/src/grpc.ts
@@ -320,17 +320,26 @@ 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);
+        const methodsToWrap = [
+          ...Object.keys(methods),
+          ...(Object.keys(methods)
+            .map(methodName => methods[methodName].originalName)
+            .filter(
+              originalName =>
+                !!originalName && client.prototype.hasOwnProperty(originalName)
+            ) as string[]),
+        ];
         shimmer.massWrap(
           client.prototype as never,
-          Object.keys(methods) as never[],
+          methodsToWrap as never[],
           // tslint:disable-next-line:no-any
           plugin._getPatchedClientMethods() as any
         );
@@ -352,12 +361,9 @@ export class GrpcPlugin extends BasePlugin<grpc> {
             parent: plugin._tracer.getCurrentSpan(),
           })
           .setAttribute(AttributeNames.COMPONENT, GrpcPlugin.component);
-        return plugin._makeGrpcClientRemoteCall(
-          original,
-          args,
-          this,
-          plugin
-        )(span);
+        return plugin._makeGrpcClientRemoteCall(original, args, this, plugin)(
+          span
+        );
       };
     };
   }