diff --git a/leak-utils/leak-utils/main.go b/leak-utils/leak-utils/main.go index 66b17a3..8daef46 100644 --- a/leak-utils/leak-utils/main.go +++ b/leak-utils/leak-utils/main.go @@ -133,6 +133,7 @@ func main() { var outputFile *string = flag.StringP("output", "o", "", "Output Parquet file") var compression *string = flag.StringP("compression", "c", "ZSTD", "Compression codec (UNCOMPRESSED, SNAPPY, GZIP, BROTLI, LZ4, ZSTD)") var noColors *bool = flag.Bool("no-colors", false, "Remove all colors") + var deleteTempFile *bool = flag.Bool("delete-temp-file", true, "Delete the temporary file") var debug *bool = flag.Bool("debug", false, "Debug mode") flag.Parse() if *inputFile == "" || *outputFile == "" { @@ -143,7 +144,7 @@ func main() { } lu.Compression = *compression lu.Debug = *debug - err := misc.JsonToParquet(lu, *inputFile, *outputFile) + err := misc.JsonToParquet(lu, *inputFile, *outputFile, *deleteTempFile) if err != nil { log.Fatal("Failed to transform JSON file", "error", err) } diff --git a/leak-utils/misc/json.go b/leak-utils/misc/json.go index 4ce662e..f8b3295 100644 --- a/leak-utils/misc/json.go +++ b/leak-utils/misc/json.go @@ -110,10 +110,12 @@ func flattenJSONFile(inputFile string, outputFile string) error { return nil } -func JsonToParquet(lu settings.LeakUtils, inputFile string, outputFile string) error { +func JsonToParquet(lu settings.LeakUtils, inputFile string, outputFile string, deleteTempFile bool) error { tmpFile := filepath.Join("/tmp", "leak-utils.flat.json") err := flattenJSONFile(inputFile, tmpFile) - defer os.Remove(tmpFile) + if deleteTempFile { + defer os.Remove(tmpFile) + } query := fmt.Sprintf(`COPY (FROM read_json('%s', union_by_name=true, ignore_errors=true)) TO '%s' (FORMAT 'parquet', COMPRESSION '%s', ROW_GROUP_SIZE 200000);`, tmpFile, outputFile, lu.Compression)