From ba6a050d489bf09dbef418ba1865f2fca770c098 Mon Sep 17 00:00:00 2001 From: Keith Chason Date: Fri, 15 Nov 2024 05:15:44 -0500 Subject: [PATCH] Batch JSONL output and add trailing commas (#5705) --- internal/runner/options.go | 13 ++++- pkg/reporting/exporters/jsonl/jsonl.go | 81 +++++++++++++++++++------- pkg/reporting/reporting.go | 2 +- 3 files changed, 70 insertions(+), 26 deletions(-) diff --git a/internal/runner/options.go b/internal/runner/options.go index 4ad62a855a..e36c248a64 100644 --- a/internal/runner/options.go +++ b/internal/runner/options.go @@ -304,10 +304,17 @@ func createReportingOptions(options *types.Options) (*reporting.Options, error) OmitRaw: options.OmitRawRequests, } } + // Combine options. if options.JSONLExport != "" { - reportingOptions.JSONLExporter = &jsonl.Options{ - File: options.JSONLExport, - OmitRaw: options.OmitRawRequests, + // Combine the CLI options with the config file options with the CLI options taking precedence + if reportingOptions.JSONLExporter != nil { + reportingOptions.JSONLExporter.File = options.JSONLExport + reportingOptions.JSONLExporter.OmitRaw = options.OmitRawRequests + } else { + reportingOptions.JSONLExporter = &jsonl.Options{ + File: options.JSONLExport, + OmitRaw: options.OmitRawRequests, + } } } diff --git a/pkg/reporting/exporters/jsonl/jsonl.go b/pkg/reporting/exporters/jsonl/jsonl.go index c02f90baba..ea7c44ef23 100644 --- a/pkg/reporting/exporters/jsonl/jsonl.go +++ b/pkg/reporting/exporters/jsonl/jsonl.go @@ -10,16 +10,21 @@ import ( ) type Exporter struct { - options *Options - mutex *sync.Mutex - rows []output.ResultEvent + options *Options + mutex *sync.Mutex + rows []output.ResultEvent + outputFile *os.File } // Options contains the configuration options for JSONL exporter client type Options struct { // File is the file to export found JSONL result to - File string `yaml:"file"` - OmitRaw bool `yaml:"omit-raw"` + File string `yaml:"file"` + // OmitRaw whether to exclude the raw request and response from the output + OmitRaw bool `yaml:"omit-raw"` + // BatchSize the number of records to keep in memory before writing them out to the JSONL file or 0 to disable + // batching (default) + BatchSize int `yaml:"batch-size"` } // New creates a new JSONL exporter integration client based on options. @@ -32,8 +37,7 @@ func New(options *Options) (*Exporter, error) { return exporter, nil } -// Export appends the passed result event to the list of objects to be exported to -// the resulting JSONL file +// Export appends the passed result event to the list of objects to be exported to the resulting JSONL file func (exporter *Exporter) Export(event *output.ResultEvent) error { exporter.mutex.Lock() defer exporter.mutex.Unlock() @@ -46,23 +50,36 @@ func (exporter *Exporter) Export(event *output.ResultEvent) error { // Add the event to the rows exporter.rows = append(exporter.rows, *event) + // If the batch size is greater than 0 and the number of rows has reached the batch, flush it to the database + if exporter.options.BatchSize > 0 && len(exporter.rows) >= exporter.options.BatchSize { + err := exporter.WriteRows() + if err != nil { + // The error is already logged, return it to bubble up to the caller + return err + } + } + return nil } -// Close writes the in-memory data to the JSONL file specified by options.JSONLExport -// and closes the exporter after operation -func (exporter *Exporter) Close() error { - exporter.mutex.Lock() - defer exporter.mutex.Unlock() - - // Open the JSONL file for writing and create it if it doesn't exist - f, err := os.OpenFile(exporter.options.File, os.O_WRONLY|os.O_CREATE, 0644) - if err != nil { - return errors.Wrap(err, "failed to create JSONL file") +// WriteRows writes all rows from the rows list to JSONL file and removes them from the list +func (exporter *Exporter) WriteRows() error { + // Open the file for writing if it's not already. + // This will recreate the file if it exists, but keep the file handle so that batched writes within the same + // execution are appended to the same file. + var err error + if exporter.outputFile == nil { + // Open the JSONL file for writing and create it if it doesn't exist + exporter.outputFile, err = os.OpenFile(exporter.options.File, os.O_WRONLY|os.O_CREATE, 0644) + if err != nil { + return errors.Wrap(err, "failed to create JSONL file") + } } - // Loop through the rows and convert each to a JSON byte array and write to file - for _, row := range exporter.rows { + // Loop through the rows and write them, removing them as they're entered + for len(exporter.rows) > 0 { + row := exporter.rows[0] + // Convert the row to JSON byte array and append a trailing newline. This is treated as a single line in JSONL obj, err := json.Marshal(row) if err != nil { @@ -70,16 +87,36 @@ func (exporter *Exporter) Close() error { } // Add a trailing newline to the JSON byte array to confirm with the JSONL format - obj = append(obj, '\n') + obj = append(obj, ',', '\n') // Attempt to append the JSON line to file specified in options.JSONLExport - if _, err = f.Write(obj); err != nil { + if _, err = exporter.outputFile.Write(obj); err != nil { return errors.Wrap(err, "failed to append JSONL line") } + + // Remove the item from the list + exporter.rows = exporter.rows[1:] + } + + return nil +} + +// Close writes the in-memory data to the JSONL file specified by options.JSONLExport and closes the exporter after +// operation +func (exporter *Exporter) Close() error { + exporter.mutex.Lock() + defer exporter.mutex.Unlock() + + // Write any remaining rows to the file + // Write all pending rows + err := exporter.WriteRows() + if err != nil { + // The error is already logged, return it to bubble up to the caller + return err } // Close the file - if err := f.Close(); err != nil { + if err := exporter.outputFile.Close(); err != nil { return errors.Wrap(err, "failed to close JSONL file") } diff --git a/pkg/reporting/reporting.go b/pkg/reporting/reporting.go index ddc5428636..20fe85f1a6 100644 --- a/pkg/reporting/reporting.go +++ b/pkg/reporting/reporting.go @@ -197,7 +197,7 @@ func New(options *Options, db string, doNotDedupe bool) (Client, error) { return client, nil } -// CreateConfigIfNotExists creates report-config if it doesn't exists +// CreateConfigIfNotExists creates report-config if it doesn't exist func CreateConfigIfNotExists() error { reportingConfig := config.DefaultConfig.GetReportingConfigFilePath()