feat(pgsql-test): add IdHash support for snapshot utilities

Author: theothersideofgod

packages/pgsql-test/__tests__/utils.test.ts

@@ -5,7 +5,8 @@ import {
   pruneIds,
   pruneIdArrays,
   pruneUUIDs,
-  pruneHashes
+  pruneHashes,
+  IdHash
 } from '../src/utils';
 
 describe('snapshot utilities', () => {
@@ -46,6 +47,7 @@ describe('snapshot utilities', () => {
       const row: Record<string, unknown> = { id: null, user_id: undefined };
       expect(pruneIds(row)).toEqual({ id: null, user_id: undefined });
     });
+
   });
 
   describe('pruneIdArrays', () => {
@@ -113,6 +115,7 @@ describe('snapshot utilities', () => {
         name: 'Alice'
       });
     });
+
   });
 
   describe('snapshot', () => {
@@ -177,5 +180,72 @@ describe('snapshot utilities', () => {
       expect(snapshot(null)).toBe(null);
       expect(snapshot(undefined)).toBe(undefined);
     });
+
+  });
+
+  describe('IdHash support', () => {
+    it('pruneIds uses IdHash to map IDs to numbered placeholders', () => {
+      const idHash: IdHash = { 'uuid-1': 1, 'uuid-2': 2 };
+      const row = { id: 'uuid-1', user_id: 'uuid-2', org_id: 'unknown', name: 'Alice' };
+      expect(pruneIds(row, idHash)).toEqual({
+        id: '[ID-1]',
+        user_id: '[ID-2]',
+        org_id: '[ID]',
+        name: 'Alice'
+      });
+    });
+
+    it('prune applies IdHash mapping when provided', () => {
+      const idHash: IdHash = { '1': 1, '2': 2 };
+      const row = { id: 1, user_id: 2, name: 'Alice' };
+      expect(prune(row, idHash)).toEqual({
+        id: '[ID-1]',
+        user_id: '[ID-2]',
+        name: 'Alice'
+      });
+    });
+
+    it('snapshot uses IdHash to preserve ID relationships', () => {
+      const idHash: IdHash = {
+        'user-uuid-1': 1,
+        'user-uuid-2': 2,
+        'post-uuid-1': 3
+      };
+      const data = [
+        { id: 'user-uuid-1', name: 'Alice', post_id: 'post-uuid-1' },
+        { id: 'user-uuid-2', name: 'Bob', post_id: 'post-uuid-1' }
+      ];
+      expect(snapshot(data, idHash)).toEqual([
+        { id: '[ID-1]', name: 'Alice', post_id: '[ID-3]' },
+        { id: '[ID-2]', name: 'Bob', post_id: '[ID-3]' }
+      ]);
+    });
+
+    it('snapshot uses IdHash with nested objects', () => {
+      const idHash: IdHash = { 'org-1': 1, 'user-1': 2 };
+      const data = {
+        org: { id: 'org-1', name: 'Acme' },
+        user: { id: 'user-1', org_id: 'org-1' }
+      };
+      expect(snapshot(data, idHash)).toEqual({
+        org: { id: '[ID-1]', name: 'Acme' },
+        user: { id: '[ID-2]', org_id: '[ID-1]' }
+      });
+    });
+
+    it('snapshot uses IdHash with string labels', () => {
+      const idHash: IdHash = {
+        'uuid-user-1': 'user1',
+        'uuid-user-2': 'user2',
+        'uuid-group-1': 'group1',
+        'uuid-group-2': 'group2'
+      };
+      const data = [
+        { actor_id: 'uuid-user-2', entity_id: 'uuid-group-2', is_admin: true }
+      ];
+      expect(snapshot(data, idHash)).toEqual([
+        { actor_id: '[ID-user2]', entity_id: '[ID-group2]', is_admin: true }
+      ]);
+    });
   });
 });

packages/pgsql-test/src/utils.ts

@@ -1,6 +1,15 @@
 const uuidRegexp = /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i;
 
-const idReplacement = (v: unknown): string | unknown => (!v ? v : '[ID]');
+// ID hash map for tracking ID relationships in snapshots
+// Values can be numbers (e.g., 1 -> [ID-1]) or strings (e.g., 'user2' -> [ID-user2])
+export type IdHash = Record<string, number | string>;
+
+const idReplacement = (v: unknown, idHash?: IdHash): string | unknown => {
+  if (!v) return v;
+  if (!idHash) return '[ID]';
+  const key = String(v);
+  return idHash[key] !== undefined ? `[ID-${idHash[key]}]` : '[ID]';
+};
 
 // Generic object type for any key-value mapping
 type AnyObject = Record<string, any>;
@@ -32,11 +41,11 @@ export const pruneDates = (row: AnyObject): AnyObject =>
     return v;
   });
 
-export const pruneIds = (row: AnyObject): AnyObject =>
+export const pruneIds = (row: AnyObject, idHash?: IdHash): AnyObject =>
   mapValues(row, (v, k) =>
     (k === 'id' || (typeof k === 'string' && k.endsWith('_id'))) &&
       (typeof v === 'string' || typeof v === 'number')
-      ? idReplacement(v)
+      ? idReplacement(v, idHash)
       : v
   );
 
@@ -95,27 +104,35 @@ export const composePruners = (...pruners: Pruner[]): Pruner =>
   (row: AnyObject): AnyObject =>
     pruners.reduce((acc, pruner) => pruner(acc), row);
 
-// Default pruners used by prune/snapshot
+// Pruner with optional IdHash support
+type PrunerWithIdHash = (row: AnyObject, idHash?: IdHash) => AnyObject;
+
+// Default pruners used by prune/snapshot (without IdHash)
 export const defaultPruners: Pruner[] = [
   pruneTokens,
   prunePeoplestamps,
   pruneDates,
   pruneIdArrays,
-  pruneIds,
   pruneUUIDs,
   pruneHashes
 ];
 
-export const prune = composePruners(...defaultPruners);
+// Compose pruners and apply pruneIds with IdHash support
+export const prune = (row: AnyObject, idHash?: IdHash): AnyObject => {
+  const pruned = composePruners(...defaultPruners)(row);
+  return pruneIds(pruned, idHash);
+};
 
 // Factory to create a snapshot function with custom pruners
 export const createSnapshot = (pruners: Pruner[]) => {
   const pruneFn = composePruners(...pruners);
-  const snap = (obj: unknown): unknown => {
+  const snap = (obj: unknown, idHash?: IdHash): unknown => {
     if (Array.isArray(obj)) {
-      return obj.map(snap);
+      return obj.map((el) => snap(el, idHash));
     } else if (obj && typeof obj === 'object') {
-      return mapValues(pruneFn(obj as AnyObject), snap);
+      const pruned = pruneFn(obj as AnyObject);
+      const prunedWithIds = pruneIds(pruned, idHash);
+      return mapValues(prunedWithIds, (v) => snap(v, idHash));
     }
     return obj;
   };