Add: The crawlFile API adds a beforeSave callback function, allowing users to customize the storage path and name of each requested file

Author: coder-hxl

src/api.ts

@@ -13,7 +13,8 @@ import {
   log,
   logError,
   logSuccess,
-  logWarn
+  logWarn,
+  mkdirDirSync
 } from './utils'
 
 import {
@@ -146,7 +147,7 @@ function loaderCommonConfig(
       priority = 0
     }
 
-    return { url, timeout, proxy, maxRetry, priority }
+    return { ...requestConfig, url, timeout, proxy, maxRetry, priority }
   })
 
   // 2.intervalTime
@@ -258,6 +259,23 @@ function loaderFileConfig(
   // 装载公共配置到 loaderConfig
   loaderCommonConfig(baseConfig, requestObjecs, loaderConfig)
 
+  // 装载单独的配置
+  loaderConfig.requestConfigs.forEach((requestConfig) => {
+    if (
+      isUndefined(requestConfig.storeDir) &&
+      !isUndefined(rawConfig.fileConfig?.storeDir)
+    ) {
+      requestConfig.storeDir = rawConfig.fileConfig!.storeDir
+    }
+
+    if (
+      isUndefined(requestConfig.extension) &&
+      !isUndefined(rawConfig.fileConfig?.extension)
+    ) {
+      requestConfig.extension = rawConfig.fileConfig!.extension
+    }
+  })
+
   return loaderConfig
 }
 
@@ -494,10 +512,6 @@ export function createCrawlFile(baseConfig: LoaderXCrawlBaseConfig) {
       config
     )
 
