feat(files): export workspace files to Google Drive

apps/sim/app/api/tools/google_drive/upload/route.ts

@@ -1,12 +1,12 @@
 import { createLogger } from '@sim/logger'
 import { getErrorMessage } from '@sim/utils/errors'
-import { generateShortId } from '@sim/utils/id'
 import { type NextRequest, NextResponse } from 'next/server'
 import { googleDriveUploadContract } from '@/lib/api/contracts/tools/google'
 import { parseRequest } from '@/lib/api/server'
 import { checkInternalAuth } from '@/lib/auth/hybrid'
 import { generateRequestId } from '@/lib/core/utils/request'
 import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
+import { DriveUploadError, uploadBufferToDrive } from '@/lib/google-drive/upload-to-drive'
 import { processSingleFileToUserFile } from '@/lib/uploads/utils/file-utils'
 import { downloadFileFromStorage } from '@/lib/uploads/utils/file-utils.server'
 import { assertToolFileAccess } from '@/app/api/files/authorization'
@@ -20,35 +20,6 @@ export const dynamic = 'force-dynamic'
 
 const logger = createLogger('GoogleDriveUploadAPI')
 
-const GOOGLE_DRIVE_API_BASE = 'https://www.googleapis.com/upload/drive/v3/files'
-
-/**
- * Build multipart upload body for Google Drive API
- */
-function buildMultipartBody(
-  metadata: Record<string, any>,
-  fileBuffer: Buffer,
-  mimeType: string,
-  boundary: string
-): string {
-  const parts: string[] = []
-
-  parts.push(`--${boundary}`)
-  parts.push('Content-Type: application/json; charset=UTF-8')
-  parts.push('')
-  parts.push(JSON.stringify(metadata))
-
-  parts.push(`--${boundary}`)
-  parts.push(`Content-Type: ${mimeType}`)
-  parts.push('Content-Transfer-Encoding: base64')
-  parts.push('')
-  parts.push(fileBuffer.toString('base64'))
-
-  parts.push(`--${boundary}--`)
-
-  return parts.join('\r\n')
-}
-
 export const POST = withRouteHandler(async (request: NextRequest) => {
   const requestId = generateRequestId()
 
@@ -159,99 +130,34 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
       }
     }
 
-    const metadata: {
-      name: string
-      mimeType: string
-      parents?: string[]
-    } = {
-      name: validatedData.fileName,
-      mimeType: requestedMimeType,
-    }
-
-    if (validatedData.folderId && validatedData.folderId.trim() !== '') {
-      metadata.parents = [validatedData.folderId.trim()]
-    }
-
-    const boundary = `boundary_${Date.now()}_${generateShortId(7)}`
-
-    const multipartBody = buildMultipartBody(metadata, fileBuffer, uploadMimeType, boundary)
-
     logger.info(`[${requestId}] Uploading to Google Drive via multipart upload`, {
       fileName: validatedData.fileName,
       size: fileBuffer.length,
       uploadMimeType,
       requestedMimeType,
     })
 
