ts-proto - npm Package Compare versions

Comparing version 1.4.0 to 1.4.1

build/main.js

@@ -11,3 +11,7 @@ "use strict";
 const dataloader = ts_poet_1.TypeNames.anyType('DataLoader=dataloader');
-function generateFile(typeMap, fileDesc) {
+function generateFile(typeMap, fileDesc, parameter) {
+    const options = { useContext: false };
+    if (parameter && parameter.includes('context=true')) {
+        options.useContext = true;
+    }
     // Google's protofiles are organized like Java, where package == the folder the file
@@ -56,7 +60,10 @@ // is in, and file == a specific service within the package. I.e. you can have multiple
     visitServices(fileDesc, serviceDesc => {
-        file = file.addInterface(generateService(typeMap, fileDesc, serviceDesc));
-        file = file.addClass(generateServiceClientImpl(typeMap, fileDesc, serviceDesc));
+        file = file.addInterface(generateService(typeMap, fileDesc, serviceDesc, options));
+        file = file.addClass(generateServiceClientImpl(typeMap, fileDesc, serviceDesc, options));
     });
     if (fileDesc.service.length > 0) {
-        file = file.addInterface(generateRpcType());
+        file = file.addInterface(generateRpcType(options));
+        if (options.useContext) {
+            file = file.addInterface(generateDataLoadersType(options));
+        }
     }
@@ -523,15 +530,29 @@ file = addLongUtilityMethod(file);
 }
-function generateService(typeMap, fileDesc, serviceDesc) {
+const contextTypeVar = ts_poet_1.TypeNames.typeVariable('Context', ts_poet_1.TypeNames.bound('DataLoaders'));
+function generateService(typeMap, fileDesc, serviceDesc, options) {
     let service = ts_poet_1.InterfaceSpec.create(serviceDesc.name).addModifiers(ts_poet_1.Modifier.EXPORT);
+    if (options.useContext) {
+        service = service.addTypeVariable(contextTypeVar);
+    }
     for (const methodDesc of serviceDesc.method) {
-        service = service.addFunction(ts_poet_1.FunctionSpec.create(methodDesc.name)
-            .addParameter('request', requestType(typeMap, methodDesc))
-            .returns(responsePromise(typeMap, methodDesc)));
-        const batchMethod = detectBatchMethod(typeMap, fileDesc, methodDesc);
-        if (batchMethod) {
-            const name = batchMethod.methodDesc.name.replace('Batch', 'Get');
-            service = service.addFunction(ts_poet_1.FunctionSpec.create(name)
-                .addParameter(utils_1.singular(batchMethod.inputFieldName), batchMethod.inputType)
-                .returns(ts_poet_1.TypeNames.PROMISE.param(batchMethod.outputType)));
+        let requestFn = ts_poet_1.FunctionSpec.create(methodDesc.name);
+        if (options.useContext) {
+            requestFn = requestFn.addParameter('ctx', ts_poet_1.TypeNames.typeVariable('Context'));
         }
+        requestFn = requestFn.addParameter('request', requestType(typeMap, methodDesc));
+        requestFn = requestFn.returns(responsePromise(typeMap, methodDesc));
+        service = service.addFunction(requestFn);
+        if (options.useContext) {
+            const batchMethod = detectBatchMethod(typeMap, fileDesc, serviceDesc, methodDesc);
+            if (batchMethod) {
+                const name = batchMethod.methodDesc.name.replace('Batch', 'Get');
+                let batchFn = ts_poet_1.FunctionSpec.create(name);
+                if (options.useContext) {
+                    batchFn = batchFn.addParameter('ctx', ts_poet_1.TypeNames.typeVariable('Context'));
+                }
+                batchFn = batchFn.addParameter(utils_1.singular(batchMethod.inputFieldName), batchMethod.inputType);
+                batchFn = batchFn.returns(ts_poet_1.TypeNames.PROMISE.param(batchMethod.outputType));
+                service = service.addFunction(batchFn);
+            }
+        }
     }
@@ -543,25 +564,50 @@ return service;
 }
-function generateServiceClientImpl(typeMap, fileDesc, serviceDesc) {
+function generateRegularRpcMethod(options, typeMap, fileDesc, serviceDesc, methodDesc) {
+    let requestFn = ts_poet_1.FunctionSpec.create(methodDesc.name);
+    if (options.useContext) {
+        requestFn = requestFn.addParameter('ctx', ts_poet_1.TypeNames.typeVariable('Context'));
+    }
+    return requestFn
+        .addParameter('request', requestType(typeMap, methodDesc))
+        .addStatement('const data = %L.encode(request).finish()', requestType(typeMap, methodDesc))
+        .addStatement('const promise = this.rpc.request(%L"%L.%L", %S, %L)', options.useContext ? 'ctx, ' : '', // sneak ctx in as the 1st parameter to our rpc call
+    fileDesc.package, serviceDesc.name, methodDesc.name, 'data')
+        .addStatement('return promise.then(data => %L.decode(new %T(data)))', responseType(typeMap, methodDesc), 'Reader@protobufjs/minimal')
+        .returns(responsePromise(typeMap, methodDesc));
+}
+function generateServiceClientImpl(typeMap, fileDesc, serviceDesc, options) {
+    // Define the FooServiceImpl class
     let client = ts_poet_1.ClassSpec.create(`${serviceDesc.name}ClientImpl`).addModifiers(ts_poet_1.Modifier.EXPORT);
+    if (options.useContext) {
+        client = client.addTypeVariable(contextTypeVar);
+        client = client.addInterface(`${serviceDesc.name}<Context>`);
+    }
+    else {
+        client = client.addInterface(serviceDesc.name);
+    }
+    // Create the constructor(rpc: Rpc)
+    const rpcType = options.useContext ? 'Rpc<Context>' : 'Rpc';
     client = client.addFunction(ts_poet_1.FunctionSpec.createConstructor()
-        .addParameter('rpc', 'Rpc')
+        .addParameter('rpc', rpcType)
         .addStatement('this.rpc = rpc'));
-    client = client.addProperty('rpc', 'Rpc', { modifiers: [ts_poet_1.Modifier.PRIVATE, ts_poet_1.Modifier.READONLY] });
+    client = client.addProperty('rpc', rpcType, { modifiers: [ts_poet_1.Modifier.PRIVATE, ts_poet_1.Modifier.READONLY] });
+    // Create a method for each FooService method
     for (const methodDesc of serviceDesc.method) {
-        // add a batch method if this fuzzy matches to a batch lookup method
-        const arrayBatchMethod = detectBatchMethod(typeMap, fileDesc, methodDesc);
-        if (arrayBatchMethod) {
-            client = generateBatchingMethod(typeMap, client, arrayBatchMethod);
+        // See if this this fuzzy matches to a batchable method
+        if (options.useContext) {
+            const batchMethod = detectBatchMethod(typeMap, fileDesc, serviceDesc, methodDesc);
+            if (batchMethod) {
+                client = client.addFunction(generateBatchingRpcMethod(typeMap, batchMethod));
+            }
         }
-        // generate the regular method
-        client = client.addFunction(ts_poet_1.FunctionSpec.create(methodDesc.name)
-            .addParameter('request', requestType(typeMap, methodDesc))
-            .addStatement('const data = %L.encode(request).finish()', requestType(typeMap, methodDesc))
-            .addStatement('const promise = this.rpc.request("%L.%L", %S, %L)', fileDesc.package, serviceDesc.name, methodDesc.name, 'data')
-            .addStatement('return promise.then(data => %L.decode(new %T(data)))', responseType(typeMap, methodDesc), 'Reader@protobufjs/minimal')
-            .returns(responsePromise(typeMap, methodDesc)));
+        if (options.useContext && methodDesc.name.match(/^Get[A-Z]/)) {
+            client = client.addFunction(generateCachingRpcMethod(typeMap, fileDesc, serviceDesc, methodDesc));
+        }
+        else {
+            client = client.addFunction(generateRegularRpcMethod(options, typeMap, fileDesc, serviceDesc, methodDesc));
+        }
     }
     return client;
 }
-function detectBatchMethod(typeMap, fileDesc, methodDesc) {
+function detectBatchMethod(typeMap, fileDesc, serviceDesc, methodDesc) {
     const nameMatches = methodDesc.name.startsWith('Batch');
@@ -584,4 +630,6 @@ const inputType = typeMap.get(methodDesc.inputType.substring(1)); // drop the `.` prefix
             }
+            const uniqueIdentifier = `${fileDesc.package}.${serviceDesc.name}.${methodDesc.name}`;
             return {
                 methodDesc,
+                uniqueIdentifier,
                 singleMethodName,
@@ -598,28 +646,49 @@ inputFieldName,
 }
-function generateBatchingMethod(typeMap, client, batchMethod) {
-    const name = batchMethod.singleMethodName.replace('Get', '');
-    const loaderFieldName = `${utils_1.lowerFirst(name)}Loader`;
-    const { methodDesc, singleMethodName, inputFieldName, inputType, outputFieldName, outputType, mapType } = batchMethod;
-    // add a dataloader field
+/** We've found a BatchXxx method, create a synthetic GetXxx method that calls it. */
+function generateBatchingRpcMethod(typeMap, batchMethod) {
+    const { methodDesc, singleMethodName, inputFieldName, inputType, outputFieldName, outputType, mapType, uniqueIdentifier } = batchMethod;
+    // Create the `(keys) => ...` lambda we'll pass to the DataLoader constructor
     let lambda = ts_poet_1.CodeBlock.lambda(inputFieldName) // e.g. keys
         .addStatement('const request = { %L }', inputFieldName);
     if (mapType) {
+        // If the return type is a map, lookup each key in the result
         lambda = lambda
-            .beginLambda('return this.%L(request).then(res =>', methodDesc.name)
-            .addStatement('return %L.map(e => res.%L[e])', inputFieldName, outputFieldName)
+            .beginLambda('return this.%L(ctx, request).then(res =>', methodDesc.name)
+            .addStatement('return %L.map(key => res.%L[key])', inputFieldName, outputFieldName)
             .endLambda(')');
     }
     else {
-        lambda = lambda.addStatement('return this.%L(request).then(res => res.%L)', methodDesc.name, outputFieldName);
+        // Otherwise assume they come back in order
+        lambda = lambda.addStatement('return this.%L(ctx, request).then(res => res.%L)', methodDesc.name, outputFieldName);
     }
-    client = client.addProperty(ts_poet_1.PropertySpec.create(loaderFieldName, dataloader.param(inputType, outputType))
-        .addModifiers(ts_poet_1.Modifier.PRIVATE)
-        .setImplicitlyTyped()
-        .initializer('new %T(%L)', dataloader.param(inputType, outputType), lambda));
-    client = client.addFunction(ts_poet_1.FunctionSpec.create(singleMethodName)
+    return ts_poet_1.FunctionSpec.create(singleMethodName)
+        .addParameter('ctx', 'Context')
         .addParameter(utils_1.singular(inputFieldName), inputType)
-        .addStatement('return this.%L.load(%L)', loaderFieldName, utils_1.singular(inputFieldName))
-        .returns(ts_poet_1.TypeNames.PROMISE.param(outputType)));
-    return client;
+        .addCode('const dl = ctx.getDataLoader(%S, () => {%>\n', uniqueIdentifier)
+        .addCode('return new %T<%T, %T>(%L, { cacheKeyFn: %T });\n', dataloader, inputType, outputType, lambda, ts_poet_1.TypeNames.anyType('hash=object-hash'))
+        .addCode('%<});\n')
+        .addStatement('return dl.load(%L)', utils_1.singular(inputFieldName))
+        .returns(ts_poet_1.TypeNames.PROMISE.param(outputType));
 }
+/** We're not going to batch, but use DataLoader for per-request caching. */
+function generateCachingRpcMethod(typeMap, fileDesc, serviceDesc, methodDesc) {
+    const inputType = requestType(typeMap, methodDesc);
+    const outputType = responseType(typeMap, methodDesc);
+    let lambda = ts_poet_1.CodeBlock.lambda('requests')
+        .beginLambda('const responses = requests.map(async request =>')
+        .addStatement('const data = %L.encode(request).finish()', inputType)
+        .addStatement('const response = await this.rpc.request(ctx, "%L.%L", %S, %L)', fileDesc.package, serviceDesc.name, methodDesc.name, 'data')
+        .addStatement('return %L.decode(new %T(response))', responseType(typeMap, methodDesc), 'Reader@protobufjs/minimal')
+        .endLambda(')')
+        .addStatement('return Promise.all(responses)');
+    const uniqueIdentifier = `${fileDesc.package}.${serviceDesc.name}.${methodDesc.name}`;
+    return ts_poet_1.FunctionSpec.create(methodDesc.name)
+        .addParameter('ctx', 'Context')
+        .addParameter('request', requestType(typeMap, methodDesc))
+        .addCode('const dl = ctx.getDataLoader(%S, () => {%>\n', uniqueIdentifier)
+        .addCode('return new %T<%T, %T>(%L, { cacheKeyFn: %T });\n', dataloader, inputType, outputType, lambda, ts_poet_1.TypeNames.anyType('hash=object-hash'))
+        .addCode('%<});\n')
+        .addStatement('return dl.load(request)')
+        .returns(ts_poet_1.TypeNames.PROMISE.param(outputType));
+}
 /**
@@ -634,10 +703,29 @@ * Creates an `Rpc.request(service, method, data)` abstraction.
  */
-function generateRpcType() {
+function generateRpcType(options) {
     const data = ts_poet_1.TypeNames.anyType('Uint8Array');
-    return ts_poet_1.InterfaceSpec.create('Rpc').addFunction(ts_poet_1.FunctionSpec.create('request')
+    let fn = ts_poet_1.FunctionSpec.create('request');
+    if (options.useContext) {
+        fn = fn.addParameter('ctx', 'Context');
+    }
+    fn = fn
         .addParameter('service', ts_poet_1.TypeNames.STRING)
         .addParameter('method', ts_poet_1.TypeNames.STRING)
         .addParameter('data', data)
-        .returns(ts_poet_1.TypeNames.PROMISE.param(data)));
+        .returns(ts_poet_1.TypeNames.PROMISE.param(data));
+    let rpc = ts_poet_1.InterfaceSpec.create('Rpc');
+    if (options.useContext) {
+        rpc = rpc.addTypeVariable(ts_poet_1.TypeNames.typeVariable('Context'));
+    }
+    rpc = rpc.addFunction(fn);
+    return rpc;
 }
+function generateDataLoadersType(options) {
+    // TODO Maybe should be a generic `Context.get<T>(id, () => T): T` method
+    let fn = ts_poet_1.FunctionSpec.create('getDataLoader')
+        .addTypeVariable(ts_poet_1.TypeNames.typeVariable('T'))
+        .addParameter('identifier', ts_poet_1.TypeNames.STRING)
+        .addParameter('cstrFn', ts_poet_1.TypeNames.lambda2([], ts_poet_1.TypeNames.typeVariable('T')))
+        .returns(ts_poet_1.TypeNames.typeVariable('T'));
+    return ts_poet_1.InterfaceSpec.create('DataLoaders').addFunction(fn);
+}
 function requestType(typeMap, methodDesc) {
@@ -644,0 +732,0 @@ return types_1.messageToTypeName(typeMap, methodDesc.inputType);

build/plugin.js

@@ -18,3 +18,3 @@ "use strict";
     const files = request.protoFile.map(file => {
-        const spec = main_1.generateFile(typeMap, file);
+        const spec = main_1.generateFile(typeMap, file, request.parameter);
         return new CodeGeneratorResponse.File({
@@ -21,0 +21,0 @@ name: spec.path,

package.json

 {
   "name": "ts-proto",
-  "version": "1.4.0",
+  "version": "1.4.1",
   "description": "",
@@ -30,7 +30,9 @@ "main": "build/plugin.js",
   "dependencies": {
+    "@types/object-hash": "^1.3.0",
     "dataloader": "^1.4.0",
+    "object-hash": "^1.3.1",
     "protobufjs": "^6.8.8",
     "sequency": "^0.19.2",
-    "ts-poet": "^0.5.4"
+    "ts-poet": "^1.0.0"
   }
 }

src/main.ts

@@ -37,3 +37,3 @@ import {
 import { asSequence } from 'sequency';
-import { lowerFirst, singular } from './utils';
+import { singular } from './utils';
 import DescriptorProto = google.protobuf.DescriptorProto;
@@ -48,3 +48,12 @@ import FieldDescriptorProto = google.protobuf.FieldDescriptorProto;
 
-export function generateFile(typeMap: TypeMap, fileDesc: FileDescriptorProto): FileSpec {
+export type Options = {
+  useContext: boolean;
+};
+
+export function generateFile(typeMap: TypeMap, fileDesc: FileDescriptorProto, parameter: string): FileSpec {
+  const options: Options = { useContext: false };
+  if (parameter && parameter.includes('context=true')) {
+    options.useContext = true;
+  }
+
   // Google's protofiles are organized like Java, where package == the folder the file
@@ -104,8 +113,11 @@ // is in, and file == a specific service within the package. I.e. you can have multiple
   visitServices(fileDesc, serviceDesc => {
-    file = file.addInterface(generateService(typeMap, fileDesc, serviceDesc));
-    file = file.addClass(generateServiceClientImpl(typeMap, fileDesc, serviceDesc));
+    file = file.addInterface(generateService(typeMap, fileDesc, serviceDesc, options));
+    file = file.addClass(generateServiceClientImpl(typeMap, fileDesc, serviceDesc, options));
   });
 
   if (fileDesc.service.length > 0) {
-    file = file.addInterface(generateRpcType());
+    file = file.addInterface(generateRpcType(options));
+    if (options.useContext) {
+      file = file.addInterface(generateDataLoadersType(options));
+    }
   }
@@ -632,23 +644,36 @@
 
+const contextTypeVar = TypeNames.typeVariable('Context', TypeNames.bound('DataLoaders'));
+
 function generateService(
   typeMap: TypeMap,
   fileDesc: FileDescriptorProto,
-  serviceDesc: ServiceDescriptorProto
+  serviceDesc: ServiceDescriptorProto,
+  options: Options
 ): InterfaceSpec {
   let service = InterfaceSpec.create(serviceDesc.name).addModifiers(Modifier.EXPORT);
+  if (options.useContext) {
+    service = service.addTypeVariable(contextTypeVar);
+  }
   for (const methodDesc of serviceDesc.method) {
-    service = service.addFunction(
-      FunctionSpec.create(methodDesc.name)
-        .addParameter('request', requestType(typeMap, methodDesc))
-        .returns(responsePromise(typeMap, methodDesc))
-    );
-    const batchMethod = detectBatchMethod(typeMap, fileDesc, methodDesc);
-    if (batchMethod) {
-      const name = batchMethod.methodDesc.name.replace('Batch', 'Get');
-      service = service.addFunction(
-        FunctionSpec.create(name)
-          .addParameter(singular(batchMethod.inputFieldName), batchMethod.inputType)
-          .returns(TypeNames.PROMISE.param(batchMethod.outputType))
-      );
+    let requestFn = FunctionSpec.create(methodDesc.name);
+    if (options.useContext) {
+      requestFn = requestFn.addParameter('ctx', TypeNames.typeVariable('Context'));
     }
+    requestFn = requestFn.addParameter('request', requestType(typeMap, methodDesc));
+    requestFn = requestFn.returns(responsePromise(typeMap, methodDesc));
+    service = service.addFunction(requestFn);
+
+    if (options.useContext) {
+      const batchMethod = detectBatchMethod(typeMap, fileDesc, serviceDesc, methodDesc);
+      if (batchMethod) {
+        const name = batchMethod.methodDesc.name.replace('Batch', 'Get');
+        let batchFn = FunctionSpec.create(name);
+        if (options.useContext) {
+          batchFn = batchFn.addParameter('ctx', TypeNames.typeVariable('Context'));
+        }
+        batchFn = batchFn.addParameter(singular(batchMethod.inputFieldName), batchMethod.inputType);
+        batchFn = batchFn.returns(TypeNames.PROMISE.param(batchMethod.outputType));
+        service = service.addFunction(batchFn);
+      }
+    }
   }
@@ -660,2 +685,4 @@ return service;
   methodDesc: MethodDescriptorProto;
+  // a ${package + service + method name} key to identify this method in caches
+  uniqueIdentifier: string;
   singleMethodName: string;
@@ -673,39 +700,71 @@ inputFieldName: string;
 
+function generateRegularRpcMethod(
+  options: Options,
+  typeMap: TypeMap,
+  fileDesc: google.protobuf.FileDescriptorProto,
+  serviceDesc: google.protobuf.ServiceDescriptorProto,
+  methodDesc: google.protobuf.MethodDescriptorProto
+) {
+  let requestFn = FunctionSpec.create(methodDesc.name);
+  if (options.useContext) {
+    requestFn = requestFn.addParameter('ctx', TypeNames.typeVariable('Context'));
+  }
+  return requestFn
+    .addParameter('request', requestType(typeMap, methodDesc))
+    .addStatement('const data = %L.encode(request).finish()', requestType(typeMap, methodDesc))
+    .addStatement(
+      'const promise = this.rpc.request(%L"%L.%L", %S, %L)',
+      options.useContext ? 'ctx, ' : '', // sneak ctx in as the 1st parameter to our rpc call
+      fileDesc.package,
+      serviceDesc.name,
+      methodDesc.name,
+      'data'
+    )
+    .addStatement(
+      'return promise.then(data => %L.decode(new %T(data)))',
+      responseType(typeMap, methodDesc),
+      'Reader@protobufjs/minimal'
+    )
+    .returns(responsePromise(typeMap, methodDesc));
+}
+
 function generateServiceClientImpl(
   typeMap: TypeMap,
   fileDesc: FileDescriptorProto,
-  serviceDesc: ServiceDescriptorProto
+  serviceDesc: ServiceDescriptorProto,
+  options: Options
 ): ClassSpec {
+  // Define the FooServiceImpl class
   let client = ClassSpec.create(`${serviceDesc.name}ClientImpl`).addModifiers(Modifier.EXPORT);
+  if (options.useContext) {
+    client = client.addTypeVariable(contextTypeVar);
+    client = client.addInterface(`${serviceDesc.name}<Context>`);
+  } else {
+    client = client.addInterface(serviceDesc.name);
+  }
+
+  // Create the constructor(rpc: Rpc)
+  const rpcType = options.useContext ? 'Rpc<Context>' : 'Rpc';
   client = client.addFunction(
     FunctionSpec.createConstructor()
-      .addParameter('rpc', 'Rpc')
+      .addParameter('rpc', rpcType)
       .addStatement('this.rpc = rpc')
   );
-  client = client.addProperty('rpc', 'Rpc', { modifiers: [Modifier.PRIVATE, Modifier.READONLY] });
+  client = client.addProperty('rpc', rpcType, { modifiers: [Modifier.PRIVATE, Modifier.READONLY] });
+
+  // Create a method for each FooService method
   for (const methodDesc of serviceDesc.method) {
-    // add a batch method if this fuzzy matches to a batch lookup method
-    const arrayBatchMethod = detectBatchMethod(typeMap, fileDesc, methodDesc);
-    if (arrayBatchMethod) {
-      client = generateBatchingMethod(typeMap, client, arrayBatchMethod);
+    // See if this this fuzzy matches to a batchable method
+    if (options.useContext) {
+      const batchMethod = detectBatchMethod(typeMap, fileDesc, serviceDesc, methodDesc);
+      if (batchMethod) {
+        client = client.addFunction(generateBatchingRpcMethod(typeMap, batchMethod));
+      }
     }
-    // generate the regular method
-    client = client.addFunction(
-      FunctionSpec.create(methodDesc.name)
-        .addParameter('request', requestType(typeMap, methodDesc))
-        .addStatement('const data = %L.encode(request).finish()', requestType(typeMap, methodDesc))
-        .addStatement(
-          'const promise = this.rpc.request("%L.%L", %S, %L)',
-          fileDesc.package,
-          serviceDesc.name,
-          methodDesc.name,
-          'data'
-        )
-        .addStatement(
-          'return promise.then(data => %L.decode(new %T(data)))',
-          responseType(typeMap, methodDesc),
-          'Reader@protobufjs/minimal'
-        )
-        .returns(responsePromise(typeMap, methodDesc))
-    );
+
+    if (options.useContext && methodDesc.name.match(/^Get[A-Z]/)) {
+      client = client.addFunction(generateCachingRpcMethod(typeMap, fileDesc, serviceDesc, methodDesc));
+    } else {
+      client = client.addFunction(generateRegularRpcMethod(options, typeMap, fileDesc, serviceDesc, methodDesc));
+    }
   }
@@ -718,2 +777,3 @@ return client;
   fileDesc: FileDescriptorProto,
+  serviceDesc: ServiceDescriptorProto,
   methodDesc: MethodDescriptorProto
@@ -738,4 +798,6 @@ ): BatchMethod | undefined {
       }
+      const uniqueIdentifier = `${fileDesc.package}.${serviceDesc.name}.${methodDesc.name}`;
       return {
         methodDesc,
+        uniqueIdentifier,
         singleMethodName,
@@ -753,32 +815,84 @@ inputFieldName,
 
-function generateBatchingMethod(typeMap: TypeMap, client: ClassSpec, batchMethod: BatchMethod): ClassSpec {
-  const name = batchMethod.singleMethodName.replace('Get', '');
-  const loaderFieldName = `${lowerFirst(name)}Loader`;
-  const { methodDesc, singleMethodName, inputFieldName, inputType, outputFieldName, outputType, mapType } = batchMethod;
-  // add a dataloader field
+/** We've found a BatchXxx method, create a synthetic GetXxx method that calls it. */
+function generateBatchingRpcMethod(typeMap: TypeMap, batchMethod: BatchMethod): FunctionSpec {
+  const {
+    methodDesc,
+    singleMethodName,
+    inputFieldName,
+    inputType,
+    outputFieldName,
+    outputType,
+    mapType,
+    uniqueIdentifier
+  } = batchMethod;
+  // Create the `(keys) => ...` lambda we'll pass to the DataLoader constructor
   let lambda = CodeBlock.lambda(inputFieldName) // e.g. keys
     .addStatement('const request = { %L }', inputFieldName);
   if (mapType) {
+    // If the return type is a map, lookup each key in the result
     lambda = lambda
-      .beginLambda('return this.%L(request).then(res =>', methodDesc.name)
-      .addStatement('return %L.map(e => res.%L[e])', inputFieldName, outputFieldName)
+      .beginLambda('return this.%L(ctx, request).then(res =>', methodDesc.name)
+      .addStatement('return %L.map(key => res.%L[key])', inputFieldName, outputFieldName)
       .endLambda(')');
   } else {
-    lambda = lambda.addStatement('return this.%L(request).then(res => res.%L)', methodDesc.name, outputFieldName);
+    // Otherwise assume they come back in order
+    lambda = lambda.addStatement('return this.%L(ctx, request).then(res => res.%L)', methodDesc.name, outputFieldName);
   }
-  client = client.addProperty(
-    PropertySpec.create(loaderFieldName, dataloader.param(inputType, outputType))
-      .addModifiers(Modifier.PRIVATE)
-      .setImplicitlyTyped()
-      .initializer('new %T(%L)', dataloader.param(inputType, outputType), lambda)
-  );
-  client = client.addFunction(
-    FunctionSpec.create(singleMethodName)
-      .addParameter(singular(inputFieldName), inputType)
-      .addStatement('return this.%L.load(%L)', loaderFieldName, singular(inputFieldName))
-      .returns(TypeNames.PROMISE.param(outputType))
-  );
-  return client;
+  return FunctionSpec.create(singleMethodName)
+    .addParameter('ctx', 'Context')
+    .addParameter(singular(inputFieldName), inputType)
+    .addCode('const dl = ctx.getDataLoader(%S, () => {%>\n', uniqueIdentifier)
+    .addCode(
+      'return new %T<%T, %T>(%L, { cacheKeyFn: %T });\n',
+      dataloader,
+      inputType,
+      outputType,
+      lambda,
+      TypeNames.anyType('hash=object-hash')
+    )
+    .addCode('%<});\n')
+    .addStatement('return dl.load(%L)', singular(inputFieldName))
+    .returns(TypeNames.PROMISE.param(outputType));
 }
 
+/** We're not going to batch, but use DataLoader for per-request caching. */
+function generateCachingRpcMethod(
+  typeMap: TypeMap,
+  fileDesc: FileDescriptorProto,
+  serviceDesc: ServiceDescriptorProto,
+  methodDesc: MethodDescriptorProto
+): FunctionSpec {
+  const inputType = requestType(typeMap, methodDesc);
+  const outputType = responseType(typeMap, methodDesc);
+  let lambda = CodeBlock.lambda('requests')
+    .beginLambda('const responses = requests.map(async request =>')
+    .addStatement('const data = %L.encode(request).finish()', inputType)
+    .addStatement(
+      'const response = await this.rpc.request(ctx, "%L.%L", %S, %L)',
+      fileDesc.package,
+      serviceDesc.name,
+      methodDesc.name,
+      'data'
+    )
+    .addStatement('return %L.decode(new %T(response))', responseType(typeMap, methodDesc), 'Reader@protobufjs/minimal')
+    .endLambda(')')
+    .addStatement('return Promise.all(responses)');
+  const uniqueIdentifier = `${fileDesc.package}.${serviceDesc.name}.${methodDesc.name}`;
+  return FunctionSpec.create(methodDesc.name)
+    .addParameter('ctx', 'Context')
+    .addParameter('request', requestType(typeMap, methodDesc))
+    .addCode('const dl = ctx.getDataLoader(%S, () => {%>\n', uniqueIdentifier)
+    .addCode(
+      'return new %T<%T, %T>(%L, { cacheKeyFn: %T });\n',
+      dataloader,
+      inputType,
+      outputType,
+      lambda,
+      TypeNames.anyType('hash=object-hash')
+    )
+    .addCode('%<});\n')
+    .addStatement('return dl.load(request)')
+    .returns(TypeNames.PROMISE.param(outputType));
+}
+
 /**
@@ -793,13 +907,31 @@ * Creates an `Rpc.request(service, method, data)` abstraction.
  */
-function generateRpcType(): InterfaceSpec {
+function generateRpcType(options: Options): InterfaceSpec {
   const data = TypeNames.anyType('Uint8Array');
-  return InterfaceSpec.create('Rpc').addFunction(
-    FunctionSpec.create('request')
-      .addParameter('service', TypeNames.STRING)
-      .addParameter('method', TypeNames.STRING)
-      .addParameter('data', data)
-      .returns(TypeNames.PROMISE.param(data))
-  );
+  let fn = FunctionSpec.create('request');
+  if (options.useContext) {
+    fn = fn.addParameter('ctx', 'Context');
+  }
+  fn = fn
+    .addParameter('service', TypeNames.STRING)
+    .addParameter('method', TypeNames.STRING)
+    .addParameter('data', data)
+    .returns(TypeNames.PROMISE.param(data));
+  let rpc = InterfaceSpec.create('Rpc');
+  if (options.useContext) {
+    rpc = rpc.addTypeVariable(TypeNames.typeVariable('Context'));
+  }
+  rpc = rpc.addFunction(fn);
+  return rpc;
 }
 
+function generateDataLoadersType(options: Options): InterfaceSpec {
+  // TODO Maybe should be a generic `Context.get<T>(id, () => T): T` method
+  let fn = FunctionSpec.create('getDataLoader')
+    .addTypeVariable(TypeNames.typeVariable('T'))
+    .addParameter('identifier', TypeNames.STRING)
+    .addParameter('cstrFn', TypeNames.lambda2([], TypeNames.typeVariable('T')))
+    .returns(TypeNames.typeVariable('T'));
+  return InterfaceSpec.create('DataLoaders').addFunction(fn);
+}
+
 function requestType(typeMap: TypeMap, methodDesc: MethodDescriptorProto): TypeName {
@@ -806,0 +938,0 @@ return messageToTypeName(typeMap, methodDesc.inputType);

src/plugin.ts

@@ -17,3 +17,3 @@ import { promisify } from 'util';
   const files = request.protoFile.map(file => {
-    const spec = generateFile(typeMap, file);
+    const spec = generateFile(typeMap, file, request.parameter);
     return new CodeGeneratorResponse.File({
@@ -20,0 +20,0 @@ name: spec.path,

tests/vector-test.ts

@@ -14,3 +14,3 @@ import { google } from '../build/pbjs';
     for (let file of request.protoFile) {
-      const spec = generateFile(typeMap, file);
+      const spec = generateFile(typeMap, file, "");
       const out = new StringBuffer();
@@ -17,0 +17,0 @@ spec.emit(out);

New alerts

License Policy Violation

License

This package is not allowed per your license policy. Review the package's license to ensure compliance.

Found 1 instance in 1 package

Fixed alerts

License Policy Violation

License

This package is not allowed per your license policy. Review the package's license to ensure compliance.

Found 1 instance in 1 package

Improved metrics

Number of package files

46

increased by4.55%

Lines of code

12764

increased by7.04%

Number of lines in readme file

115

increased by18.56%

Worsened metrics

Total package byte prevSize

834045

decreased by-0.23%

Dependency count

6

increased by50%

Dependency changes

• + Added@types/object-hash@^1.3.0

• + Addedobject-hash@^1.3.1

• + Added@types/object-hash@1.3.4(transitive)

• + Addedobject-hash@1.3.1(transitive)

• + Addedts-poet@1.0.1(transitive)

• - Removedts-poet@0.5.5(transitive)

• Updatedts-poet@^1.0.0