github.com/influxdata/influxdb/v2@v2.7.6/pkg/tar/stream.go (about) 1 package tar 2 3 import ( 4 "archive/tar" 5 "fmt" 6 "io" 7 "os" 8 "path/filepath" 9 "strings" 10 "time" 11 12 "github.com/influxdata/influxdb/v2/pkg/file" 13 ) 14 15 // Stream is a convenience function for creating a tar of a shard dir. It walks over the directory and subdirs, 16 // possibly writing each file to a tar writer stream. By default StreamFile is used, which will result in all files 17 // being written. A custom writeFunc can be passed so that each file may be written, modified+written, or skipped 18 // depending on the custom logic. 19 func Stream(w io.Writer, dir, relativePath string, writeFunc func(f os.FileInfo, shardRelativePath, fullPath string, tw *tar.Writer) error) error { 20 tw := tar.NewWriter(w) 21 defer tw.Close() 22 23 if writeFunc == nil { 24 writeFunc = StreamFile 25 } 26 27 return filepath.WalkDir(dir, func(path string, entry os.DirEntry, err error) error { 28 if err != nil { 29 return err 30 } 31 32 // Skip adding an entry for the root dir 33 if dir == path && entry.IsDir() { 34 return nil 35 } 36 37 // Figure out the full relative path including any sub-dirs 38 subDir, _ := filepath.Split(path) 39 subDir, err = filepath.Rel(dir, subDir) 40 if err != nil { 41 return err 42 } 43 f, err := entry.Info() 44 if err != nil { 45 return err 46 } 47 48 return writeFunc(f, filepath.Join(relativePath, subDir), path, tw) 49 }) 50 } 51 52 // Generates a filtering function for Stream that checks an incoming file, and only writes the file to the stream if 53 // its mod time is later than since. Example: to tar only files newer than a certain datetime, use 54 // tar.Stream(w, dir, relativePath, SinceFilterTarFile(datetime)) 55 func SinceFilterTarFile(since time.Time) func(f os.FileInfo, shardRelativePath, fullPath string, tw *tar.Writer) error { 56 return func(f os.FileInfo, shardRelativePath, fullPath string, tw *tar.Writer) error { 57 if f.ModTime().After(since) { 58 return StreamFile(f, shardRelativePath, fullPath, tw) 59 } 60 return nil 61 } 62 } 63 64 // stream a single file to tw, extending the header name using the shardRelativePath 65 func StreamFile(f os.FileInfo, shardRelativePath, fullPath string, tw *tar.Writer) error { 66 return StreamRenameFile(f, f.Name(), shardRelativePath, fullPath, tw) 67 } 68 69 // Stream a single file to tw, using tarHeaderFileName instead of the actual filename 70 // e.g., when we want to write a *.tmp file using the original file's non-tmp name. 71 func StreamRenameFile(f os.FileInfo, tarHeaderFileName, relativePath, fullPath string, tw *tar.Writer) error { 72 h, err := tar.FileInfoHeader(f, f.Name()) 73 if err != nil { 74 return err 75 } 76 h.Name = filepath.ToSlash(filepath.Join(relativePath, tarHeaderFileName)) 77 78 if err := tw.WriteHeader(h); err != nil { 79 return err 80 } 81 82 if !f.Mode().IsRegular() { 83 return nil 84 } 85 86 fr, err := os.Open(fullPath) 87 if err != nil { 88 return err 89 } 90 91 defer fr.Close() 92 93 _, err = io.CopyN(tw, fr, h.Size) 94 95 return err 96 } 97 98 // Restore reads a tar archive from r and extracts all of its files into dir, 99 // using only the base name of each file. 100 func Restore(r io.Reader, dir string) error { 101 tr := tar.NewReader(r) 102 for { 103 if err := extractFile(tr, dir); err == io.EOF { 104 break 105 } else if err != nil { 106 return err 107 } 108 } 109 110 return file.SyncDir(dir) 111 } 112 113 // extractFile copies the next file from tr into dir, using the file's base name. 114 func extractFile(tr *tar.Reader, dir string) error { 115 // Read next archive file. 116 hdr, err := tr.Next() 117 if err != nil { 118 return err 119 } 120 121 // The hdr.Name is the relative path of the file from the root data dir. 122 // e.g (db/rp/1/xxxxx.tsm or db/rp/1/index/xxxxxx.tsi) 123 sections := strings.Split(filepath.FromSlash(hdr.Name), string(filepath.Separator)) 124 if len(sections) < 3 { 125 return fmt.Errorf("invalid archive path: %s", hdr.Name) 126 } 127 128 relativePath := filepath.Join(sections[3:]...) 129 130 subDir, _ := filepath.Split(relativePath) 131 // If this is a directory entry (usually just `index` for tsi), create it an move on. 132 if hdr.Typeflag == tar.TypeDir { 133 return os.MkdirAll(filepath.Join(dir, subDir), os.FileMode(hdr.Mode).Perm()) 134 } 135 136 // Make sure the dir we need to write into exists. It should, but just double check in 137 // case we get a slightly invalid tarball. 138 if subDir != "" { 139 if err := os.MkdirAll(filepath.Join(dir, subDir), 0755); err != nil { 140 return err 141 } 142 } 143 144 destPath := filepath.Join(dir, relativePath) 145 tmp := destPath + ".tmp" 146 147 // Create new file on disk. 148 f, err := os.OpenFile(tmp, os.O_CREATE|os.O_RDWR, os.FileMode(hdr.Mode).Perm()) 149 if err != nil { 150 return err 151 } 152 defer f.Close() 153 154 // Copy from archive to the file. 155 if _, err := io.CopyN(f, tr, hdr.Size); err != nil { 156 return err 157 } 158 159 // Sync to disk & close. 160 if err := f.Sync(); err != nil { 161 return err 162 } 163 164 if err := f.Close(); err != nil { 165 return err 166 } 167 168 return file.RenameFile(tmp, destPath) 169 }