github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/dm/pkg/utils/file.go (about)

     1  // Copyright 2019 PingCAP, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package utils
    15  
    16  import (
    17  	"io"
    18  	"os"
    19  	"path"
    20  	"path/filepath"
    21  	"strconv"
    22  	"strings"
    23  
    24  	"github.com/docker/go-units"
    25  	"github.com/pingcap/tiflow/dm/pkg/log"
    26  	"github.com/pingcap/tiflow/dm/pkg/terror"
    27  	"go.uber.org/zap"
    28  )
    29  
    30  // IsFileExists checks if file exists.
    31  func IsFileExists(name string) bool {
    32  	f, err := os.Stat(name)
    33  	if err != nil {
    34  		if os.IsNotExist(err) {
    35  			return false
    36  		}
    37  	}
    38  
    39  	if f.IsDir() {
    40  		return false
    41  	}
    42  
    43  	return true
    44  }
    45  
    46  // IsDirExists checks if dir exists.
    47  func IsDirExists(name string) bool {
    48  	f, err := os.Stat(name)
    49  	if err != nil {
    50  		if os.IsNotExist(err) {
    51  			return false
    52  		}
    53  	}
    54  
    55  	if !f.IsDir() {
    56  		return false
    57  	}
    58  
    59  	return true
    60  }
    61  
    62  // GetFileSize return the size of the file.
    63  // NOTE: do not support to get the size of the directory now.
    64  func GetFileSize(file string) (int64, error) {
    65  	if !IsFileExists(file) {
    66  		return 0, terror.ErrGetFileSize.Generate(file)
    67  	}
    68  
    69  	stat, err := os.Stat(file)
    70  	if err != nil {
    71  		return 0, terror.ErrGetFileSize.Delegate(err, file)
    72  	}
    73  	return stat.Size(), nil
    74  }
    75  
    76  // ParseFileSize parses the size in MiB from input.
    77  func ParseFileSize(fileSizeStr string, defaultSize uint64) (uint64, error) {
    78  	var fileSize uint64
    79  	if len(fileSizeStr) == 0 {
    80  		fileSize = defaultSize
    81  	} else if fileSizeMB, err := strconv.ParseUint(fileSizeStr, 10, 64); err == nil {
    82  		fileSize = fileSizeMB * units.MiB
    83  	} else if size, err := units.RAMInBytes(fileSizeStr); err == nil {
    84  		fileSize = uint64(size)
    85  	} else {
    86  		return 0, err
    87  	}
    88  	return fileSize, nil
    89  }
    90  
    91  // WriteFileAtomic writes file to temp and atomically move when everything else succeeds.
    92  func WriteFileAtomic(filename string, data []byte, perm os.FileMode) error {
    93  	dir, name := path.Dir(filename), path.Base(filename)
    94  	f, err := os.CreateTemp(dir, name)
    95  	if err != nil {
    96  		return err
    97  	}
    98  	n, err := f.Write(data)
    99  	f.Close()
   100  	if err == nil && n < len(data) {
   101  		err = io.ErrShortWrite
   102  	} else {
   103  		err = os.Chmod(f.Name(), perm)
   104  	}
   105  	if err != nil {
   106  		err2 := os.Remove(f.Name())
   107  		log.L().Warn("failed to remove the temporary file",
   108  			zap.String("filename", f.Name()),
   109  			zap.Error(err2))
   110  		return err
   111  	}
   112  	return os.Rename(f.Name(), filename)
   113  }
   114  
   115  // CollectDirFiles gets files in path.
   116  func CollectDirFiles(path string) (map[string]struct{}, error) {
   117  	files := make(map[string]struct{})
   118  	err := filepath.Walk(path, func(_ string, f os.FileInfo, err error) error {
   119  		if err != nil {
   120  			return err
   121  		}
   122  
   123  		if f == nil {
   124  			return nil
   125  		}
   126  
   127  		if f.IsDir() {
   128  			return nil
   129  		}
   130  
   131  		name := strings.TrimSpace(f.Name())
   132  		files[name] = struct{}{}
   133  		return nil
   134  	})
   135  
   136  	return files, err
   137  }
   138  
   139  // GetDBFromDumpFilename extracts db name from dump filename.
   140  func GetDBFromDumpFilename(filename string) (db string, ok bool) {
   141  	if !strings.HasSuffix(filename, "-schema-create.sql") {
   142  		return "", false
   143  	}
   144  
   145  	idx := strings.LastIndex(filename, "-schema-create.sql")
   146  	return filename[:idx], true
   147  }
   148  
   149  // GetTableFromDumpFilename extracts db and table name from dump filename.
   150  func GetTableFromDumpFilename(filename string) (db, table string, ok bool) {
   151  	if !strings.HasSuffix(filename, "-schema.sql") {
   152  		return "", "", false
   153  	}
   154  
   155  	idx := strings.LastIndex(filename, "-schema.sql")
   156  	name := filename[:idx]
   157  	fields := strings.Split(name, ".")
   158  	if len(fields) != 2 {
   159  		return "", "", false
   160  	}
   161  	return fields[0], fields[1], true
   162  }