From c0a72123bd45b6e47c48658f28122a4532883d90 Mon Sep 17 00:00:00 2001
From: Catalin Pit <catalinpit@gmail.com>
Date: Wed, 23 Jul 2025 14:42:16 +0300
Subject: [PATCH] chore: implement feedback

---
 .../attachment/set-document-attachments.ts    | 46 +++++++----------
 .../attachment/set-template-attachments.ts    | 50 ++++++++-----------
 .../trpc/server/attachment-router/router.ts   |  3 ++
 3 files changed, 44 insertions(+), 55 deletions(-)

diff --git a/packages/lib/server-only/attachment/set-document-attachments.ts b/packages/lib/server-only/attachment/set-document-attachments.ts
index cdbae58d4..3f72404c9 100644
--- a/packages/lib/server-only/attachment/set-document-attachments.ts
+++ b/packages/lib/server-only/attachment/set-document-attachments.ts
@@ -7,11 +7,13 @@ import { AppError } from '../../errors/app-error';
 import { AppErrorCode } from '../../errors/app-error';
 import { DOCUMENT_AUDIT_LOG_TYPE } from '../../types/document-audit-logs';
 import { createDocumentAuditLogData } from '../../utils/document-audit-logs';
+import { buildTeamWhereQuery } from '../../utils/teams';
 
 export type CreateAttachmentsOptions = {
   documentId: number;
   attachments: Pick<Attachment, 'id' | 'label' | 'url' | 'type'>[];
   user: Pick<User, 'id' | 'email' | 'name'>;
+  teamId: number;
   requestMetadata: RequestMetadata;
 };
 
@@ -19,11 +21,13 @@ export const setDocumentAttachments = async ({
   documentId,
   attachments,
   user,
+  teamId,
   requestMetadata,
 }: CreateAttachmentsOptions) => {
   const document = await prisma.document.findUnique({
     where: {
       id: documentId,
+      team: buildTeamWhereQuery({ teamId, userId: user.id }),
     },
   });
 
@@ -53,33 +57,21 @@ export const setDocumentAttachments = async ({
   const upsertedAttachments: Attachment[] = [];
 
   for (const attachment of attachments) {
-    if (attachment.id) {
-      const updated = await prisma.attachment.upsert({
-        where: { id: attachment.id },
-        update: {
-          label: attachment.label,
-          url: attachment.url,
-          type: attachment.type,
-        },
-        create: {
-          label: attachment.label,
-          url: attachment.url,
-          type: attachment.type,
-          documentId,
-        },
-      });
-      upsertedAttachments.push(updated);
-    } else {
-      const created = await prisma.attachment.create({
-        data: {
-          label: attachment.label,
-          url: attachment.url,
-          type: attachment.type,
-          documentId,
-        },
-      });
-      upsertedAttachments.push(created);
-    }
+    const updated = await prisma.attachment.upsert({
+      where: { id: attachment.id, documentId: document.id },
+      update: {
+        label: attachment.label,
+        url: attachment.url,
+        type: attachment.type,
+      },
+      create: {
+        label: attachment.label,
+        url: attachment.url,
+        type: attachment.type,
+        documentId,
+      },
+    });
+    upsertedAttachments.push(updated);
   }
 
   const isAttachmentsSame = upsertedAttachments.every((attachment) => {
diff --git a/packages/lib/server-only/attachment/set-template-attachments.ts b/packages/lib/server-only/attachment/set-template-attachments.ts
index 3edd0730b..68c69fd49 100644
--- a/packages/lib/server-only/attachment/set-template-attachments.ts
+++ b/packages/lib/server-only/attachment/set-template-attachments.ts
@@ -4,19 +4,25 @@ import { prisma } from '@documenso/prisma';
 
 import { AppError } from '../../errors/app-error';
 import { AppErrorCode } from '../../errors/app-error';
+import { buildTeamWhereQuery } from '../../utils/teams';
 
 export type CreateAttachmentsOptions = {
   templateId: number;
   attachments: Pick<Attachment, 'id' | 'label' | 'url' | 'type'>[];
+  userId: number;
+  teamId: number;
 };
 
 export const setTemplateAttachments = async ({
   templateId,
   attachments,
+  userId,
+  teamId,
 }: CreateAttachmentsOptions) => {
   const template = await prisma.template.findUnique({
     where: {
       id: templateId,
+      team: buildTeamWhereQuery({ teamId, userId }),
     },
   });
 
@@ -46,34 +52,22 @@ export const setTemplateAttachments = async ({
   const upsertedAttachments: Attachment[] = [];
 
   for (const attachment of attachments) {
-    if (attachment.id) {
-      const updated = await prisma.attachment.upsert({
-        where: { id: attachment.id },
-        update: {
-          label: attachment.label,
-          url: attachment.url,
-          type: attachment.type,
-          templateId,
-        },
-        create: {
-          label: attachment.label,
-          url: attachment.url,
-          type: attachment.type,
-          templateId,
-        },
-      });
-      upsertedAttachments.push(updated);
-    } else {
-      const created = await prisma.attachment.create({
-        data: {
-          label: attachment.label,
-          url: attachment.url,
-          type: attachment.type,
-          templateId,
-        },
-      });
-      upsertedAttachments.push(created);
-    }
+    const updated = await prisma.attachment.upsert({
+      where: { id: attachment.id, templateId: template.id },
+      update: {
+        label: attachment.label,
+        url: attachment.url,
+        type: attachment.type,
+        templateId,
+      },
+      create: {
+        label: attachment.label,
+        url: attachment.url,
+        type: attachment.type,
+        templateId,
+      },
+    });
+    upsertedAttachments.push(updated);
   }
 
   return upsertedAttachments;
diff --git a/packages/trpc/server/attachment-router/router.ts b/packages/trpc/server/attachment-router/router.ts
index b7698563f..b5c31ebf8 100644
--- a/packages/trpc/server/attachment-router/router.ts
+++ b/packages/trpc/server/attachment-router/router.ts
@@ -47,6 +47,7 @@ export const attachmentRouter = router({
         documentId,
         attachments,
         user: ctx.user,
+        teamId: ctx.teamId,
         requestMetadata: ctx.metadata.requestMetadata,
       });
 
@@ -81,6 +82,8 @@ export const attachmentRouter = router({
 
       const updatedAttachments = await setTemplateAttachments({
         templateId,
+        userId: ctx.user.id,
+        teamId: ctx.teamId,
         attachments,
       });