Fixed pongo transaction cache handling

Author: oskardudycz

src/packages/pongo/src/core/cache/collectionCache.int.spec.ts

@@ -231,4 +231,267 @@ describe('pongoCollection cache integration', () => {
       expect(doc?.name).toBe('Karl');
     });
   });
+
+  describe('cache type resolution cascade', () => {
+    it('collection identity-map overrides client in-memory', async () => {
+      client = pongoClient({
+        driver: sqlite3Driver,
+        connectionString: memoryConnectionString(),
+        cache: { type: 'in-memory', max: 5 },
+      });
+      await client.connect();
+
+      const col = client
+        .db('db')
+        .collection<User>('users', { cache: { type: 'identity-map' } });
+
+      const ids: string[] = [];
+      for (let i = 0; i < 10; i++) {
+        const { insertedId } = await col.insertOne({ name: `User${i}` });
+        ids.push(insertedId!);
+      }
+
+      for (const id of ids) {
+        const doc = await col.findOne({ _id: id });
+        expect(doc).not.toBeNull();
+      }
+    });
+
+    it('client identity-map is used by collection with no override', async () => {
+      client = pongoClient({
+        driver: sqlite3Driver,
+        connectionString: memoryConnectionString(),
+        cache: { type: 'identity-map' },
+      });
+      await client.connect();
+
+      const col = client.db('db').collection<User>('users');
+      const ids: string[] = [];
+      for (let i = 0; i < 2000; i++) {
+        const { insertedId } = await col.insertOne({ name: `User${i}` });
+        ids.push(insertedId!);
+      }
+
+      for (const id of ids) {
+        const doc = await col.findOne({ _id: id });
+        expect(doc).not.toBeNull();
+      }
+    });
+
+    it('collection disabled overrides client identity-map', async () => {
+      client = pongoClient({
+        driver: sqlite3Driver,
+        connectionString: memoryConnectionString(),
+        cache: { type: 'identity-map' },
+      });
+      await client.connect();
+
+      const db = client.db('db');
+      const col = db.collection<User>('users', { cache: 'disabled' });
+      const noCache = db.collection<User>('users', { cache: 'disabled' });
+
+      const { insertedId } = await col.insertOne({ name: 'Leo' });
+      await noCache.updateOne(
+        { _id: insertedId! },
+        { $set: { name: 'LeoUpdated' } },
+      );
+
+      const doc = await col.findOne({ _id: insertedId! });
+      expect(doc?.name).toBe('LeoUpdated');
+    });
+
+    it('two collections can use different cache types on the same db', async () => {
+      client = pongoClient({
+        driver: sqlite3Driver,
+        connectionString: memoryConnectionString(),
+      });
+      await client.connect();
+
+      const db = client.db('db');
+      const lruCol = db.collection<User>('users', {
+        cache: { type: 'in-memory', max: 5 },
+      });
+      const idMapCol = db.collection<User>('orders', {
+        cache: { type: 'identity-map' },
+      });
+
+      const lruIds: string[] = [];
+      const idMapIds: string[] = [];
+
+      for (let i = 0; i < 10; i++) {
+        const { insertedId: u } = await lruCol.insertOne({ name: `User${i}` });
+        lruIds.push(u!);
+        const { insertedId: o } = await idMapCol.insertOne({
+          name: `Order${i}`,
+        });
+        idMapIds.push(o!);
+      }
+
+      // LRU with max 5: oldest 5 should be evicted from cache (reads go to DB)
+      // identity-map: all 10 present in cache
+      for (const id of idMapIds) {
+        const doc = await idMapCol.findOne({ _id: id });
+        expect(doc).not.toBeNull();
+      }
+
+      // At least the most-recent LRU entries should still be in cache
+      const recent = lruIds.slice(-5);
+      for (const id of recent) {
+        const doc = await lruCol.findOne({ _id: id });
+        expect(doc).not.toBeNull();
+      }
+    });
+  });
+
+  describe('transaction cache integration', () => {
+    beforeEach(async () => {
+      client = pongoClient({
+        driver: sqlite3Driver,
+        connectionString: memoryConnectionString(),
+      });
+      await client.connect();
+    });
+
+    it('insertOne within transaction does NOT populate collection cache until commit', async () => {
+      const { cache, spies } = spyCache('col');
+      const col = client.db('db').collection<User>('users', { cache });
+
+      const session = client.startSession();
+      session.startTransaction();
+      await col.insertOne({ name: 'Alice' }, { session });
+
+      expect(spies.set).not.toHaveBeenCalled();
+
+      await session.commitTransaction();
+      await session.endSession();
+
+      expect(spies.set).toHaveBeenCalled();
+    });
+
+    it('findOne within transaction returns doc from transaction cache', async () => {
+      const { cache, spies } = spyCache('col');
+      const col = client.db('db').collection<User>('users', { cache });
+
+      const session = client.startSession();
+      session.startTransaction();
+      const { insertedId } = await col.insertOne(
+        { name: 'Carol' },
+        { session },
+      );
+
+      const doc = await col.findOne({ _id: insertedId! }, { session });
+      expect(doc?.name).toBe('Carol');
+      // should have hit tx cache, not collection cache get
+      expect(spies.get).not.toHaveBeenCalled();
+
+      await session.abortTransaction();
+      await session.endSession();
+    });
+
+    it('findOne within transaction falls through to collection cache on tx miss', async () => {
+      const { cache } = spyCache('col');
+      const col = client.db('db').collection<User>('users', { cache });
+
+      const { insertedId } = await col.insertOne({ name: 'Dave' });
+
+      const session = client.startSession();
+      session.startTransaction();
+      const doc = await col.findOne({ _id: insertedId! }, { session });
+      expect(doc?.name).toBe('Dave');
+
+      await session.abortTransaction();
+      await session.endSession();
+    });
+
+    it('rollback discards buffered cache ops — collection cache stays clean', async () => {
+      const { cache, spies } = spyCache('col');
+      const col = client.db('db').collection<User>('users', { cache });
+
+      const session = client.startSession();
+      session.startTransaction();
+      await col.insertOne({ name: 'Bob' }, { session });
+      await session.abortTransaction();
+      await session.endSession();
+
+      expect(spies.set).not.toHaveBeenCalled();
+    });
+
+    it('deleteOne within transaction does NOT evict from collection cache until commit', async () => {
+      const { cache, spies } = spyCache('col');
+      const col = client.db('db').collection<User>('users', { cache });
+
+      const { insertedId } = await col.insertOne({ name: 'Eve' });
+      spies.delete.mockClear();
+
+      const session = client.startSession();
+      session.startTransaction();
+      await col.deleteOne({ _id: insertedId! }, { session });
+
+      expect(spies.delete).not.toHaveBeenCalled();
+
+      await session.commitTransaction();
+      await session.endSession();
+
+      expect(spies.delete).toHaveBeenCalled();
+    });
+
+    it('updateOne within transaction does NOT evict from collection cache until commit', async () => {
+      const { cache, spies } = spyCache('col');
+      const col = client.db('db').collection<User>('users', { cache });
+
+      const { insertedId } = await col.insertOne({ name: 'Frank' });
+      spies.delete.mockClear();
+
+      const session = client.startSession();
+      session.startTransaction();
+      await col.updateOne(
+        { _id: insertedId! },
+        { $set: { name: 'Updated' } },
+        { session },
+      );
+
+      expect(spies.delete).not.toHaveBeenCalled();
+
+      // reading without session should still return cached 'Frank'
+      const doc = await col.findOne({ _id: insertedId! });
+      expect(doc?.name).toBe('Frank');
+
+      await session.commitTransaction();
+      await session.endSession();
+
+      expect(spies.delete).toHaveBeenCalled();
+    });
+
+    it('replaceOne within transaction buffers cache set until commit', async () => {
+      const { cache, spies } = spyCache('col');
+      const col = client.db('db').collection<User>('users', { cache });
+
+      const { insertedId } = await col.insertOne({ name: 'Grace' });
+      spies.set.mockClear();
+
+      const session = client.startSession();
+      session.startTransaction();
+      await col.replaceOne(
+        { _id: insertedId! },
+        { name: 'Replaced' },
+        { session },
+      );
+
+      expect(spies.set).not.toHaveBeenCalled();
+
+      await session.commitTransaction();
+      await session.endSession();
+
+      expect(spies.set).toHaveBeenCalled();
+    });
+
+    it('without transaction, operations update collection cache directly', async () => {
+      const { cache, spies } = spyCache('col');
+      const col = client.db('db').collection<User>('users', { cache });
+
+      await col.insertOne({ name: 'Hank' });
+
+      expect(spies.set).toHaveBeenCalled();
+    });
+  });
 });

