fix: allow IN queries on __key__

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'],
       });