feat(storage): Add support for createMany, updateMany and upsert

packages/uploads/src/__tests__/queryExtensions.test.ts

@@ -250,4 +250,270 @@ describe('Query extensions', () => {
       expect(fs.unlink).toHaveBeenCalledWith('im-a-invalid-path')
     })
   })
+
+  describe('upsert', () => {
+    it('will remove old files and save new ones on upsert, if it exists [UPDATE]', async () => {
+      const ogDumbo = await prismaClient.dumbo.create({
+        data: {
+          firstUpload: '/tmp/oldFirst.txt',
+          secondUpload: '/tmp/oldSecond.txt',
+        },
+      })
+
+      const updatedDumbo = await prismaClient.dumbo.upsert({
+        update: {
+          firstUpload: '/tmp/newFirst.txt',
+        },
+        create: {
+          // won't be used
+          firstUpload: 'x',
+          secondUpload: 'x',
+        },
+        where: {
+          id: ogDumbo.id,
+        },
+      })
+
+      expect(updatedDumbo.firstUpload).toBe('/tmp/newFirst.txt')
+      expect(updatedDumbo.secondUpload).toBe('/tmp/oldSecond.txt')
+      expect(fs.unlink).toHaveBeenCalledOnce()
+      expect(fs.unlink).toHaveBeenCalledWith('/tmp/oldFirst.txt')
+    })
+
+    it('will create a new record (findOrCreate)', async () => {
+      const newDumbo = await prismaClient.dumbo.upsert({
+        create: {
+          firstUpload: '/tmp/first.txt',
+          secondUpload: '/bazinga/second.txt',
+        },
+        update: {},
+        where: {
+          id: 444444444,
+        },
+      })
+
+      expect(newDumbo.firstUpload).toBe('/tmp/first.txt')
+      expect(newDumbo.secondUpload).toBe('/bazinga/second.txt')
+    })
+
+    it('will remove processed files if upsert CREATION fails (findOrCreate)', async () => {
+      // This is essentially findOrCreate, because update is empty
+      try {
+        await prismaClient.dumbo.upsert({
+          create: {
+            firstUpload: '/tmp/first.txt',
+            secondUpload: '/bazinga/second.txt',
+            // @ts-expect-error Checking the error here
+            id: 'this-is-the-incorrect-type',
+          },
+        })
+      } catch {
+        expect(fs.unlink).toHaveBeenNthCalledWith(1, '/tmp/first.txt')
+        expect(fs.unlink).toHaveBeenNthCalledWith(2, '/bazinga/second.txt')
+      }
+
+      expect.assertions(2)
+    })
+
+    it('will remove processed files if upsert UPDATE fails', async () => {
+      // Bit of a contrived case... why would you ever have different values for update and create...
+
+      const ogDumbo = await prismaClient.dumbo.create({
+        data: {
+          firstUpload: '/tmp/oldFirst.txt',
+          secondUpload: '/tmp/oldSecond.txt',
+        },
+      })
+
+      try {
+        await prismaClient.dumbo.upsert({
+          where: {
+            id: ogDumbo.id,
+          },
+          update: {
+            firstUpload: '/tmp/newFirst.txt',
+            secondUpload: '/tmp/newSecond.txt',
+            // @ts-expect-error Intentionally causing an error
+            id: 'this-should-cause-an-error',
+          },
+          create: {
+            firstUpload: '/tmp/createFirst.txt',
+            secondUpload: '/tmp/createSecond.txt',
+          },
+        })
+      } catch (error) {
+        expect(fs.unlink).toHaveBeenCalledTimes(2)
+        expect(fs.unlink).not.toHaveBeenCalledWith('/tmp/createFirst.txt')
+        expect(fs.unlink).not.toHaveBeenCalledWith('/tmp/createSecond.txt')
+        expect(fs.unlink).toHaveBeenNthCalledWith(1, '/tmp/newFirst.txt')
+        expect(fs.unlink).toHaveBeenNthCalledWith(2, '/tmp/newSecond.txt')
+        expect(error).toBeDefined()
+      }
+
+      // Verify the original files weren't deleted
+      const unchangedDumbo = await prismaClient.dumbo.findUnique({
+        where: { id: ogDumbo.id },
+      })
+      expect(unchangedDumbo?.firstUpload).toBe('/tmp/oldFirst.txt')
+      expect(unchangedDumbo?.secondUpload).toBe('/tmp/oldSecond.txt')
+
+      expect.assertions(8)
+    })
+
+  })
+
+  describe('createMany', () => {
+    it('createMany will remove files if all the create fails', async () => {
+      try {
+        await prismaClient.dumbo.createMany({
+          data: [
+            {
+              firstUpload: '/one/first.txt',
+              secondUpload: '/one/second.txt',
+              // @ts-expect-error Intentional
+              id: 'break',
+            },
+            {
+              firstUpload: '/two/first.txt',
+              secondUpload: '/two/second.txt',
+              // @ts-expect-error Intentional
+              id: 'break2',
+            },
+          ],
+        })
+      } catch {
+        expect(fs.unlink).toHaveBeenCalledTimes(4)
+        expect(fs.unlink).toHaveBeenNthCalledWith(1, '/one/first.txt')
+        expect(fs.unlink).toHaveBeenNthCalledWith(2, '/one/second.txt')
+        expect(fs.unlink).toHaveBeenNthCalledWith(3, '/two/first.txt')
+        expect(fs.unlink).toHaveBeenNthCalledWith(4, '/two/second.txt')
+      }
+
+      expect.assertions(5)
+    })
+
+    it('createMany will remove all files, even if one of them errors', async () => {
+      try {
+        await prismaClient.dumbo.createMany({
+          data: [
+            // This one is correct, but createMany fails together
+            // so all the files should be removed!
+            {
+              firstUpload: '/one/first.txt',
+              secondUpload: '/one/second.txt',
+              id: 9158125,
+            },
+            {
+              firstUpload: '/two/first.txt',
+              secondUpload: '/two/second.txt',
+              // @ts-expect-error Intentional
+              id: 'break2',
+            },
+          ],
+        })
+      } catch (e) {
+        // This one doesn't actually get created!
+        expect(
+          prismaClient.dumbo.findUnique({ where: { id: 9158125 } }),
+        ).resolves.toBeNull()
+
+        expect(fs.unlink).toHaveBeenCalledTimes(4)
+        expect(fs.unlink).toHaveBeenNthCalledWith(1, '/one/first.txt')
+        expect(fs.unlink).toHaveBeenNthCalledWith(2, '/one/second.txt')
+        expect(fs.unlink).toHaveBeenNthCalledWith(3, '/two/first.txt')
+        expect(fs.unlink).toHaveBeenNthCalledWith(4, '/two/second.txt')
+      }
+
+      expect.assertions(6)
+    })
+  })
+
+  describe('updateMany', () => {
+    it('will remove old files and save new ones on update, if they exist', async () => {
+      const ogDumbo1 = await prismaClient.dumbo.create({
+        data: {
+          firstUpload: '/FINDME/oldFirst1.txt',
+          secondUpload: '/FINDME/oldSecond1.txt',
+        },
+      })
+
+      const ogDumbo2 = await prismaClient.dumbo.create({
+        data: {
+          firstUpload: '/FINDME/oldFirst2.txt',
+          secondUpload: '/FINDME/oldSecond2.txt',
+        },
+      })
+
+      const updatedDumbos = await prismaClient.dumbo.updateMany({
+        data: {
+          firstUpload: '/REPLACED/newFirst.txt',
+          secondUpload: '/REPLACED/newSecond.txt',
+        },
+        where: {
+          firstUpload: {
+            contains: 'FINDME',
+          },
+        },
+      })
+
+      expect(updatedDumbos.count).toBe(2)
+
+      const updatedDumbo1 = await prismaClient.dumbo.findFirstOrThrow({
+        where: {
+          id: ogDumbo1.id,
+        },
+      })
+
+      const updatedDumbo2 = await prismaClient.dumbo.findFirstOrThrow({
+        where: {
+          id: ogDumbo2.id,
+        },
+      })
+
+      // Still performs the update
+      expect(updatedDumbo1.firstUpload).toBe('/REPLACED/newFirst.txt')
+      expect(updatedDumbo1.secondUpload).toBe('/REPLACED/newSecond.txt')
+      expect(updatedDumbo2.firstUpload).toBe('/REPLACED/newFirst.txt')
+      expect(updatedDumbo2.secondUpload).toBe('/REPLACED/newSecond.txt')
+
+      // Then deletes the old files
+      expect(fs.unlink).toHaveBeenCalledTimes(4)
+      expect(fs.unlink).toHaveBeenNthCalledWith(1, '/FINDME/oldFirst1.txt')
+      expect(fs.unlink).toHaveBeenNthCalledWith(2, '/FINDME/oldSecond1.txt')
+      expect(fs.unlink).toHaveBeenNthCalledWith(3, '/FINDME/oldFirst2.txt')
+      expect(fs.unlink).toHaveBeenNthCalledWith(4, '/FINDME/oldSecond2.txt')
+    })
+
+    it('will __not__ remove files if the update fails', async () => {
+      const ogDumbo1 = await prismaClient.dumbo.create({
+        data: {
+          firstUpload: '/tmp/oldFirst1.txt',
+          secondUpload: '/tmp/oldSecond1.txt',
+        },
+      })
+
+      const ogDumbo2 = await prismaClient.dumbo.create({
+        data: {
+          firstUpload: '/tmp/oldFirst2.txt',
+          secondUpload: '/tmp/oldSecond2.txt',
+        },
+      })
+
+      const failedUpdatePromise = prismaClient.dumbo.updateMany({
+        data: {
+          // @ts-expect-error Intentional
+          id: 'this-is-the-incorrect-type',
+        },
+        where: {
+          OR: [{ id: ogDumbo1.id }, { id: ogDumbo2.id }],
+        },
+      })
+
+      // Id is invalid, so the update should fail
+      await expect(failedUpdatePromise).rejects.toThrowError()
+
+      // The old files should NOT be deleted
+      expect(fs.unlink).not.toHaveBeenCalled()
+    })
+  })
 })

