@fragno-dev/db 0.1.6 → 0.1.8

package/src/adapters/drizzle/drizzle-uow-compiler.test.ts

@@ -8,11 +8,12 @@ import {
   type AnySchema,
 } from "../../schema/create";
 import { createDrizzleUOWCompiler } from "./drizzle-uow-compiler";
-import type { DrizzleConfig } from "./drizzle-adapter";
 import { drizzle } from "drizzle-orm/pglite";
 import type { DBType } from "./shared";
 import { UnitOfWork, type UOWDecoder } from "../../query/unit-of-work";
 import { writeAndLoadSchema } from "./test-utils";
+import type { ConnectionPool } from "../../shared/connection-pool";
+import { createDrizzleConnectionPool } from "./drizzle-connection-pool";
 
 /**
  * Integration tests for Drizzle UOW compiler and executor.
@@ -48,7 +49,7 @@ describe("drizzle-uow-compiler", () => {
   });
 
   let db: DBType;
-  let config: DrizzleConfig;
+  let pool: ConnectionPool<DBType>;
 
   beforeAll(async () => {
     // Write schema to file and dynamically import it
@@ -63,10 +64,8 @@ describe("drizzle-uow-compiler", () => {
       schema: schemaModule,
     }) as unknown as DBType;
 
-    config = {
-      db,
-      provider: "postgresql",
-    };
+    // Wrap in connection pool
+    pool = createDrizzleConnectionPool(db);
 
     return async () => {
       await cleanup();
@@ -74,7 +73,7 @@ describe("drizzle-uow-compiler", () => {
   });
 
   function createTestUOWWithSchema<const T extends AnySchema>(schema: T) {
-    const compiler = createDrizzleUOWCompiler(schema, config);
+    const compiler = createDrizzleUOWCompiler(schema, pool, "postgresql");
     const mockExecutor = {
       executeRetrievalPhase: async () => [],
       executeMutationPhase: async () => ({ success: true, createdInternalIds: [] }),
@@ -89,7 +88,7 @@ describe("drizzle-uow-compiler", () => {
   }
 
   function createTestUOW(name?: string) {
-    const compiler = createDrizzleUOWCompiler(testSchema, config);
+    const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
     const mockExecutor = {
       executeRetrievalPhase: async () => [],
       executeMutationPhase: async () => ({ success: true, createdInternalIds: [] }),
@@ -104,7 +103,7 @@ describe("drizzle-uow-compiler", () => {
   }
 
   it("should create a compiler with the correct structure", () => {
-    const compiler = createDrizzleUOWCompiler(testSchema, config);
+    const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
 
     expect(compiler).toBeDefined();
     expect(compiler.compileRetrievalOperation).toBeInstanceOf(Function);
@@ -118,7 +117,7 @@ describe("drizzle-uow-compiler", () => {
         b.whereIndex("idx_email", (eb) => eb("email", "=", "test@example.com")),
       );
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
 
       expect(compiled.retrievalBatch).toHaveLength(1);
@@ -134,7 +133,7 @@ describe("drizzle-uow-compiler", () => {
         b.whereIndex("idx_name", (eb) => eb("name", "=", "Alice")).select(["id", "name"]),
       );
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
 
       expect(compiled.retrievalBatch).toHaveLength(1);
@@ -148,7 +147,7 @@ describe("drizzle-uow-compiler", () => {
       const uow = createTestUOW();
       uow.find("users", (b) => b.whereIndex("primary").pageSize(10));
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
 
       expect(compiled.retrievalBatch).toHaveLength(1);
@@ -162,7 +161,7 @@ describe("drizzle-uow-compiler", () => {
       const uow = createTestUOW();
       uow.find("users", (b) => b.whereIndex("primary").orderByIndex("primary", "desc"));
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
 
       expect(compiled.retrievalBatch).toHaveLength(1);
@@ -175,7 +174,7 @@ describe("drizzle-uow-compiler", () => {
       const uow = createTestUOW();
       uow.find("users", (b) => b.whereIndex("idx_name").orderByIndex("idx_name", "desc"));
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
 
       expect(compiled.retrievalBatch).toHaveLength(1);
@@ -191,7 +190,7 @@ describe("drizzle-uow-compiler", () => {
       );
       uow.find("posts", (b) => b.whereIndex("idx_title", (eb) => eb("title", "contains", "test")));
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
 
       expect(compiled.retrievalBatch).toHaveLength(2);
@@ -207,7 +206,7 @@ describe("drizzle-uow-compiler", () => {
       const uow = createTestUOW();
       uow.find("users", (b) => b.whereIndex("primary").selectCount());
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
 
       expect(compiled.retrievalBatch).toHaveLength(1);
@@ -222,7 +221,7 @@ describe("drizzle-uow-compiler", () => {
         b.whereIndex("idx_name", (eb) => eb("name", "starts with", "John")).selectCount(),
       );
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
 
       expect(compiled.retrievalBatch).toHaveLength(1);
@@ -239,7 +238,7 @@ describe("drizzle-uow-compiler", () => {
         b.whereIndex("idx_name").orderByIndex("idx_name", "asc").after(cursor).pageSize(10),
       );
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
 
       expect(compiled.retrievalBatch).toHaveLength(1);
@@ -256,7 +255,7 @@ describe("drizzle-uow-compiler", () => {
         b.whereIndex("idx_name").orderByIndex("idx_name", "desc").before(cursor).pageSize(10),
       );
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
 
       expect(compiled.retrievalBatch).toHaveLength(1);
@@ -277,7 +276,7 @@ describe("drizzle-uow-compiler", () => {
           .pageSize(5),
       );
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
 
       expect(compiled.retrievalBatch).toHaveLength(1);
@@ -295,7 +294,7 @@ describe("drizzle-uow-compiler", () => {
           .join((jb) => jb.author()),
       );
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
 
       expect(compiled.retrievalBatch).toHaveLength(1);
@@ -318,7 +317,7 @@ describe("drizzle-uow-compiler", () => {
           ),
       );
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
 
       expect(compiled.retrievalBatch).toHaveLength(1);
@@ -336,7 +335,7 @@ describe("drizzle-uow-compiler", () => {
           .join((jb) => jb.author((builder) => builder.orderByIndex("idx_name", "desc"))),
       );
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
 
       expect(compiled.retrievalBatch).toHaveLength(1);
@@ -352,7 +351,7 @@ describe("drizzle-uow-compiler", () => {
         b.whereIndex("primary").join((jb) => jb.author((builder) => builder.pageSize(5))),
       );
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
 
       expect(compiled.retrievalBatch).toHaveLength(1);
@@ -373,7 +372,7 @@ describe("drizzle-uow-compiler", () => {
         age: 30,
       });
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
       const [batch] = compiled.mutationBatch;
       assert(batch);
@@ -389,6 +388,131 @@ describe("drizzle-uow-compiler", () => {
       ]);
     });
 
+    it("should compile create operation with external id string for reference column", () => {
+      const uow = createTestUOW();
+      // Create a post with userId as just an external id string
+      uow.create("posts", {
+        title: "Test Post",
+        content: "Post content",
+        userId: "user_external_id_123",
+        viewCount: 5,
+      });
+
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
+      const compiled = uow.compile(compiler);
+      const [batch] = compiled.mutationBatch;
+      assert(batch);
+      expect(batch.expectedAffectedRows).toBeNull();
+      expect(batch.query.sql).toMatchInlineSnapshot(
+        `"insert into "posts" ("id", "title", "content", "userId", "viewCount", "_internalId", "_version") values ($1, $2, $3, (select "_internalId" from "users" where "id" = $4 limit 1), $5, default, default)"`,
+      );
+      expect(batch.query.params).toMatchObject([
+        expect.any(String), // auto-generated post ID
+        "Test Post",
+        "Post content",
+        "user_external_id_123", // external id string
+        5, // viewCount
+      ]);
+    });
+
+    it("should compile create operation with bigint for reference column (no subquery)", () => {
+      const uow = createTestUOW();
+      // Create a post with userId as a bigint directly (internal ID)
+      uow.create("posts", {
+        title: "Direct ID Post",
+        content: "Content with direct bigint",
+        userId: 12345n,
+      });
+
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
+      const compiled = uow.compile(compiler);
+      const [batch] = compiled.mutationBatch;
+      assert(batch);
+      expect(batch.expectedAffectedRows).toBeNull();
+      // Should NOT have a subquery when using bigint directly
+      expect(batch.query.sql).not.toMatch(/\(select.*from.*users/i);
+      expect(batch.query.sql).toMatchInlineSnapshot(
+        `"insert into "posts" ("id", "title", "content", "userId", "viewCount", "_internalId", "_version") values ($1, $2, $3, $4, default, default, default)"`,
+      );
+      expect(batch.query.params).toMatchObject([
+        expect.any(String), // auto-generated post ID
+        "Direct ID Post",
+        "Content with direct bigint",
+        12345n, // bigint internal ID directly
+      ]);
+    });
+
+    it("should compile create operation with FragnoId object for reference column", () => {
+      const uow = createTestUOW();
+      const userId = FragnoId.fromExternal("user_ext_789", 0);
+      // Create a post with userId as a FragnoId object
+      uow.create("posts", {
+        title: "Post with FragnoId",
+        content: "Content",
+        userId,
+      });
+
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
+      const compiled = uow.compile(compiler);
+      const [batch] = compiled.mutationBatch;
+      assert(batch);
+      expect(batch.expectedAffectedRows).toBeNull();
+      // FragnoId should use internal ID directly (no subquery needed if available)
+      // But since we don't have the internal ID populated in the test, it should serialize
+      expect(batch.query.sql).toMatchInlineSnapshot(
+        `"insert into "posts" ("id", "title", "content", "userId", "viewCount", "_internalId", "_version") values ($1, $2, $3, $4, default, default, default)"`,
+      );
+    });
+
+    it("should compile update operation with external id string for reference column", () => {
+      const uow = createTestUOW();
+      const postId = FragnoId.fromExternal("post123", 0);
+      uow.update("posts", postId, (b) =>
+        b.set({
+          userId: "new_user_external_id_456",
+        }),
+      );
+
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
+      const compiled = uow.compile(compiler);
+      const [batch] = compiled.mutationBatch;
+      assert(batch);
+      expect(batch.expectedAffectedRows).toBeNull();
+      // Should generate a subquery for the string external ID in UPDATE
+      expect(batch.query.sql).toMatchInlineSnapshot(
+        `"update "posts" set "userId" = (select "_internalId" from "users" where "id" = $1 limit 1), "_version" = COALESCE(_version, 0) + 1 where "posts"."id" = $2"`,
+      );
+      expect(batch.query.params).toMatchObject([
+        "new_user_external_id_456", // external id string
+        "post123", // post external id
+      ]);
+    });
+
+    it("should compile update operation with bigint for reference column (no subquery)", () => {
+      const uow = createTestUOW();
+      const postId = FragnoId.fromExternal("post456", 0);
+      uow.update("posts", postId, (b) =>
+        b.set({
+          userId: 99999n,
+        }),
+      );
+
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
+      const compiled = uow.compile(compiler);
+      const [batch] = compiled.mutationBatch;
+      assert(batch);
+      expect(batch.expectedAffectedRows).toBeNull();
+      // Should NOT have a subquery when using bigint directly
+      expect(batch.query.sql).not.toMatch(/\(select.*from.*users/i);
+      expect(batch.query.sql).toMatchInlineSnapshot(
+        `"update "posts" set "userId" = $1, "_version" = COALESCE(_version, 0) + 1 where "posts"."id" = $2"`,
+      );
+      expect(batch.query.params).toMatchObject([
+        99999n, // bigint internal ID directly
+        "post456", // post external id
+      ]);
+    });
+
     it("should compile update operation with ID", () => {
       const uow = createTestUOW();
       const userId = FragnoId.fromExternal("user123", 0);
@@ -399,7 +523,7 @@ describe("drizzle-uow-compiler", () => {
         }),
       );
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
       const [batch] = compiled.mutationBatch;
       assert(batch);
@@ -415,7 +539,7 @@ describe("drizzle-uow-compiler", () => {
       const userId = FragnoId.fromExternal("user123", 5);
       uow.update("users", userId, (b) => b.set({ age: 18 }).check());
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
       const [batch] = compiled.mutationBatch;
       assert(batch);
@@ -431,7 +555,7 @@ describe("drizzle-uow-compiler", () => {
       const userId = FragnoId.fromExternal("user123", 0);
       uow.delete("users", userId);
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
       const [batch] = compiled.mutationBatch;
       assert(batch);
@@ -447,7 +571,7 @@ describe("drizzle-uow-compiler", () => {
       const userId = FragnoId.fromExternal("user123", 3);
       uow.delete("users", userId, (b) => b.check());
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
       const [batch] = compiled.mutationBatch;
       assert(batch);
@@ -467,7 +591,7 @@ describe("drizzle-uow-compiler", () => {
         }),
       );
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
       const [batch] = compiled.mutationBatch;
       assert(batch);
@@ -482,7 +606,7 @@ describe("drizzle-uow-compiler", () => {
       const uow = createTestUOW();
       uow.delete("users", "user123");
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
       const [batch] = compiled.mutationBatch;
       assert(batch);
@@ -522,7 +646,7 @@ describe("drizzle-uow-compiler", () => {
       const userId = FragnoId.fromExternal("user456", 0);
       uow.delete("posts", userId);
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
       const [createBatch, updateBatch, deleteBatch] = compiled.mutationBatch;
 
@@ -561,7 +685,7 @@ describe("drizzle-uow-compiler", () => {
       const userId = FragnoId.fromExternal("user123", 3);
       uow.update("users", userId, (b) => b.set({ age: 31 }).check());
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
 
       expect(compiled.name).toBe("update-user-balance");
@@ -593,7 +717,7 @@ describe("drizzle-uow-compiler", () => {
         ),
       );
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
 
       expect(compiled.retrievalBatch).toHaveLength(1);
@@ -607,7 +731,7 @@ describe("drizzle-uow-compiler", () => {
       const uow = createTestUOW();
       uow.find("users", (b) => b.whereIndex("primary", () => false));
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
 
       // When condition is false, the operation should return null and not be added to batch
@@ -618,7 +742,7 @@ describe("drizzle-uow-compiler", () => {
       const uow = createTestUOW();
       uow.find("users", (b) => b.whereIndex("primary", () => true));
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
 
       expect(compiled.retrievalBatch).toHaveLength(1);
@@ -635,7 +759,7 @@ describe("drizzle-uow-compiler", () => {
       const userId = FragnoId.fromExternal("user123", 5);
       uow.update("users", userId, (b) => b.set({ age: 31 }).check());
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
       const [batch] = compiled.mutationBatch;
       assert(batch);
@@ -652,7 +776,7 @@ describe("drizzle-uow-compiler", () => {
       const userId = FragnoId.fromExternal("user456", 3);
       uow.delete("users", userId, (b) => b.check());
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
       const [batch] = compiled.mutationBatch;
       assert(batch);
@@ -672,7 +796,7 @@ describe("drizzle-uow-compiler", () => {
       uow.update("users", userId, (b) => b.set({ age: 30 }).check());
       uow.update("posts", postId, (b) => b.set({ viewCount: 100 }).check());
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
       const [userBatch, postBatch] = compiled.mutationBatch;
 
@@ -699,7 +823,7 @@ describe("drizzle-uow-compiler", () => {
       const userId = FragnoId.fromExternal("user1", 0);
       uow.update("users", userId, (b) => b.set({ age: 25 }));
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
       const [batch] = compiled.mutationBatch;
       assert(batch);
@@ -715,7 +839,7 @@ describe("drizzle-uow-compiler", () => {
     it("should handle UOW with no operations", () => {
       const uow = createTestUOW();
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
 
       expect(compiled.retrievalBatch).toHaveLength(0);
@@ -726,7 +850,7 @@ describe("drizzle-uow-compiler", () => {
       const uow = createTestUOW();
       uow.find("users");
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
 
       expect(compiled.retrievalBatch).toHaveLength(1);
@@ -740,7 +864,7 @@ describe("drizzle-uow-compiler", () => {
         email: "test@example.com",
       });
 
-      const compiler = createDrizzleUOWCompiler(testSchema, config);
+      const compiler = createDrizzleUOWCompiler(testSchema, pool, "postgresql");
       const compiled = uow.compile(compiler);
 
       expect(compiled.retrievalBatch).toHaveLength(0);
@@ -770,7 +894,7 @@ describe("drizzle-uow-compiler", () => {
     });
 
     let defaultsDb: DBType;
-    let defaultsConfig: DrizzleConfig;
+    let defaultsPool: ConnectionPool<DBType>;
 
     beforeAll(async () => {
       // Write schema to file and dynamically import it
@@ -785,10 +909,8 @@ describe("drizzle-uow-compiler", () => {
         schema: schemaModule,
       }) as unknown as DBType;
 
-      defaultsConfig = {
-        db: defaultsDb,
-        provider: "postgresql",
-      };
+      // Wrap in connection pool
+      defaultsPool = createDrizzleConnectionPool(defaultsDb);
 
       return async () => {
         await cleanup();
@@ -802,7 +924,7 @@ describe("drizzle-uow-compiler", () => {
         message: "Test log",
       });
 
-      const compiler = createDrizzleUOWCompiler(defaultsSchema, defaultsConfig);
+      const compiler = createDrizzleUOWCompiler(defaultsSchema, defaultsPool, "postgresql");
       const compiled = uow.compile(compiler);
       const [batch] = compiled.mutationBatch;
       assert(batch);
@@ -832,7 +954,7 @@ describe("drizzle-uow-compiler", () => {
         status: "active",
       });
 
-      const compiler = createDrizzleUOWCompiler(defaultsSchema, defaultsConfig);
+      const compiler = createDrizzleUOWCompiler(defaultsSchema, defaultsPool, "postgresql");
       const compiled = uow.compile(compiler);
       const [batch] = compiled.mutationBatch;
       assert(batch);
@@ -854,7 +976,7 @@ describe("drizzle-uow-compiler", () => {
         counter: 999, // override the function default
       });
 
-      const compiler = createDrizzleUOWCompiler(defaultsSchema, defaultsConfig);
+      const compiler = createDrizzleUOWCompiler(defaultsSchema, defaultsPool, "postgresql");
       const compiled = uow.compile(compiler);
       const [batch] = compiled.mutationBatch;
       assert(batch);
@@ -875,7 +997,7 @@ describe("drizzle-uow-compiler", () => {
       uow.create("logs", { message: "Log 2" });
       uow.create("logs", { message: "Log 3" });
 
-      const compiler = createDrizzleUOWCompiler(defaultsSchema, defaultsConfig);
+      const compiler = createDrizzleUOWCompiler(defaultsSchema, defaultsPool, "postgresql");
       const compiled = uow.compile(compiler);
 
       expect(compiled.mutationBatch).toHaveLength(3);
@@ -930,7 +1052,7 @@ describe("drizzle-uow-compiler", () => {
     });
 
     let nestedDb: DBType;
-    let nestedConfig: DrizzleConfig;
+    let nestedPool: ConnectionPool<DBType>;
 
     beforeAll(async () => {
       // Write schema to file and dynamically import it
@@ -945,10 +1067,8 @@ describe("drizzle-uow-compiler", () => {
         schema: schemaModule,
       }) as unknown as DBType;
 
-      nestedConfig = {
-        db: nestedDb,
-        provider: "postgresql",
-      };
+      // Wrap in connection pool
+      nestedPool = createDrizzleConnectionPool(nestedDb);
 
       return async () => {
         await cleanup();
@@ -956,7 +1076,7 @@ describe("drizzle-uow-compiler", () => {
     }, 20000);
 
     function createNestedUOW(name?: string) {
-      const compiler = createDrizzleUOWCompiler(nestedSchema, nestedConfig);
+      const compiler = createDrizzleUOWCompiler(nestedSchema, nestedPool, "postgresql");
       const mockExecutor = {
         executeRetrievalPhase: async () => [],
         executeMutationPhase: async () => ({ success: true, createdInternalIds: [] }),
@@ -984,7 +1104,7 @@ describe("drizzle-uow-compiler", () => {
           ),
       );
 
-      const compiler = createDrizzleUOWCompiler(nestedSchema, nestedConfig);
+      const compiler = createDrizzleUOWCompiler(nestedSchema, nestedPool, "postgresql");
       const compiled = uow.compile(compiler);
 
       expect(compiled.retrievalBatch).toHaveLength(1);
@@ -1022,7 +1142,7 @@ describe("drizzle-uow-compiler", () => {
           ),
       );
 
-      const compiler = createDrizzleUOWCompiler(nestedSchema, nestedConfig);
+      const compiler = createDrizzleUOWCompiler(nestedSchema, nestedPool, "postgresql");
       const compiled = uow.compile(compiler);
 
       expect(compiled.retrievalBatch).toHaveLength(1);
@@ -1055,7 +1175,7 @@ describe("drizzle-uow-compiler", () => {
           ),
       );
 
-      const compiler = createDrizzleUOWCompiler(nestedSchema, nestedConfig);
+      const compiler = createDrizzleUOWCompiler(nestedSchema, nestedPool, "postgresql");
       const compiled = uow.compile(compiler);
 
       expect(compiled.retrievalBatch).toHaveLength(1);
@@ -1077,7 +1197,7 @@ describe("drizzle-uow-compiler", () => {
         ),
       );
 
-      const compiler = createDrizzleUOWCompiler(nestedSchema, nestedConfig);
+      const compiler = createDrizzleUOWCompiler(nestedSchema, nestedPool, "postgresql");
       const compiled = uow.compile(compiler);
 
       expect(compiled.retrievalBatch).toHaveLength(1);