github.com/stellar/stellar-etl@v1.0.1-0.20240312145900-4874b6bf2b89/cmd/command_utils.go (about)

     1  package cmd
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/json"
     6  	"fmt"
     7  	"os"
     8  	"path/filepath"
     9  )
    10  
    11  type CloudStorage interface {
    12  	UploadTo(credentialsPath, bucket, path string) error
    13  }
    14  
    15  func createOutputFile(filepath string) error {
    16  	var _, err = os.Stat(filepath)
    17  	if os.IsNotExist(err) {
    18  		var _, err = os.Create(filepath)
    19  		if err != nil {
    20  			return err
    21  		}
    22  	}
    23  
    24  	return nil
    25  }
    26  
    27  func mustOutFile(path string) *os.File {
    28  	absolutePath, err := filepath.Abs(path)
    29  	if err != nil {
    30  		cmdLogger.Fatal("could not get absolute filepath: ", err)
    31  	}
    32  
    33  	err = os.MkdirAll(filepath.Dir(path), os.ModePerm)
    34  	if err != nil {
    35  		cmdLogger.Fatalf("could not create directory %s: %s", path, err)
    36  	}
    37  
    38  	err = createOutputFile(absolutePath)
    39  	if err != nil {
    40  		cmdLogger.Fatal("could not create output file: ", err)
    41  	}
    42  
    43  	outFile, err := os.OpenFile(absolutePath, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0644)
    44  	if err != nil {
    45  		cmdLogger.Fatal("error in opening output file: ", err)
    46  	}
    47  
    48  	return outFile
    49  }
    50  
    51  func exportEntry(entry interface{}, outFile *os.File, extra map[string]string) (int, error) {
    52  	// This extra marshalling/unmarshalling is silly, but it's required to properly handle the null.[String|Int*] types, and add the extra fields.
    53  	m, err := json.Marshal(entry)
    54  	if err != nil {
    55  		cmdLogger.Errorf("Error marshalling %+v: %v ", entry, err)
    56  	}
    57  	i := map[string]interface{}{}
    58  	// Use a decoder here so that 'UseNumber' ensures large ints are properly decoded
    59  	decoder := json.NewDecoder(bytes.NewReader(m))
    60  	decoder.UseNumber()
    61  	err = decoder.Decode(&i)
    62  	if err != nil {
    63  		cmdLogger.Errorf("Error unmarshalling %+v: %v ", i, err)
    64  	}
    65  	for k, v := range extra {
    66  		i[k] = v
    67  	}
    68  
    69  	marshalled, err := json.Marshal(i)
    70  	if err != nil {
    71  		return 0, fmt.Errorf("could not json encode %+v: %s", entry, err)
    72  	}
    73  	cmdLogger.Debugf("Writing entry to %s", outFile.Name())
    74  	numBytes, err := outFile.Write(marshalled)
    75  	if err != nil {
    76  		cmdLogger.Errorf("Error writing %+v to file: %s", entry, err)
    77  	}
    78  	newLineNumBytes, err := outFile.WriteString("\n")
    79  	if err != nil {
    80  		cmdLogger.Errorf("Error writing new line to file %s: %s", outFile.Name(), err)
    81  	}
    82  	return numBytes + newLineNumBytes, nil
    83  }
    84  
    85  // Prints the number of attempted, failed, and successful transformations as a JSON object
    86  func printTransformStats(attempts, failures int) {
    87  	resultsMap := map[string]int{
    88  		"attempted_transforms":  attempts,
    89  		"failed_transforms":     failures,
    90  		"successful_transforms": attempts - failures,
    91  	}
    92  
    93  	results, err := json.Marshal(resultsMap)
    94  	if err != nil {
    95  		cmdLogger.Fatal("Could not marshal results: ", err)
    96  	}
    97  
    98  	cmdLogger.Info(string(results))
    99  }
   100  
   101  func exportFilename(start, end uint32, dataType string) string {
   102  	return fmt.Sprintf("%d-%d-%s.txt", start, end-1, dataType)
   103  }
   104  
   105  func deleteLocalFiles(path string) {
   106  	err := os.RemoveAll(path)
   107  	if err != nil {
   108  		cmdLogger.Errorf("Unable to remove %s: %s", path, err)
   109  		return
   110  	}
   111  	cmdLogger.Infof("Successfully deleted %s", path)
   112  }
   113  
   114  func maybeUpload(cloudCredentials, cloudStorageBucket, cloudProvider, path string) {
   115  	if cloudProvider == "" {
   116  		cmdLogger.Info("No cloud provider specified for upload. Skipping upload.")
   117  		return
   118  	}
   119  
   120  	if len(cloudStorageBucket) == 0 {
   121  		cmdLogger.Error("No bucket specified")
   122  		return
   123  	}
   124  
   125  	var cloudStorage CloudStorage
   126  	switch cloudProvider {
   127  	case "gcp":
   128  		cloudStorage = newGCS(cloudCredentials, cloudStorageBucket)
   129  		err := cloudStorage.UploadTo(cloudCredentials, cloudStorageBucket, path)
   130  		if err != nil {
   131  			cmdLogger.Errorf("Unable to upload output to GCS: %s", err)
   132  			return
   133  		}
   134  	default:
   135  		cmdLogger.Error("Unknown cloud provider")
   136  	}
   137  }