packages/uploads/src/prismaExtension.ts

@@ -93,6 +93,21 @@ export const createUploadsExtension = <MNames extends ModelNames = ModelNames>(
           throw e
         }
       },
+      async createMany({ query, args }) {
+        try {
+          const result = await query(args)
+          return result
+        } catch (e) {
+          const createDatas = args.data as []
+
+          // If the create fails, we need to delete the uploaded files
+          for await (const createData of createDatas) {
+            await removeUploadedFiles(uploadFields, createData)
+          }
+
+          throw e
+        }
+      },
       async update({ query, model, args }) {
         // Check if any of the uploadFields are present in args.data
         // We only want to process fields that are being updated
@@ -136,6 +151,91 @@ export const createUploadsExtension = <MNames extends ModelNames = ModelNames>(
           }
         }
       },
+      async updateMany({ query, model, args }) {
+        // Check if any of the uploadFields are present in args.data
+        // We only want to process fields that are being updated
+        const uploadFieldsToUpdate = uploadFields.filter(
+          (field) =>
+            // All of this non-sense is to make typescript happy. I'm not sure how data could be anything but an object
+            typeof args.data === 'object' &&
+            args.data !== null &&
+            field in args.data,
+        )
+
+        if (uploadFieldsToUpdate.length == 0) {
+          return query(args)
+        } else {
+          // MULTIPLE!
+          const originalRecords = await prismaInstance[
+            model as ModelNames
+            // @ts-expect-error TS in strict mode will error due to union type. We cannot narrow it down here.
+          ].findMany({
+            where: args.where,
+            // @TODO: should we select here to reduce the amount of data we're handling
+          })
+
+          try {
+            const result = await query(args)
+
+            // Remove the uploaded files from each of the original records
+            for await (const originalRecord of originalRecords) {
+              await removeUploadedFiles(uploadFieldsToUpdate, originalRecord)
+            }
+
+            return result
+          } catch (e) {
+            // If the update many fails, we need to delete the newly uploaded files
+            // but not the ones that already exist!
+            await removeUploadedFiles(
+              uploadFieldsToUpdate,
+              args.data as Record<string, string>,
+            )
+            throw e
+          }
+        }
+      },
+      async upsert({ query, model, args }) {
+        // null if we don't know yet
+        let isUpdate: boolean | null = null
+        const uploadFieldsToUpdate = uploadFields.filter(
+          (field) =>
+            typeof args.update === 'object' &&
+            args.update !== null &&
+            field in args.update,
+        )
+
+        try {
+          // We only need to check for existing records if we're updating
+          let existingRecord: Record<string, string> | undefined
+          if (args.update) {
+            existingRecord = await prismaInstance[
+              model as ModelNames
+              // @ts-expect-error TS in strict mode will error due to union type. We cannot narrow it down here.
+            ].findUnique({
+              where: args.where,
+            })
+            isUpdate = !!existingRecord
+          }
+
+          const result = await query(args)
+
+          if (isUpdate && existingRecord) {
+            // If the record existed, remove old uploaded files
+            await removeUploadedFiles(uploadFieldsToUpdate, existingRecord)
+          }
+
+          return result
+        } catch (e) {
+          // If the upsert fails, we need to delete any newly uploaded files
+          await removeUploadedFiles(
+            // Only delete files we're updating on update
+            isUpdate ? uploadFieldsToUpdate : uploadFields,
+            (isUpdate ? args.update : args.create) as Record<string, string>,
+          )
+
+          throw e
+        }
+      },
 
       async delete({ query, args }) {
         const deleteResult = await query(args)
@@ -215,6 +315,11 @@ export const createUploadsExtension = <MNames extends ModelNames = ModelNames>(
     })
   })
 
+  /**
+   * This function deletes files from the storage adapter, but importantly,
+   * it does NOT throw, because if the file is already gone, that's fine,
+   * no need to stop the actual db operation
+   */
   async function removeUploadedFiles(
     fieldsToDelete: string[],
     data: Record<string, string>,