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 }