fix: allow IN queries on __key__ (#1085)

This change removes special handling encoding property filters when the property provided is `__key__`. In this special case, changes are made to signatures to cause typescript compilation errors when values associated with the `__key__` property are not keys or arrays of keys. This is because we only expect and support such values when `__key__` is provided in place of property so it is better if the user understands which values are supported at compile time. Special care with extra tests was taken to ensure that when removing this logic, tests for all supported use cases were still passing and every new use case worked as intended. Changes to make code more idiomatic were also introduced.

src/filter.ts

@@ -14,6 +14,7 @@
 
 import {Operator, Filter as IFilter} from './query';
 import {entity} from './entity';
+import Key = entity.Key;
 
 const OP_TO_OPERATOR = new Map([
   ['=', 'EQUAL'],
@@ -56,42 +57,34 @@ export abstract class EntityFilter {
   abstract toProto(): any;
 }
 
+export type AllowedFilterValueType<T> = T extends '__key__'
+  ? Key | Key[]
+  : unknown;
+
 /**
  * A PropertyFilter is a filter that gets applied to a query directly.
  *
  * @see {@link https://cloud.google.com/datastore/docs/concepts/queries#property_filters| Property filters Reference}
  *
  * @class
  */
-export class PropertyFilter extends EntityFilter implements IFilter {
-  name: string;
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  val: any;
-  op: Operator;
-
+export class PropertyFilter<T extends string>
+  extends EntityFilter
+  implements IFilter
+{
   /**
    * Build a Property Filter object.
    *
    * @param {string} Property
    * @param {Operator} operator
    * @param {any} val
    */
-  constructor(property: string, operator: Operator, val: any) {
+  constructor(
+    public name: T,
+    public op: Operator,
+    public val: AllowedFilterValueType<T>
+  ) {
     super();
-    this.name = property;
-    this.op = operator;
-    this.val = val;
-  }
-
-  private encodedValue(): any {
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    let value: any = {};
-    if (this.name === '__key__') {
-      value.keyValue = entity.keyToKeyProto(this.val);
-    } else {
-      value = entity.encodeValue(this.val, this.name);
-    }
-    return value;
   }
 
   /**
@@ -100,18 +93,13 @@ export class PropertyFilter extends EntityFilter implements IFilter {
    */
   // eslint-disable-next-line
   toProto(): any {
-    const value = new PropertyFilter(
-      this.name,
-      this.op,
-      this.val
-    ).encodedValue();
     return {
       propertyFilter: {
         property: {
           name: this.name,
         },
         op: OP_TO_OPERATOR.get(this.op),
-        value,
+        value: entity.encodeValue(this.val, this.name),
       },
     };
   }
@@ -156,5 +144,5 @@ class CompositeFilter extends EntityFilter {
 }
 
 export function isFilter(filter: any): filter is EntityFilter {
-  return (filter as EntityFilter).toProto !== undefined;
+  return filter instanceof EntityFilter;
 }

src/query.ts

@@ -18,7 +18,7 @@ import arrify = require('arrify');
 import {Key} from 'readline';
 import {Datastore} from '.';
 import {Entity} from './entity';
-import {EntityFilter, isFilter} from './filter';
+import {EntityFilter, isFilter, AllowedFilterValueType} from './filter';
 import {Transaction} from './transaction';
 import {CallOptions} from 'google-gax';
 import {RunQueryStreamOptions} from '../src/request';
@@ -207,12 +207,20 @@ class Query {
    * const keyQuery = query.filter('__key__', key);
    * ```
    */
-  filter(propertyOrFilter: string | EntityFilter, value?: {} | null): Query;
-  filter(propertyOrFilter: string, operator: Operator, value: {} | null): Query;
-  filter(
-    propertyOrFilter: string | EntityFilter,
-    operatorOrValue?: Operator,
-    value?: {} | null
+  filter(filter: EntityFilter): Query;
+  filter<T extends string>(
+    property: T,
+    value: AllowedFilterValueType<T>
+  ): Query;
+  filter<T extends string>(
+    property: T,
+    operator: Operator,
+    value: AllowedFilterValueType<T>
+  ): Query;
+  filter<T extends string>(
+    propertyOrFilter: T | EntityFilter,
+    operatorOrValue?: Operator | AllowedFilterValueType<T>,
+    value?: AllowedFilterValueType<T>
   ): Query {
     if (isFilter(propertyOrFilter)) {
       this.entityFilters.push(propertyOrFilter);
@@ -223,7 +231,7 @@ class Query {
       );
       let operator = operatorOrValue as Operator;
       if (arguments.length === 2) {
-        value = operatorOrValue as {};
+        value = operatorOrValue as AllowedFilterValueType<T>;
         operator = '=';
       }
 

system-test/datastore.ts

@@ -22,6 +22,8 @@ import {google} from '../protos/protos';
 import {Storage} from '@google-cloud/storage';
 import {AggregateField} from '../src/aggregate';
 import {PropertyFilter, EntityFilter, and, or} from '../src/filter';
+import {entity} from '../src/entity';
+import KEY_SYMBOL = entity.KEY_SYMBOL;
 
 describe('Datastore', () => {
   const testKinds: string[] = [];
@@ -794,6 +796,38 @@ describe('Datastore', () => {
       assert.strictEqual(entities!.length, 3);
     });
 
+    it('should filter queries with __key__ and IN', async () => {
+      const key1 = datastore.key(['Book', 'GoT', 'Character', 'Rickard']);
+      const key2 = datastore.key([
+        'Book',
+        'GoT',
+        'Character',
+        'Rickard',
+        'Character',
+        'Eddard',
+        'Character',
+        'Sansa',
+      ]);
+      const key3 = datastore.key([
+        'Book',
+        'GoT',
+        'Character',
+        'Rickard',
+        'Character',
+        'Eddard',
+      ]);
+      const value = [key1, key2, key3];
+      const q = datastore
+        .createQuery('Character')
+        .hasAncestor(ancestor)
+        .filter('__key__', 'IN', value);
+      const [entities] = await datastore.runQuery(q);
+      assert.strictEqual(entities!.length, 3);
+      assert.deepStrictEqual(entities[0][KEY_SYMBOL], key1);
+      assert.deepStrictEqual(entities[1][KEY_SYMBOL], key3);
+      assert.deepStrictEqual(entities[2][KEY_SYMBOL], key2);
+    });
+
     it('should filter queries with NOT_IN', async () => {
       const q = datastore
         .createQuery('Character')

test/entity.ts

@@ -17,7 +17,7 @@ import {beforeEach, afterEach, describe, it} from 'mocha';
 import * as extend from 'extend';
 import * as sinon from 'sinon';
 import {Datastore} from '../src';
-import {Entity} from '../src/entity';
+import {Entity, entity as globalEntity} from '../src/entity';
 import {IntegerTypeCastOptions} from '../src/query';
 import {PropertyFilter, EntityFilter, and} from '../src/filter';
 
@@ -1873,7 +1873,7 @@ describe('entity', () => {
     };
 
     it('should support all configurations of a query', () => {
-      const ancestorKey = new entity.Key({
+      const ancestorKey = new globalEntity.Key({
         path: ['Kind2', 'somename'],
       });
 
@@ -1894,8 +1894,62 @@ describe('entity', () => {
       assert.deepStrictEqual(entity.queryToQueryProto(query), queryProto);
     });
 
+    it('should support using __key__ with array as value', () => {
+      const keyWithInQuery = {
+        distinctOn: [],
+        filter: {
+          compositeFilter: {
+            filters: [
+              {
+                propertyFilter: {
+                  op: 'IN',
+                  property: {
+                    name: '__key__',
+                  },
+                  value: {
+                    arrayValue: {
+                      values: [
+                        {
+                          keyValue: {
+                            path: [
+                              {
+                                kind: 'Kind1',
+                                name: 'key1',
+                              },
+                            ],
+                          },
+                        },
+                      ],
+                    },
+                  },
+                },
+              },
+            ],
+            op: 'AND',
+          },
+        },
+        kind: [
+          {
+            name: 'Kind1',
+          },
+        ],
+        order: [],
+        projection: [],
+      };
+
+      const ds = new Datastore({projectId: 'project-id'});
+
+      const query = ds
+        .createQuery('Kind1')
+        .filter('__key__', 'IN', [
+          new globalEntity.Key({path: ['Kind1', 'key1']}),
+        ]);
+
+      assert.deepStrictEqual(entity.queryToQueryProto(query), keyWithInQuery);
+    });
+
     it('should support the filter method with Filter objects', () => {
-      const ancestorKey = new entity.Key({
+      const ancestorKey = new globalEntity.Key({
         path: ['Kind2', 'somename'],
       });
 
@@ -1916,7 +1970,7 @@ describe('entity', () => {
     });
 
     it('should support the filter method with AND', () => {
-      const ancestorKey = new entity.Key({
+      const ancestorKey = new globalEntity.Key({
         path: ['Kind2', 'somename'],
       });