Add some hydration tests.

Author: rkistner

packages/sync-rules/test/src/parameter_queries.test.ts

@@ -1,7 +1,19 @@
-import { describe, expect, test } from 'vitest';
-import { UnscopedParameterLookup, SqlParameterQuery, SourceTableInterface } from '../../src/index.js';
+import { beforeEach, describe, expect, test } from 'vitest';
+import {
+  UnscopedParameterLookup,
+  SqlParameterQuery,
+  SourceTableInterface,
+  debugHydratedMergedSource,
+  BucketParameterQuerier,
+  QuerierError,
+  GetQuerierOptions,
+  RequestParameters,
+  ScopedParameterLookup,
+  mergeParameterIndexLookupCreators
+} from '../../src/index.js';
 import { StaticSqlParameterQuery } from '../../src/StaticSqlParameterQuery.js';
 import { BASIC_SCHEMA, EMPTY_DATA_SOURCE, normalizeTokenParameters, PARSE_OPTIONS } from './util.js';
+import { HydrationState } from '../../src/HydrationState.js';
 
 describe('parameter queries', () => {
   const table = (name: string): SourceTableInterface => ({
@@ -855,6 +867,85 @@ describe('parameter queries', () => {
     expect(query.getLookups(requestParams)).toEqual([UnscopedParameterLookup.normalized(['user1'])]);
   });
 
+  describe('custom hydrationState', function () {
+    const hydrationState: HydrationState = {
+      getBucketSourceScope(source) {
+        return { bucketPrefix: `${source.uniqueName}-test` };
+      },
+      getParameterIndexLookupScope(source) {
+        return {
+          lookupName: `${source.defaultLookupScope.lookupName}.test`,
+          queryId: `${source.defaultLookupScope.queryId}.test`
+        };
+      }
+    };
+
+    let query: SqlParameterQuery;
+
+    beforeEach(() => {
+      const sql = 'SELECT id as group_id FROM groups WHERE token_parameters.user_id IN groups.user_ids';
+      query = SqlParameterQuery.fromSql(
+        'mybucket',
+        sql,
+        PARSE_OPTIONS,
+        'myquery',
+        EMPTY_DATA_SOURCE
+      ) as SqlParameterQuery;
+
+      expect(query.errors).toEqual([]);
+    });
+
+    test('for lookups', function () {
+      const merged = mergeParameterIndexLookupCreators(hydrationState, [query]);
+      const result = merged.evaluateParameterRow(TABLE_GROUPS, {
+        id: 'group1',
+        user_ids: JSON.stringify(['test-user', 'other-user'])
+      });
+      expect(result).toEqual([
+        {
+          lookup: ScopedParameterLookup.direct({ lookupName: 'mybucket.test', queryId: 'myquery.test' }, ['test-user']),
+          bucketParameters: [{ group_id: 'group1' }]
+        },
+        {
+          lookup: ScopedParameterLookup.direct({ lookupName: 'mybucket.test', queryId: 'myquery.test' }, [
+            'other-user'
+          ]),
+          bucketParameters: [{ group_id: 'group1' }]
+        }
+      ]);
+    });
+
+    test('for queries', function () {
+      const hydrated = query.createParameterQuerierSource({ hydrationState });
+
+      const queriers: BucketParameterQuerier[] = [];
+      const errors: QuerierError[] = [];
+      const pending = { queriers, errors };
+
+      const querierOptions: GetQuerierOptions = {
+        hasDefaultStreams: true,
+        globalParameters: new RequestParameters(
+          {
+            sub: 'test-user'
+          },
+          {}
+        ),
+        streams: {}
+      };
+
+      hydrated.pushBucketParameterQueriers(pending, querierOptions);
+
+      expect(errors).toEqual([]);
+      expect(queriers.length).toBe(1);
+
+      const querier = queriers[0];
+
+      expect(querier.parameterQueryLookups).toEqual([
+        ScopedParameterLookup.direct({ lookupName: 'mybucket.test', queryId: 'myquery.test' }, ['test-user'])
+      ]);
+    });
+  });
+
   test('invalid OR in parameter queries', () => {
     // Supporting this case is more tricky. We can do this by effectively denormalizing the OR clause
     // into separate queries, but it's a significant change. For now, developers should do that manually.

packages/sync-rules/test/src/static_parameter_queries.test.ts

@@ -1,6 +1,13 @@
 import { describe, expect, test } from 'vitest';
-import { BucketDataScope } from '../../src/HydrationState.js';
-import { RequestParameters, SqlParameterQuery } from '../../src/index.js';
+import { BucketDataScope, HydrationState } from '../../src/HydrationState.js';
+import {
+  BucketParameterQuerier,
+  GetQuerierOptions,
+  QuerierError,
+  RequestParameters,
+  ScopedParameterLookup,
+  SqlParameterQuery
+} from '../../src/index.js';
 import { StaticSqlParameterQuery } from '../../src/StaticSqlParameterQuery.js';
 import { EMPTY_DATA_SOURCE, normalizeTokenParameters, PARSE_OPTIONS } from './util.js';
 
@@ -365,4 +372,63 @@ describe('static parameter queries', () => {
       "select request.parameters() ->> 'project_id' as project_id where request.jwt() ->> 'role' = 'authenticated'"
     );
   });
+
+  test('custom hydrationState for buckets', function () {
+    const sql = 'SELECT token_parameters.user_id';
+    const query = SqlParameterQuery.fromSql(
+      'mybucket',
+      sql,
+      PARSE_OPTIONS,
+      '1',
+      EMPTY_DATA_SOURCE
+    ) as StaticSqlParameterQuery;
+
+    expect(query.errors).toEqual([]);
+
+    const hydrationState: HydrationState = {
+      getBucketSourceScope(source) {
+        return { bucketPrefix: `${source.uniqueName}-test` };
+      },
+      getParameterIndexLookupScope(source) {
+        return {
+          lookupName: `${source.defaultLookupScope.lookupName}.test`,
+          queryId: `${source.defaultLookupScope.queryId}.test`
+        };
+      }
+    };
+
+    // Internal API
+    const hydrated = query.createParameterQuerierSource({ hydrationState });
+
+    const queriers: BucketParameterQuerier[] = [];
+    const errors: QuerierError[] = [];
+    const pending = { queriers, errors };
+
+    const querierOptions: GetQuerierOptions = {
+      hasDefaultStreams: true,
+      globalParameters: new RequestParameters(
+        {
+          sub: 'test-user'
+        },
+        {}
+      ),
+      streams: {}
+    };
+
+    hydrated.pushBucketParameterQueriers(pending, querierOptions);
+
+    expect(errors).toEqual([]);
+    expect(queriers).toMatchObject([
+      {
+        staticBuckets: [
+          {
+            bucket: 'mybucket-test["test-user"]',
+            definition: 'mybucket',
+            inclusion_reasons: ['default'],
+            priority: 3
+          }
+        ]
+      }
+    ]);
+  });
 });