-    if (!fs.existsSync(fileConfig.storeDir)) {
-      fs.mkdirSync(fileConfig.storeDir)
-    }
-
     const controllerRes = await controller(
       'file',
       baseConfig.mode,
@@ -517,7 +531,8 @@ export function createCrawlFile(baseConfig: LoaderXCrawlBaseConfig) {
         maxRetry,
         crawlCount,
         errorQueue,
-        crawlSingleRes
+        crawlSingleRes,
+        requestConfig
       } = item
 
       const crawlRes: CrawlFileSingleRes = {
@@ -532,14 +547,47 @@ export function createCrawlFile(baseConfig: LoaderXCrawlBaseConfig) {
 
       if (isSuccess && crawlSingleRes) {
         const mimeType = crawlSingleRes.headers['content-type'] ?? ''
-        const fileExtension = fileConfig.extension ?? mimeType.split('/').pop()
-        const fileName = new Date().getTime().toString()
-        const filePath = path.resolve(
-          fileConfig.storeDir,
-          `${fileName}.${fileExtension}`
-        )
+        let fileName = ''
+        let fileExtension = ''
 
-        const saveFileItem = writeFile(filePath, crawlSingleRes.data)
+        if (!isUndefined(requestConfig.fileName)) {
+          fileName = requestConfig.fileName
+        } else {
+          fileName = new Date().getTime().toString()
+        }
+
+        if (!isUndefined(requestConfig.extension)) {
+          fileExtension = requestConfig.extension
+        } else {
+          fileExtension = '.' + mimeType.split('/').pop()
+        }
+
+        if (
+          !isUndefined(requestConfig.storeDir) &&
+          !fs.existsSync(requestConfig.storeDir)
+        ) {
+          mkdirDirSync(requestConfig.storeDir)
+        }
+
+        const storePath = requestConfig.storeDir ?? __dirname
+        const filePath = path.resolve(storePath, fileName + fileExtension)
+
+        // 在保存前的回调
+        let data = crawlSingleRes.data
+        if (fileConfig?.beforeSave) {
+          const newData = fileConfig.beforeSave({
+            id,
+            fileName,
+            filePath,
+            data
+          })
+
+          if (newData) {
+            data = newData
+          }
+        }
+
+        const saveFileItem = writeFile(filePath, data)
           .catch((err) => {
             const message = `File save error at id ${id}: ${err.message}`
             const valueOf = () => id
@@ -551,9 +599,18 @@ export function createCrawlFile(baseConfig: LoaderXCrawlBaseConfig) {
           .then((isError) => {
             const size = crawlSingleRes.data.length
             const isSuccess = !isError
-            const fileInfo = { isSuccess, fileName, mimeType, size, filePath }
 
-            crawlRes.data = { ...crawlSingleRes, data: fileInfo }
+            crawlRes.data = {
+              ...crawlSingleRes,
+              data: {
+                isSuccess,
+                fileName,
+                fileExtension,
+                mimeType,
+                size,
+                filePath
+              }
+            }
 
             if (callback) {
               callback(crawlRes)

src/types/api.ts

@@ -94,6 +94,9 @@ export interface FileRequestConfig {
   proxy?: string
   maxRetry?: number
   priority?: number
+  storeDir?: string
+  fileName?: string
+  extension?: string
 }
 
 // CrawlConfig
@@ -134,9 +137,15 @@ export interface CrawlFileConfig<R extends CrawlFileRequestConfig> {
   timeout?: number
   intervalTime?: IntervalTime
   maxRetry?: number
-  fileConfig: {
-    storeDir: string
+  fileConfig?: {
+    storeDir?: string
     extension?: string
+    beforeSave?: (info: {
+      id: number
+      fileName: string
+      filePath: string
+      data: Buffer
+    }) => Buffer | void
   }
 }
 
@@ -179,6 +188,7 @@ export interface CrawlFileSingleRes extends CrawlCommonRes {
     data: {
       isSuccess: boolean
       fileName: string
+      fileExtension: string
       mimeType: string
       size: number
       filePath: string

src/utils.ts

@@ -1,3 +1,5 @@
+import fs from 'node:fs'
+import path from 'node:path'
 import chalk from 'chalk'
 
 export function sleep(timeout: number) {
@@ -14,6 +16,20 @@ export function random(max: number, min = 0) {
   return res
 }
 
+export function mkdirDirSync(dir: string) {
+  const dirSplit = path.resolve(dir).split(path.sep)
+
+  dirSplit.reduce((prev, item, index) => {
+    const currentDir = index !== 0 ? path.join(prev, item) : item
+
+    if (!fs.existsSync(currentDir)) {
+      fs.mkdirSync(currentDir)
+    }
+
+    return currentDir
+  }, '')
+}
+
 export const log = console.log
 export const logNumber = chalk.hex('#a57fff')
 export const logSuccess = chalk.green

test/environment/crawlFile.test.ts

@@ -103,6 +103,40 @@ async function loaderAPIConfig() {
   return res.reduce((prev, item) => prev && item.isSuccess, true)
 }
 
+/* 3.Store Config */
+async function storeConfig() {
+  const testXCrawl = xCrawl({
+    baseUrl:
+      'https://raw.githubusercontent.com/coder-hxl/airbnb-upload/master/area',
+    proxy: 'http://localhost:14892'
+  })
+
+  const record: string[] = []
+  const res = await testXCrawl.crawlFile({
+    requestConfig: [
+      { url: '/4401.jpg', fileName: '4401' },
+      { url: '/4403.jpg', fileName: '4403' }
+    ],
+    fileConfig: {
+      storeDir: path.resolve(__dirname, './upload'),
+      extension: '.jpg',
+      beforeSave(info) {
+        record.push(info.fileName)
+      }
+    }
+  })
+
+  let isSuccess = true
+  res.forEach((item) => {
+    if (isSuccess) {
+      const hasName = record.includes(item.data?.data.fileName ?? '')
+      isSuccess = item.isSuccess && hasName
+    }
+  })
+
+  return isSuccess
+}
+
 /* 1.Written */
 test('crawlFile - writtenString', async () => {
   console.log(
@@ -147,3 +181,11 @@ test('crawlFile - loaderAPIConfig', async () => {
   )
   await expect(loaderAPIConfig()).resolves.toBe(true)
 })
+
+/* 2.Store Config */
+test('crawlFile - loaderAPIConfig', async () => {
+  console.log(
+    chalk.bgGreen('================ crawlFile - storeConfig ================')
+  )
+  await expect(storeConfig()).resolves.toBe(true)
+})