@fragno-dev/db 0.1.11 → 0.1.13

package/src/adapters/drizzle/drizzle-query.test.ts

@@ -1,6 +1,6 @@
 import { drizzle } from "drizzle-orm/pglite";
-import { beforeAll, beforeEach, describe, expect, it } from "vitest";
-import { column, idColumn, referenceColumn, schema } from "../../schema/create";
+import { beforeAll, beforeEach, describe, expect, expectTypeOf, it } from "vitest";
+import { column, FragnoId, idColumn, referenceColumn, schema } from "../../schema/create";
 import type { DBType } from "./shared";
 import { writeAndLoadSchema } from "./test-utils";
 import { fromDrizzle } from "./drizzle-query";
@@ -134,10 +134,14 @@ describe("drizzle-query", () => {
     it("should find with select subset of columns", async () => {
       const someExternalId = "some-external-id";
 
-      await orm.findFirst("user", (b) =>
+      const res = await orm.findFirst("user", (b) =>
         b.whereIndex("primary", (eb) => eb("id", "=", someExternalId)).select(["id", "email"]),
       );
 
+      if (res) {
+        expectTypeOf(res.email).toEqualTypeOf<string>();
+      }
+
       const [query] = queries;
       expect(query.sql).toMatchInlineSnapshot(
         `"select "id", "email", "_internalId", "_version" from "user" "user" where "user"."id" = $1 limit $2"`,
@@ -194,7 +198,7 @@ describe("drizzle-query", () => {
     });
 
     it("should find with select subset", async () => {
-      await orm.find("user", (b) => b.whereIndex("primary").select(["id", "email"]));
+      const _res = await orm.find("user", (b) => b.whereIndex("primary").select(["id", "email"]));
 
       const [query] = queries;
       expect(query.sql).toMatchInlineSnapshot(
@@ -331,6 +335,33 @@ describe("drizzle-query", () => {
       );
       expect(query.params).toEqual([newExpiresAt.toISOString(), sessionId]);
     });
+
+    it("should update with version check using FragnoId", async () => {
+      const userId = FragnoId.fromExternal("user-123", 5);
+
+      await orm.update("user", userId, (b) =>
+        b
+          .set({
+            email: "checked@example.com",
+          })
+          .check(),
+      );
+
+      // Verify the SQL query includes version check in WHERE clause
+      const [query] = queries;
+      expect(query.sql).toMatchInlineSnapshot(
+        `"update "user" set "email" = $1, "_version" = COALESCE(_version, 0) + 1 where ("user"."id" = $2 and "user"."_version" = $3)"`,
+      );
+      expect(query.params).toEqual(["checked@example.com", "user-123", 5]);
+    });
+
+    it("should throw when trying to check() with string ID", async () => {
+      await expect(
+        orm.update("user", "user-123", (b) => b.set({ email: "test@example.com" }).check()),
+      ).rejects.toThrow(
+        'Cannot use check() with a string ID on table "user". Version checking requires a FragnoId with version information.',
+      );
+    });
   });
 
   describe("updateMany", () => {
@@ -378,6 +409,25 @@ describe("drizzle-query", () => {
       expect(query.sql).toMatchInlineSnapshot(`"delete from "session" where "session"."id" = $1"`);
       expect(query.params).toEqual([sessionId]);
     });
+
+    it("should delete with version check using FragnoId", async () => {
+      const userId = FragnoId.fromExternal("user-789", 3);
+
+      await orm.delete("user", userId, (b) => b.check());
+
+      // Verify the SQL query includes version check in WHERE clause
+      const [query] = queries;
+      expect(query.sql).toMatchInlineSnapshot(
+        `"delete from "user" where ("user"."id" = $1 and "user"."_version" = $2)"`,
+      );
+      expect(query.params).toEqual(["user-789", 3]);
+    });
+
+    it("should throw when trying to check() with string ID on delete", async () => {
+      await expect(orm.delete("user", "user-123", (b) => b.check())).rejects.toThrow(
+        'Cannot use check() with a string ID on table "user". Version checking requires a FragnoId with version information.',
+      );
+    });
   });
 
   describe("deleteMany", () => {

package/src/adapters/drizzle/drizzle-query.ts

@@ -1,12 +1,14 @@
 import type { AbstractQuery } from "../../query/query";
-import type { AnySchema } from "../../schema/create";
-import type { CompiledMutation, UOWExecutor } from "../../query/unit-of-work";
+import type { AnySchema, AnyTable } from "../../schema/create";
+import type { CompiledMutation, UOWExecutor, ValidIndexName } from "../../query/unit-of-work";
 import { createDrizzleUOWCompiler, type DrizzleCompiledQuery } from "./drizzle-uow-compiler";
 import { executeDrizzleRetrievalPhase, executeDrizzleMutationPhase } from "./drizzle-uow-executor";
 import { UnitOfWork } from "../../query/unit-of-work";
 import { parseDrizzle, type DrizzleResult, type TableNameMapper, type DBType } from "./shared";
 import { createDrizzleUOWDecoder } from "./drizzle-uow-decoder";
 import type { ConnectionPool } from "../../shared/connection-pool";
+import type { TableToUpdateValues } from "../../query/query";
+import type { CursorResult } from "../../query/cursor";
 
 /**
  * Configuration options for creating a Drizzle Unit of Work
@@ -24,6 +26,37 @@ export interface DrizzleUOWConfig {
   dryRun?: boolean;
 }
 
+/**
+ * Special builder for updateMany operations that captures configuration
+ */
+class UpdateManySpecialBuilder<TTable extends AnyTable> {
+  #indexName?: string;
+  #condition?: unknown;
+  #setValues?: TableToUpdateValues<TTable>;
+
+  whereIndex<TIndexName extends ValidIndexName<TTable>>(
+    indexName: TIndexName,
+    condition?: unknown,
+  ): this {
+    this.#indexName = indexName as string;
+    this.#condition = condition;
+    return this;
+  }
+
+  set(values: TableToUpdateValues<TTable>): this {
+    this.#setValues = values;
+    return this;
+  }
+
+  getConfig() {
+    return {
+      indexName: this.#indexName,
+      condition: this.#condition,
+      setValues: this.#setValues,
+    };
+  }
+}
+
 /**
  * Creates a Drizzle-based query engine for the given schema.
  *
@@ -114,15 +147,28 @@ export function fromDrizzle<T extends AnySchema>(
 
   return {
     async find(tableName, builderFn) {
-      const uow = createUOW({ config: uowConfig }).find(tableName, builderFn);
+      // Safe: builderFn returns a FindBuilder (or void), which matches UnitOfWork signature
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      const uow = createUOW({ config: uowConfig }).find(tableName, builderFn as any);
       const [result] = await uow.executeRetrieve();
       return result;
     },
 
+    async findWithCursor(tableName, builderFn) {
+      // Safe: builderFn returns a FindBuilder, which matches UnitOfWork signature
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      const uow = createUOW({ config: uowConfig }).findWithCursor(tableName, builderFn as any);
+      const [result] = await uow.executeRetrieve();
+      return result as CursorResult<unknown>;
+    },
+
     async findFirst(tableName, builderFn) {
       const uow = createUOW({ config: uowConfig });
       if (builderFn) {
-        uow.find(tableName, (b) => builderFn(b as never).pageSize(1));
+        uow.find(tableName, (b) => {
+          builderFn(b);
+          return b.pageSize(1);
+        });
       } else {
         uow.find(tableName, (b) => b.whereIndex("primary").pageSize(1));
       }
@@ -133,7 +179,7 @@ export function fromDrizzle<T extends AnySchema>(
 
     async create(tableName, values) {
       const uow = createUOW({ config: uowConfig });
-      uow.create(tableName as string, values as never);
+      uow.create(tableName, values);
       const { success } = await uow.executeMutations();
       if (!success) {
         throw new Error("Failed to create record");
@@ -150,7 +196,7 @@ export function fromDrizzle<T extends AnySchema>(
     async createMany(tableName, valuesArray) {
       const uow = createUOW({ config: uowConfig });
       for (const values of valuesArray) {
-        uow.create(tableName as string, values as never);
+        uow.create(tableName, values);
       }
       const { success } = await uow.executeMutations();
       if (!success) {
@@ -162,7 +208,7 @@ export function fromDrizzle<T extends AnySchema>(
 
     async update(tableName, id, builderFn) {
       const uow = createUOW({ config: uowConfig });
-      uow.update(tableName as string, id, builderFn as never);
+      uow.update(tableName, id, builderFn);
       const { success } = await uow.executeMutations();
       if (!success) {
         throw new Error("Failed to update record (version conflict or record not found)");
@@ -170,25 +216,17 @@ export function fromDrizzle<T extends AnySchema>(
     },
 
     async updateMany(tableName, builderFn) {
-      // FIXME: This is not correct
-
-      let whereConfig: { indexName?: string; condition?: unknown } = {};
-      let setValues: unknown;
-
-      const specialBuilder = {
-        whereIndex(indexName: string, condition?: unknown) {
-          whereConfig = { indexName, condition };
-          return this;
-        },
-        set(values: unknown) {
-          setValues = values;
-          return this;
-        },
-      };
+      const table = schema.tables[tableName];
+      if (!table) {
+        throw new Error(`Table ${tableName} not found in schema`);
+      }
 
+      const specialBuilder = new UpdateManySpecialBuilder<typeof table>();
       builderFn(specialBuilder);
 
-      if (!whereConfig.indexName) {
+      const { indexName, condition, setValues } = specialBuilder.getConfig();
+
+      if (!indexName) {
         throw new Error("whereIndex() must be called in updateMany");
       }
       if (!setValues) {
@@ -197,24 +235,22 @@ export function fromDrizzle<T extends AnySchema>(
 
       const findUow = createUOW({ config: uowConfig });
       findUow.find(tableName, (b) => {
-        if (whereConfig.condition) {
-          return b.whereIndex(whereConfig.indexName as never, whereConfig.condition as never);
+        if (condition) {
+          // Safe: condition is captured from whereIndex call with proper typing
+          // eslint-disable-next-line @typescript-eslint/no-explicit-any
+          return b.whereIndex(indexName as ValidIndexName<typeof table>, condition as any);
         }
-        return b.whereIndex(whereConfig.indexName as never);
+        return b.whereIndex(indexName as ValidIndexName<typeof table>);
       });
-      const findResults = await findUow.executeRetrieve();
-      const records = (findResults as unknown as [unknown])[0];
+      const [records]: unknown[][] = await findUow.executeRetrieve();
 
-      // @ts-expect-error - Type narrowing doesn't work through unknown cast
       if (!records || records.length === 0) {
         return;
       }
 
       const updateUow = createUOW({ config: uowConfig });
-      for (const record of records as never as Array<{ id: unknown }>) {
-        updateUow.update(tableName as string, record.id as string, (b) =>
-          b.set(setValues as never),
-        );
+      for (const record of records as Array<{ id: unknown }>) {
+        updateUow.update(tableName, record.id as string, (b) => b.set(setValues));
       }
       const { success } = await updateUow.executeMutations();
       if (!success) {
@@ -224,7 +260,7 @@ export function fromDrizzle<T extends AnySchema>(
 
     async delete(tableName, id, builderFn) {
       const uow = createUOW({ config: uowConfig });
-      uow.delete(tableName as string, id, builderFn as never);
+      uow.delete(tableName, id, builderFn);
       const { success } = await uow.executeMutations();
       if (!success) {
         throw new Error("Failed to delete record (version conflict or record not found)");
@@ -232,39 +268,17 @@ export function fromDrizzle<T extends AnySchema>(
     },
 
     async deleteMany(tableName, builderFn) {
-      let whereConfig: { indexName?: string; condition?: unknown } = {};
-
-      const specialBuilder = {
-        whereIndex(indexName: string, condition?: unknown) {
-          whereConfig = { indexName, condition };
-          return this;
-        },
-      };
-
-      builderFn(specialBuilder as never);
-
-      if (!whereConfig.indexName) {
-        throw new Error("whereIndex() must be called in deleteMany");
-      }
-
       const findUow = createUOW({ config: uowConfig });
-      findUow.find(tableName as string, (b) => {
-        if (whereConfig.condition) {
-          return b.whereIndex(whereConfig.indexName as never, whereConfig.condition as never);
-        }
-        return b.whereIndex(whereConfig.indexName as never);
-      });
-      const findResults2 = await findUow.executeRetrieve();
-      const records = (findResults2 as unknown as [unknown])[0];
+      findUow.find(tableName, builderFn);
+      const [records]: unknown[][] = await findUow.executeRetrieve();
 
-      // @ts-expect-error - Type narrowing doesn't work through unknown cast
       if (!records || records.length === 0) {
         return;
       }
 
       const deleteUow = createUOW({ config: uowConfig });
-      for (const record of records as never as Array<{ id: unknown }>) {
-        deleteUow.delete(tableName as string, record.id as string);
+      for (const record of records as Array<{ id: unknown }>) {
+        deleteUow.delete(tableName, record.id as string);
       }
       const { success } = await deleteUow.executeMutations();
       if (!success) {

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

@@ -14,6 +14,7 @@ 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";
+import { Cursor } from "../../query/cursor";
 
 /**
  * Integration tests for Drizzle UOW compiler and executor.
@@ -233,7 +234,12 @@ describe("drizzle-uow-compiler", () => {
 
     it("should compile find operation with cursor pagination using after", () => {
       const uow = createTestUOW();
-      const cursor = "eyJpbmRleFZhbHVlcyI6eyJuYW1lIjoiQWxpY2UifSwiZGlyZWN0aW9uIjoiZm9yd2FyZCJ9"; // {"indexValues":{"name":"Alice"},"direction":"forward"}
+      const cursor = new Cursor({
+        indexName: "idx_name",
+        orderDirection: "asc",
+        pageSize: 10,
+        indexValues: { name: "Alice" },
+      });
       uow.find("users", (b) =>
         b.whereIndex("idx_name").orderByIndex("idx_name", "asc").after(cursor).pageSize(10),
       );
@@ -250,7 +256,12 @@ describe("drizzle-uow-compiler", () => {
 
     it("should compile find operation with cursor pagination using before", () => {
       const uow = createTestUOW();
-      const cursor = "eyJpbmRleFZhbHVlcyI6eyJuYW1lIjoiQm9iIn0sImRpcmVjdGlvbiI6ImJhY2t3YXJkIn0="; // {"indexValues":{"name":"Bob"},"direction":"backward"}
+      const cursor = new Cursor({
+        indexName: "idx_name",
+        orderDirection: "desc",
+        pageSize: 10,
+        indexValues: { name: "Bob" },
+      });
       uow.find("users", (b) =>
         b.whereIndex("idx_name").orderByIndex("idx_name", "desc").before(cursor).pageSize(10),
       );
@@ -267,7 +278,12 @@ describe("drizzle-uow-compiler", () => {
 
     it("should compile find operation with cursor pagination and additional where conditions", () => {
       const uow = createTestUOW();
-      const cursor = "eyJpbmRleFZhbHVlcyI6eyJuYW1lIjoiQWxpY2UifSwiZGlyZWN0aW9uIjoiZm9yd2FyZCJ9";
+      const cursor = new Cursor({
+        indexName: "idx_name",
+        orderDirection: "asc",
+        pageSize: 5,
+        indexValues: { name: "Alice" },
+      });
       uow.find("users", (b) =>
         b
           .whereIndex("idx_name", (eb) => eb("name", "starts with", "John"))
@@ -457,10 +473,9 @@ describe("drizzle-uow-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
+      // FragnoId should generate a subquery to lookup the internal ID from external ID
       expect(batch.query.sql).toMatchInlineSnapshot(
-        `"insert into "posts" ("id", "title", "content", "userId", "viewCount", "_internalId", "_version") values ($1, $2, $3, $4, default, default, default)"`,
+        `"insert into "posts" ("id", "title", "content", "userId", "viewCount", "_internalId", "_version") values ($1, $2, $3, (select "_internalId" from "users" where "id" = $4 limit 1), default, default, default)"`,
       );
     });
 
@@ -1303,5 +1318,66 @@ describe("drizzle-uow-compiler", () => {
       );
       expect(compiled.retrievalBatch[0].params).toEqual([1, sessionId]);
     });
+
+    it("should support creating and using ID in same UOW", () => {
+      const uow = createAuthUOW("create-user-and-session");
+
+      // Create user and capture the returned ID
+      const userId = uow.create("user", {
+        email: "test@example.com",
+        passwordHash: "hashed_password",
+      });
+
+      // Use the returned FragnoId directly to create a session
+      // The compiler should extract externalId and generate a subquery
+      uow.create("session", {
+        userId: userId,
+        expiresAt: new Date("2025-12-31"),
+      });
+
+      const compiler = createDrizzleUOWCompiler(authSchema, authPool, "postgresql");
+      const compiled = uow.compile(compiler);
+
+      // Should have no retrieval operations
+      expect(compiled.retrievalBatch).toHaveLength(0);
+
+      // Should have 2 mutation operations (create user, create session)
+      expect(compiled.mutationBatch).toHaveLength(2);
+
+      const [userCreate, sessionCreate] = compiled.mutationBatch;
+      assert(userCreate);
+      assert(sessionCreate);
+
+      // Verify user create SQL
+      expect(userCreate.query.sql).toMatchInlineSnapshot(
+        `"insert into "user" ("id", "email", "passwordHash", "createdAt", "_internalId", "_version") values ($1, $2, $3, $4, default, default)"`,
+      );
+      expect(userCreate.query.params).toMatchObject([
+        userId.externalId, // The generated ID
+        "test@example.com",
+        "hashed_password",
+        expect.any(String), // timestamp
+      ]);
+      expect(userCreate.expectedAffectedRows).toBeNull();
+
+      // Verify session create SQL - FragnoId generates subquery to lookup internal ID
+      expect(sessionCreate.query.sql).toMatchInlineSnapshot(
+        `"insert into "session" ("id", "userId", "expiresAt", "createdAt", "_internalId", "_version") values ($1, (select "_internalId" from "user" where "id" = $2 limit 1), $3, $4, default, default)"`,
+      );
+      expect(sessionCreate.query.params).toMatchObject([
+        expect.any(String), // generated session ID
+        userId.externalId, // FragnoId's externalId is used in the subquery
+        expect.any(String), // expiresAt timestamp
+        expect.any(String), // createdAt timestamp
+      ]);
+      expect(sessionCreate.expectedAffectedRows).toBeNull();
+
+      // Verify the returned FragnoId has the expected structure
+      expect(userId).toMatchObject({
+        externalId: expect.any(String),
+        version: 0,
+        internalId: undefined,
+      });
+    });
   });
 });

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

@@ -423,8 +423,9 @@ export function createDrizzleUOWCompiler<TSchema extends AnySchema>(
           // Add cursor-based pagination conditions
           if ((after || before) && indexColumns.length > 0) {
             const cursor = after || before;
-            const cursorData = decodeCursor(cursor!);
-            const serializedValues = serializeCursorValues(cursorData, indexColumns, provider);
+            // Decode cursor if it's a string, otherwise use it as-is
+            const cursorObj = typeof cursor === "string" ? decodeCursor(cursor!) : cursor!;
+            const serializedValues = serializeCursorValues(cursorObj, indexColumns, provider);
 
             // Build tuple comparison for cursor pagination
             // For "after" with "asc": (col1, col2, ...) > (val1, val2, ...)

package/src/adapters/drizzle/drizzle-uow-decoder.ts

@@ -4,6 +4,7 @@ import type { RetrievalOperation, UOWDecoder } from "../../query/unit-of-work";
 import { decodeResult } from "../../query/result-transform";
 import { getOrderedJoinColumns } from "./join-column-utils";
 import type { DrizzleResult } from "./shared";
+import { createCursorFromRecord, Cursor, type CursorResult } from "../../query/cursor";
 
 /**
  * Join information with nested join support
@@ -173,10 +174,48 @@ export function createDrizzleUOWDecoder<TSchema extends AnySchema>(
       }
 
       // Handle find operations - decode each row
-      return result.rows.map((row) => {
+      const decodedRows = result.rows.map((row) => {
         const transformedRow = transformJoinArraysToObjects(row, op, provider);
         return decodeResult(transformedRow, op.table, provider);
       });
+
+      // If cursor generation is requested, wrap in CursorResult
+      if (op.withCursor) {
+        let cursor: Cursor | undefined;
+
+        // Generate cursor from last item if results exist
+        if (decodedRows.length > 0 && op.options.orderByIndex && op.options.pageSize) {
+          const lastItem = decodedRows[decodedRows.length - 1];
+          const indexName = op.options.orderByIndex.indexName;
+
+          // Get index columns
+          let indexColumns;
+          if (indexName === "_primary") {
+            indexColumns = [op.table.getIdColumn()];
+          } else {
+            const index = op.table.indexes[indexName];
+            if (index) {
+              indexColumns = index.columns;
+            }
+          }
+
+          if (indexColumns && lastItem) {
+            cursor = createCursorFromRecord(lastItem as Record<string, unknown>, indexColumns, {
+              indexName: op.options.orderByIndex.indexName,
+              orderDirection: op.options.orderByIndex.direction,
+              pageSize: op.options.pageSize,
+            });
+          }
+        }
+
+        const cursorResult: CursorResult<unknown> = {
+          items: decodedRows,
+          cursor,
+        };
+        return cursorResult;
+      }
+
+      return decodedRows;
     });
   };
 }