xixu-me/gdrive-upload-robot
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor upload logic and remove upload tests

Browse Source 
Refactored file upload logic to pass file size directly and streamline handling of large files with a new resumableUpload function. Removed the upload.test.js file.
This commit is contained in:
xixu-me committed 2025-07-04 13:02:24 +08:00
1 parent cb5b329a5b
commit eac9e037e2
2 files changed
+40 -213
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/index.js
+40 -141
View File
@@ -15,17 +15,19 @@ export default {
async function handleDocument(message, env) { async function handleDocument(message, env) {
const fileId = message.document.file_id; const fileId = message.document.file_id;
const fileName = message.document.file_name; const fileName = message.document.file_name;
await uploadToDrive(fileId, fileName, env); const fileSize = message.document.file_size;
await uploadToDrive(fileId, fileName, fileSize, env);
} }
async function handlePhoto(message, env) { async function handlePhoto(message, env) {
const photo = message.photo[message.photo.length - 1]; // Get the highest resolution photo const photo = message.photo[message.photo.length - 1]; // Get the highest resolution photo
const fileId = photo.file_id; const fileId = photo.file_id;
const fileName = `${fileId}.jpg`; const fileName = `${fileId}.jpg`;
await uploadToDrive(fileId, fileName, env); const fileSize = photo.file_size;
await uploadToDrive(fileId, fileName, fileSize, env);
} }
async function uploadToDrive(fileId, fileName, env) { async function uploadToDrive(fileId, fileName, fileSize, env) {
const botToken = env.TELEGRAM_BOT_TOKEN; const botToken = env.TELEGRAM_BOT_TOKEN;
const driveFolderId = env.GOOGLE_DRIVE_FOLDER_ID; const driveFolderId = env.GOOGLE_DRIVE_FOLDER_ID;
const adminChatId = env.ADMIN_CHAT_ID; const adminChatId = env.ADMIN_CHAT_ID;
@@ -42,27 +44,20 @@ async function uploadToDrive(fileId, fileName, env) {
const filePath = fileDetails.result.file_path; const filePath = fileDetails.result.file_path;
const fileUrl = `https://api.telegram.org/file/bot${botToken}/${filePath}`; const fileUrl = `https://api.telegram.org/file/bot${botToken}/${filePath}`;
// Download file from Telegram
const fileResponse = await fetch(fileUrl);
const fileData = await fileResponse.arrayBuffer();
const fileSize = fileData.byteLength;
// Authenticate with Google Drive // Authenticate with Google Drive
const jwtToken = await getGoogleAuthToken(googleCredentials); const jwtToken = await getGoogleAuthToken(googleCredentials);
// Check if file is larger than 20MB (20 * 1024 * 1024 = 20971520 bytes) // For files larger than 20MB, use resumable upload
const MAX_SIMPLE_UPLOAD_SIZE = 20971520; if (fileSize > 20 * 1024 * 1024) {
await resumableUpload(fileUrl, fileName, driveFolderId, jwtToken, adminChatId, botToken);
if (fileSize > MAX_SIMPLE_UPLOAD_SIZE) { return;
await sendMessage(adminChatId, `File "${fileName}" is ${Math.round(fileSize / 1024 / 1024)}MB, using resumable upload...`, botToken);
await uploadLargeFile(fileData, fileName, driveFolderId, jwtToken, adminChatId, botToken);
} else {
await uploadSmallFile(fileData, fileName, driveFolderId, jwtToken, adminChatId, botToken);
} }
}
async function uploadSmallFile(fileData, fileName, driveFolderId, jwtToken, adminChatId, botToken) { // Download file from Telegram
// Upload file to Google Drive using multipart upload const fileResponse = await fetch(fileUrl);
const fileData = await fileResponse.arrayBuffer();
// Upload file to Google Drive using multipart upload for smaller files
const uploadUrl = 'https://www.googleapis.com/upload/drive/v3/files?uploadType=multipart'; const uploadUrl = 'https://www.googleapis.com/upload/drive/v3/files?uploadType=multipart';
const metadata = { const metadata = {
name: fileName, name: fileName,
@@ -89,160 +84,64 @@ async function uploadSmallFile(fileData, fileName, driveFolderId, jwtToken, admi
} }
} }
async function uploadLargeFile(fileData, fileName, driveFolderId, jwtToken, adminChatId, botToken) { async function resumableUpload(fileUrl, fileName, driveFolderId, jwtToken, adminChatId, botToken) {
const metadata = { const metadata = {
name: fileName, name: fileName,
parents: [driveFolderId], parents: [driveFolderId],
}; };
// Step 1: Initiate resumable upload const initiationResponse = await fetch('https://www.googleapis.com/upload/drive/v3/files?uploadType=resumable', {
const initiateUrl = 'https://www.googleapis.com/upload/drive/v3/files?uploadType=resumable';
const initiateResponse = await fetch(initiateUrl, {
method: 'POST', method: 'POST',
headers: { headers: {
Authorization: `Bearer ${jwtToken}`, Authorization: `Bearer ${jwtToken}`,
'Content-Type': 'application/json', 'Content-Type': 'application/json',
'X-Upload-Content-Type': 'application/octet-stream',
'X-Upload-Content-Length': fileData.byteLength.toString(),
}, },
body: JSON.stringify(metadata), body: JSON.stringify(metadata),
}); });
if (!initiateResponse.ok) { if (!initiationResponse.ok) {
const error = await initiateResponse.text(); await sendMessage(adminChatId, `Error initiating resumable upload: ${await initiationResponse.text()}`, botToken);
await sendMessage(adminChatId, `Error initiating resumable upload: ${error}`, botToken);
return; return;
} }
const resumableUploadUrl = initiateResponse.headers.get('Location'); const location = initiationResponse.headers.get('Location');
if (!resumableUploadUrl) {
await sendMessage(adminChatId, `Error: No resumable upload URL received`, botToken); const fileResponse = await fetch(fileUrl);
return; const reader = fileResponse.body.getReader();
let uploaded = 0;
const chunkSize = 256 * 1024 * 10; // 2.56MB chunks
const totalSize = Number(fileResponse.headers.get('Content-Length'));
while (true) {
const { done, value } = await reader.read();
if (done) {
break;
} }
// Step 2: Check if any bytes have already been uploaded const chunk = value;
let uploadedBytes = await getUploadStatus(resumableUploadUrl, jwtToken); const uploadResponse = await fetch(location, {
if (uploadedBytes > 0) {
await sendMessage(adminChatId, `Resuming upload from ${Math.round(uploadedBytes / 1024 / 1024)}MB`, botToken);
}
// Step 3: Upload file in chunks
const CHUNK_SIZE = 10 * 1024 * 1024; // 10MB chunks
const MAX_RETRIES = 3;
const totalSize = fileData.byteLength;
while (uploadedBytes < totalSize) {
const start = uploadedBytes;
const end = Math.min(start + CHUNK_SIZE, totalSize);
const chunk = fileData.slice(start, end);
const chunkSize = end - start;
let retryCount = 0;
let success = false;
while (retryCount < MAX_RETRIES && !success) {
try {
const chunkResponse = await fetch(resumableUploadUrl, {
method: 'PUT', method: 'PUT',
headers: { headers: {
Authorization: `Bearer ${jwtToken}`, 'Content-Range': `bytes ${uploaded}-${uploaded + chunk.length - 1}/${totalSize}`,
'Content-Range': `bytes ${start}-${end - 1}/${totalSize}`,
'Content-Length': chunkSize.toString(),
}, },
body: chunk, body: chunk,
}); });
if (chunkResponse.status === 308) { if (uploadResponse.status === 308) {
// Chunk uploaded successfully, continue with next chunk // Resume Incomplete
const rangeHeader = chunkResponse.headers.get('Range'); uploaded += chunk.length;
if (rangeHeader) { } else if (uploadResponse.ok) {
const rangeMatch = rangeHeader.match(/bytes=0-(\d+)/); const uploadResult = await uploadResponse.json();
if (rangeMatch) {
uploadedBytes = parseInt(rangeMatch[1]) + 1;
} else {
uploadedBytes = end;
}
} else {
uploadedBytes = end;
}
// Send progress update (only every 20% to avoid spam)
const progress = Math.round((uploadedBytes / totalSize) * 100);
if (progress % 20 === 0 || uploadedBytes === end) {
await sendMessage(
adminChatId,
`Upload progress: ${progress}% (${Math.round(uploadedBytes / 1024 / 1024)}MB / ${Math.round(totalSize / 1024 / 1024)}MB)`,
botToken
);
}
success = true;
} else if (chunkResponse.status === 200 || chunkResponse.status === 201) {
// Upload completed
const uploadResult = await chunkResponse.json();
if (uploadResult.id) { if (uploadResult.id) {
await sendMessage(adminChatId, `File "${fileName}" uploaded to Google Drive successfully!`, botToken); await sendMessage(adminChatId, `File "${fileName}" uploaded to Google Drive successfully!`, botToken);
} else { } else {
await sendMessage(adminChatId, `Error completing upload: ${JSON.stringify(uploadResult)}`, botToken); await sendMessage(adminChatId, `Error completing resumable upload: ${JSON.stringify(uploadResult.error)}`, botToken);
} }
return; return;
} else if (chunkResponse.status === 404) {
// Session expired, need to restart
await sendMessage(adminChatId, `Upload session expired. Please try uploading the file again.`, botToken);
return;
} else if (chunkResponse.status >= 500) {
// Server error, retry
retryCount++;
if (retryCount < MAX_RETRIES) {
await new Promise((resolve) => setTimeout(resolve, 1000 * retryCount)); // Exponential backoff
await sendMessage(adminChatId, `Server error, retrying chunk upload... (${retryCount}/${MAX_RETRIES})`, botToken);
}
} else { } else {
// Other error, don't retry await sendMessage(adminChatId, `Error during resumable upload chunk: ${await uploadResponse.text()}`, botToken);
const error = await chunkResponse.text();
await sendMessage(adminChatId, `Error uploading chunk: ${error}`, botToken);
return; return;
} }
} catch (error) {
retryCount++;
if (retryCount < MAX_RETRIES) {
await new Promise((resolve) => setTimeout(resolve, 1000 * retryCount)); // Exponential backoff
await sendMessage(adminChatId, `Network error, retrying chunk upload... (${retryCount}/${MAX_RETRIES})`, botToken);
} else {
await sendMessage(adminChatId, `Network error uploading chunk after ${MAX_RETRIES} retries: ${error.message}`, botToken);
return;
}
}
}
if (!success) {
await sendMessage(adminChatId, `Failed to upload chunk after ${MAX_RETRIES} retries`, botToken);
return;
}
}
}
async function getUploadStatus(resumableUploadUrl, jwtToken) {
try {
const statusResponse = await fetch(resumableUploadUrl, {
method: 'PUT',
headers: {
Authorization: `Bearer ${jwtToken}`,
'Content-Range': 'bytes */*',
},
});
if (statusResponse.status === 308) {
const rangeHeader = statusResponse.headers.get('Range');
if (rangeHeader) {
const rangeMatch = rangeHeader.match(/bytes=0-(\d+)/);
if (rangeMatch) {
return parseInt(rangeMatch[1]) + 1;
}
}
}
return 0;
} catch (error) {
return 0;
} }
} }
test/upload.test.js
-72
View File
@@ -1,72 +0,0 @@
import { beforeEach, describe, expect, it, vi } from 'vitest';
describe('File Upload Tests', () => {
beforeEach(() => {
// Mock global fetch
global.fetch = vi.fn();
global.FormData = vi.fn(() => ({
append: vi.fn(),
}));
global.Blob = vi.fn();
global.crypto = {
subtle: {
importKey: vi.fn(),
sign: vi.fn(),
},
};
});
it('should use regular upload for files under 20MB', async () => {
// Mock a 10MB file
const mockFileData = new ArrayBuffer(10 * 1024 * 1024);
// Mock fetch responses
global.fetch
.mockResolvedValueOnce({
json: () => Promise.resolve({ ok: true, result: { file_path: 'test/file.jpg' } }),
})
.mockResolvedValueOnce({
arrayBuffer: () => Promise.resolve(mockFileData),
})
.mockResolvedValueOnce({
json: () => Promise.resolve({ access_token: 'test_token' }),
})
.mockResolvedValueOnce({
json: () => Promise.resolve({ id: 'test_file_id' }),
});
// This would test the uploadToDrive function
// Since we can't easily import the function due to the export default structure,
// this serves as a placeholder for testing structure
expect(mockFileData.byteLength).toBeLessThan(20 * 1024 * 1024);
});
it('should use resumable upload for files over 20MB', async () => {
// Mock a 25MB file
const mockFileData = new ArrayBuffer(25 * 1024 * 1024);
expect(mockFileData.byteLength).toBeGreaterThan(20 * 1024 * 1024);
});
it('should handle chunk sizes correctly', () => {
const CHUNK_SIZE = 10 * 1024 * 1024;
const totalSize = 25 * 1024 * 1024;
// Calculate expected number of chunks
const expectedChunks = Math.ceil(totalSize / CHUNK_SIZE);
expect(expectedChunks).toBe(3);
// Test chunk boundaries
const chunks = [];
for (let i = 0; i < totalSize; i += CHUNK_SIZE) {
const start = i;
const end = Math.min(i + CHUNK_SIZE, totalSize);
chunks.push({ start, end, size: end - start });
}
expect(chunks).toHaveLength(3);
expect(chunks[0].size).toBe(CHUNK_SIZE);
expect(chunks[1].size).toBe(CHUNK_SIZE);
expect(chunks[2].size).toBe(totalSize - 2 * CHUNK_SIZE);
});
});