Skip to content

Commit

Permalink
fix(grpc): patch original client methods (open-telemetry#631)
Browse files Browse the repository at this point in the history
* fix(grpc): patch original client methods

* fix: review comments and build

* fix: add test
  • Loading branch information
mayurkale22 authored Dec 20, 2019
1 parent bc583b8 commit 587a5f5
Show file tree
Hide file tree
Showing 2 changed files with 46 additions and 3 deletions.
22 changes: 19 additions & 3 deletions packages/opentelemetry-plugin-grpc/src/grpc.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
);
Expand All @@ -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) => {
Expand Down
27 changes: 27 additions & 0 deletions packages/opentelemetry-plugin-grpc/test/grpc.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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[]
Expand Down Expand Up @@ -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',
Expand Down

0 comments on commit 587a5f5

Please sign in to comment.