github.com/cloudreve/Cloudreve/v3@v3.0.0-20240224133659-3edb00a6484c/pkg/task/import.go (about)

     1  package task
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"path"
     7  
     8  	model "github.com/cloudreve/Cloudreve/v3/models"
     9  	"github.com/cloudreve/Cloudreve/v3/pkg/filesystem"
    10  	"github.com/cloudreve/Cloudreve/v3/pkg/filesystem/fsctx"
    11  	"github.com/cloudreve/Cloudreve/v3/pkg/util"
    12  )
    13  
    14  // ImportTask 导入务
    15  type ImportTask struct {
    16  	User      *model.User
    17  	TaskModel *model.Task
    18  	TaskProps ImportProps
    19  	Err       *JobError
    20  }
    21  
    22  // ImportProps 导入任务属性
    23  type ImportProps struct {
    24  	PolicyID  uint   `json:"policy_id"`    // 存储策略ID
    25  	Src       string `json:"src"`          // 原始路径
    26  	Recursive bool   `json:"is_recursive"` // 是否递归导入
    27  	Dst       string `json:"dst"`          // 目的目录
    28  }
    29  
    30  // Props 获取任务属性
    31  func (job *ImportTask) Props() string {
    32  	res, _ := json.Marshal(job.TaskProps)
    33  	return string(res)
    34  }
    35  
    36  // Type 获取任务状态
    37  func (job *ImportTask) Type() int {
    38  	return ImportTaskType
    39  }
    40  
    41  // Creator 获取创建者ID
    42  func (job *ImportTask) Creator() uint {
    43  	return job.User.ID
    44  }
    45  
    46  // Model 获取任务的数据库模型
    47  func (job *ImportTask) Model() *model.Task {
    48  	return job.TaskModel
    49  }
    50  
    51  // SetStatus 设定状态
    52  func (job *ImportTask) SetStatus(status int) {
    53  	job.TaskModel.SetStatus(status)
    54  }
    55  
    56  // SetError 设定任务失败信息
    57  func (job *ImportTask) SetError(err *JobError) {
    58  	job.Err = err
    59  	res, _ := json.Marshal(job.Err)
    60  	job.TaskModel.SetError(string(res))
    61  }
    62  
    63  // SetErrorMsg 设定任务失败信息
    64  func (job *ImportTask) SetErrorMsg(msg string, err error) {
    65  	jobErr := &JobError{Msg: msg}
    66  	if err != nil {
    67  		jobErr.Error = err.Error()
    68  	}
    69  	job.SetError(jobErr)
    70  }
    71  
    72  // GetError 返回任务失败信息
    73  func (job *ImportTask) GetError() *JobError {
    74  	return job.Err
    75  }
    76  
    77  // Do 开始执行任务
    78  func (job *ImportTask) Do() {
    79  	ctx := context.Background()
    80  
    81  	// 查找存储策略
    82  	policy, err := model.GetPolicyByID(job.TaskProps.PolicyID)
    83  	if err != nil {
    84  		job.SetErrorMsg("Policy not exist.", err)
    85  		return
    86  	}
    87  
    88  	// 创建文件系统
    89  	job.User.Policy = policy
    90  	fs, err := filesystem.NewFileSystem(job.User)
    91  	if err != nil {
    92  		job.SetErrorMsg(err.Error(), nil)
    93  		return
    94  	}
    95  	defer fs.Recycle()
    96  
    97  	fs.Policy = &policy
    98  	if err := fs.DispatchHandler(); err != nil {
    99  		job.SetErrorMsg("Failed to dispatch policy.", err)
   100  		return
   101  	}
   102  
   103  	// 注册钩子
   104  	fs.Use("BeforeAddFile", filesystem.HookValidateFile)
   105  	fs.Use("BeforeAddFile", filesystem.HookValidateCapacity)
   106  
   107  	// 列取目录、对象
   108  	job.TaskModel.SetProgress(ListingProgress)
   109  	coxIgnoreConflict := context.WithValue(context.Background(), fsctx.IgnoreDirectoryConflictCtx,
   110  		true)
   111  	objects, err := fs.Handler.List(ctx, job.TaskProps.Src, job.TaskProps.Recursive)
   112  	if err != nil {
   113  		job.SetErrorMsg("Failed to list files.", err)
   114  		return
   115  	}
   116  
   117  	job.TaskModel.SetProgress(InsertingProgress)
   118  
   119  	// 虚拟目录路径与folder对象ID的对应
   120  	pathCache := make(map[string]*model.Folder, len(objects))
   121  
   122  	// 插入目录记录到用户文件系统
   123  	for _, object := range objects {
   124  		if object.IsDir {
   125  			// 创建目录
   126  			virtualPath := path.Join(job.TaskProps.Dst, object.RelativePath)
   127  			folder, err := fs.CreateDirectory(coxIgnoreConflict, virtualPath)
   128  			if err != nil {
   129  				util.Log().Warning("Importing task cannot create user directory %q: %s", virtualPath, err)
   130  			} else if folder.ID > 0 {
   131  				pathCache[virtualPath] = folder
   132  			}
   133  		}
   134  	}
   135  
   136  	// 插入文件记录到用户文件系统
   137  	for _, object := range objects {
   138  		if !object.IsDir {
   139  			// 创建文件信息
   140  			virtualPath := path.Dir(path.Join(job.TaskProps.Dst, object.RelativePath))
   141  			fileHeader := fsctx.FileStream{
   142  				Size:        object.Size,
   143  				VirtualPath: virtualPath,
   144  				Name:        object.Name,
   145  				SavePath:    object.Source,
   146  			}
   147  
   148  			// 查找父目录
   149  			parentFolder := &model.Folder{}
   150  			if parent, ok := pathCache[virtualPath]; ok {
   151  				parentFolder = parent
   152  			} else {
   153  				folder, err := fs.CreateDirectory(context.Background(), virtualPath)
   154  				if err != nil {
   155  					util.Log().Warning("Importing task cannot create user directory %q: %s",
   156  						virtualPath, err)
   157  					continue
   158  				}
   159  				parentFolder = folder
   160  
   161  			}
   162  
   163  			// 插入文件记录
   164  			_, err := fs.AddFile(context.Background(), parentFolder, &fileHeader)
   165  			if err != nil {
   166  				util.Log().Warning("Importing task cannot insert user file %q: %s",
   167  					object.RelativePath, err)
   168  				if err == filesystem.ErrInsufficientCapacity {
   169  					job.SetErrorMsg("Insufficient storage capacity.", err)
   170  					return
   171  				}
   172  			}
   173  
   174  		}
   175  	}
   176  }
   177  
   178  // NewImportTask 新建导入任务
   179  func NewImportTask(user, policy uint, src, dst string, recursive bool) (Job, error) {
   180  	creator, err := model.GetActiveUserByID(user)
   181  	if err != nil {
   182  		return nil, err
   183  	}
   184  
   185  	newTask := &ImportTask{
   186  		User: &creator,
   187  		TaskProps: ImportProps{
   188  			PolicyID:  policy,
   189  			Recursive: recursive,
   190  			Src:       src,
   191  			Dst:       dst,
   192  		},
   193  	}
   194  
   195  	record, err := Record(newTask)
   196  	if err != nil {
   197  		return nil, err
   198  	}
   199  	newTask.TaskModel = record
   200  
   201  	return newTask, nil
   202  }
   203  
   204  // NewImportTaskFromModel 从数据库记录中恢复导入任务
   205  func NewImportTaskFromModel(task *model.Task) (Job, error) {
   206  	user, err := model.GetActiveUserByID(task.UserID)
   207  	if err != nil {
   208  		return nil, err
   209  	}
   210  	newTask := &ImportTask{
   211  		User:      &user,
   212  		TaskModel: task,
   213  	}
   214  
   215  	err = json.Unmarshal([]byte(task.Props), &newTask.TaskProps)
   216  	if err != nil {
   217  		return nil, err
   218  	}
   219  
   220  	return newTask, nil
   221  }