[0.2.0]: custom_uniqueKey field (#14)

* Add additional test and add complexity to example

* Move findRandom functions to separate file

* Change from 'defineExtension' to 'getExtensionContext'

* Rename helper functions to match main functions

* Introduce the custom_uniqueKey param

* Bump version

* Update copyright

LICENSE

@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2023 Nicolas Keil
+Copyright (c) 2024 Nicolas Keil
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

README.md

@@ -4,18 +4,16 @@
 
 ## Features
 
-- Random Row Retrieval: easily retrieve a random row from your database table using `table.findRandom()`.
-- Random Row Subset retrieval: use `table.findManyRandom()` to query for a customized random row subset of an arbitrary `findMany()` query.
+- Random row retrieval: easily retrieve a random row from your database table using `findRandom()`.
+- Random multi-row retrieval: use `findManyRandom()` to query for a random subset of a `findMany()` query.
 
 ## Installation
 
-Install the Prisma Random Query Extension using your favorite npm package manager:
-
 ```bash
-npm install prisma-extension-random       # npm
-yarn add prisma-extension-random          # yarn
-bun add prisma-extension-random           # bun
-pnpm add prisma-extension-random          # pnpm
+npm install prisma-extension-random
+yarn add prisma-extension-random
+bun add prisma-extension-random
+pnpm add prisma-extension-random
 ```
 
 ## Usage
@@ -49,6 +47,16 @@ const movies = await prisma.movie.findManyRandom(5, {
 });
 ```
 
+Note: when using models with a non-standard id field, you must specify the name of that field using `custom_uniqueKey` as below. If not specified, the name is assumed to be "id".
+
+```typescript
+// Assuming the User table's `@id` column is called 'email'
+const users = await prisma.user.findManyRandom(5, {
+  select: { name: true },
+  custom_uniqueKey: 'email',
+});
+```
+
 ## Contributing
 
 Issues and pull requests are welcome. I'll review them as soon as possible!

example/db.test.ts

@@ -73,3 +73,11 @@ test('findManyRandom', async () => {
   const users2 = await prisma.user.findManyRandom(POPULATION + 10009);
   assert.lengthOf(users2, POPULATION);
 });
+
+test('findRandom filtered', async () => {
+  const users1 = await prisma.user.findRandom({
+    where: { firstName: 'User0' },
+  });
+  assert.isNotNull(users1);
+  assert.equal(users1?.firstName, 'User0');
+});

example/db.ts

@@ -1,4 +1,6 @@
 import { PrismaClient } from '@prisma/client';
+
 import prismaRandom from '../dist/index.js';
+// import prismaRandom from '../src/index.js';
 
 export const prisma = new PrismaClient().$extends(prismaRandom());

example/index.ts

@@ -1,17 +1,21 @@
 import { prisma } from './db.js';
 
 const main = async () => {
-  const user = await prisma.user.findRandom();
+  const user = await prisma.user.findRandom({
+    where: { id: { gt: 2 } },
+    select: { id: true, firstName: true },
+  });
 
   const post = await prisma.post.findManyRandom(10, {
-    select: { id: true, title: true },
+    select: { title: true },
     where: {
       OR: [
         { title: { contains: 'prisma' } },
         { content: { contains: 'prisma' } },
       ],
       published: true,
     },
+    custom_uniqueKey: 'id',
   });
 
   console.log({ user, post });

package.json

@@ -1,6 +1,6 @@
 {
   "name": "prisma-extension-random",
-  "version": "0.1.7",
+  "version": "0.2.0",
   "type": "module",
   "license": "MIT",
   "author": "Nicolas Keil <nkeil.dev@gmail.com>",

src/helpers.ts

@@ -0,0 +1,58 @@
+import { Prisma, User } from '@prisma/client';
+
+// "User" types are used to simplify type inference
+
+export const $findRandom = async (
+  context: Prisma.UserDelegate,
+  args?: Prisma.UserFindFirstArgs,
+): Promise<any> => {
+  const numRows = await context.count({
+    where: args?.where,
+  });
+  return await context.findFirst({
+    ...args,
+    skip: Math.max(0, Math.floor(Math.random() * numRows)),
+  });
+};
+
+export const $findManyRandom = async (
+  context: Prisma.UserDelegate,
+  num: number,
+  args?: {
+    select?: Prisma.UserFindFirstArgs['select'];
+    where?: Prisma.UserFindFirstArgs['where'];
+    custom_uniqueKey?: 'id';
+  },
+): Promise<any> => {
+  const uniqueKey = args?.custom_uniqueKey ?? 'id';
+
+  const rows = [];
+  const rowIds: User['id'][] = [];
+
+  const select = args?.select ?? {};
+  select[uniqueKey] = true;
+
+  const where = args?.where ?? {};
+  where[uniqueKey] = { notIn: rowIds };
+
+  let numRows = await context.count({ where });
+  for (let i = 0; i < num && numRows > 0; ++i) {
+    const row = await context.findFirst({
+      select,
+      where,
+      skip: Math.max(0, Math.floor(Math.random() * numRows)),
+    });
+
+    if (!row) {
+      console.error(
+        `get random row failed. Where clause: ${JSON.stringify(where)}`,
+      );
+      break;
+    }
+    rows.push(row);
+    rowIds.push(row[uniqueKey]);
+    numRows--;
+  }
+
+  return rows;
+};

src/index.ts

@@ -1,85 +1,54 @@
 import { Prisma } from '@prisma/client';
+import { $findManyRandom, $findRandom } from './helpers.js';
 
 type Args = {};
 
 export default (_extensionArgs?: Args) =>
-  Prisma.defineExtension((prisma) => {
-    return prisma.$extends({
-      name: 'prisma-extension-random',
-      model: {
-        $allModels: {
-          async findRandom<T, A>(
-            this: T,
-            args?: Prisma.Exact<A, Prisma.Args<T, 'findFirst'>> & object,
-          ) {
-            const context = Prisma.getExtensionContext(this);
-
-            const numRows = (await (context as any).count({
-              where: (args as { where?: object } | undefined)?.where,
-            })) as number;
-            return (await (context as any).findFirst({
-              ...args,
-              skip: Math.max(0, Math.floor(Math.random() * numRows)),
-            })) as Prisma.Result<T, A, 'findFirst'>;
-          },
+  Prisma.getExtensionContext({
+    name: 'prisma-extension-random',
+    model: {
+      $allModels: {
+        async findRandom<T, A>(
+          this: T,
+          args?: Prisma.Exact<A, Prisma.Args<T, 'findFirst'>> & object,
+        ) {
+          const context = Prisma.getExtensionContext(this);
+
+          return (await $findRandom(
+            context as any,
+            args as any,
+          )) as Prisma.Result<T, A, 'findFirst'>;
+        },
 
-          async findManyRandom<T, TWhere, TSelect>(
-            this: T,
-            num: number,
-            args?: {
-              where?: Prisma.Exact<
-                TWhere,
-                Prisma.Args<T, 'findFirst'>['where']
-              >;
-              select?: Prisma.Exact<
-                TSelect,
-                Prisma.Args<T, 'findFirst'>['select'] & { id: true }
-              >;
-            },
-          ) {
-            const context = Prisma.getExtensionContext(this);
-            type FindFirstResult = Prisma.Result<
-              T,
-              { where: TWhere; select: TSelect },
-              'findFirst'
+        async findManyRandom<T, TWhere, TSelect, TUnique extends string>(
+          this: T,
+          num: number,
+          args?: {
+            where?: Prisma.Exact<TWhere, Prisma.Args<T, 'findFirst'>['where']>;
+            select?: Prisma.Exact<
+              TSelect,
+              Prisma.Args<T, 'findFirst'>['select']
             >;
-
-            const select = args?.select ?? { id: true as const };
-            let where = args?.where ?? {};
-
-            let numRows = (await (context as any).count({ where })) as number;
-
-            const rows: Array<NonNullable<FindFirstResult>> = [];
-            const rowIds: string[] = [];
-
-            where = {
-              ...where,
-              id: { notIn: rowIds },
-            };
-
-            for (let i = 0; i < num && numRows > 0; ++i) {
-              const row = (await (context as any).findFirst({
-                select,
-                where,
-                skip: Math.max(0, Math.floor(Math.random() * numRows)),
-              })) as FindFirstResult;
-
-              if (!row) {
-                console.error(
-                  `get random row failed. Where clause: ${JSON.stringify(
-                    where,
-                  )}`,
-                );
-                break;
-              }
-              rows.push(row);
-              rowIds.push((row as unknown as { id: string }).id);
-              numRows--;
-            }
-
-            return rows;
+            custom_uniqueKey?: TUnique; // TODO: add intellisense?
           },
+        ) {
+          const context = Prisma.getExtensionContext(this);
+          type ExtendedSelect = TSelect & Record<TUnique, true>;
+
+          return (await $findManyRandom(
+            context as any,
+            num,
+            args as any,
+          )) as Array<
+            NonNullable<
+              Prisma.Result<
+                T,
+                { where: TWhere; select: ExtendedSelect },
+                'findFirst'
+              >
+            >
+          >;
         },
       },
-    });
+    },
   });