-    const uploadResponse = await fetch(
-      `${GOOGLE_DRIVE_API_BASE}?uploadType=multipart&supportsAllDrives=true`,
-      {
-        method: 'POST',
-        headers: {
-          Authorization: `Bearer ${validatedData.accessToken}`,
-          'Content-Type': `multipart/related; boundary=${boundary}`,
-          'Content-Length': Buffer.byteLength(multipartBody, 'utf-8').toString(),
-        },
-        body: multipartBody,
-      }
-    )
-
-    if (!uploadResponse.ok) {
-      const errorText = await uploadResponse.text()
-      logger.error(`[${requestId}] Google Drive API error:`, {
-        status: uploadResponse.status,
-        statusText: uploadResponse.statusText,
-        error: errorText,
+    let finalFile
+    try {
+      finalFile = await uploadBufferToDrive({
+        accessToken: validatedData.accessToken,
+        name: validatedData.fileName,
+        mimeType: requestedMimeType,
+        uploadMimeType,
+        buffer: fileBuffer,
+        folderId: validatedData.folderId ?? undefined,
       })
-      return NextResponse.json(
-        {
-          success: false,
-          error: `Google Drive API error: ${uploadResponse.statusText}`,
-        },
-        { status: uploadResponse.status }
-      )
-    }
-
-    const uploadData = await uploadResponse.json()
-    const fileId = uploadData.id
-
-    logger.info(`[${requestId}] File uploaded successfully`, { fileId })
-
-    if (GOOGLE_WORKSPACE_MIME_TYPES.includes(requestedMimeType)) {
-      logger.info(`[${requestId}] Updating file name to ensure it persists after conversion`)
-
-      const updateNameResponse = await fetch(
-        `https://www.googleapis.com/drive/v3/files/${fileId}?supportsAllDrives=true`,
-        {
-          method: 'PATCH',
-          headers: {
-            Authorization: `Bearer ${validatedData.accessToken}`,
-            'Content-Type': 'application/json',
-          },
-          body: JSON.stringify({
-            name: validatedData.fileName,
-          }),
-        }
-      )
-
-      if (!updateNameResponse.ok) {
-        logger.warn(
-          `[${requestId}] Failed to update filename after conversion, but content was uploaded`
-        )
+    } catch (error) {
+      if (error instanceof DriveUploadError) {
+        logger.error(`[${requestId}] Google Drive API error:`, {
+          status: error.status,
+          error: error.message,
+        })
+        return NextResponse.json({ success: false, error: error.message }, { status: error.status })
       }
+      throw error
     }
 
-    const finalFileResponse = await fetch(
-      `https://www.googleapis.com/drive/v3/files/${fileId}?supportsAllDrives=true&fields=id,name,mimeType,webViewLink,webContentLink,size,createdTime,modifiedTime,parents`,
-      {
-        headers: {
-          Authorization: `Bearer ${validatedData.accessToken}`,
-        },
-      }
-    )
-
-    const finalFile = await finalFileResponse.json()
-
     logger.info(`[${requestId}] Upload complete`, {
       fileId: finalFile.id,
       fileName: finalFile.name,

apps/sim/app/api/workspaces/[id]/files/export-to-drive/route.test.ts

@@ -0,0 +1,158 @@
+/**
+ * @vitest-environment node
+ */
+import { authMockFns } from '@sim/testing'
+import { NextRequest } from 'next/server'
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+const {
+  mockListWorkspaceFiles,
+  mockFetchWorkspaceFileBuffer,
+  mockRefreshAccessTokenIfNeeded,
+  mockUploadBufferToDrive,
+  mockVerifyWorkspaceMembership,
+} = vi.hoisted(() => ({
+  mockListWorkspaceFiles: vi.fn(),
+  mockFetchWorkspaceFileBuffer: vi.fn(),
+  mockRefreshAccessTokenIfNeeded: vi.fn(),
+  mockUploadBufferToDrive: vi.fn(),
+  mockVerifyWorkspaceMembership: vi.fn(),
+}))
+
+vi.mock('@/lib/uploads/contexts/workspace', () => ({
+  listWorkspaceFiles: mockListWorkspaceFiles,
+  fetchWorkspaceFileBuffer: mockFetchWorkspaceFileBuffer,
+}))
+
+vi.mock('@/app/api/auth/oauth/utils', () => ({
+  refreshAccessTokenIfNeeded: mockRefreshAccessTokenIfNeeded,
+}))
+
+vi.mock('@/lib/google-drive/upload-to-drive', () => ({
+  uploadBufferToDrive: mockUploadBufferToDrive,
+}))
+
+vi.mock('@/app/api/workflows/utils', () => ({
+  verifyWorkspaceMembership: mockVerifyWorkspaceMembership,
+}))
+
+const WS = '7727ef3f-8cf6-4686-b063-2bb006a10785'
+
+import { POST } from '@/app/api/workspaces/[id]/files/export-to-drive/route'
+
+const params = (id = WS) => ({ params: Promise.resolve({ id }) })
+
+const makeRequest = (body: unknown) =>
+  new NextRequest(`http://localhost/api/workspaces/${WS}/files/export-to-drive`, {
+    method: 'POST',
+    headers: { 'Content-Type': 'application/json' },
+    body: JSON.stringify(body),
+  })
+
+const fileRecord = (id: string, name: string, type = 'application/pdf') => ({
+  id,
+  name,
+  type,
+  key: `workspace/${WS}/${id}-${name}`,
+  storageContext: 'workspace',
+})
+
+const validBody = { fileIds: ['file-1', 'file-2'], credentialId: 'cred-1' }
+
+describe('POST /api/workspaces/[id]/files/export-to-drive', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    authMockFns.mockGetSession.mockResolvedValue({ user: { id: 'user-1' } })
+    mockVerifyWorkspaceMembership.mockResolvedValue('read')
+    mockRefreshAccessTokenIfNeeded.mockResolvedValue('access-token')
+    mockListWorkspaceFiles.mockResolvedValue([
+      fileRecord('file-1', 'a.pdf'),
+      fileRecord('file-2', 'b.pdf'),
+    ])
+    mockFetchWorkspaceFileBuffer.mockResolvedValue(Buffer.from('content'))
+    mockUploadBufferToDrive.mockImplementation(({ name }: { name: string }) =>
+      Promise.resolve({ id: `drive-${name}`, name, webViewLink: `https://drive/${name}` })
+    )
+  })
+
+  it('returns 401 when unauthenticated', async () => {
+    authMockFns.mockGetSession.mockResolvedValueOnce(null)
+    const res = await POST(makeRequest(validBody), params())
+    expect(res.status).toBe(401)
+    expect(mockUploadBufferToDrive).not.toHaveBeenCalled()
+  })
+
+  it('returns 403 when the user is not a workspace member', async () => {
+    mockVerifyWorkspaceMembership.mockResolvedValueOnce(null)
+    const res = await POST(makeRequest(validBody), params())
+    expect(res.status).toBe(403)
+    expect(mockUploadBufferToDrive).not.toHaveBeenCalled()
+  })
+
+  it('returns 400 when the body is invalid (no files selected)', async () => {
+    const res = await POST(makeRequest({ fileIds: [], credentialId: 'cred-1' }), params())
+    expect(res.status).toBe(400)
+  })
+
+  it('returns 400 when the Google Drive token cannot be resolved', async () => {
+    mockRefreshAccessTokenIfNeeded.mockResolvedValueOnce(null)
+    const res = await POST(makeRequest(validBody), params())
+    const body = await res.json()
+    expect(res.status).toBe(400)
+    expect(body.error).toContain('not connected')
+    expect(mockUploadBufferToDrive).not.toHaveBeenCalled()
+  })
+
+  it('returns 400 when none of the requested files exist', async () => {
+    mockListWorkspaceFiles.mockResolvedValueOnce([fileRecord('other', 'x.pdf')])
+    const res = await POST(makeRequest(validBody), params())
+    expect(res.status).toBe(400)
+    expect(mockUploadBufferToDrive).not.toHaveBeenCalled()
+  })
+
+  it('exports all matching files and returns success', async () => {
+    const res = await POST(makeRequest(validBody), params())
+    const body = await res.json()
+    expect(res.status).toBe(200)
+    expect(body.success).toBe(true)
+    expect(body.exported).toHaveLength(2)
+    expect(body.failed).toHaveLength(0)
+    expect(body.exported[0]).toMatchObject({ fileId: 'file-1', driveFileId: 'drive-a.pdf' })
+    expect(mockUploadBufferToDrive).toHaveBeenCalledTimes(2)
+  })
+
+  it('reports partial failure without aborting the batch', async () => {
+    mockUploadBufferToDrive
+      .mockResolvedValueOnce({ id: 'drive-a', name: 'a.pdf', webViewLink: 'https://drive/a' })
+      .mockRejectedValueOnce(new Error('Drive quota exceeded'))
+    const res = await POST(makeRequest(validBody), params())
+    const body = await res.json()
+    expect(res.status).toBe(200)
+    expect(body.success).toBe(false)
+    expect(body.exported).toHaveLength(1)
+    expect(body.failed).toHaveLength(1)
+    expect(body.failed[0]).toMatchObject({ fileId: 'file-2', error: 'Drive quota exceeded' })
+  })
+
+  it('reports requested files that no longer exist as failures', async () => {
+    mockListWorkspaceFiles.mockResolvedValueOnce([fileRecord('file-1', 'a.pdf')])
+    const res = await POST(
+      makeRequest({ fileIds: ['file-1', 'file-gone'], credentialId: 'cred-1' }),
+      params()
+    )
+    const body = await res.json()
+    expect(res.status).toBe(200)
+    expect(body.success).toBe(false)
+    expect(body.exported).toHaveLength(1)
+    expect(body.failed).toEqual([{ fileId: 'file-gone', error: 'File not found' }])
+    expect(mockUploadBufferToDrive).toHaveBeenCalledTimes(1)
+  })
+
+  it('only exports files that match the requested ids', async () => {
+    await POST(makeRequest({ fileIds: ['file-1'], credentialId: 'cred-1' }), params())
+    expect(mockUploadBufferToDrive).toHaveBeenCalledTimes(1)
+    expect(mockUploadBufferToDrive).toHaveBeenCalledWith(
+      expect.objectContaining({ name: 'a.pdf', accessToken: 'access-token' })
+    )
+  })
+})