diff --git a/CHANGELOG.md b/CHANGELOG.md index 854503f..1a5be1e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,9 +18,14 @@ and this project adheres to override automatic selection. JSON records use `time`, `level`, and `message` fields with additional context, including `error` for failures. - Per-column CSV formats for fixed-precision floats and custom boolean labels. +- Cooperative cancellation via `Run` and CLI handling of SIGINT and SIGTERM, + plus SIGHUP on Unix. Cancellation cleans up temporary files and preserves + existing output until publication starts. ### Changed +- The Go library's `Run` now requires a `context.Context` as its first argument. + Pass `context.Background()` when cancellation is not needed. - Filesystem paths now interpret `${name}` as a parameter, even when no variables are supplied. Undefined parameters and malformed placeholders are errors; escape a literal `${` as `$${`. @@ -42,6 +47,8 @@ and this project adheres to ### Fixed +- Nested 128-bit unsigned integers in maps and arrays are now serialized as + exact JSON numbers in CSV and Parquet string columns. - Extra command-line arguments are now rejected instead of silently ignored. - CSV output now includes the configured header when no data rows are emitted, including when an IP family has no records. This applies when `include_header` diff --git a/cancellation_test.go b/cancellation_test.go new file mode 100644 index 0000000..bee2bca --- /dev/null +++ b/cancellation_test.go @@ -0,0 +1,250 @@ +package mmdbconvert + +import ( + "bytes" + "context" + "errors" + "fmt" + "net/netip" + "os" + "testing" + + "github.com/maxmind/mmdbwriter/v2/mmdbtype" + "github.com/stretchr/testify/require" + + "github.com/maxmind/mmdbconvert/internal/merger" + "github.com/maxmind/mmdbconvert/internal/writer" +) + +func TestRun_Canceled(t *testing.T) { + for _, format := range []string{"csv", "parquet", "mmdb"} { + t.Run(format, func(t *testing.T) { + configPath, _, paths := outputTestConfig(t, format, false, `["country", "iso_code"]`) + require.NoError(t, os.WriteFile(paths[0], []byte("previous output"), 0o600)) + ctx, cancel := context.WithCancel(t.Context()) + cancel() + require.ErrorIs(t, Run(ctx, Options{ConfigPath: configPath}), context.Canceled) + assertFileContent(t, paths[0], "previous output") + assertOutputDirectory(t, paths) + }) + } +} + +func TestConvert_CancellationCleansOutputs(t *testing.T) { + for _, format := range []string{"csv", "parquet", "mmdb"} { + for _, split := range []bool{false, true} { + if split && format == "mmdb" { + continue + } + for _, stage := range []string{"merge", "flush", "after flush"} { + t.Run(fmt.Sprintf("%s/split=%t/%s", format, split, stage), func(t *testing.T) { + _, cfg, paths := outputTestConfig(t, format, split, `["country", "iso_code"]`) + for _, path := range paths { + require.NoError(t, os.WriteFile(path, []byte("previous output"), 0o600)) + } + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + readers := openOutputTestReaders(t, cfg) + rowWriter, outputs, err := prepareRowWriter(ctx, cfg, readers) + require.NoError(t, err) + wrapped := &cancelRowWriter{RowWriter: rowWriter, cancel: cancel, stage: stage} + var observed merger.RowWriter = wrapped + _, supportsRanges := rowWriter.(merger.RangeRowWriter) + if supportsRanges { + observed = &cancelRangeRowWriter{cancelRowWriter: wrapped} + } + err = convert(ctx, cfg, readers, observed, outputs) + require.ErrorIs(t, err, context.Canceled) + if supportsRanges { + require.Positive(t, wrapped.ranges) + require.Zero(t, wrapped.rows) + } else { + require.Positive(t, wrapped.rows) + } + if stage == "flush" { + // The real serializer must stop; convert's post-flush check + // alone would allow a canceled flush to write the full file. + require.ErrorIs(t, wrapped.flushErr, context.Canceled) + } + for _, path := range paths { + assertFileContent(t, path, "previous output") + } + assertOutputDirectory(t, paths) + }) + } + } + } +} + +func TestConvert_CancellationDuringFinalAccumulatorFlush(t *testing.T) { + _, cfg, paths := outputTestConfig(t, "mmdb", false, `["value"]`) + require.NoError(t, os.WriteFile(paths[0], []byte("previous output"), 0o600)) + // A single data range stays in the accumulator until Merge's final Flush. + cfg.Databases[0].Path = createCSVTestDatabase(t, 4, []string{"1.0.0.0/24"}) + + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + readers := openOutputTestReaders(t, cfg) + rowWriter, outputs, err := prepareRowWriter(ctx, cfg, readers) + require.NoError(t, err) + wrapped := &cancelRowWriter{RowWriter: rowWriter, cancel: cancel, stage: "merge"} + err = convert(ctx, cfg, readers, &cancelRangeRowWriter{cancelRowWriter: wrapped}, outputs) + require.ErrorIs(t, err, context.Canceled) + require.Equal(t, 1, wrapped.ranges) + require.Zero(t, wrapped.flushes, "cancellation must prevent format serialization") + assertFileContent(t, paths[0], "previous output") + assertOutputDirectory(t, paths) +} + +type cancelRowWriter struct { + merger.RowWriter + + cancel context.CancelFunc + stage string + rows int + ranges int + flushes int + flushErr error +} + +func (w *cancelRowWriter) WriteRow(prefix netip.Prefix, data []mmdbtype.DataType) error { + w.rows++ + if err := w.RowWriter.WriteRow(prefix, data); err != nil { + return err + } + if w.stage == "merge" { + w.cancel() + } + return nil +} + +func (w *cancelRowWriter) Flush() error { + w.flushes++ + if w.stage == "flush" { + w.cancel() + } + w.flushErr = w.RowWriter.(interface{ Flush() error }).Flush() + if w.flushErr != nil { + return w.flushErr + } + if w.stage == "after flush" { + w.cancel() + } + return nil +} + +// Preserve the prepared writer's optional range capability in the observer. +type cancelRangeRowWriter struct { + *cancelRowWriter +} + +func (w *cancelRangeRowWriter) WriteRange(start, end netip.Addr, data []mmdbtype.DataType) error { + w.ranges++ + if err := w.RowWriter.(merger.RangeRowWriter).WriteRange(start, end, data); err != nil { + return err + } + if w.stage == "merge" { + w.cancel() + } + return nil +} + +func TestContextWriter_Cancellation(t *testing.T) { + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + var buf bytes.Buffer + w := contextWriter{ctx: ctx, writer: &buf} + n, err := w.Write([]byte("first")) + require.NoError(t, err) + require.Equal(t, 5, n) + cancel() + n, err = w.Write([]byte("second")) + require.ErrorIs(t, err, context.Canceled) + require.Zero(t, n) + require.Equal(t, "first", buf.String()) +} + +func TestConvert_CancellationAfterOutputWrite(t *testing.T) { + for _, format := range []string{"csv", "parquet", "mmdb"} { + t.Run(format, func(t *testing.T) { + _, cfg, paths := outputTestConfig(t, format, false, `["country", "iso_code"]`) + require.NoError(t, os.WriteFile(paths[0], []byte("previous output"), 0o600)) + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + file, err := preparePendingOutput(paths[0]) + require.NoError(t, err) + defer func() { require.NoError(t, file.Cleanup()) }() + // Cancel after an underlying write succeeds. Even if serialization + // finishes in that write, the publication check must preserve the old output. + output := &cancelWriteOutput{pendingOutput: file, cancel: cancel} + readers := openOutputTestReaders(t, cfg) + rowWriter, err := newRowWriter( + contextWriter{ctx: ctx, writer: output}, + cfg, + readers, + writer.IPVersionAny, + ) + require.NoError(t, err) + err = convert(ctx, cfg, readers, rowWriter, []pendingOutput{file}) + require.ErrorIs(t, err, context.Canceled) + require.Positive(t, output.bytesWritten) + assertFileContent(t, paths[0], "previous output") + assertOutputDirectory(t, paths) + }) + } +} + +type cancelWriteOutput struct { + pendingOutput + + cancel context.CancelFunc + bytesWritten int +} + +func (o *cancelWriteOutput) Write(p []byte) (int, error) { + n, err := o.pendingOutput.Write(p) + o.bytesWritten += n + o.cancel() + return n, err +} + +func TestConvert_CancellationPreservesCleanupErrors(t *testing.T) { + _, cfg, _ := outputTestConfig(t, "csv", false, `["country", "iso_code"]`) + ctx, cancel := context.WithCancel(t.Context()) + cancel() + cleanupErr := errors.New("cleanup failure") + output := &cleanupErrorOutput{cleanupErr: cleanupErr} + err := convert( + ctx, + cfg, + openOutputTestReaders(t, cfg), + &errorRowWriter{}, + []pendingOutput{output}, + ) + require.ErrorIs(t, err, context.Canceled) + require.ErrorIs(t, err, cleanupErr) + require.Equal(t, 1, output.cleanups) + require.Zero(t, output.commits) +} + +func TestPublication_CompletesAfterCancellation(t *testing.T) { + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + first := &cancelCommitOutput{cancel: cancel} + second := &cleanupErrorOutput{} + require.NoError(t, flushAndCommit(ctx, &errorRowWriter{}, []pendingOutput{first, second})) + require.ErrorIs(t, ctx.Err(), context.Canceled) + require.Equal(t, 1, first.commits) + require.Equal(t, 1, second.commits) +} + +type cancelCommitOutput struct { + cleanupErrorOutput + + cancel context.CancelFunc +} + +func (o *cancelCommitOutput) Commit() error { + o.cancel() + return o.cleanupErrorOutput.Commit() +} diff --git a/cmd/mmdbconvert/main.go b/cmd/mmdbconvert/main.go index cfcb0dd..3675251 100644 --- a/cmd/mmdbconvert/main.go +++ b/cmd/mmdbconvert/main.go @@ -2,6 +2,7 @@ package main import ( + "context" "errors" "flag" "fmt" @@ -28,10 +29,23 @@ const ( var version = unknownVersion func main() { - os.Exit(runCLI(os.Args[1:], os.Stdout, os.Stderr, term.IsTerminal(int(os.Stderr.Fd())))) + os.Exit( + runCLI( + context.Background(), + os.Args[1:], + os.Stdout, + os.Stderr, + term.IsTerminal(int(os.Stderr.Fd())), + ), + ) } -func runCLI(args []string, stdout, stderr io.Writer, stderrIsTerminal bool) int { +func runCLI( + ctx context.Context, + args []string, + stdout, stderr io.Writer, + stderrIsTerminal bool, +) int { // Define command-line flags var ( configPath string @@ -142,89 +156,14 @@ func runCLI(args []string, stdout, stderr io.Writer, stderrIsTerminal bool) int } logger = logger.With("config_path", configPath, "disable_cache", disableCache) - // Start CPU profiling if requested - var cpuProfileFile *os.File - if cpuprofile != "" { - // #nosec G304 -- cpuprofile path comes from trusted command-line flag - f, err := os.Create(cpuprofile) - if err != nil { - logger.Error("Creating CPU profile", "error", err) - return 1 - } - cpuProfileFile = f - if err := pprof.StartCPUProfile(f); err != nil { - logger.Error("Starting CPU profile", "error", err) - f.Close() - return 1 - } - } - - // Run the conversion - runErr := run(mmdbconvert.Options{ + opts := mmdbconvert.Options{ ConfigPath: configPath, Variables: variables, DisableCache: disableCache, - }, logger) - if runErr != nil { - logger.Error("Converting databases", "error", runErr) - } - - // Stop CPU profiling and close file before potentially exiting - if cpuProfileFile != nil { - pprof.StopCPUProfile() - cpuProfileFile.Close() } - - // Write memory profile if requested - if memprofile != "" { - // #nosec G304 -- memprofile path comes from trusted command-line flag - f, err := os.Create(memprofile) - if err != nil { - logger.Error("Creating memory profile", "error", err) - return 1 - } - if err := pprof.WriteHeapProfile(f); err != nil { - f.Close() - logger.Error("Writing memory profile", "error", err) - return 1 - } - f.Close() - } - - if runErr != nil { - return 1 - } - return 0 -} - -func extraConfigArgsMessage(args []string) string { - for _, arg := range args { - if strings.HasPrefix(arg, "-") && arg != "-" { - return "Flags must precede a positional config path; use --config to specify the path alongside flags" - } - } - return "Only one config file path may be specified" -} - -// run performs the main conversion process. -func run(opts mmdbconvert.Options, logger *slog.Logger) error { - startTime := time.Now() - - logger.Info("Starting mmdbconvert") - logger.Info("Loading configuration") - logger.Info("Merging databases and writing output") - - err := mmdbconvert.Run(opts) - if err != nil { - return err - } - - logger.Info( - "Successfully completed", - slog.Duration("elapsed", time.Since(startTime).Round(time.Millisecond)), - ) - - return nil + return runWithSignals(ctx, logger, func(ctx context.Context) int { + return runConversion(ctx, opts, logger, cpuprofile, memprofile) + }) } func usage(w io.Writer) { @@ -289,3 +228,90 @@ MORE INFORMATION: `, ) } + +func extraConfigArgsMessage(args []string) string { + for _, arg := range args { + if strings.HasPrefix(arg, "-") && arg != "-" { + return "Flags must precede a positional config path; use --config to specify the path alongside flags" + } + } + return "Only one config file path may be specified" +} + +func runConversion( + ctx context.Context, + opts mmdbconvert.Options, + logger *slog.Logger, + cpuprofile, memprofile string, +) int { + // Start CPU profiling if requested + var cpuProfileFile *os.File + if cpuprofile != "" { + // #nosec G304 -- cpuprofile path comes from trusted command-line flag + f, err := os.Create(cpuprofile) + if err != nil { + logger.Error("Creating CPU profile", "error", err) + return 1 + } + cpuProfileFile = f + if err := pprof.StartCPUProfile(f); err != nil { + logger.Error("Starting CPU profile", "error", err) + f.Close() + return 1 + } + } + + // Run the conversion + runErr := run(ctx, opts, logger) + if runErr != nil { + logger.Error("Converting databases", "error", runErr) + } + + // Stop CPU profiling and close file before potentially exiting + if cpuProfileFile != nil { + pprof.StopCPUProfile() + cpuProfileFile.Close() + } + + // Write memory profile if requested + if memprofile != "" { + // #nosec G304 -- memprofile path comes from trusted command-line flag + f, err := os.Create(memprofile) + if err != nil { + logger.Error("Creating memory profile", "error", err) + return 1 + } + if err := pprof.WriteHeapProfile(f); err != nil { + f.Close() + logger.Error("Writing memory profile", "error", err) + return 1 + } + f.Close() + } + + if runErr != nil { + return 1 + } + return 0 +} + +// run performs the main conversion process. +func run(ctx context.Context, opts mmdbconvert.Options, logger *slog.Logger) error { + startTime := time.Now() + + logger.Info("Starting mmdbconvert") + logger.Info("Loading configuration") + logger.Info("Merging databases and writing output") + + err := mmdbconvert.Run(ctx, opts) + if err != nil { + return err + } + + logger.Info( + "Successfully completed", + slog.Duration("elapsed", time.Since(startTime).Round(time.Millisecond)), + ) + + return nil +} diff --git a/cmd/mmdbconvert/main_linux_test.go b/cmd/mmdbconvert/main_linux_test.go index 84588a2..d932db0 100644 --- a/cmd/mmdbconvert/main_linux_test.go +++ b/cmd/mmdbconvert/main_linux_test.go @@ -21,7 +21,7 @@ func TestRunCLIMemoryProfileWriteError(t *testing.T) { if info.Mode()&os.ModeCharDevice == 0 { t.Skip("/dev/full is not a character device") } - configPath, _ := writeTestConfig(t, testDatabasePath(t)) + configPath, outputPath := writeTestConfig(t, testDatabasePath(t)) for _, failConversion := range []bool{false, true} { t.Run(fmt.Sprintf("conversion_failure=%t", failConversion), func(t *testing.T) { path := configPath @@ -30,7 +30,7 @@ func TestRunCLIMemoryProfileWriteError(t *testing.T) { } var stdout, stderr bytes.Buffer args := []string{"--quiet", "--memprofile", "/dev/full", path} - assert.Equal(t, 1, runCLI(args, &stdout, &stderr, false)) + assert.Equal(t, 1, runCLI(t.Context(), args, &stdout, &stderr, false)) assert.Empty(t, stdout.String()) records := decodeLogRecords(t, stderr.String()) @@ -40,6 +40,10 @@ func TestRunCLIMemoryProfileWriteError(t *testing.T) { assert.Contains(t, records[0]["error"], path) } else { require.Len(t, records, 1) + // A profiling failure can follow successful output publication. + output, statErr := os.Stat(outputPath) + require.NoError(t, statErr) + require.Positive(t, output.Size()) } for _, record := range records { assert.Equal(t, "ERROR", record["level"]) diff --git a/cmd/mmdbconvert/main_test.go b/cmd/mmdbconvert/main_test.go index caed97b..aced9c3 100644 --- a/cmd/mmdbconvert/main_test.go +++ b/cmd/mmdbconvert/main_test.go @@ -2,7 +2,9 @@ package main import ( "bytes" + "context" "encoding/json" + "errors" "fmt" "io" "os" @@ -16,6 +18,44 @@ import ( "github.com/stretchr/testify/require" ) +func TestRunCLICanceled(t *testing.T) { + for _, format := range []string{"json", "text"} { + t.Run(format, func(t *testing.T) { + configPath, outputPath := writeTestConfig(t, testDatabasePath(t)) + require.NoError(t, os.WriteFile(outputPath, []byte("previous output"), 0o600)) + ctx, cancel := context.WithCancelCause(t.Context()) + cancel(errors.New("caller canceled conversion")) + var stdout, stderr bytes.Buffer + require.Equal( + t, + 1, + runCLI( + ctx, + []string{"--log-format=" + format, configPath}, + &stdout, + &stderr, + false, + ), + ) + require.Empty(t, stdout.String()) + require.Contains(t, stderr.String(), "context canceled") + require.Contains(t, stderr.String(), "caller canceled conversion") + require.NotContains(t, stderr.String(), "Signal received") + if format == "json" { + records := decodeLogRecords(t, stderr.String()) + warning := records[len(records)-1] + require.Equal(t, "Conversion context canceled", warning["message"]) + require.Contains(t, warning["error"], "caller canceled conversion") + require.NotContains(t, warning, "cause") + } + require.NotContains(t, stderr.String(), "Successfully completed") + data, err := os.ReadFile(filepath.Clean(outputPath)) + require.NoError(t, err) + require.Equal(t, "previous output", string(data)) + }) + } +} + func TestRunCLIFormats(t *testing.T) { tests := []struct { name string @@ -78,7 +118,7 @@ func TestRunCLIFormats(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { var stdout, stderr bytes.Buffer - assert.Equal(t, tt.code, runCLI(tt.args, &stdout, &stderr, tt.terminal)) + assert.Equal(t, tt.code, runCLI(t.Context(), tt.args, &stdout, &stderr, tt.terminal)) assert.Empty(t, stdout.String()) if tt.json { records := decodeLogRecords(t, stderr.String()) @@ -107,7 +147,7 @@ func TestRunCLISuccess(t *testing.T) { configPath, } var stdout, stderr bytes.Buffer - assert.Equal(t, 0, runCLI(args, &stdout, &stderr, false)) + assert.Equal(t, 0, runCLI(t.Context(), args, &stdout, &stderr, false)) assert.Empty(t, stdout.String()) output, err := os.ReadFile(filepath.Clean(outputPath)) require.NoError(t, err) @@ -234,7 +274,7 @@ func TestRunCLIArgumentErrors(t *testing.T) { if code == 0 { code = 2 } - assert.Equal(t, code, runCLI(args, &stdout, &stderr, format == "text")) + assert.Equal(t, code, runCLI(t.Context(), args, &stdout, &stderr, format == "text")) assert.Empty(t, stdout.String()) message := tt.message if message == "" { @@ -308,7 +348,7 @@ func TestRunCLIOperationalErrors(t *testing.T) { []string{"--log-format=" + format, fmt.Sprintf("--quiet=%t", quiet)}, tt.args...) var stdout, stderr bytes.Buffer - assert.Equal(t, 1, runCLI(args, &stdout, &stderr, false)) + assert.Equal(t, 1, runCLI(t.Context(), args, &stdout, &stderr, false)) assert.Empty(t, stdout.String()) assert.NotContains(t, stderr.String(), "USAGE:") if format == "text" { @@ -342,7 +382,7 @@ func TestRunCLIQuietErrorContext(t *testing.T) { configPath := filepath.Join(t.TempDir(), "missing.toml") args := []string{"--quiet", fmt.Sprintf("--disable-cache=%t", disableCache), configPath} var stdout, stderr bytes.Buffer - assert.Equal(t, 1, runCLI(args, &stdout, &stderr, false)) + assert.Equal(t, 1, runCLI(t.Context(), args, &stdout, &stderr, false)) records := decodeLogRecords(t, stderr.String()) require.Len(t, records, 1) assert.Equal(t, "ERROR", records[0]["level"]) @@ -359,7 +399,7 @@ func TestRunCLIHelpAndVersion(t *testing.T) { t.Run(format+"/"+flagName, func(t *testing.T) { var stdout, stderr bytes.Buffer args := []string{"--quiet", "--log-format=" + format, flagName} - assert.Equal(t, 0, runCLI(args, &stdout, &stderr, false)) + assert.Equal(t, 0, runCLI(t.Context(), args, &stdout, &stderr, false)) if flagName == "--version" { assert.Equal(t, "mmdbconvert version "+version+"\n", stdout.String()) assert.Empty(t, stderr.String()) @@ -380,7 +420,7 @@ func TestRunCLIHelpWithInvalidFormat(t *testing.T) { for _, args := range [][]string{{"--log-format=xml", help}, {help, "--log-format=xml"}} { t.Run(strings.Join(args, " "), func(t *testing.T) { var stdout, stderr bytes.Buffer - assert.Equal(t, 0, runCLI(args, &stdout, &stderr, false)) + assert.Equal(t, 0, runCLI(t.Context(), args, &stdout, &stderr, false)) assert.Empty(t, stdout.String()) assert.Contains(t, stderr.String(), "USAGE:") assert.Contains(t, stderr.String(), "--log-format") @@ -395,7 +435,7 @@ func TestRunCLIVersionWriteError(t *testing.T) { require.NoError(t, err) require.NoError(t, stdout.Close()) var stderr bytes.Buffer - assert.Equal(t, 1, runCLI([]string{"--version"}, stdout, &stderr, false)) + assert.Equal(t, 1, runCLI(t.Context(), []string{"--version"}, stdout, &stderr, false)) records := decodeLogRecords(t, stderr.String()) require.Len(t, records, 1) assert.Equal(t, "ERROR", records[0]["level"]) @@ -417,7 +457,7 @@ func TestRunCLIProfiling(t *testing.T) { } var stdout, stderr bytes.Buffer args := []string{"--quiet", "--cpuprofile", cpuPath, "--memprofile", memoryPath, path} - assert.Equal(t, wantCode, runCLI(args, &stdout, &stderr, false)) + assert.Equal(t, wantCode, runCLI(t.Context(), args, &stdout, &stderr, false)) assert.Empty(t, stdout.String()) if failConversion { records := decodeLogRecords(t, stderr.String()) @@ -446,7 +486,7 @@ func TestRunCLIConversionAndMemoryProfileErrors(t *testing.T) { memoryPath := t.TempDir() var stdout, stderr bytes.Buffer args := []string{"--quiet", "--cpuprofile", cpuPath, "--memprofile", memoryPath, configPath} - assert.Equal(t, 1, runCLI(args, &stdout, &stderr, false)) + assert.Equal(t, 1, runCLI(t.Context(), args, &stdout, &stderr, false)) assert.Empty(t, stdout.String()) records := decodeLogRecords(t, stderr.String()) require.Len(t, records, 2) @@ -470,7 +510,7 @@ func TestRunCLICPUProfileAlreadyStarted(t *testing.T) { cpuPath := filepath.Join(t.TempDir(), "cpu.prof") var stdout, stderr bytes.Buffer args := []string{"--quiet", "--cpuprofile", cpuPath, "config.toml"} - assert.Equal(t, 1, runCLI(args, &stdout, &stderr, false)) + assert.Equal(t, 1, runCLI(t.Context(), args, &stdout, &stderr, false)) assert.Empty(t, stdout.String()) records := decodeLogRecords(t, stderr.String()) require.Len(t, records, 1) diff --git a/cmd/mmdbconvert/signals.go b/cmd/mmdbconvert/signals.go new file mode 100644 index 0000000..9a6fef1 --- /dev/null +++ b/cmd/mmdbconvert/signals.go @@ -0,0 +1,91 @@ +package main + +import ( + "context" + "errors" + "fmt" + "log/slog" + "os" + "os/signal" +) + +// runWithSignals owns signal handling only for the lifetime of the CLI. +func runWithSignals(ctx context.Context, logger *slog.Logger, run func(context.Context) int) int { + signals := make(chan os.Signal, 1) + if watched := terminationSignals(); len(watched) != 0 { + signal.Notify(signals, watched...) + } + code, cause := runWithSignalChannel(ctx, run, signals, logger) + if interruption, ok := errors.AsType[*signalError](cause); ok { + message := "Signal received" + if code == 0 { + message = "Signal received; conversion completed and outputs were published" + } + logger.Warn(message, "error", cause) + return exitAfterSignal(interruption.signal, logger) + } + if cause != nil && code != 0 { + logger.Warn("Conversion context canceled", "error", cause) + } + return code +} + +// runWithSignalChannel consumes the first signal and owns stopping the channel. +func runWithSignalChannel( + parent context.Context, + run func(context.Context) int, + signals chan os.Signal, + logger *slog.Logger, +) (int, error) { + ctx, cancel := context.WithCancelCause(parent) + defer cancel(nil) + finished := make(chan struct{}) + done := make(chan struct{}) + var interruption *signalError + go func(finished <-chan struct{}) { + defer close(done) + for { + select { + case sig, ok := <-signals: + if !ok { + return + } + if interruption != nil { + //revive:disable-next-line:deep-exit A second signal must bypass blocked cleanup. + os.Exit(exitAfterSignal(sig, logger)) + } + // Restore default handling before cancellation becomes visible, + // so a second signal can terminate blocked cleanup. + signal.Stop(signals) + interruption = &signalError{signal: sig} + cancel(interruption) + case <-finished: + // Stop guarantees no more sends. Drain queued signals before + // normal completion can cancel the context without a cause. + signal.Stop(signals) + close(signals) + finished = nil + } + } + }(finished) + code := run(ctx) + close(finished) + <-done + // Parent cancellation may have won the context's cause. A received signal + // still determines process termination, independently of the callback's cause. + if interruption != nil { + return code, fmt.Errorf("running conversion: %w", interruption) + } + if cause := context.Cause(ctx); cause != nil { + return code, fmt.Errorf("running conversion: %w", cause) + } + return code, nil +} + +type signalError struct { + signal os.Signal +} + +func (e *signalError) Error() string { + return fmt.Sprintf("%s signal received", e.signal) +} diff --git a/cmd/mmdbconvert/signals_other.go b/cmd/mmdbconvert/signals_other.go new file mode 100644 index 0000000..8ac7acc --- /dev/null +++ b/cmd/mmdbconvert/signals_other.go @@ -0,0 +1,18 @@ +//go:build !unix && !windows + +package main + +import ( + "log/slog" + "os" +) + +// Other platforms retain their native signal handling. Callers can still cancel +// conversions through their context. +func terminationSignals() []os.Signal { + return nil +} + +func exitAfterSignal(_ os.Signal, _ *slog.Logger) int { + return 1 +} diff --git a/cmd/mmdbconvert/signals_test.go b/cmd/mmdbconvert/signals_test.go new file mode 100644 index 0000000..a0f2581 --- /dev/null +++ b/cmd/mmdbconvert/signals_test.go @@ -0,0 +1,88 @@ +package main + +import ( + "bytes" + "context" + "errors" + "fmt" + "io" + "log/slog" + "os" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestRunWithSignalChannelCompletion(t *testing.T) { + for _, interrupted := range []bool{false, true} { + name := "normal" + if interrupted { + name = "queued interrupt" + } + t.Run(name, func(t *testing.T) { + signals := make(chan os.Signal, 1) + code, cause := runWithSignalChannel(t.Context(), func(context.Context) int { + if interrupted { + signals <- os.Interrupt + } + return 0 + }, signals, newLogger(io.Discard, logFormatJSON, slog.LevelWarn)) + require.Zero(t, code) + if interrupted { + var interruption *signalError + require.ErrorAs(t, cause, &interruption) + require.Equal(t, os.Interrupt, interruption.signal) + } else { + require.NoError(t, cause) + } + }) + } +} + +func TestRunWithSignalChannelCanceledParent(t *testing.T) { + for _, code := range []int{0, 1} { + t.Run(fmt.Sprintf("code=%d", code), func(t *testing.T) { + ctx, cancel := context.WithCancelCause(t.Context()) + defer cancel(nil) + parentCause := errors.New("parent stopped conversion") + cancel(parentCause) + signals := make(chan os.Signal, 1) + signals <- os.Interrupt + got, cause := runWithSignalChannel(ctx, func(ctx context.Context) int { + require.ErrorIs(t, context.Cause(ctx), parentCause) + return code + }, signals, newLogger(io.Discard, logFormatJSON, slog.LevelWarn)) + require.Equal(t, code, got) + var interruption *signalError + require.ErrorAs(t, cause, &interruption) + require.Equal(t, os.Interrupt, interruption.signal) + require.ErrorIs(t, context.Cause(ctx), parentCause) + }) + } +} + +func TestRunWithSignalsParentCancellation(t *testing.T) { + for _, code := range []int{0, 1} { + t.Run(fmt.Sprintf("code=%d", code), func(t *testing.T) { + ctx, cancel := context.WithCancelCause(t.Context()) + defer cancel(nil) + parentCause := errors.New("parent stopped conversion") + var stderr bytes.Buffer + logger := newLogger(&stderr, logFormatJSON, slog.LevelWarn) + got := runWithSignals(ctx, logger, func(context.Context) int { + cancel(parentCause) + return code + }) + require.Equal(t, code, got) + if code == 0 { + require.Empty(t, stderr.String()) + return + } + records := decodeLogRecords(t, stderr.String()) + require.Len(t, records, 1) + require.Equal(t, "WARN", records[0]["level"]) + require.Equal(t, "Conversion context canceled", records[0]["message"]) + require.Contains(t, records[0]["error"], parentCause.Error()) + }) + } +} diff --git a/cmd/mmdbconvert/signals_unix.go b/cmd/mmdbconvert/signals_unix.go new file mode 100644 index 0000000..e694f81 --- /dev/null +++ b/cmd/mmdbconvert/signals_unix.go @@ -0,0 +1,37 @@ +//go:build unix + +package main + +import ( + "log/slog" + "os" + "os/signal" + "syscall" + "time" +) + +func terminationSignals() []os.Signal { + var signals []os.Signal + // Honor the parent's dispositions, including nohup and background jobs. + for _, sig := range []os.Signal{os.Interrupt, syscall.SIGTERM, syscall.SIGHUP} { + if !signal.Ignored(sig) { + signals = append(signals, sig) + } + } + return signals +} + +func exitAfterSignal(sig os.Signal, logger *slog.Logger) int { + // Shells distinguish signal termination from a normal exit with code 130. + // Re-raise after cleanup so an interrupt also stops a waiting shell loop. + signal.Reset(sig) + if err := syscall.Kill(os.Getpid(), sig.(syscall.Signal)); err != nil { + logger.Error("Re-raising signal", "signal", sig, "error", err) + } else { + // Delivery is asynchronous. Normally the signal terminates us immediately; + // retain a bounded fallback if the process unexpectedly survives it. + time.Sleep(time.Second) + logger.Warn("Signal did not terminate process", "signal", sig) + } + return 128 + int(sig.(syscall.Signal)) +} diff --git a/cmd/mmdbconvert/signals_unix_test.go b/cmd/mmdbconvert/signals_unix_test.go new file mode 100644 index 0000000..505236c --- /dev/null +++ b/cmd/mmdbconvert/signals_unix_test.go @@ -0,0 +1,290 @@ +//go:build unix + +package main + +import ( + "bufio" + "bytes" + "context" + "errors" + "fmt" + "io" + "log/slog" + "os" + "os/exec" + "os/signal" + "sync" + "syscall" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestCLISignals(t *testing.T) { + for _, tt := range []struct { + name string + first os.Signal + second os.Signal + completed bool + }{ + {name: "interrupt", first: os.Interrupt}, + {name: "terminated", first: syscall.SIGTERM}, + {name: "hangup", first: syscall.SIGHUP}, + {name: "interrupt with publication", first: os.Interrupt, completed: true}, + {name: "terminated with publication", first: syscall.SIGTERM, completed: true}, + {name: "hangup with publication", first: syscall.SIGHUP, completed: true}, + {name: "interrupt twice", first: os.Interrupt, second: os.Interrupt}, + {name: "terminated twice", first: syscall.SIGTERM, second: syscall.SIGTERM}, + {name: "hangup twice", first: syscall.SIGHUP, second: syscall.SIGHUP}, + {name: "interrupt then terminate", first: os.Interrupt, second: syscall.SIGTERM}, + {name: "terminate then interrupt", first: syscall.SIGTERM, second: os.Interrupt}, + } { + t.Run(tt.name, func(t *testing.T) { + skipIgnoredSignal(t, tt.first) + if tt.second != nil { + skipIgnoredSignal(t, tt.second) + } + child := startSignalChild(t, signalChildOptions{ + blockCleanup: tt.second != nil, + completed: tt.completed, + }) + require.NoError(t, child.cmd.Process.Signal(tt.first)) + last := tt.first + if tt.second != nil { + child.requireStage(t, "canceled") + require.NoError(t, child.cmd.Process.Signal(tt.second)) + last = tt.second + } + status, stderr := child.wait(t) + require.Equal(t, last, status.Signal()) + if tt.second != nil { + require.NotContains(t, stderr, "cleanup complete") + return + } + require.Contains(t, stderr, "cleanup complete") + records := decodeLogRecords(t, stderr) + require.Len(t, records, 2) + warning := records[1] + message := "Signal received" + if tt.completed { + message = "Signal received; conversion completed and outputs were published" + } + require.Equal(t, message, warning["message"]) + require.Contains(t, warning["error"], tt.first.String()+" signal received") + require.NotContains(t, warning, "cause") + }) + } +} + +type signalChildOptions struct { + blockCleanup bool + shellLoop bool + completed bool +} + +type signalChild struct { + //nolint:containedctx // This test fixture owns the subprocess's bounded lifetime. + ctx context.Context + cmd *exec.Cmd + stdout *bufio.Reader + stderr bytes.Buffer + waited bool + stopHelpers func() +} + +func startSignalChild(t *testing.T, opts signalChildOptions) *signalChild { + t.Helper() + ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second) + t.Cleanup(cancel) + binary, err := os.Executable() + require.NoError(t, err) + // #nosec G204 -- run only this test binary's synchronized signal helper. + cmd := exec.CommandContext(ctx, binary, "-test.run=^TestCLISignalHelper$") + if opts.shellLoop { + // #nosec G204 -- run the fixed loop with this test binary as a positional argument. + cmd = exec.CommandContext(ctx, requireBash(t), "-c", `for n in 1 2; do "$@"; done`, + "signal-test", binary, "-test.run=^TestCLISignalHelper$") + cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} + } + cmd.Env = append(os.Environ(), "MMDBCONVERT_SIGNAL_TEST_HELPER=1") + if opts.blockCleanup { + cmd.Env = append(cmd.Env, "MMDBCONVERT_SIGNAL_TEST_BLOCK=1") + } + if !opts.completed { + cmd.Env = append(cmd.Env, "MMDBCONVERT_SIGNAL_TEST_FAILED=1") + } + child := &signalChild{ctx: ctx, cmd: cmd, stopHelpers: func() {}} + closeInput := func() {} + if opts.shellLoop { + input, lifetime, err := os.Pipe() + require.NoError(t, err) + closeInput = sync.OnceFunc(func() { assert.NoError(t, input.Close()) }) + t.Cleanup(closeInput) + child.stopHelpers = sync.OnceFunc(func() { assert.NoError(t, lifetime.Close()) }) + t.Cleanup(child.stopHelpers) + // Cmd.Wait may reap Bash before a helper releases its inherited stderr. + // Keep this independent of os/exec's context watcher, which stops at reap. + stop := context.AfterFunc(ctx, child.stopHelpers) + t.Cleanup(func() { stop() }) + cmd.Stdin = input + cmd.Env = append(cmd.Env, "MMDBCONVERT_SIGNAL_TEST_LIFETIME=1") + } + cmd.Stderr = &child.stderr + stdout, err := cmd.StdoutPipe() + require.NoError(t, err) + child.stdout = bufio.NewReader(stdout) + require.NoError(t, cmd.Start()) + t.Cleanup(func() { + child.stopHelpers() + cancel() + if !child.waited { + child.waited = true + // Early assertion failures may leave a child to kill and reap. + var exitErr *exec.ExitError + if err := cmd.Wait(); err != nil && !errors.As(err, &exitErr) { + assert.ErrorIs(t, err, context.Canceled) + } + } + }) + closeInput() + // The child announces readiness only after signal handling is installed. + child.requireStage(t, "ready") + return child +} + +func (child *signalChild) requireStage(t *testing.T, stage string) { + t.Helper() + line, err := child.stdout.ReadString('\n') + require.NoError(t, err) + require.Equal(t, stage+"\n", line) +} + +func (child *signalChild) wait(t *testing.T) (syscall.WaitStatus, string) { + t.Helper() + require.False(t, child.waited, "child must be reaped exactly once") + child.waited = true + err := child.cmd.Wait() + // Wait completes the stderr copy before the buffer is read. + stderr := child.stderr.String() + var exitErr *exec.ExitError + require.ErrorAs(t, err, &exitErr, stderr) + require.NoError(t, child.ctx.Err(), "child must exit through a signal, not timeout") + status := exitErr.Sys().(syscall.WaitStatus) + require.True(t, status.Signaled(), stderr) + return status, stderr +} + +func TestCLISignalHelper(t *testing.T) { + if os.Getenv("MMDBCONVERT_SIGNAL_TEST_HELPER") != "1" { + return + } + if os.Getenv("MMDBCONVERT_SIGNAL_TEST_LIFETIME") == "1" { + go func() { + // EOF ends every helper, including a second loop iteration orphaned + // after Bash exits. A read failure must also stop the helper. + _, err := io.Copy(io.Discard, os.Stdin) + assert.NoError(t, err) + //revive:disable-next-line:deep-exit The parent is tearing down the subprocess fixture. + os.Exit(1) + }() + } + var ignored []os.Signal + switch os.Getenv("MMDBCONVERT_SIGNAL_TEST_IGNORED") { + case "INT": + ignored = []os.Signal{os.Interrupt} + case "HUP": + ignored = []os.Signal{syscall.SIGHUP} + case "ALL": + ignored = []os.Signal{os.Interrupt, syscall.SIGTERM, syscall.SIGHUP} + // Unlike INT/HUP, Go may install a handler over inherited ignored TERM. + // Ignore it explicitly to exercise having no signals to subscribe to. + signal.Ignore(ignored...) + } + logger := newLogger(os.Stderr, logFormatJSON, slog.LevelWarn) + code := runWithSignals(t.Context(), logger, func(ctx context.Context) int { + if len(ignored) != 0 { + for _, sig := range ignored { + require.True(t, signal.Ignored(sig), "%s must remain ignored", sig) + } + return 0 + } + defer logger.Warn("cleanup complete") + _, err := fmt.Fprintln(os.Stdout, "ready") + require.NoError(t, err) + <-ctx.Done() + require.ErrorIs(t, ctx.Err(), context.Canceled) + if os.Getenv("MMDBCONVERT_SIGNAL_TEST_BLOCK") == "1" { + _, err = fmt.Fprintln(os.Stdout, "canceled") + require.NoError(t, err) + select {} + } + // Publication may complete successfully even after a signal arrives. + if os.Getenv("MMDBCONVERT_SIGNAL_TEST_FAILED") == "1" { + return 1 + } + return 0 + }) + //revive:disable-next-line:deep-exit Subprocess helper must expose the CLI's exit status. + os.Exit(code) +} + +func TestCLIPreservesIgnoredSignals(t *testing.T) { + bash := requireBash(t) + binary, err := os.Executable() + require.NoError(t, err) + for _, ignored := range []string{"INT", "HUP", "ALL"} { + t.Run(ignored, func(t *testing.T) { + trap := ignored + if ignored == "ALL" { + trap = "INT" + } + // #nosec G204 -- inherit ignored signals before executing this test binary. + cmd := exec.CommandContext(t.Context(), bash, "-c", `trap '' "$1"; shift; exec "$@"`, + "signal-test", trap, binary, "-test.run=^TestCLISignalHelper$") + cmd.Env = append( + os.Environ(), + "MMDBCONVERT_SIGNAL_TEST_HELPER=1", + "MMDBCONVERT_SIGNAL_TEST_IGNORED="+ignored, + ) + output, err := cmd.CombinedOutput() + require.NoError(t, err, string(output)) + }) + } +} + +func TestCLIInterruptStopsShellLoop(t *testing.T) { + skipIgnoredSignal(t, os.Interrupt) + child := startSignalChild(t, signalChildOptions{shellLoop: true, completed: true}) + require.NoError(t, syscall.Kill(-child.cmd.Process.Pid, syscall.SIGINT)) + status, stderr := child.wait(t) + require.Contains(t, stderr, "cleanup complete") + require.Equal(t, syscall.SIGINT, status.Signal()) +} + +func TestCLIShellLoopCleanup(t *testing.T) { + child := startSignalChild(t, signalChildOptions{shellLoop: true}) + // Killing only Bash leaves its helper holding stderr open. Closing the + // inherited input must stop that helper so Wait can finish before timeout. + require.NoError(t, child.cmd.Process.Kill()) + child.stopHelpers() + status, _ := child.wait(t) + require.Equal(t, syscall.SIGKILL, status.Signal()) +} + +func skipIgnoredSignal(t *testing.T, sig os.Signal) { + t.Helper() + if signal.Ignored(sig) { + t.Skipf("%s is inherited as ignored", sig) + } +} + +func requireBash(t *testing.T) string { + t.Helper() + path, err := exec.LookPath("bash") + if err != nil { + t.Skipf("bash is unavailable: %v", err) + } + return path +} diff --git a/cmd/mmdbconvert/signals_windows.go b/cmd/mmdbconvert/signals_windows.go new file mode 100644 index 0000000..6e1bf24 --- /dev/null +++ b/cmd/mmdbconvert/signals_windows.go @@ -0,0 +1,17 @@ +package main + +import ( + "log/slog" + "os" + "syscall" +) + +func terminationSignals() []os.Signal { + return []os.Signal{os.Interrupt, syscall.SIGTERM} +} + +func exitAfterSignal(_ os.Signal, _ *slog.Logger) int { + // Preserve the native console termination status, STATUS_CONTROL_C_EXIT + // (0xC000013A), expressed as a signed exit code for os.Exit. + return -1073741510 +} diff --git a/cmd/mmdbconvert/signals_windows_test.go b/cmd/mmdbconvert/signals_windows_test.go new file mode 100644 index 0000000..3920848 --- /dev/null +++ b/cmd/mmdbconvert/signals_windows_test.go @@ -0,0 +1,30 @@ +package main + +import ( + "io" + "log/slog" + "os" + "os/exec" + "syscall" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestCLIWindowsSignalExit(t *testing.T) { + if os.Getenv("MMDBCONVERT_WINDOWS_SIGNAL_HELPER") == "1" { + logger := newLogger(io.Discard, logFormatJSON, slog.LevelWarn) + //revive:disable-next-line:deep-exit Subprocess helper exposes the native console exit status. + os.Exit(exitAfterSignal(os.Interrupt, logger)) + } + binary, err := os.Executable() + require.NoError(t, err) + // #nosec G204 -- run only this test binary's exit-status helper. + cmd := exec.CommandContext(t.Context(), binary, "-test.run=^TestCLIWindowsSignalExit$") + cmd.Env = append(os.Environ(), "MMDBCONVERT_WINDOWS_SIGNAL_HELPER=1") + err = cmd.Run() + var exitErr *exec.ExitError + require.ErrorAs(t, err, &exitErr) + status := exitErr.Sys().(syscall.WaitStatus) + require.Equal(t, uint32(0xC000013A), status.ExitCode) +} diff --git a/cmd/mmdbconvert/variables_test.go b/cmd/mmdbconvert/variables_test.go index 9ba82d2..3ea149f 100644 --- a/cmd/mmdbconvert/variables_test.go +++ b/cmd/mmdbconvert/variables_test.go @@ -35,7 +35,7 @@ func TestRunCLIPathVariables(t *testing.T) { args = append([]string{"--config", configPath}, args...) } var stdout, stderr bytes.Buffer - require.Equal(t, 0, runCLI(args, &stdout, &stderr, false), stderr.String()) + require.Equal(t, 0, runCLI(t.Context(), args, &stdout, &stderr, false), stderr.String()) assert.Empty(t, stdout.String()) assert.Empty(t, stderr.String()) for _, path := range []string{ipv4, ipv6} { @@ -114,7 +114,7 @@ func TestRunCLIPathVariableErrors(t *testing.T) { args := append([]string{"--log-format=" + format}, baseArgs...) args = append(args, tt.args...) var stdout, stderr bytes.Buffer - assert.Equal(t, tt.code, runCLI(args, &stdout, &stderr, false)) + assert.Equal(t, tt.code, runCLI(t.Context(), args, &stdout, &stderr, false)) assert.Empty(t, stdout.String()) if format == "json" { records := decodeLogRecords(t, stderr.String()) @@ -143,7 +143,7 @@ func TestRunCLIPathVariableErrors(t *testing.T) { func TestRunCLIPathVariablesRequired(t *testing.T) { var stdout, stderr bytes.Buffer args := []string{"--quiet", "--log-format=json", "--config", writePathVariableConfig(t)} - assert.Equal(t, 1, runCLI(args, &stdout, &stderr, false)) + assert.Equal(t, 1, runCLI(t.Context(), args, &stdout, &stderr, false)) records := decodeLogRecords(t, stderr.String()) require.Len(t, records, 1) assert.Contains(t, records[0]["error"], `output.ipv4_file: undefined variable "ipv4"`) diff --git a/docs/config.md b/docs/config.md index 7db2268..d074405 100644 --- a/docs/config.md +++ b/docs/config.md @@ -139,6 +139,34 @@ Split IPv4/IPv6 files are replaced separately, so a failure can leave only one updated. Use a fresh output directory for each build if both files must be published together. Abrupt termination can leave temporary files behind. +Go callers pass a context to `Run(ctx, opts)`, which does not install signal +handlers. Use `context.Background()` when cancellation is not needed. +Cancellation removes staged files and preserves existing outputs before +publication starts. When cancellation stops conversion, the returned error +matches `context.Canceled` or `context.DeadlineExceeded` through `errors.Is`. +Callers can inspect a custom cancellation cause with `context.Cause(ctx)`. + +Cancellation is cooperative: cleanup may wait for an operation already in +progress. Once publication starts, it finishes before exiting. + +On Unix and Windows, the CLI handles SIGINT and SIGTERM (plus SIGHUP on Unix) by +canceling conversion and cleaning up temporary files. On Unix, inherited ignored +SIGINT and SIGHUP stay ignored, preserving background-job and `nohup` behavior. +Other platforms retain their native signal handling. + +After cleanup, the CLI re-raises the handled signal on Unix so an interrupt +stops waiting shell loops, even if the output was published. If re-raising fails +or the process does not terminate within one second, the CLI logs the condition +and exits with `128 + signal number`. On Windows it uses the native console +termination status `STATUS_CONTROL_C_EXIT` (`0xC000013A`, or `-1073741510` as a +signed value). + +A second handled signal forces termination without waiting for cleanup. A +wrapper that forwards a signal the CLI also receives directly can trigger this +with a single Ctrl-C. Windows also limits cleanup time for console close, +logoff, and shutdown events. These forced terminations can leave temporary files +behind. + #### CSV Options When `format = "csv"`, you can specify CSV-specific options: diff --git a/format_integration_test.go b/format_integration_test.go index 54393f2..85088f0 100644 --- a/format_integration_test.go +++ b/format_integration_test.go @@ -72,7 +72,7 @@ func TestRun_FormatsAfterMergingAndFiltering(t *testing.T) { require.NoError(t, err) configPath := filepath.Join(dir, "config.toml") require.NoError(t, os.WriteFile(configPath, data, 0o600)) - require.NoError(t, Run(Options{ConfigPath: configPath})) + require.NoError(t, Run(t.Context(), Options{ConfigPath: configPath})) assertFileContent(t, outputPath, tt.want) }) } @@ -109,7 +109,7 @@ func TestRun_FormatTypeMismatchPreservesOutput(t *testing.T) { for _, path := range paths { require.NoError(t, os.WriteFile(path, []byte("previous output"), 0o600)) } - err = Run(Options{ConfigPath: configPath}) + err = Run(t.Context(), Options{ConfigPath: configPath}) require.ErrorContains(t, err, "column 'country_code'") require.ErrorContains(t, err, tt.want) for _, path := range paths { diff --git a/geoip_csv_test.go b/geoip_csv_test.go index 4ccbd75..dd4f34f 100644 --- a/geoip_csv_test.go +++ b/geoip_csv_test.go @@ -24,7 +24,7 @@ func TestRun_CityCSVParity(t *testing.T) { require.NoError(t, err) configPath := filepath.Join(t.TempDir(), "city.toml") require.NoError(t, os.WriteFile(configPath, data, 0o600)) - require.NoError(t, Run(Options{ConfigPath: configPath})) + require.NoError(t, Run(t.Context(), Options{ConfigPath: configPath})) for _, tt := range []struct { version string diff --git a/go.mod b/go.mod index 6728629..9bfeda3 100644 --- a/go.mod +++ b/go.mod @@ -3,7 +3,7 @@ module github.com/maxmind/mmdbconvert go 1.27.0 require ( - github.com/maxmind/mmdbwriter/v2 v2.0.0-20260923191345-8f49d9bce747 + github.com/maxmind/mmdbwriter/v2 v2.0.0-20260928214421-cc514d224474 github.com/oschwald/maxminddb-golang/v2 v2.6.0 github.com/parquet-go/parquet-go v0.32.0 github.com/pelletier/go-toml/v2 v2.4.3 diff --git a/go.sum b/go.sum index fd0d2eb..d5d875e 100644 --- a/go.sum +++ b/go.sum @@ -14,8 +14,8 @@ github.com/hexops/gotextdiff v1.0.3 h1:gitA9+qJrrTCsiCl7+kh75nPqQt1cx4ZkudSTLoUq github.com/hexops/gotextdiff v1.0.3/go.mod h1:pSWU5MAI3yDq+fZBTazCSJysOMbxWL1BSow5/V2vxeg= github.com/klauspost/compress v1.18.2 h1:iiPHWW0YrcFgpBYhsA6D1+fqHssJscY/Tm/y2Uqnapk= github.com/klauspost/compress v1.18.2/go.mod h1:R0h/fSBs8DE4ENlcrlib3PsXS61voFxhIs2DeRhCvJ4= -github.com/maxmind/mmdbwriter/v2 v2.0.0-20260923191345-8f49d9bce747 h1:FSl/wHLKIt3uJhFl50atRbUhO0HjRtRGlaNq5YRM2wg= -github.com/maxmind/mmdbwriter/v2 v2.0.0-20260923191345-8f49d9bce747/go.mod h1:hPQyr0pfn16zAHNQsUS9w4SHjBZ/Y8+Jfmf04SS1K8A= +github.com/maxmind/mmdbwriter/v2 v2.0.0-20260928214421-cc514d224474 h1:g1xsksNM6TU0WRdd+EexImVGW1ffFy4XFVLI6Rx6LKE= +github.com/maxmind/mmdbwriter/v2 v2.0.0-20260928214421-cc514d224474/go.mod h1:hPQyr0pfn16zAHNQsUS9w4SHjBZ/Y8+Jfmf04SS1K8A= github.com/oschwald/maxminddb-golang/v2 v2.6.0 h1:pRlHCdJmc+4uxMOSthmKDt5HOw3JTX8TJZlhyP5ew0w= github.com/oschwald/maxminddb-golang/v2 v2.6.0/go.mod h1:sjqpB3z2BZrMduDp9TAUTCkZDoT3nDhixUc4Dge2qRQ= github.com/parquet-go/bitpack v1.0.0 h1:AUqzlKzPPXf2bCdjfj4sTeacrUwsT7NlcYDMUQxPcQA= diff --git a/internal/merger/merger.go b/internal/merger/merger.go index 4a96f43..42e04d8 100644 --- a/internal/merger/merger.go +++ b/internal/merger/merger.go @@ -7,6 +7,7 @@ package merger import ( + "context" "errors" "fmt" "net/netip" @@ -187,12 +188,15 @@ func NewMerger(readers *mmdb.Readers, cfg *config.Config, writer RowWriter) (*Me // Merge performs the streaming merge of all databases. // It uses nested NetworksWithin iteration to find the smallest overlapping // networks across all databases, then extracts data and streams to accumulator. -func (m *Merger) Merge() error { +func (m *Merger) Merge(ctx context.Context) error { // readersList and dbNamesList are already built in NewMerger() firstReader := m.readersList[0] // Iterate all networks in the first database for result := range firstReader.Networks(maxminddb.IncludeNetworksWithoutData()) { + if err := ctx.Err(); err != nil { + return fmt.Errorf("iterating first database: %w", err) + } if err := result.Err(); err != nil { return fmt.Errorf("iterating first database: %w", err) } @@ -212,11 +216,14 @@ func (m *Merger) Merge() error { m.resultsBuffer[0] = result // Process this network through remaining databases starting at index 1 - if err := m.processNetwork(prefix, 1); err != nil { + if err := m.processNetwork(ctx, prefix, 1); err != nil { return err } } + if err := ctx.Err(); err != nil { + return fmt.Errorf("flushing accumulator: %w", err) + } // Flush any remaining accumulated data if err := m.acc.Flush(); err != nil { return fmt.Errorf("flushing accumulator: %w", err) @@ -233,6 +240,7 @@ func (m *Merger) Merge() error { // - effectivePrefix is the smallest network across all databases so far. // - With IncludeNetworksWithoutData, we always get at least one Result per database. func (m *Merger) processNetwork( + ctx context.Context, effectivePrefix netip.Prefix, dbIndex int, ) error { @@ -247,6 +255,9 @@ func (m *Merger) processNetwork( // Iterate networks within effectivePrefix in this database // With IncludeNetworksWithoutData, this ALWAYS yields at least one Result for result := range currentReader.NetworksWithin(effectivePrefix, maxminddb.IncludeNetworksWithoutData()) { + if err := ctx.Err(); err != nil { + return fmt.Errorf("iterating database within %s: %w", effectivePrefix, err) + } if err := result.Err(); err != nil { return fmt.Errorf("iterating database within %s: %w", effectivePrefix, err) } @@ -262,7 +273,7 @@ func (m *Merger) processNetwork( // Recurse with the smallest prefix // NOTE: smallest may be smaller than result.Prefix() - that's OK! // The result contains data for a broader network that covers smallest. - if err := m.processNetwork(smallest, dbIndex+1); err != nil { + if err := m.processNetwork(ctx, smallest, dbIndex+1); err != nil { return err } diff --git a/internal/merger/merger_bench_test.go b/internal/merger/merger_bench_test.go index 8e68577..55d9c28 100644 --- a/internal/merger/merger_bench_test.go +++ b/internal/merger/merger_bench_test.go @@ -54,7 +54,7 @@ func BenchmarkMergerFullMerge(b *testing.B) { require.NoError(b, err) b.StartTimer() - err = merger.Merge() + err = merger.Merge(b.Context()) b.StopTimer() require.NoError(b, err) diff --git a/internal/merger/merger_test.go b/internal/merger/merger_test.go index c32c0dd..fa42bc5 100644 --- a/internal/merger/merger_test.go +++ b/internal/merger/merger_test.go @@ -1,10 +1,16 @@ package merger import ( + "bytes" + "context" "errors" + "fmt" "net/netip" + "os" + "path/filepath" "testing" + "github.com/maxmind/mmdbwriter/v2" "github.com/maxmind/mmdbwriter/v2/mmdbtype" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -48,7 +54,7 @@ func TestMerger_SingleDatabase(t *testing.T) { // Create merger and run merger, err := NewMerger(readers, cfg, writer) require.NoError(t, err) - err = merger.Merge() + err = merger.Merge(t.Context()) require.NoError(t, err) // Should have written some rows @@ -61,6 +67,61 @@ func TestMerger_SingleDatabase(t *testing.T) { } } +func TestMerger_Cancellation(t *testing.T) { + // An empty first database makes its entire address space one outer network, + // so the multi-database case must cancel within the nested City iteration. + empty, err := mmdbwriter.New(mmdbwriter.Options{IPVersion: 6, IncludeReservedNetworks: true}) + require.NoError(t, err) + var buf bytes.Buffer + _, err = empty.WriteTo(&buf) + require.NoError(t, err) + emptyPath := filepath.Join(t.TempDir(), "empty.mmdb") + require.NoError(t, os.WriteFile(emptyPath, buf.Bytes(), 0o600)) + for _, multiple := range []bool{false, true} { + t.Run(fmt.Sprintf("multiple=%t", multiple), func(t *testing.T) { + readers, err := mmdb.OpenDatabases( + map[string]string{"city": cityTestDB, "empty": emptyPath}, + ) + require.NoError(t, err) + defer readers.Close() + cfg := &config.Config{Columns: []config.Column{ + {Name: "country", Database: "city", Path: config.Path{"country", "iso_code"}}, + }} + if multiple { + cfg.Columns = append([]config.Column{{ + Name: "empty", Database: "empty", Path: config.Path{"value"}, + }}, cfg.Columns...) + } + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + writer := &cancelWriter{cancel: cancel} + m, err := NewMerger(readers, cfg, writer) + require.NoError(t, err) + require.ErrorIs(t, m.Merge(ctx), context.Canceled) + require.Equal(t, 1, writer.rows, "merging must stop before writing another range") + }) + } +} + +type cancelWriter struct { + cancel context.CancelFunc + rows int +} + +var _ RangeRowWriter = (*cancelWriter)(nil) + +func (w *cancelWriter) WriteRow(netip.Prefix, []mmdbtype.DataType) error { + w.rows++ + w.cancel() + return nil +} + +func (w *cancelWriter) WriteRange(netip.Addr, netip.Addr, []mmdbtype.DataType) error { + w.rows++ + w.cancel() + return nil +} + func TestSimpleReaderIteration(t *testing.T) { reader, err := mmdb.Open(cityTestDB) require.NoError(t, err) @@ -113,7 +174,7 @@ func TestMerger_MultipleDatabases(t *testing.T) { // Create merger and run merger, err := NewMerger(readers, cfg, writer) require.NoError(t, err) - err = merger.Merge() + err = merger.Merge(t.Context()) require.NoError(t, err) // Should have written some rows @@ -164,7 +225,7 @@ func TestMerger_EmitsNetworksPresentOnlyInLaterDatabase(t *testing.T) { merger, err := NewMerger(readers, cfg, writer) require.NoError(t, err) - err = merger.Merge() + err = merger.Merge(t.Context()) require.ErrorIs(t, err, errStopIteration) assert.True(t, writer.found, "expected to detect coverage for 214.0.0.1") } @@ -202,7 +263,7 @@ func TestMerger_AdjacentNetworkMerging(t *testing.T) { // Create merger and run merger, err := NewMerger(readers, cfg, writer) require.NoError(t, err) - err = merger.Merge() + err = merger.Merge(t.Context()) require.NoError(t, err) // The merger should have consolidated some networks @@ -325,7 +386,7 @@ func TestMerger_NilValues(t *testing.T) { merger, err := NewMerger(readers, cfg, writer) require.NoError(t, err) - err = merger.Merge() + err = merger.Merge(t.Context()) require.NoError(t, err) // Should have some rows @@ -418,7 +479,7 @@ func TestMerger_ResultAlignment(t *testing.T) { assert.Equal(t, 0, merger.extractors[0].dbIndex, "city column should map to index 0") assert.Equal(t, 1, merger.extractors[1].dbIndex, "anon column should map to index 1") - err = merger.Merge() + err = merger.Merge(t.Context()) require.NoError(t, err) // Should have successfully merged data from both databases @@ -449,7 +510,7 @@ func TestMerger_BroaderDatabase(t *testing.T) { merger, err := NewMerger(readers, cfg, writer) require.NoError(t, err) - err = merger.Merge() + err = merger.Merge(t.Context()) require.NoError(t, err) // Verify we got results and they contain data from both databases @@ -495,7 +556,7 @@ func TestMerger_MissingData(t *testing.T) { merger, err := NewMerger(readers, cfg, writer) require.NoError(t, err) - err = merger.Merge() + err = merger.Merge(t.Context()) require.NoError(t, err) // Should have rows @@ -562,7 +623,7 @@ func TestMerger_NoRedundantLookups(t *testing.T) { assert.Equal(t, 1, merger.extractors[3].dbIndex, "anon column should map to index 1") // Run the merge - err = merger.Merge() + err = merger.Merge(t.Context()) require.NoError(t, err) // Verify we got results with data from multiple columns diff --git a/internal/writer/csv.go b/internal/writer/csv.go index d23797f..274283f 100644 --- a/internal/writer/csv.go +++ b/internal/writer/csv.go @@ -422,8 +422,8 @@ func convertToString(value any) (string, error) { return strconv.FormatUint(uint64(v), 10), nil case mmdbtype.Uint64: return strconv.FormatUint(uint64(v), 10), nil - case *mmdbtype.Uint128: - return (*big.Int)(v).String(), nil + case mmdbtype.Uint128: + return formatUint128(v), nil case mmdbtype.Float32: return strconv.FormatFloat(float64(v), 'g', -1, 32), nil case mmdbtype.Float64: @@ -431,13 +431,13 @@ func convertToString(value any) (string, error) { case mmdbtype.Bytes: return hex.EncodeToString([]byte(v)), nil case mmdbtype.Map: - b, err := json.Marshal(v) + b, err := json.Marshal(normalizeMMDBJSON(v)) if err != nil { return "", fmt.Errorf("marshaling map to JSON: %w", err) } return string(b), nil case mmdbtype.Slice: - b, err := json.Marshal(v) + b, err := json.Marshal(normalizeMMDBJSON(v)) if err != nil { return "", fmt.Errorf("marshaling slice to JSON: %w", err) } @@ -447,3 +447,10 @@ func convertToString(value any) (string, error) { return fmt.Sprintf("%v", v), nil } } + +func formatUint128(value mmdbtype.Uint128) string { + if value.High == 0 { + return strconv.FormatUint(value.Low, 10) + } + return value.BigInt().String() +} diff --git a/internal/writer/csv_test.go b/internal/writer/csv_test.go index 8fa56c0..70349b0 100644 --- a/internal/writer/csv_test.go +++ b/internal/writer/csv_test.go @@ -4,6 +4,7 @@ import ( "bytes" "encoding/csv" "errors" + "math" "math/big" "net/netip" "strconv" @@ -600,6 +601,10 @@ func TestConvertToString(t *testing.T) { {"uint16", mmdbtype.Uint16(42), "42"}, {"uint32", mmdbtype.Uint32(42), "42"}, {"uint64", mmdbtype.Uint64(42), "42"}, + {"uint128 zero", mmdbtype.Uint128{}, "0"}, + {"uint128 low word", mmdbtype.Uint128{Low: 42}, "42"}, + {"uint128 uint64 boundary", mmdbtype.Uint128{Low: math.MaxUint64}, "18446744073709551615"}, + {"uint128 high word", mmdbtype.Uint128{High: 1, Low: 1}, "18446744073709551617"}, {"float64", mmdbtype.Float64(3.14), "3.14"}, {"bool true", mmdbtype.Bool(true), "1"}, {"bool false", mmdbtype.Bool(false), "0"}, diff --git a/internal/writer/format.go b/internal/writer/format.go index f734b11..21d14e0 100644 --- a/internal/writer/format.go +++ b/internal/writer/format.go @@ -71,7 +71,7 @@ func mmdbTypeName(value mmdbtype.DataType) string { return "uint32" case mmdbtype.Uint64: return "uint64" - case *mmdbtype.Uint128: + case mmdbtype.Uint128: return "uint128" default: return "unknown" diff --git a/internal/writer/format_test.go b/internal/writer/format_test.go index 0d740be..1ad27c7 100644 --- a/internal/writer/format_test.go +++ b/internal/writer/format_test.go @@ -5,7 +5,6 @@ import ( "encoding/csv" "fmt" "math" - "math/big" "net/netip" "testing" @@ -154,7 +153,7 @@ func TestCSVWriter_FormatTypeMismatch(t *testing.T) { {mmdbtype.Uint16(1), "uint16"}, {mmdbtype.Uint32(1), "uint32"}, {mmdbtype.Uint64(1), "uint64"}, - {(*mmdbtype.Uint128)(big.NewInt(1)), "uint128"}, + {mmdbtype.Uint128{Low: 1}, "uint128"}, {mmdbtype.String("1.2"), "string"}, {mmdbtype.Map{"value": mmdbtype.Float64(1.2)}, "map"}, {mmdbtype.Slice{mmdbtype.Float64(1.2)}, "array"}, diff --git a/internal/writer/integration_test.go b/internal/writer/integration_test.go index a923da4..1c40bc6 100644 --- a/internal/writer/integration_test.go +++ b/internal/writer/integration_test.go @@ -64,7 +64,7 @@ func TestEndToEnd_MMDBExport(t *testing.T) { require.NoError(t, err) m, err := merger.NewMerger(readers, cfg, writer) require.NoError(t, err) - require.NoError(t, m.Merge()) + require.NoError(t, m.Merge(t.Context())) require.NoError(t, writer.Flush()) output, err := maxminddb.OpenBytes(buf.Bytes()) @@ -154,7 +154,7 @@ func TestEndToEnd_CSVExport(t *testing.T) { // Create merger and run m, err := merger.NewMerger(readers, cfg, csvWriter) require.NoError(t, err) - err = m.Merge() + err = m.Merge(t.Context()) require.NoError(t, err) // Flush CSV @@ -223,7 +223,7 @@ func TestEndToEnd_CSVExport_MultipleNetworkColumns(t *testing.T) { m, err := merger.NewMerger(readers, cfg, csvWriter) require.NoError(t, err) - err = m.Merge() + err = m.Merge(t.Context()) require.NoError(t, err) err = csvWriter.Flush() @@ -287,7 +287,7 @@ func TestEndToEnd_CSVExport_MultipleColumns(t *testing.T) { m, err := merger.NewMerger(readers, cfg, csvWriter) require.NoError(t, err) - err = m.Merge() + err = m.Merge(t.Context()) require.NoError(t, err) err = csvWriter.Flush() @@ -358,7 +358,7 @@ func TestEndToEnd_CSVExport_NilValues(t *testing.T) { m, err := merger.NewMerger(readers, cfg, csvWriter) require.NoError(t, err) - err = m.Merge() + err = m.Merge(t.Context()) require.NoError(t, err) err = csvWriter.Flush() diff --git a/internal/writer/json.go b/internal/writer/json.go new file mode 100644 index 0000000..01c81b7 --- /dev/null +++ b/internal/writer/json.go @@ -0,0 +1,66 @@ +package writer + +import ( + "encoding/json" + "reflect" + + "github.com/maxmind/mmdbwriter/v2/mmdbtype" +) + +func normalizeMMDBJSON(value mmdbtype.DataType) any { + return jsonValue(value, map[jsonContainer]any{}) +} + +type jsonContainer struct { + pointer uintptr + length int + kind reflect.Kind +} + +// jsonValue replaces Uint128 leaves with exact JSON numbers. Memoized copies +// preserve sharing and cycles without changing the input; json.Marshal retains +// responsibility for rejecting cycles and unsupported scalar values. +func jsonValue(value mmdbtype.DataType, seen map[jsonContainer]any) any { + switch v := value.(type) { + case mmdbtype.Uint128: + return json.Number(formatUint128(v)) + case mmdbtype.Map: + if v == nil { + return nil + } + key := jsonContainer{ + pointer: reflect.ValueOf(v).Pointer(), + length: len(v), + kind: reflect.Map, + } + if previous, ok := seen[key]; ok { + return previous + } + result := make(map[string]any, len(v)) + seen[key] = result + for name, child := range v { + result[string(name)] = jsonValue(child, seen) + } + return result + case mmdbtype.Slice: + if v == nil { + return nil + } + key := jsonContainer{ + pointer: reflect.ValueOf(v).Pointer(), + length: len(v), + kind: reflect.Slice, + } + if previous, ok := seen[key]; ok { + return previous + } + result := make([]any, len(v)) + seen[key] = result + for i, child := range v { + result[i] = jsonValue(child, seen) + } + return result + default: + return value + } +} diff --git a/internal/writer/json_test.go b/internal/writer/json_test.go new file mode 100644 index 0000000..94fb878 --- /dev/null +++ b/internal/writer/json_test.go @@ -0,0 +1,87 @@ +package writer + +import ( + "encoding/json" + "math" + "testing" + + "github.com/maxmind/mmdbwriter/v2/mmdbtype" + "github.com/stretchr/testify/require" +) + +func TestConvertToString_JSON(t *testing.T) { + sharedMap := mmdbtype.Map{"n": mmdbtype.Uint128{Low: 42}} + sharedSlice := mmdbtype.Slice{mmdbtype.Uint128{Low: 1}, mmdbtype.Uint128{High: 1, Low: 1}} + for _, tt := range []struct { + name string + value mmdbtype.DataType + want string + }{ + {"map", sharedMap, `{"n":42}`}, + {"slice", mmdbtype.Slice{mmdbtype.Uint128{}, mmdbtype.Uint128{High: 1, Low: 1}}, `[0,18446744073709551617]`}, + { + "mixed nesting", + mmdbtype.Map{"values": mmdbtype.Slice{mmdbtype.Map{"n": mmdbtype.Uint128{High: math.MaxUint64, Low: math.MaxUint64}}}}, + `{"values":[{"n":340282366920938463463374607431768211455}]}`, + }, + { + "other scalars", + mmdbtype.Slice{mmdbtype.Bool(true), mmdbtype.Bytes{0xaa, 0xbb}, mmdbtype.Int32(-1), mmdbtype.Float64(1.5), mmdbtype.String("text"), nil}, + `[true,"qrs=",-1,1.5,"text",null]`, + }, + {"nil map", mmdbtype.Map(nil), `null`}, + {"empty map", mmdbtype.Map{}, `{}`}, + {"nil slice", mmdbtype.Slice(nil), `null`}, + {"empty slice", mmdbtype.Slice{}, `[]`}, + {"shared map", mmdbtype.Slice{sharedMap, sharedMap}, `[{"n":42},{"n":42}]`}, + {"shared slice", mmdbtype.Slice{sharedSlice, sharedSlice}, `[[1,18446744073709551617],[1,18446744073709551617]]`}, + {"slice prefixes", mmdbtype.Slice{sharedSlice[:1], sharedSlice}, `[[1],[1,18446744073709551617]]`}, + } { + t.Run(tt.name, func(t *testing.T) { + got, err := convertToString(tt.value) + require.NoError(t, err) + // Exact text preserves integer precision that JSON float decoding loses. + require.Equal(t, tt.want, got) + }) + } + // Normalizing a value must not replace leaves in the original containers. + require.Equal(t, mmdbtype.Uint128{Low: 42}, sharedMap["n"]) + require.Equal(t, mmdbtype.Uint128{High: 1, Low: 1}, sharedSlice[1]) +} + +func TestConvertToString_JSONErrors(t *testing.T) { + cyclicMap := mmdbtype.Map{} + cyclicMap["self"] = cyclicMap + cyclicSlice := make(mmdbtype.Slice, 1) + cyclicSlice[0] = cyclicSlice + for _, tt := range []struct { + name string + value mmdbtype.DataType + context string + }{ + {"map cycle", cyclicMap, "marshaling map to JSON"}, + {"slice cycle", cyclicSlice, "marshaling slice to JSON"}, + {"nested NaN", mmdbtype.Map{"values": mmdbtype.Slice{mmdbtype.Float64(math.NaN())}}, "marshaling map to JSON"}, + {"nested infinity", mmdbtype.Slice{mmdbtype.Map{"n": mmdbtype.Float64(math.Inf(1))}}, "marshaling slice to JSON"}, + } { + t.Run(tt.name, func(t *testing.T) { + got, err := convertToString(tt.value) + require.Empty(t, got) + require.ErrorContains(t, err, tt.context) + var unsupported *json.UnsupportedValueError + require.ErrorAs(t, err, &unsupported) + }) + } +} + +func TestConvertToParquetType_JSON(t *testing.T) { + value := mmdbtype.Map{"n": mmdbtype.Slice{mmdbtype.Uint128{High: 1, Low: 1}}} + for _, typeHint := range []string{"", "string"} { + t.Run("type="+typeHint, func(t *testing.T) { + got, err := convertToParquetType(value, typeHint) + require.NoError(t, err) + //nolint:testifylint // JSONEq loses integer precision through float64 decoding. + require.Equal(t, `{"n":[18446744073709551617]}`, got) + }) + } +} diff --git a/internal/writer/parquet.go b/internal/writer/parquet.go index cf433f8..e010ef4 100644 --- a/internal/writer/parquet.go +++ b/internal/writer/parquet.go @@ -4,7 +4,7 @@ import ( "errors" "fmt" "io" - "math/big" + "math" "net/netip" "github.com/maxmind/mmdbwriter/v2/mmdbtype" @@ -394,12 +394,11 @@ func convertToParquetType(value any, typeHint string) (any, error) { return nil, fmt.Errorf("uint64 value %d overflows int64", v) } return int64(v), nil - case *mmdbtype.Uint128: - i := (*big.Int)(v) - if !i.IsInt64() { - return nil, fmt.Errorf("uint128 value %s overflows int64", i.String()) + case mmdbtype.Uint128: + if v.High != 0 || v.Low > math.MaxInt64 { + return nil, fmt.Errorf("uint128 value %s overflows int64", formatUint128(v)) } - return i.Int64(), nil + return int64(v.Low), nil default: return nil, fmt.Errorf("cannot convert %T to int64", value) } @@ -418,8 +417,12 @@ func convertToParquetType(value any, typeHint string) (any, error) { return float64(v), nil case mmdbtype.Uint64: return float64(v), nil - case *mmdbtype.Uint128: - i := (*big.Int)(v) + case mmdbtype.Uint128: + if v.High == 0 { + return float64(v.Low), nil + } + // Converting the two words separately can round twice. + i := v.BigInt() f, _ := i.Float64() return f, nil default: diff --git a/internal/writer/parquet_test.go b/internal/writer/parquet_test.go index 7706e88..51f11b8 100644 --- a/internal/writer/parquet_test.go +++ b/internal/writer/parquet_test.go @@ -4,6 +4,7 @@ import ( "bytes" "fmt" "io" + "math" "net/netip" "testing" @@ -495,6 +496,24 @@ func TestConvertToParquetType(t *testing.T) { {"int32 to int64", mmdbtype.Int32(42), "int64", int64(42), false}, {"uint16 to int64", mmdbtype.Uint16(42), "int64", int64(42), false}, {"uint32 to int64", mmdbtype.Uint32(42), "int64", int64(42), false}, + {"uint128 to int64", mmdbtype.Uint128{Low: 42}, "int64", int64(42), false}, + { + "uint128 int64 boundary", + mmdbtype.Uint128{Low: 1<<63 - 1}, + "int64", + int64(1<<63 - 1), + false, + }, + {"uint128 low word overflow", mmdbtype.Uint128{Low: 1 << 63}, "int64", nil, true}, + {"uint128 high word overflow", mmdbtype.Uint128{High: 1}, "int64", nil, true}, + {"uint128 to float64", mmdbtype.Uint128{High: 1}, "float64", float64(1 << 64), false}, + { + "uint128 to string", + mmdbtype.Uint128{High: 1, Low: 1}, + "string", + "18446744073709551617", + false, + }, {"float32 to float64", mmdbtype.Float32(3.14), "float64", float64(float32(3.14)), false}, {"uint32 to float64", mmdbtype.Uint32(42), "float64", float64(42), false}, {"bool", mmdbtype.Bool(true), "bool", true, false}, @@ -516,6 +535,25 @@ func TestConvertToParquetType(t *testing.T) { } } +func TestConvertToParquetType_Uint128FloatRounding(t *testing.T) { + for _, value := range []mmdbtype.Uint128{ + {}, + {Low: 42}, + {Low: math.MaxUint64}, + {High: 1}, + {High: 1<<53 + 1, Low: 1 << 63}, + {High: math.MaxUint64, Low: math.MaxUint64}, + } { + t.Run(value.BigInt().String(), func(t *testing.T) { + want, _ := value.BigInt().Float64() + got, err := convertToParquetType(value, "float64") + require.NoError(t, err) + // Compare bits to catch one-ULP rounding differences. + require.Equal(t, math.Float64bits(want), math.Float64bits(got.(float64))) + }) + } +} + func TestParquetWriter_NetworkBucket_IPv4(t *testing.T) { tests := []struct { name string diff --git a/mmdbconvert.go b/mmdbconvert.go index 0457b9c..7dd2d0c 100644 --- a/mmdbconvert.go +++ b/mmdbconvert.go @@ -3,6 +3,7 @@ package mmdbconvert import ( + "context" "errors" "fmt" @@ -28,8 +29,17 @@ type Options struct { DisableCache bool } -// Run performs the MMDB conversion using the specified options. -func Run(opts Options) error { +// Run performs the conversion, stopping and cleaning staged outputs when +// ctx is canceled before publication. Once publication starts, it runs to completion. +// Cancellation is cooperative and may wait for an operation already in progress. +// Run does not install signal handlers. +// When cancellation stops conversion, the returned error matches context.Canceled +// or context.DeadlineExceeded through errors.Is. Custom causes remain available +// through context.Cause(ctx). +func Run(ctx context.Context, opts Options) error { + if err := ctx.Err(); err != nil { + return fmt.Errorf("starting conversion: %w", err) + } if opts.ConfigPath == "" { return errors.New("config path is required") } @@ -58,15 +68,16 @@ func Run(opts Options) error { return fmt.Errorf("validating network columns: %w", err) } - rowWriter, outputs, err := prepareRowWriter(cfg, readers) + rowWriter, outputs, err := prepareRowWriter(ctx, cfg, readers) if err != nil { return err } - return convert(cfg, readers, rowWriter, outputs) + return convert(ctx, cfg, readers, rowWriter, outputs) } // convert owns the prepared outputs through conversion, publication, and cleanup. func convert( + ctx context.Context, cfg *config.Config, readers *mmdb.Readers, rowWriter merger.RowWriter, @@ -80,11 +91,11 @@ func convert( if err != nil { return fmt.Errorf("creating merger: %w", err) } - if err := m.Merge(); err != nil { + if err := m.Merge(ctx); err != nil { return fmt.Errorf("merging databases: %w", err) } - return flushAndCommit(rowWriter, outputs) + return flushAndCommit(ctx, rowWriter, outputs) } func detectIPVersionFromDatabases(cfg *config.Config, readers *mmdb.Readers) (int, error) { diff --git a/mmdbconvert_test.go b/mmdbconvert_test.go index 06971ea..26fab76 100644 --- a/mmdbconvert_test.go +++ b/mmdbconvert_test.go @@ -44,7 +44,7 @@ path = ["country", "iso_code"] err = os.WriteFile(configFile, []byte(configContent), 0o600) require.NoError(t, err) - err = Run(Options{ConfigPath: configFile}) + err = Run(t.Context(), Options{ConfigPath: configFile}) require.NoError(t, err) // Verify output file was created @@ -86,7 +86,7 @@ path = ["country", "iso_code"] require.NoError(t, err) // Run with DisableCache option - err = Run(Options{ + err = Run(t.Context(), Options{ ConfigPath: configFile, DisableCache: true, }) @@ -99,13 +99,13 @@ path = ["country", "iso_code"] } func TestRun_MissingConfigPath(t *testing.T) { - err := Run(Options{}) + err := Run(t.Context(), Options{}) require.Error(t, err) assert.Contains(t, err.Error(), "config path is required") } func TestRun_NonexistentConfigFile(t *testing.T) { - err := Run(Options{ConfigPath: "/nonexistent/config.toml"}) + err := Run(t.Context(), Options{ConfigPath: "/nonexistent/config.toml"}) require.Error(t, err) assert.Contains(t, err.Error(), "loading config") } @@ -117,7 +117,7 @@ func TestRun_InvalidConfig(t *testing.T) { err := os.WriteFile(configFile, []byte("invalid toml [[["), 0o600) require.NoError(t, err) - err = Run(Options{ConfigPath: configFile}) + err = Run(t.Context(), Options{ConfigPath: configFile}) require.Error(t, err) assert.Contains(t, err.Error(), "loading config") } @@ -144,7 +144,7 @@ path = ["country", "iso_code"] err := os.WriteFile(configFile, []byte(configContent), 0o600) require.NoError(t, err) - err = Run(Options{ConfigPath: configFile}) + err = Run(t.Context(), Options{ConfigPath: configFile}) require.Error(t, err) assert.Contains(t, err.Error(), "opening databases") } @@ -179,7 +179,7 @@ path = ["country", "iso_code"] err = os.WriteFile(configFile, []byte(configContent), 0o600) require.NoError(t, err) - err = Run(Options{ConfigPath: configFile}) + err = Run(t.Context(), Options{ConfigPath: configFile}) require.NoError(t, err) // Verify output file was created @@ -216,7 +216,7 @@ path = ["country", "iso_code"] err = os.WriteFile(configFile, []byte(configContent), 0o600) require.NoError(t, err) - err = Run(Options{ConfigPath: configFile}) + err = Run(t.Context(), Options{ConfigPath: configFile}) require.NoError(t, err) // Verify both output files were created @@ -339,7 +339,7 @@ path = [%q] `, outputPaths, csvConfig, tomlPath(databasePath), tt.field) configFile := filepath.Join(tmpDir, "config.toml") require.NoError(t, os.WriteFile(configFile, []byte(configContent), 0o600)) - require.NoError(t, Run(Options{ConfigPath: configFile})) + require.NoError(t, Run(t.Context(), Options{ConfigPath: configFile})) for name, want := range wantFiles { content, err := os.ReadFile(filepath.Clean(filepath.Join(tmpDir, name))) diff --git a/output.go b/output.go index 987cd55..e9e1ee7 100644 --- a/output.go +++ b/output.go @@ -1,6 +1,7 @@ package mmdbconvert import ( + "context" "errors" "fmt" "io" @@ -23,6 +24,7 @@ type pendingOutput interface { } func prepareRowWriter( + ctx context.Context, cfg *config.Config, readers *mmdb.Readers, ) (merger.RowWriter, []pendingOutput, error) { @@ -40,6 +42,12 @@ func prepareRowWriter( var outputs []pendingOutput var writers []merger.RowWriter for i, destination := range destinations { + if err := ctx.Err(); err != nil { + return nil, nil, errors.Join( + fmt.Errorf("preparing outputs: %w", err), + cleanupOutputs(outputs), + ) + } file, err := newPendingOutput(destination) if err != nil { return nil, nil, errors.Join(err, cleanupOutputs(outputs)) @@ -53,7 +61,12 @@ func prepareRowWriter( } } outputs = append(outputs, output) - rowWriter, err := newRowWriter(file, cfg, readers, versions[i]) + rowWriter, err := newRowWriter( + contextWriter{ctx: ctx, writer: file}, + cfg, + readers, + versions[i], + ) if err != nil { return nil, nil, errors.Join( fmt.Errorf("preparing output %s: %w", destination.path, err), @@ -200,12 +213,45 @@ func newRowWriter( } } -func flushAndCommit(rowWriter merger.RowWriter, outputs []pendingOutput) error { +// contextWriter checks cancellation at serialization write boundaries. It never +// closes the underlying output; convert retains cleanup ownership. +// +//nolint:containedctx // io.Writer cannot take a context argument; this adapter belongs to one conversion. +type contextWriter struct { + ctx context.Context + writer io.Writer +} + +func (w contextWriter) Write(p []byte) (int, error) { + if err := w.ctx.Err(); err != nil { + return 0, fmt.Errorf("writing output: %w", err) + } + n, err := w.writer.Write(p) + if err != nil { + return n, fmt.Errorf("writing output: %w", err) + } + return n, nil +} + +func flushAndCommit( + ctx context.Context, + rowWriter merger.RowWriter, + outputs []pendingOutput, +) error { + // The accumulator's final write can finish after cancellation. Avoid starting + // format serialization, which may do substantial work before its first write. + if err := ctx.Err(); err != nil { + return fmt.Errorf("flushing output: %w", err) + } if flusher, ok := rowWriter.(interface{ Flush() error }); ok { if err := flusher.Flush(); err != nil { return fmt.Errorf("flushing output: %w", err) } } + if err := ctx.Err(); err != nil { + return fmt.Errorf("publishing outputs: %w", err) + } + // Once publication starts, finish the commit sequence even if canceled. // All writers must flush before any output is published. for _, file := range outputs { if err := file.Commit(); err != nil { diff --git a/output_test.go b/output_test.go index 014b09d..487f75c 100644 --- a/output_test.go +++ b/output_test.go @@ -38,7 +38,7 @@ func TestRun_OutputPublication(t *testing.T) { for _, path := range paths { require.NoError(t, os.WriteFile(path, []byte("previous output"), 0o600)) } - require.NoError(t, Run(Options{ConfigPath: configPath})) + require.NoError(t, Run(t.Context(), Options{ConfigPath: configPath})) for _, path := range paths { data, err := os.ReadFile(filepath.Clean(path)) require.NoError(t, err) @@ -88,7 +88,11 @@ func TestRun_ConversionFailurePreservesOutputs(t *testing.T) { for _, path := range paths { require.NoError(t, os.WriteFile(path, []byte("previous output"), 0o600)) } - require.ErrorContains(t, Run(Options{ConfigPath: configPath}), "merging databases") + require.ErrorContains( + t, + Run(t.Context(), Options{ConfigPath: configPath}), + "merging databases", + ) for _, path := range paths { assertFileContent(t, path, "previous output") } @@ -109,7 +113,7 @@ func TestRun_RejectsDuplicateSplitPathsBeforeOpeningFiles(t *testing.T) { require.NoError(t, err) require.NoError(t, os.WriteFile(configPath, data, 0o600)) require.NoError(t, os.WriteFile(paths[0], []byte("previous output"), 0o600)) - err = Run(Options{ConfigPath: configPath}) + err = Run(t.Context(), Options{ConfigPath: configPath}) require.ErrorContains( t, err, @@ -128,7 +132,7 @@ func TestPrepareRowWriter_SecondOutputFailure(t *testing.T) { require.NoError(t, os.WriteFile(paths[0], []byte("previous output"), 0o600)) cfg.Output.IPv6File = filepath.Join(filepath.Dir(paths[1]), "missing", "out") readers := openOutputTestReaders(t, cfg) - _, _, err := prepareRowWriter(cfg, readers) + _, _, err := prepareRowWriter(t.Context(), cfg, readers) require.ErrorContains(t, err, cfg.Output.IPv6File) assertFileContent(t, paths[0], "previous output") assertOutputDirectory(t, paths[:1]) @@ -190,7 +194,7 @@ func TestRun_RejectsDirectoryBeforeConversion(t *testing.T) { } dir := paths[len(paths)-1] require.NoError(t, os.Mkdir(dir, 0o700)) - err := Run(Options{ConfigPath: configPath}) + err := Run(t.Context(), Options{ConfigPath: configPath}) require.ErrorContains(t, err, "not a regular file") require.Contains(t, filepath.ToSlash(err.Error()), filepath.ToSlash(dir)) require.DirExists(t, dir) @@ -234,7 +238,7 @@ func TestOutput_FlushFailure(t *testing.T) { []mmdbtype.DataType{mmdbtype.String("GB")}, ), ) - err = flushAndCommit(rowWriter, []pendingOutput{file}) + err = flushAndCommit(t.Context(), rowWriter, []pendingOutput{file}) require.ErrorIs(t, err, writeErr) require.NoError(t, file.Cleanup()) assertFileContent(t, paths[0], "previous output") @@ -266,7 +270,7 @@ func TestOutput_SplitFlushFailure(t *testing.T) { ), ) } - require.ErrorIs(t, flushAndCommit(split, files), writeErr) + require.ErrorIs(t, flushAndCommit(t.Context(), split, files), writeErr) require.NoError(t, cleanupOutputs(files)) for _, path := range paths { assertFileContent(t, path, "previous output") @@ -355,7 +359,7 @@ func TestOutput_CommitFailure(t *testing.T) { []mmdbtype.DataType{mmdbtype.String("GB")}, ), ) - require.ErrorContains(t, flushAndCommit(split, files), paths[1]) + require.ErrorContains(t, flushAndCommit(t.Context(), split, files), paths[1]) require.NoError(t, cleanupOutputs(files)) assertFileContent(t, paths[0], "network,country_code\n2.0.0.0/24,GB\n") require.DirExists(t, paths[1]) @@ -395,7 +399,7 @@ func TestConvert_PreservesPrimaryAndCleanupErrors(t *testing.T) { case "publication": first.commitErr = primaryErr } - err := convert(cfg, readers, rowWriter, []pendingOutput{first, second}) + err := convert(t.Context(), cfg, readers, rowWriter, []pendingOutput{first, second}) if stage != "cleanup only" { require.ErrorIs(t, err, primaryErr) } @@ -593,7 +597,7 @@ func TestRun_FilesystemAliases(t *testing.T) { require.NoError(t, os.WriteFile(configPath, data, 0o600)) require.ErrorContains( t, - Run(Options{ConfigPath: configPath}), + Run(t.Context(), Options{ConfigPath: configPath}), "refer to the same file", ) if existing { diff --git a/output_unix_test.go b/output_unix_test.go index 42fe326..ef8cf38 100644 --- a/output_unix_test.go +++ b/output_unix_test.go @@ -160,7 +160,7 @@ func TestRun_SplitParentAliases(t *testing.T) { require.NoError(t, err) require.NoError(t, os.WriteFile(configPath, data, 0o600)) require.NoError(t, os.WriteFile(paths[0], []byte("previous output"), 0o600)) - err = Run(Options{ConfigPath: configPath}) + err = Run(t.Context(), Options{ConfigPath: configPath}) if !distinct { require.ErrorContains(t, err, "ignoring case") assertFileContent(t, paths[0], "previous output") diff --git a/variables_test.go b/variables_test.go index ee271ed..4335389 100644 --- a/variables_test.go +++ b/variables_test.go @@ -45,7 +45,10 @@ path = ["value"] require.NoError(t, os.WriteFile(configPath, []byte(content), 0o600)) variables := map[string]string{"input": databasePath, "output": outputPath} original := maps.Clone(variables) - require.NoError(t, Run(Options{ConfigPath: configPath, Variables: variables})) + require.NoError( + t, + Run(t.Context(), Options{ConfigPath: configPath, Variables: variables}), + ) assert.Equal(t, original, variables) switch format { @@ -102,7 +105,7 @@ database = "source" path = ["value"] ` require.NoError(t, os.WriteFile(configPath, []byte(content), 0o600)) - require.NoError(t, Run(Options{ + require.NoError(t, Run(t.Context(), Options{ ConfigPath: configPath, Variables: map[string]string{"input": relativeInput, "output_dir": "."}, })) @@ -157,6 +160,7 @@ path = %q variables := map[string]string{"output": outputPath} maps.Copy(variables, tt.variables) err := Run( + t.Context(), Options{ConfigPath: configPath, Variables: variables}, ) require.ErrorContains(