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 }