github.com/cloudberrydb/gpbackup@v1.0.3-0.20240118031043-5410fd45eed6/utils/io.go (about)

     1  package utils
     2  
     3  /*
     4   * This file contains structs and functions related to interacting with files
     5   * and directories, both locally and remotely over SSH.
     6   */
     7  
     8  import (
     9  	"fmt"
    10  	"io"
    11  	"io/ioutil"
    12  	"os"
    13  
    14  	"github.com/cloudberrydb/gp-common-go-libs/gplog"
    15  )
    16  
    17  /*
    18   * Generic file/directory manipulation functions
    19   */
    20  
    21  func MustPrintf(file io.Writer, s string, v ...interface{}) uint64 {
    22  	bytesWritten, err := fmt.Fprintf(file, s, v...)
    23  	if err != nil {
    24  		gplog.Fatal(err, "Unable to write to file")
    25  	}
    26  	return uint64(bytesWritten)
    27  }
    28  
    29  func MustPrintln(file io.Writer, v ...interface{}) uint64 {
    30  	bytesWritten, err := fmt.Fprintln(file, v...)
    31  	if err != nil {
    32  		gplog.Fatal(err, "Unable to write to file")
    33  	}
    34  	return uint64(bytesWritten)
    35  }
    36  
    37  /*
    38   * Structs and functions for file readers/writers that track bytes read/written
    39   */
    40  
    41  type FileWithByteCount struct {
    42  	Filename  string
    43  	Writer    io.Writer
    44  	File      *os.File
    45  	ByteCount uint64
    46  }
    47  
    48  func NewFileWithByteCount(writer io.Writer) *FileWithByteCount {
    49  	return &FileWithByteCount{"", writer, nil, 0}
    50  }
    51  
    52  func NewFileWithByteCountFromFile(filename string) *FileWithByteCount {
    53  	file, err := OpenFileForWrite(filename)
    54  	gplog.FatalOnError(err)
    55  	return &FileWithByteCount{filename, file, file, 0}
    56  }
    57  
    58  func (file *FileWithByteCount) Close() {
    59  	if file.File != nil {
    60  		err := file.File.Sync()
    61  		gplog.FatalOnError(err)
    62  		err = file.File.Close()
    63  		gplog.FatalOnError(err)
    64  		if file.Filename != "" {
    65  			err := os.Chmod(file.Filename, 0444)
    66  			gplog.FatalOnError(err)
    67  		}
    68  	}
    69  }
    70  
    71  func (file *FileWithByteCount) MustPrintln(v ...interface{}) {
    72  	bytesWritten, err := fmt.Fprintln(file.Writer, v...)
    73  	gplog.FatalOnError(err, "Unable to write to file")
    74  	file.ByteCount += uint64(bytesWritten)
    75  }
    76  
    77  func (file *FileWithByteCount) MustPrintf(s string, v ...interface{}) {
    78  	bytesWritten, err := fmt.Fprintf(file.Writer, s, v...)
    79  	gplog.FatalOnError(err, "Unable to write to file")
    80  	file.ByteCount += uint64(bytesWritten)
    81  }
    82  
    83  func (file *FileWithByteCount) MustPrint(s string) {
    84  	bytesWritten, err := fmt.Fprint(file.Writer, s)
    85  	gplog.FatalOnError(err, "Unable to write to file")
    86  	file.ByteCount += uint64(bytesWritten)
    87  }
    88  
    89  func CopyFile(src, dest string) error {
    90  	info, err := os.Stat(src)
    91  	if err == nil {
    92  		var content []byte
    93  		content, err = ioutil.ReadFile(src)
    94  		if err != nil {
    95  			gplog.Error(fmt.Sprintf("Error: %v, encountered when reading file: %s", err, src))
    96  			return err
    97  		}
    98  		return ioutil.WriteFile(dest, content, info.Mode())
    99  	}
   100  	gplog.Error(fmt.Sprintf("Error: %v, encountered when trying to stat file: %s", err, src))
   101  	return err
   102  }