refactor(file-uploader): enhance file handling and improve type safety

- Update getFileSize utility to format file sizes consistently.
- Refactor traverseFileEntry to use specific FileSystemEntry types for better type safety.
- Modify useFileUpload hook to accept an optional allowedExtensions parameter for custom file type filtering.
- Update tests to reflect changes in file type handling and ensure comprehensive coverage.
This commit is contained in:
CodingOnStar
2026-01-29 16:56:15 +08:00
parent f88afd72ab
commit 1ab9e0b8a3
9 changed files with 130 additions and 280 deletions

View File

@ -216,13 +216,22 @@ describe('image-uploader utils', () => {
type FileCallback = (file: MockFile) => void
type EntriesCallback = (entries: FileSystemEntry[]) => void
// Helper to create mock FileSystemEntry with required properties
const createMockEntry = (props: {
isFile: boolean
isDirectory: boolean
name?: string
file?: (callback: FileCallback) => void
createReader?: () => { readEntries: (callback: EntriesCallback) => void }
}): FileSystemEntry => props as unknown as FileSystemEntry
it('should resolve with file array for file entry', async () => {
const mockFile: MockFile = { name: 'test.png' }
const mockEntry = {
const mockEntry = createMockEntry({
isFile: true,
isDirectory: false,
file: (callback: FileCallback) => callback(mockFile),
}
})
const result = await traverseFileEntry(mockEntry)
expect(result).toHaveLength(1)
@ -232,11 +241,11 @@ describe('image-uploader utils', () => {
it('should resolve with file array with prefix for nested file', async () => {
const mockFile: MockFile = { name: 'test.png' }
const mockEntry = {
const mockEntry = createMockEntry({
isFile: true,
isDirectory: false,
file: (callback: FileCallback) => callback(mockFile),
}
})
const result = await traverseFileEntry(mockEntry, 'folder/')
expect(result).toHaveLength(1)
@ -244,24 +253,24 @@ describe('image-uploader utils', () => {
})
it('should resolve empty array for unknown entry type', async () => {
const mockEntry = {
const mockEntry = createMockEntry({
isFile: false,
isDirectory: false,
}
})
const result = await traverseFileEntry(mockEntry)
expect(result).toEqual([])
})
it('should handle directory with no files', async () => {
const mockEntry = {
const mockEntry = createMockEntry({
isFile: false,
isDirectory: true,
name: 'empty-folder',
createReader: () => ({
readEntries: (callback: EntriesCallback) => callback([]),
}),
}
})
const result = await traverseFileEntry(mockEntry)
expect(result).toEqual([])
@ -271,20 +280,20 @@ describe('image-uploader utils', () => {
const mockFile1: MockFile = { name: 'file1.png' }
const mockFile2: MockFile = { name: 'file2.png' }
const mockFileEntry1 = {
const mockFileEntry1 = createMockEntry({
isFile: true,
isDirectory: false,
file: (callback: FileCallback) => callback(mockFile1),
}
})
const mockFileEntry2 = {
const mockFileEntry2 = createMockEntry({
isFile: true,
isDirectory: false,
file: (callback: FileCallback) => callback(mockFile2),
}
})
let readCount = 0
const mockEntry = {
const mockEntry = createMockEntry({
isFile: false,
isDirectory: true,
name: 'folder',
@ -292,14 +301,14 @@ describe('image-uploader utils', () => {
readEntries: (callback: EntriesCallback) => {
if (readCount === 0) {
readCount++
callback([mockFileEntry1, mockFileEntry2] as unknown as FileSystemEntry[])
callback([mockFileEntry1, mockFileEntry2])
}
else {
callback([])
}
},
}),
}
})
const result = await traverseFileEntry(mockEntry)
expect(result).toHaveLength(2)

View File

@ -14,28 +14,21 @@ export const getFileType = (currentFile: File) => {
return arr[arr.length - 1]
}
export const getFileSize = (size: number): string => {
if (size / 1024 < 10)
return `${(size / 1024).toFixed(2)}KB`
return `${(size / 1024 / 1024).toFixed(2)}MB`
}
type FileWithPath = {
relativePath?: string
} & File
export const traverseFileEntry = (entry: any, prefix = ''): Promise<FileWithPath[]> => {
export const traverseFileEntry = (entry: FileSystemEntry, prefix = ''): Promise<FileWithPath[]> => {
return new Promise((resolve) => {
if (entry.isFile) {
entry.file((file: FileWithPath) => {
(entry as FileSystemFileEntry).file((file: FileWithPath) => {
file.relativePath = `${prefix}${file.name}`
resolve([file])
})
}
else if (entry.isDirectory) {
const reader = entry.createReader()
const entries: any[] = []
const reader = (entry as FileSystemDirectoryEntry).createReader()
const entries: FileSystemEntry[] = []
const read = () => {
reader.readEntries(async (results: FileSystemEntry[]) => {
if (!results.length) {

View File

@ -178,7 +178,7 @@ export const useDSLImport = ({
if (pipeline_id)
await handleCheckPluginDependencies(pipeline_id, true)
push(`datasets/${dataset_id}/pipeline`)
push(`/datasets/${dataset_id}/pipeline`)
}
else if (status === DSLImportStatus.FAILED) {
notify({ type: 'error', message: t('creation.errorTip', { ns: 'datasetPipeline' }) })

View File

@ -1,5 +1,5 @@
import type { ReactNode } from 'react'
import type { CustomFile as File, FileItem } from '@/models/datasets'
import type { CustomFile, FileItem } from '@/models/datasets'
import { act, render, renderHook, waitFor } from '@testing-library/react'
import { beforeEach, describe, expect, it, vi } from 'vitest'
import { ToastContext } from '@/app/components/base/toast'
@ -114,7 +114,7 @@ describe('useFileUpload', () => {
() => useFileUpload({
...defaultOptions,
supportBatchUpload: false,
fileList: [{ fileID: 'file-1', file: {} as File, progress: 100 }],
fileList: [{ fileID: 'file-1', file: {} as CustomFile, progress: 100 }],
}),
{ wrapper: createWrapper() },
)
@ -201,7 +201,7 @@ describe('useFileUpload', () => {
{ wrapper: createWrapper() },
)
const mockFile = { id: 'file-123', name: 'test.pdf', size: 1024 } as File
const mockFile = { id: 'file-123', name: 'test.pdf', size: 1024 } as CustomFile
act(() => {
result.current.handlePreview(mockFile)
@ -217,7 +217,7 @@ describe('useFileUpload', () => {
{ wrapper: createWrapper() },
)
const mockFile = { name: 'test.pdf', size: 1024 } as File
const mockFile = { name: 'test.pdf', size: 1024 } as CustomFile
act(() => {
result.current.handlePreview(mockFile)
@ -862,7 +862,7 @@ describe('useFileUpload', () => {
it('should reject when total files exceed limit', () => {
const existingFiles: FileItem[] = Array.from({ length: 8 }, (_, i) => ({
fileID: `existing-${i}`,
file: { name: `existing-${i}.pdf`, size: 1024 } as File,
file: { name: `existing-${i}.pdf`, size: 1024 } as CustomFile,
progress: 100,
}))

View File

@ -27,6 +27,11 @@ export type UseFileUploadOptions = {
onFileListUpdate?: (files: FileItem[]) => void
onPreview: (file: File) => void
supportBatchUpload?: boolean
/**
* Optional list of allowed file extensions. If not provided, fetches from API.
* Pass this when you need custom extension filtering instead of using the global config.
*/
allowedExtensions?: string[]
}
export type UseFileUploadReturn = {
@ -62,6 +67,7 @@ export const useFileUpload = ({
onFileListUpdate,
onPreview,
supportBatchUpload = false,
allowedExtensions,
}: UseFileUploadOptions): UseFileUploadReturn => {
const { t } = useTranslation()
const { notify } = useContext(ToastContext)
@ -77,9 +83,10 @@ export const useFileUpload = ({
const { data: fileUploadConfigResponse } = useFileUploadConfig()
const { data: supportFileTypesResponse } = useFileSupportTypes()
// Use provided allowedExtensions or fetch from API
const supportTypes = useMemo(
() => supportFileTypesResponse?.allowed_extensions || [],
[supportFileTypesResponse?.allowed_extensions],
() => allowedExtensions ?? supportFileTypesResponse?.allowed_extensions ?? [],
[allowedExtensions, supportFileTypesResponse?.allowed_extensions],
)
const supportTypesShowNames = useMemo(() => {

View File

@ -3,10 +3,11 @@ import { RiDeleteBinLine, RiErrorWarningFill } from '@remixicon/react'
import dynamic from 'next/dynamic'
import { useMemo } from 'react'
import DocumentFileIcon from '@/app/components/datasets/common/document-file-icon'
import { getFileSize, getFileType } from '@/app/components/datasets/common/image-uploader/utils'
import { getFileType } from '@/app/components/datasets/common/image-uploader/utils'
import useTheme from '@/hooks/use-theme'
import { Theme } from '@/types/app'
import { cn } from '@/utils/classnames'
import { formatFileSize } from '@/utils/format'
import { PROGRESS_ERROR } from '../constants'
const SimplePieChart = dynamic(() => import('@/app/components/base/simple-pie-chart'), { ssr: false })
@ -60,7 +61,7 @@ const FileListItem = ({
<div className="w-full truncate text-2xs leading-3 text-text-tertiary">
<span className="uppercase">{getFileType(fileItem.file)}</span>
<span className="px-1 text-text-quaternary">·</span>
<span>{getFileSize(fileItem.file.size)}</span>
<span>{formatFileSize(fileItem.file.size)}</span>
</div>
</div>
<div className="flex w-16 shrink-0 items-center justify-end gap-1 pr-3">

View File

@ -23,6 +23,14 @@ vi.mock('@/app/components/base/file-uploader/utils', () => ({
getFileUploadErrorMessage: (e: Error, defaultMsg: string) => e.message || defaultMsg,
}))
// Mock format utils used by the shared hook
vi.mock('@/utils/format', () => ({
getFileExtension: (filename: string) => {
const parts = filename.split('.')
return parts[parts.length - 1] || ''
},
}))
// Mock react-i18next
vi.mock('react-i18next', () => ({
useTranslation: () => ({
@ -70,6 +78,12 @@ vi.mock('@/service/use-common', () => ({
file_upload_limit: 10,
},
})),
// Required by the shared useFileUpload hook
useFileSupportTypes: vi.fn(() => ({
data: {
allowed_extensions: ['pdf', 'docx', 'txt'],
},
})),
}))
// Mock upload service
@ -629,8 +643,17 @@ describe('useLocalFileUpload', () => {
const mockFile = new File(['content'], 'test.pdf', { type: 'application/pdf' })
await act(async () => {
const dropEvent = new Event('drop', { bubbles: true, cancelable: true }) as Event & { dataTransfer: { files: File[] } | null }
dropEvent.dataTransfer = { files: [mockFile] }
const dropEvent = new Event('drop', { bubbles: true, cancelable: true }) as Event & {
dataTransfer: { items: DataTransferItem[], files: File[] } | null
}
// Mock dataTransfer with items array (used by the shared hook for directory traversal)
dropEvent.dataTransfer = {
items: [{
kind: 'file',
getAsFile: () => mockFile,
}] as unknown as DataTransferItem[],
files: [mockFile],
}
dropzone.dispatchEvent(dropEvent)
})
@ -679,8 +702,17 @@ describe('useLocalFileUpload', () => {
]
await act(async () => {
const dropEvent = new Event('drop', { bubbles: true, cancelable: true }) as Event & { dataTransfer: { files: File[] } | null }
dropEvent.dataTransfer = { files }
const dropEvent = new Event('drop', { bubbles: true, cancelable: true }) as Event & {
dataTransfer: { items: DataTransferItem[], files: File[] } | null
}
// Mock dataTransfer with items array (used by the shared hook for directory traversal)
dropEvent.dataTransfer = {
items: files.map(f => ({
kind: 'file',
getAsFile: () => f,
})) as unknown as DataTransferItem[],
files,
}
dropzone.dispatchEvent(dropEvent)
})

View File

@ -1,271 +1,84 @@
import type { CustomFile as File, FileItem } from '@/models/datasets'
import { produce } from 'immer'
import { useCallback, useEffect, useMemo, useRef, useState } from 'react'
import { useTranslation } from 'react-i18next'
import { useContext } from 'use-context-selector'
import { getFileUploadErrorMessage } from '@/app/components/base/file-uploader/utils'
import { ToastContext } from '@/app/components/base/toast'
import { IS_CE_EDITION } from '@/config'
import { useLocale } from '@/context/i18n'
import { LanguagesSupported } from '@/i18n-config/language'
import { upload } from '@/service/base'
import { useFileUploadConfig } from '@/service/use-common'
import { useCallback, useRef } from 'react'
import { useFileUpload } from '@/app/components/datasets/create/file-uploader/hooks/use-file-upload'
import { useDataSourceStore, useDataSourceStoreWithSelector } from '../../store'
import { PROGRESS_ERROR, PROGRESS_NOT_STARTED } from '../constants'
export type UseLocalFileUploadOptions = {
allowedExtensions: string[]
supportBatchUpload?: boolean
}
type FileUploadConfig = {
file_size_limit: number
batch_count_limit: number
file_upload_limit: number
}
/**
* Hook for handling local file uploads in the create-from-pipeline flow.
* This is a thin wrapper around the generic useFileUpload hook that provides
* Zustand store integration for state management.
*/
export const useLocalFileUpload = ({
allowedExtensions,
supportBatchUpload = true,
}: UseLocalFileUploadOptions) => {
const { t } = useTranslation()
const { notify } = useContext(ToastContext)
const locale = useLocale()
const localFileList = useDataSourceStoreWithSelector(state => state.localFileList)
const dataSourceStore = useDataSourceStore()
const [dragging, setDragging] = useState(false)
const dropRef = useRef<HTMLDivElement>(null)
const dragRef = useRef<HTMLDivElement>(null)
const fileUploaderRef = useRef<HTMLInputElement>(null)
const fileListRef = useRef<FileItem[]>([])
const hideUpload = !supportBatchUpload && localFileList.length > 0
// Sync fileListRef with localFileList for internal tracking
fileListRef.current = localFileList
const { data: fileUploadConfigResponse } = useFileUploadConfig()
const prepareFileList = useCallback((files: FileItem[]) => {
const { setLocalFileList } = dataSourceStore.getState()
setLocalFileList(files)
fileListRef.current = files
}, [dataSourceStore])
const supportTypesShowNames = useMemo(() => {
const extensionMap: { [key: string]: string } = {
md: 'markdown',
pptx: 'pptx',
htm: 'html',
xlsx: 'xlsx',
docx: 'docx',
}
return allowedExtensions
.map(item => extensionMap[item] || item)
.map(item => item.toLowerCase())
.filter((item, index, self) => self.indexOf(item) === index)
.map(item => item.toUpperCase())
.join(locale !== LanguagesSupported[1] ? ', ' : '、 ')
}, [locale, allowedExtensions])
const acceptTypes = useMemo(
() => allowedExtensions.map((ext: string) => `.${ext}`),
[allowedExtensions],
)
const fileUploadConfig: FileUploadConfig = useMemo(() => ({
file_size_limit: fileUploadConfigResponse?.file_size_limit ?? 15,
batch_count_limit: supportBatchUpload ? (fileUploadConfigResponse?.batch_count_limit ?? 5) : 1,
file_upload_limit: supportBatchUpload ? (fileUploadConfigResponse?.file_upload_limit ?? 5) : 1,
}), [fileUploadConfigResponse, supportBatchUpload])
const updateFile = useCallback((fileItem: FileItem, progress: number, list: FileItem[]) => {
const onFileUpdate = useCallback((fileItem: FileItem, progress: number, list: FileItem[]) => {
const { setLocalFileList } = dataSourceStore.getState()
const newList = produce(list, (draft) => {
const targetIndex = draft.findIndex(file => file.fileID === fileItem.fileID)
draft[targetIndex] = {
...draft[targetIndex],
progress,
if (targetIndex !== -1) {
draft[targetIndex] = {
...draft[targetIndex],
...fileItem,
progress,
}
}
})
setLocalFileList(newList)
}, [dataSourceStore])
const updateFileList = useCallback((preparedFiles: FileItem[]) => {
const onFileListUpdate = useCallback((files: FileItem[]) => {
const { setLocalFileList } = dataSourceStore.getState()
setLocalFileList(preparedFiles)
setLocalFileList(files)
fileListRef.current = files
}, [dataSourceStore])
const handlePreview = useCallback((file: File) => {
const onPreview = useCallback((file: File) => {
const { setCurrentLocalFile } = dataSourceStore.getState()
if (file.id)
setCurrentLocalFile(file)
setCurrentLocalFile(file)
}, [dataSourceStore])
const getFileType = (currentFile: File) => {
if (!currentFile)
return ''
const arr = currentFile.name.split('.')
return arr[arr.length - 1]
}
const isValid = useCallback((file: File) => {
const { size } = file
const ext = `.${getFileType(file)}`
const isValidType = acceptTypes.includes(ext.toLowerCase())
if (!isValidType)
notify({ type: 'error', message: t('stepOne.uploader.validation.typeError', { ns: 'datasetCreation' }) })
const isValidSize = size <= fileUploadConfig.file_size_limit * 1024 * 1024
if (!isValidSize)
notify({ type: 'error', message: t('stepOne.uploader.validation.size', { ns: 'datasetCreation', size: fileUploadConfig.file_size_limit }) })
return isValidType && isValidSize
}, [notify, t, acceptTypes, fileUploadConfig.file_size_limit])
type UploadResult = Awaited<ReturnType<typeof upload>>
const fileUpload = useCallback(async (fileItem: FileItem): Promise<FileItem> => {
const formData = new FormData()
formData.append('file', fileItem.file)
const onProgress = (e: ProgressEvent) => {
if (e.lengthComputable) {
const percent = Math.floor(e.loaded / e.total * 100)
updateFile(fileItem, percent, fileListRef.current)
}
}
return upload({
xhr: new XMLHttpRequest(),
data: formData,
onprogress: onProgress,
}, false, undefined, '?source=datasets')
.then((res: UploadResult) => {
const updatedFile = Object.assign({}, fileItem.file, {
id: res.id,
...(res as Partial<File>),
}) as File
const completeFile: FileItem = {
fileID: fileItem.fileID,
file: updatedFile,
progress: PROGRESS_NOT_STARTED,
}
const index = fileListRef.current.findIndex(item => item.fileID === fileItem.fileID)
fileListRef.current[index] = completeFile
updateFile(completeFile, 100, fileListRef.current)
return Promise.resolve({ ...completeFile })
})
.catch((e) => {
const errorMessage = getFileUploadErrorMessage(e, t('stepOne.uploader.failed', { ns: 'datasetCreation' }), t)
notify({ type: 'error', message: errorMessage })
updateFile(fileItem, PROGRESS_ERROR, fileListRef.current)
return Promise.resolve({ ...fileItem })
})
}, [fileListRef, notify, updateFile, t])
const uploadBatchFiles = useCallback((bFiles: FileItem[]) => {
bFiles.forEach(bf => (bf.progress = 0))
return Promise.all(bFiles.map(fileUpload))
}, [fileUpload])
const uploadMultipleFiles = useCallback(async (files: FileItem[]) => {
const batchCountLimit = fileUploadConfig.batch_count_limit
const length = files.length
let start = 0
let end = 0
while (start < length) {
if (start + batchCountLimit > length)
end = length
else
end = start + batchCountLimit
const bFiles = files.slice(start, end)
await uploadBatchFiles(bFiles)
start = end
}
}, [fileUploadConfig, uploadBatchFiles])
const initialUpload = useCallback((files: File[]) => {
const filesCountLimit = fileUploadConfig.file_upload_limit
if (!files.length)
return false
if (files.length + localFileList.length > filesCountLimit && !IS_CE_EDITION) {
notify({ type: 'error', message: t('stepOne.uploader.validation.filesNumber', { ns: 'datasetCreation', filesNumber: filesCountLimit }) })
return false
}
const preparedFiles = files.map((file, index) => ({
fileID: `file${index}-${Date.now()}`,
file,
progress: PROGRESS_NOT_STARTED,
}))
const newFiles = [...fileListRef.current, ...preparedFiles]
updateFileList(newFiles)
fileListRef.current = newFiles
uploadMultipleFiles(preparedFiles)
}, [fileUploadConfig.file_upload_limit, localFileList.length, updateFileList, uploadMultipleFiles, notify, t])
const handleDragEnter = useCallback((e: DragEvent) => {
e.preventDefault()
e.stopPropagation()
if (e.target !== dragRef.current)
setDragging(true)
}, [])
const handleDragOver = useCallback((e: DragEvent) => {
e.preventDefault()
e.stopPropagation()
}, [])
const handleDragLeave = useCallback((e: DragEvent) => {
e.preventDefault()
e.stopPropagation()
if (e.target === dragRef.current)
setDragging(false)
}, [])
const handleDrop = useCallback((e: DragEvent) => {
e.preventDefault()
e.stopPropagation()
setDragging(false)
if (!e.dataTransfer)
return
let files = Array.from(e.dataTransfer.files) as File[]
if (!supportBatchUpload)
files = files.slice(0, 1)
const validFiles = files.filter(isValid)
initialUpload(validFiles)
}, [initialUpload, isValid, supportBatchUpload])
const selectHandle = useCallback(() => {
if (fileUploaderRef.current)
fileUploaderRef.current.click()
}, [])
const removeFile = useCallback((fileID: string) => {
if (fileUploaderRef.current)
fileUploaderRef.current.value = ''
fileListRef.current = fileListRef.current.filter(item => item.fileID !== fileID)
updateFileList([...fileListRef.current])
}, [updateFileList])
const fileChangeHandle = useCallback((e: React.ChangeEvent<HTMLInputElement>) => {
let files = Array.from(e.target.files ?? []) as File[]
files = files.slice(0, fileUploadConfig.batch_count_limit)
initialUpload(files.filter(isValid))
}, [isValid, initialUpload, fileUploadConfig.batch_count_limit])
useEffect(() => {
const dropElement = dropRef.current
dropElement?.addEventListener('dragenter', handleDragEnter)
dropElement?.addEventListener('dragover', handleDragOver)
dropElement?.addEventListener('dragleave', handleDragLeave)
dropElement?.addEventListener('drop', handleDrop)
return () => {
dropElement?.removeEventListener('dragenter', handleDragEnter)
dropElement?.removeEventListener('dragover', handleDragOver)
dropElement?.removeEventListener('dragleave', handleDragLeave)
dropElement?.removeEventListener('drop', handleDrop)
}
}, [handleDragEnter, handleDragOver, handleDragLeave, handleDrop])
const {
dropRef,
dragRef,
fileUploaderRef,
dragging,
fileUploadConfig,
acceptTypes,
supportTypesShowNames,
hideUpload,
selectHandle,
fileChangeHandle,
removeFile,
handlePreview,
} = useFileUpload({
fileList: localFileList,
prepareFileList,
onFileUpdate,
onFileListUpdate,
onPreview,
supportBatchUpload,
allowedExtensions,
})
return {
// Refs

View File

@ -1576,11 +1576,6 @@
"count": 3
}
},
"app/components/datasets/common/image-uploader/utils.ts": {
"ts/no-explicit-any": {
"count": 2
}
},
"app/components/datasets/common/retrieval-method-config/index.spec.tsx": {
"ts/no-explicit-any": {
"count": 1