src/packages/pongo/src/core/cache/pongoCache.ts

@@ -1,5 +1,5 @@
 import { pongoCacheWrapper } from './cacheWrapper';
-import { noopCacheProvider, lruCache } from './providers';
+import { noopCacheProvider, lruCache, identityMapCache } from './providers';
 import type { CacheConfig, CacheSettings, PongoCache } from './types';
 
 const DEFAULT_CONFIG: CacheSettings = { type: 'in-memory', max: 1000 };
@@ -14,6 +14,8 @@ export const pongoCache = (
 
   const config = options ?? DEFAULT_CONFIG;
 
+  if (config.type === 'identity-map') return identityMapCache();
+
   const raw = lruCache({
     ...(config.max !== undefined ? { max: config.max } : {}),
     ...(config.ttl !== undefined ? { ttl: config.ttl } : {}),

src/packages/pongo/src/core/cache/pongoCache.unit.spec.ts

@@ -0,0 +1,76 @@
+import { describe, expect, it } from 'vitest';
+import { identityMapCache } from './providers';
+import { pongoCache } from './pongoCache';
+
+describe('pongoCache factory', () => {
+  describe("'identity-map' type", () => {
+    it('creates an identity-map cache when type is identity-map', async () => {
+      const cache = pongoCache({ type: 'identity-map' });
+
+      expect(cache.cacheType).toBe('pongo:cache:identity-map');
+
+      await cache.set('db:col:1', { _id: '1', name: 'Alice' });
+      const result = await cache.get('db:col:1');
+      expect(result).toEqual({ _id: '1', name: 'Alice' });
+    });
+
+    it('identity-map cache has no max size eviction', async () => {
+      const cache = pongoCache({ type: 'identity-map' });
+      const count = 2000;
+
+      for (let i = 0; i < count; i++) {
+        await cache.set(`db:col:${i}`, { _id: String(i), n: i });
+      }
+
+      for (let i = 0; i < count; i++) {
+        const doc = await cache.get(`db:col:${i}`);
+        expect(doc).not.toBeNull();
+      }
+    });
+  });
+
+  describe("'in-memory' type", () => {
+    it('creates an LRU cache when type is in-memory', () => {
+      const cache = pongoCache({ type: 'in-memory', max: 100 });
+      expect(cache.cacheType).toBe('pongo:cache:lru');
+    });
+
+    it('defaults to in-memory when no config provided', async () => {
+      const cache = pongoCache(undefined);
+
+      await cache.set('db:col:1', { _id: '1', name: 'Bob' });
+      const result = await cache.get('db:col:1');
+      expect(result).toEqual({ _id: '1', name: 'Bob' });
+    });
+
+    it('in-memory LRU evicts when over max', async () => {
+      const cache = pongoCache({ type: 'in-memory', max: 3 });
+
+      for (let i = 0; i < 4; i++) {
+        await cache.set(`db:col:${i}`, { _id: String(i) });
+      }
+
+      // Oldest entry should be evicted
+      const first = await cache.get('db:col:0');
+      expect(first).toBeNull();
+    });
+  });
+
+  describe("'disabled'", () => {
+    it('returns noop cache when disabled', async () => {
+      const cache = pongoCache('disabled');
+
+      await cache.set('db:col:1', { _id: '1', name: 'Carol' });
+      const result = await cache.get('db:col:1');
+      expect(result).toBeNull();
+    });
+  });
+
+  describe('pre-built PongoCache passthrough', () => {
+    it('passes through a pre-built PongoCache instance', () => {
+      const raw = identityMapCache();
+      const cache = pongoCache(raw);
+      expect(cache).toBe(raw);
+    });
+  });
+});

src/packages/pongo/src/core/cache/types.ts

@@ -45,7 +45,7 @@ export type CacheHooks = {
   onError?(error: unknown, operation: string): void;
 };
 
-export type CacheType = 'in-memory';
+export type CacheType = 'in-memory' | 'identity-map';
 
 export type CacheSettings = {
   type: CacheType;

src/packages/pongo/src/core/collection/filters/filters.ts

@@ -21,12 +21,13 @@ export const idFromFilter = <T>(
   return typeof f['_id'] === 'string' ? f['_id'] : null;
 };
 
-export const idsFromInFilter = <T>(
+export const getIdsFromIdOnlyFilter = <T>(
   filter: PongoFilter<T> | SQL | undefined,
 ): string[] | null => {
   const f = asPlainObjectWithSingleKey(filter, '_id');
   if (!f) return null;
   const idVal = f['_id'];
+  if (typeof idVal === 'string') return [idVal];
   if (!idVal || typeof idVal !== 'object' || !('$in' in idVal)) return null;
   const ids = (idVal as PlainObject)['$in'];
   if (!Array.isArray(ids) || ids.some((i) => typeof i !== 'string'))