From 0c4055fc6fd9fc2d6d584914bcd060681579b03a Mon Sep 17 00:00:00 2001 From: Angus Peart Date: Thu, 25 Sep 2025 16:05:45 -0700 Subject: [PATCH 1/2] Add --quiet mode and ProgressWriter interface --- main.go | 17 +- pmtiles/cluster.go | 15 +- pmtiles/convert.go | 55 +++-- pmtiles/edit.go | 62 ++++-- pmtiles/extract.go | 35 +-- pmtiles/makesync.go | 14 +- pmtiles/progress.go | 146 +++++++++++++ pmtiles/progress_test.go | 452 +++++++++++++++++++++++++++++++++++++++ pmtiles/sync.go | 89 +++++--- pmtiles/upload.go | 11 +- 10 files changed, 801 insertions(+), 95 deletions(-) create mode 100644 pmtiles/progress.go create mode 100644 pmtiles/progress_test.go diff --git a/main.go b/main.go index a259821..f0a72cf 100644 --- a/main.go +++ b/main.go @@ -27,6 +27,8 @@ var ( ) var cli struct { + Quiet bool `help:"Suppress verbose output and progress bars"` + Show struct { Path string `arg:""` Bucket string `help:"Remote bucket"` @@ -127,6 +129,9 @@ func main() { logger := log.New(os.Stdout, "", log.Ldate|log.Ltime|log.Lshortfile) ctx := kong.Parse(&cli) + // Set quiet mode globally based on the --quiet flag + pmtiles.SetQuietMode(cli.Quiet) + switch ctx.Command() { case "show ": err := pmtiles.Show(logger, os.Stdout, cli.Show.Bucket, cli.Show.Path, cli.Show.HeaderJson, cli.Show.Metadata, cli.Show.Tilejson, cli.Show.PublicURL, false, 0, 0, 0) @@ -153,14 +158,20 @@ func main() { mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { start := time.Now() statusCode := server.ServeHTTP(w, r) - logger.Printf("served %d %s in %s", statusCode, url.PathEscape(r.URL.Path), time.Since(start)) + if !cli.Quiet { + logger.Printf("served %d %s in %s", statusCode, url.PathEscape(r.URL.Path), time.Since(start)) + } }) - logger.Printf("Serving %s %s on port %d and interface %s with Access-Control-Allow-Origin: %s\n", cli.Serve.Bucket, cli.Serve.Path, cli.Serve.Port, cli.Serve.Interface, cli.Serve.Cors) + if !cli.Quiet { + logger.Printf("Serving %s %s on port %d and interface %s with Access-Control-Allow-Origin: %s\n", cli.Serve.Bucket, cli.Serve.Path, cli.Serve.Port, cli.Serve.Interface, cli.Serve.Cors) + } if cli.Serve.AdminPort > 0 { go func() { adminPort := strconv.Itoa(cli.Serve.AdminPort) - logger.Printf("Serving /metrics on port %s and interface %s\n", adminPort, cli.Serve.Interface) + if !cli.Quiet { + logger.Printf("Serving /metrics on port %s and interface %s\n", adminPort, cli.Serve.Interface) + } adminMux := http.NewServeMux() adminMux.Handle("/metrics", promhttp.Handler()) logger.Fatal(startHTTPServer(cli.Serve.Interface+":"+adminPort, adminMux)) diff --git a/pmtiles/cluster.go b/pmtiles/cluster.go index e10f3f8..59d4c40 100644 --- a/pmtiles/cluster.go +++ b/pmtiles/cluster.go @@ -2,7 +2,6 @@ package pmtiles import ( "fmt" - "github.com/schollz/progressbar/v3" "io" "log" "os" @@ -41,7 +40,11 @@ func Cluster(logger *log.Logger, InputPMTiles string, deduplicate bool) error { return err } - bar := progressbar.Default(int64(header.TileEntriesCount)) + var progress Progress + progressWriter := getProgressWriter() + if progressWriter != nil { + progress = progressWriter.NewCountProgress(int64(header.TileEntriesCount), "") + } err = IterateEntries(header, func(offset uint64, length uint64) ([]byte, error) { @@ -52,7 +55,9 @@ func Cluster(logger *log.Logger, InputPMTiles string, deduplicate bool) error { if isNew, newData := resolver.AddTileIsNew(e.TileID, data, e.RunLength); isNew { tmpfile.Write(newData) } - bar.Add(1) + if progress != nil { + progress.Add(1) + } }) if err != nil { @@ -66,6 +71,8 @@ func Cluster(logger *log.Logger, InputPMTiles string, deduplicate bool) error { if err != nil { return err } - fmt.Printf("total directory size %d (%f%% of original)\n", newHeader.RootLength+newHeader.LeafDirectoryLength, float64(newHeader.RootLength+newHeader.LeafDirectoryLength)/float64(header.RootLength+header.LeafDirectoryLength)*100) + if !quietMode { + fmt.Printf("total directory size %d (%f%% of original)\n", newHeader.RootLength+newHeader.LeafDirectoryLength, float64(newHeader.RootLength+newHeader.LeafDirectoryLength)/float64(header.RootLength+header.LeafDirectoryLength)*100) + } return nil } diff --git a/pmtiles/convert.go b/pmtiles/convert.go index 22a2195..330dc76 100644 --- a/pmtiles/convert.go +++ b/pmtiles/convert.go @@ -16,10 +16,11 @@ import ( "time" "github.com/RoaringBitmap/roaring/roaring64" - "github.com/schollz/progressbar/v3" "zombiezen.com/go/sqlite" ) +var quietMode bool + type offsetLen struct { Offset uint64 Length uint32 @@ -156,7 +157,9 @@ func convertMbtiles(logger *log.Logger, input string, output string, deduplicate return fmt.Errorf("Failed to convert MBTiles to header JSON, %w", err) } - logger.Println("Pass 1: Assembling TileID set") + if !quietMode { + logger.Println("Pass 1: Assembling TileID set") + } // assemble a sorted set of all TileIds tileset := roaring64.New() { @@ -187,10 +190,16 @@ func convertMbtiles(logger *log.Logger, input string, output string, deduplicate return fmt.Errorf("no tiles in MBTiles archive") } - logger.Println("Pass 2: writing tiles") + if !quietMode { + logger.Println("Pass 2: writing tiles") + } resolve := newResolver(deduplicate, header.TileType == Mvt) { - bar := progressbar.Default(int64(tileset.GetCardinality())) + var progress Progress + progressWriter := getProgressWriter() + if progressWriter != nil { + progress = progressWriter.NewCountProgress(int64(tileset.GetCardinality()), "") + } i := tileset.Iterator() stmt := conn.Prep("SELECT tile_data FROM tiles WHERE zoom_level = ? AND tile_column = ? AND tile_row = ?") @@ -229,21 +238,27 @@ func convertMbtiles(logger *log.Logger, input string, output string, deduplicate stmt.ClearBindings() stmt.Reset() - bar.Add(1) + if progress != nil { + progress.Add(1) + } } } _, err = finalize(logger, resolve, header, tmpfile, output, jsonMetadata) if err != nil { return err } - logger.Println("Finished in ", time.Since(start)) + if !quietMode { + logger.Println("Finished in ", time.Since(start)) + } return nil } func finalize(logger *log.Logger, resolve *resolver, header HeaderV3, tmpfile *os.File, output string, jsonMetadata map[string]interface{}) (HeaderV3, error) { - logger.Println("# of addressed tiles: ", resolve.AddressedTiles) - logger.Println("# of tile entries (after RLE): ", len(resolve.Entries)) - logger.Println("# of tile contents: ", resolve.NumContents()) + if !quietMode { + logger.Println("# of addressed tiles: ", resolve.AddressedTiles) + logger.Println("# of tile entries (after RLE): ", len(resolve.Entries)) + logger.Println("# of tile contents: ", resolve.NumContents()) + } header.AddressedTilesCount = resolve.AddressedTiles header.TileEntriesCount = uint64(len(resolve.Entries)) @@ -258,16 +273,18 @@ func finalize(logger *log.Logger, resolve *resolver, header HeaderV3, tmpfile *o rootBytes, leavesBytes, numLeaves := optimizeDirectories(resolve.Entries, 16384-HeaderV3LenBytes, Gzip) - if numLeaves > 0 { - logger.Println("Root dir bytes: ", len(rootBytes)) - logger.Println("Leaves dir bytes: ", len(leavesBytes)) - logger.Println("Num leaf dirs: ", numLeaves) - logger.Println("Total dir bytes: ", len(rootBytes)+len(leavesBytes)) - logger.Println("Average leaf dir bytes: ", len(leavesBytes)/numLeaves) - logger.Printf("Average bytes per addressed tile: %.2f\n", float64(len(rootBytes)+len(leavesBytes))/float64(resolve.AddressedTiles)) - } else { - logger.Println("Total dir bytes: ", len(rootBytes)) - logger.Printf("Average bytes per addressed tile: %.2f\n", float64(len(rootBytes))/float64(resolve.AddressedTiles)) + if !quietMode { + if numLeaves > 0 { + logger.Println("Root dir bytes: ", len(rootBytes)) + logger.Println("Leaves dir bytes: ", len(leavesBytes)) + logger.Println("Num leaf dirs: ", numLeaves) + logger.Println("Total dir bytes: ", len(rootBytes)+len(leavesBytes)) + logger.Println("Average leaf dir bytes: ", len(leavesBytes)/numLeaves) + logger.Printf("Average bytes per addressed tile: %.2f\n", float64(len(rootBytes)+len(leavesBytes))/float64(resolve.AddressedTiles)) + } else { + logger.Println("Total dir bytes: ", len(rootBytes)) + logger.Printf("Average bytes per addressed tile: %.2f\n", float64(len(rootBytes))/float64(resolve.AddressedTiles)) + } } metadataBytes, err := SerializeMetadata(jsonMetadata, Gzip) diff --git a/pmtiles/edit.go b/pmtiles/edit.go index 8de5b73..401e330 100644 --- a/pmtiles/edit.go +++ b/pmtiles/edit.go @@ -4,7 +4,6 @@ import ( "bytes" "encoding/json" "fmt" - "github.com/schollz/progressbar/v3" "io" "io/ioutil" "log" @@ -110,31 +109,64 @@ func Edit(_ *log.Logger, inputArchive string, newHeaderJSONFile string, newMetad newHeader.LeafDirectoryOffset = newHeader.MetadataOffset + newHeader.MetadataLength newHeader.TileDataOffset = newHeader.LeafDirectoryOffset + newHeader.LeafDirectoryLength - bar := progressbar.DefaultBytes( - int64(HeaderV3LenBytes+newHeader.RootLength+uint64(len(metadataBytes))+newHeader.LeafDirectoryLength+newHeader.TileDataLength), - "writing file", - ) - buf = SerializeHeader(newHeader) - io.Copy(io.MultiWriter(outfile, bar), bytes.NewReader(buf)) + + var progress Progress + progressWriter := getProgressWriter() + if progressWriter != nil { + progress = progressWriter.NewBytesProgress( + int64(HeaderV3LenBytes+newHeader.RootLength+uint64(len(metadataBytes))+newHeader.LeafDirectoryLength+newHeader.TileDataLength), + "writing file", + ) + } + + if progress != nil { + io.Copy(io.MultiWriter(outfile, progress), bytes.NewReader(buf)) + } else { + io.Copy(outfile, bytes.NewReader(buf)) + } rootSection := io.NewSectionReader(file, int64(oldHeader.RootOffset), int64(oldHeader.RootLength)) - if _, err := io.Copy(io.MultiWriter(outfile, bar), rootSection); err != nil { - return err + if progress != nil { + if _, err := io.Copy(io.MultiWriter(outfile, progress), rootSection); err != nil { + return err + } + } else { + if _, err := io.Copy(outfile, rootSection); err != nil { + return err + } } - if _, err := io.Copy(io.MultiWriter(outfile, bar), bytes.NewReader(metadataBytes)); err != nil { - return err + if progress != nil { + if _, err := io.Copy(io.MultiWriter(outfile, progress), bytes.NewReader(metadataBytes)); err != nil { + return err + } + } else { + if _, err := io.Copy(outfile, bytes.NewReader(metadataBytes)); err != nil { + return err + } } leafSection := io.NewSectionReader(file, int64(oldHeader.LeafDirectoryOffset), int64(oldHeader.LeafDirectoryLength)) - if _, err := io.Copy(io.MultiWriter(outfile, bar), leafSection); err != nil { - return err + if progress != nil { + if _, err := io.Copy(io.MultiWriter(outfile, progress), leafSection); err != nil { + return err + } + } else { + if _, err := io.Copy(outfile, leafSection); err != nil { + return err + } } tileSection := io.NewSectionReader(file, int64(oldHeader.TileDataOffset), int64(oldHeader.TileDataLength)) - if _, err := io.Copy(io.MultiWriter(outfile, bar), tileSection); err != nil { - return err + if progress != nil { + if _, err := io.Copy(io.MultiWriter(outfile, progress), tileSection); err != nil { + return err + } + } else { + if _, err := io.Copy(outfile, tileSection); err != nil { + return err + } } // explicitly close in order to rename diff --git a/pmtiles/extract.go b/pmtiles/extract.go index abe55ba..b195fb4 100644 --- a/pmtiles/extract.go +++ b/pmtiles/extract.go @@ -8,7 +8,6 @@ import ( "github.com/RoaringBitmap/roaring/roaring64" "github.com/dustin/go-humanize" "github.com/paulmach/orb" - "github.com/schollz/progressbar/v3" "golang.org/x/sync/errgroup" "io" "io/ioutil" @@ -487,10 +486,11 @@ func Extract(_ *log.Logger, bucketURL string, key string, minzoom int8, maxzoom return err } - bar := progressbar.DefaultBytes( - int64(totalBytes), - "fetching chunks", - ) + var progress Progress + progressWriter := getProgressWriter() + if progressWriter != nil { + progress = progressWriter.NewBytesProgress(int64(totalBytes), "fetching chunks") + } var mu sync.Mutex @@ -502,15 +502,26 @@ func Extract(_ *log.Logger, bucketURL string, key string, minzoom int8, maxzoom offsetWriter := io.NewOffsetWriter(outfile, int64(header.TileDataOffset)+int64(or.Rng.DstOffset)) for _, cd := range or.CopyDiscards { + if progress != nil { + _, err := io.CopyN(io.MultiWriter(offsetWriter, progress), tileReader, int64(cd.Wanted)) + if err != nil { + return err + } - _, err := io.CopyN(io.MultiWriter(offsetWriter, bar), tileReader, int64(cd.Wanted)) - if err != nil { - return err - } + _, err = io.CopyN(progress, tileReader, int64(cd.Discard)) + if err != nil { + return err + } + } else { + _, err := io.CopyN(offsetWriter, tileReader, int64(cd.Wanted)) + if err != nil { + return err + } - _, err = io.CopyN(bar, tileReader, int64(cd.Discard)) - if err != nil { - return err + _, err = io.CopyN(io.Discard, tileReader, int64(cd.Discard)) + if err != nil { + return err + } } } tileReader.Close() diff --git a/pmtiles/makesync.go b/pmtiles/makesync.go index 6ab268d..2e5f8a1 100644 --- a/pmtiles/makesync.go +++ b/pmtiles/makesync.go @@ -6,7 +6,6 @@ import ( "encoding/json" "fmt" "github.com/cespare/xxhash/v2" - "github.com/schollz/progressbar/v3" "golang.org/x/sync/errgroup" "io" "log" @@ -106,10 +105,11 @@ func Makesync(logger *log.Logger, cliVersion string, fileName string, blockSizeK defer output.Close() - bar := progressbar.Default( - int64(header.TileEntriesCount), - "writing syncfile", - ) + var progress Progress + progressWriter := getProgressWriter() + if progressWriter != nil { + progress = progressWriter.NewCountProgress(int64(header.TileEntriesCount), "writing syncfile") + } var current syncBlock @@ -151,7 +151,9 @@ func Makesync(logger *log.Logger, cliVersion string, fileName string, blockSizeK return io.ReadAll(io.NewSectionReader(file, int64(offset), int64(length))) }, func(e EntryV3) { - bar.Add(1) + if progress != nil { + progress.Add(1) + } if current.Length == 0 { current.Start = e.TileID current.Offset = e.Offset diff --git a/pmtiles/progress.go b/pmtiles/progress.go new file mode 100644 index 0000000..6c9e7a8 --- /dev/null +++ b/pmtiles/progress.go @@ -0,0 +1,146 @@ +package pmtiles + +import ( + "io" + "sync" + + "github.com/schollz/progressbar/v3" +) + +// ProgressWriter is an interface for progress reporting during pmtiles operations. +// It supports both count-based and byte-based progress tracking. +type ProgressWriter interface { + // NewCountProgress creates a progress tracker for count-based operations + NewCountProgress(total int64, description string) Progress + // NewBytesProgress creates a progress tracker for byte-based operations + NewBytesProgress(total int64, description string) Progress +} + +// Progress represents an active progress tracker that can be updated and written to +type Progress interface { + io.Writer + // Add increments the progress by the specified amount + Add(num int) + // Close finalizes the progress tracker + Close() error +} + +var ( + // Global progress writer, protected by mutex + progressWriterMu sync.RWMutex + progressWriter ProgressWriter +) + +func init() { + // Initialize with default progress writer + progressWriter = &defaultProgressWriter{} +} + +// SetProgressWriter sets a custom progress writer for all pmtiles operations. +// Pass nil to disable all progress reporting. +func SetProgressWriter(pw ProgressWriter) { + progressWriterMu.Lock() + defer progressWriterMu.Unlock() + if pw == nil { + progressWriter = &quietProgressWriter{} + } else { + progressWriter = pw + } +} + +// getProgressWriter returns the current progress writer +func getProgressWriter() ProgressWriter { + progressWriterMu.RLock() + defer progressWriterMu.RUnlock() + return progressWriter +} + +// SetQuietMode enables or disables quiet mode for all pmtiles operations. +// When quiet mode is enabled, progress bars and verbose logging are suppressed. +// This function is maintained for backward compatibility. +func SetQuietMode(quiet bool) { + progressWriterMu.Lock() + defer progressWriterMu.Unlock() + quietMode = quiet + if quiet { + progressWriter = &quietProgressWriter{} + } else { + progressWriter = &defaultProgressWriter{} + } +} + +// IsQuietMode returns the current quiet mode setting. +func IsQuietMode() bool { + return quietMode +} + +// defaultProgressWriter implements ProgressWriter using the schollz/progressbar library +type defaultProgressWriter struct{} + +func (d *defaultProgressWriter) NewCountProgress(total int64, description string) Progress { + if quietMode { + return &quietProgress{} + } + bar := progressbar.Default(total, description) + return &progressBarWrapper{bar: bar} +} + +func (d *defaultProgressWriter) NewBytesProgress(total int64, description string) Progress { + if quietMode { + return &quietProgress{} + } + bar := progressbar.DefaultBytes(total, description) + return &progressBarWrapper{bar: bar} +} + +// progressBarWrapper wraps schollz/progressbar to implement our Progress interface +type progressBarWrapper struct { + bar *progressbar.ProgressBar +} + +func (p *progressBarWrapper) Write(data []byte) (int, error) { + if p.bar == nil { + return len(data), nil + } + return p.bar.Write(data) +} + +func (p *progressBarWrapper) Add(num int) { + if p.bar != nil { + p.bar.Add(num) + } +} + +func (p *progressBarWrapper) Close() error { + if p.bar != nil { + return p.bar.Close() + } + return nil +} + +// quietProgressWriter implements ProgressWriter with no-op operations +type quietProgressWriter struct{} + +func (q *quietProgressWriter) NewCountProgress(total int64, description string) Progress { + return &quietProgress{} +} + +func (q *quietProgressWriter) NewBytesProgress(total int64, description string) Progress { + return &quietProgress{} +} + +// quietProgress is a no-op implementation of Progress +type quietProgress struct{} + +func (q *quietProgress) Write(data []byte) (int, error) { + return len(data), nil +} + +func (q *quietProgress) Add(num int) { + // no-op - intentionally empty + _ = num // suppress unused parameter warning +} + +func (q *quietProgress) Close() error { + return nil +} diff --git a/pmtiles/progress_test.go b/pmtiles/progress_test.go new file mode 100644 index 0000000..829e342 --- /dev/null +++ b/pmtiles/progress_test.go @@ -0,0 +1,452 @@ +package pmtiles + +import ( + "bytes" + "sync" + "testing" +) + +// Mock progress writer for testing +type mockProgressWriter struct { + countProgressCalls []mockProgressCall + bytesProgressCalls []mockProgressCall + mu sync.Mutex +} + +type mockProgressCall struct { + total int64 + description string +} + +func (m *mockProgressWriter) NewCountProgress(total int64, description string) Progress { + m.mu.Lock() + defer m.mu.Unlock() + m.countProgressCalls = append(m.countProgressCalls, mockProgressCall{total, description}) + return &mockProgress{total: total, description: description} +} + +func (m *mockProgressWriter) NewBytesProgress(total int64, description string) Progress { + m.mu.Lock() + defer m.mu.Unlock() + m.bytesProgressCalls = append(m.bytesProgressCalls, mockProgressCall{total, description}) + return &mockProgress{total: total, description: description} +} + +func (m *mockProgressWriter) getCountProgressCalls() []mockProgressCall { + m.mu.Lock() + defer m.mu.Unlock() + return append([]mockProgressCall{}, m.countProgressCalls...) +} + +func (m *mockProgressWriter) getBytesProgressCalls() []mockProgressCall { + m.mu.Lock() + defer m.mu.Unlock() + return append([]mockProgressCall{}, m.bytesProgressCalls...) +} + +func (m *mockProgressWriter) reset() { + m.mu.Lock() + defer m.mu.Unlock() + m.countProgressCalls = nil + m.bytesProgressCalls = nil +} + +// Mock progress implementation +type mockProgress struct { + total int64 + current int64 + description string + closed bool + addCalls []int + writeCalls [][]byte + mu sync.Mutex +} + +func (p *mockProgress) Write(data []byte) (int, error) { + p.mu.Lock() + defer p.mu.Unlock() + p.writeCalls = append(p.writeCalls, append([]byte{}, data...)) + p.current += int64(len(data)) + return len(data), nil +} + +func (p *mockProgress) Add(num int) { + p.mu.Lock() + defer p.mu.Unlock() + p.addCalls = append(p.addCalls, num) + p.current += int64(num) +} + +func (p *mockProgress) Close() error { + p.mu.Lock() + defer p.mu.Unlock() + p.closed = true + return nil +} + +func (p *mockProgress) getCurrent() int64 { + p.mu.Lock() + defer p.mu.Unlock() + return p.current +} + +func (p *mockProgress) isClosed() bool { + p.mu.Lock() + defer p.mu.Unlock() + return p.closed +} + +func (p *mockProgress) getAddCalls() []int { + p.mu.Lock() + defer p.mu.Unlock() + return append([]int{}, p.addCalls...) +} + +func (p *mockProgress) getWriteCalls() [][]byte { + p.mu.Lock() + defer p.mu.Unlock() + result := make([][]byte, len(p.writeCalls)) + for i, call := range p.writeCalls { + result[i] = append([]byte{}, call...) + } + return result +} + +// Helper function to reset progress writer state +func resetProgressWriter() { + progressWriterMu.Lock() + defer progressWriterMu.Unlock() + progressWriter = &defaultProgressWriter{} + quietMode = false +} + +func TestSetProgressWriter(t *testing.T) { + defer resetProgressWriter() + + mock := &mockProgressWriter{} + SetProgressWriter(mock) + + writer := getProgressWriter() + if writer != mock { + t.Errorf("Expected progress writer to be mock, got %T", writer) + } +} + +func TestSetProgressWriterNil(t *testing.T) { + defer resetProgressWriter() + + SetProgressWriter(nil) + + writer := getProgressWriter() + if _, ok := writer.(*quietProgressWriter); !ok { + t.Errorf("Expected progress writer to be quietProgressWriter when set to nil, got %T", writer) + } +} + +func TestSetQuietMode(t *testing.T) { + defer resetProgressWriter() + + // Test enabling quiet mode + SetQuietMode(true) + if !IsQuietMode() { + t.Error("Expected IsQuietMode() to return true") + } + writer := getProgressWriter() + if _, ok := writer.(*quietProgressWriter); !ok { + t.Errorf("Expected progress writer to be quietProgressWriter in quiet mode, got %T", writer) + } + + // Test disabling quiet mode + SetQuietMode(false) + if IsQuietMode() { + t.Error("Expected IsQuietMode() to return false") + } + writer = getProgressWriter() + if _, ok := writer.(*defaultProgressWriter); !ok { + t.Errorf("Expected progress writer to be defaultProgressWriter when quiet mode disabled, got %T", writer) + } +} + +func TestDefaultProgressWriter(t *testing.T) { + defer resetProgressWriter() + + // Test with quiet mode disabled + SetQuietMode(false) + writer := &defaultProgressWriter{} + + // Test NewCountProgress + progress := writer.NewCountProgress(100, "test count") + if progress == nil { + t.Error("Expected non-nil progress from NewCountProgress") + } + + // Test NewBytesProgress + progress = writer.NewBytesProgress(1024, "test bytes") + if progress == nil { + t.Error("Expected non-nil progress from NewBytesProgress") + } +} + +func TestDefaultProgressWriterQuietMode(t *testing.T) { + defer resetProgressWriter() + + // Test with quiet mode enabled + SetQuietMode(true) + writer := &defaultProgressWriter{} + + // Test NewCountProgress in quiet mode + progress := writer.NewCountProgress(100, "test count") + if _, ok := progress.(*quietProgress); !ok { + t.Errorf("Expected quietProgress in quiet mode, got %T", progress) + } + + // Test NewBytesProgress in quiet mode + progress = writer.NewBytesProgress(1024, "test bytes") + if _, ok := progress.(*quietProgress); !ok { + t.Errorf("Expected quietProgress in quiet mode, got %T", progress) + } +} + +func TestProgressBarWrapper(t *testing.T) { + defer resetProgressWriter() + + SetQuietMode(false) + writer := &defaultProgressWriter{} + progress := writer.NewCountProgress(100, "test") + + // Test Write + data := []byte("test data") + n, err := progress.Write(data) + if err != nil { + t.Errorf("Unexpected error from Write: %v", err) + } + if n != len(data) { + t.Errorf("Expected Write to return %d, got %d", len(data), n) + } + + // Test Add + progress.Add(10) + + // Test Close + if err := progress.Close(); err != nil { + t.Errorf("Unexpected error from Close: %v", err) + } +} + +func TestProgressBarWrapperNilBar(t *testing.T) { + wrapper := &progressBarWrapper{bar: nil} + + // Test Write with nil bar + data := []byte("test") + n, err := wrapper.Write(data) + if err != nil { + t.Errorf("Unexpected error from Write with nil bar: %v", err) + } + if n != len(data) { + t.Errorf("Expected Write to return %d with nil bar, got %d", len(data), n) + } + + // Test Add with nil bar (should not panic) + wrapper.Add(5) + + // Test Close with nil bar + if err := wrapper.Close(); err != nil { + t.Errorf("Unexpected error from Close with nil bar: %v", err) + } +} + +func TestQuietProgressWriter(t *testing.T) { + writer := &quietProgressWriter{} + + // Test NewCountProgress + progress := writer.NewCountProgress(100, "test count") + if _, ok := progress.(*quietProgress); !ok { + t.Errorf("Expected quietProgress, got %T", progress) + } + + // Test that Add method is callable on the returned progress (covers the missing line) + progress.Add(50) + + // Test NewBytesProgress + progress = writer.NewBytesProgress(1024, "test bytes") + if _, ok := progress.(*quietProgress); !ok { + t.Errorf("Expected quietProgress, got %T", progress) + } + + // Test that Add method is callable on this progress too + progress.Add(256) +} + +func TestQuietProgress(t *testing.T) { + progress := &quietProgress{} + + // Test Write + data := []byte("test data") + n, err := progress.Write(data) + if err != nil { + t.Errorf("Unexpected error from quiet Write: %v", err) + } + if n != len(data) { + t.Errorf("Expected quiet Write to return %d, got %d", len(data), n) + } + + // Test Add (should not panic and be no-op) + progress.Add(10) + progress.Add(25) + + // Test Close + if err := progress.Close(); err != nil { + t.Errorf("Unexpected error from quiet Close: %v", err) + } +} + +func TestQuietProgressDirect(t *testing.T) { + // Directly test quietProgress.Add to ensure coverage + progress := &quietProgress{} + progress.Add(100) // This should execute the no-op function +} + +func TestMockProgressWriter(t *testing.T) { + defer resetProgressWriter() + + mock := &mockProgressWriter{} + SetProgressWriter(mock) + + writer := getProgressWriter() + + // Test NewCountProgress + progress1 := writer.NewCountProgress(100, "count test") + calls := mock.getCountProgressCalls() + if len(calls) != 1 { + t.Errorf("Expected 1 count progress call, got %d", len(calls)) + } + if calls[0].total != 100 || calls[0].description != "count test" { + t.Errorf("Unexpected count progress call: total=%d, description=%s", calls[0].total, calls[0].description) + } + + // Test NewBytesProgress + progress2 := writer.NewBytesProgress(1024, "bytes test") + byteCalls := mock.getBytesProgressCalls() + if len(byteCalls) != 1 { + t.Errorf("Expected 1 bytes progress call, got %d", len(byteCalls)) + } + if byteCalls[0].total != 1024 || byteCalls[0].description != "bytes test" { + t.Errorf("Unexpected bytes progress call: total=%d, description=%s", byteCalls[0].total, byteCalls[0].description) + } + + // Test progress operations + progress1.Add(25) + progress1.Add(35) + mockProgress1 := progress1.(*mockProgress) + addCalls := mockProgress1.getAddCalls() + if len(addCalls) != 2 || addCalls[0] != 25 || addCalls[1] != 35 { + t.Errorf("Unexpected Add calls: %v", addCalls) + } + if mockProgress1.getCurrent() != 60 { + t.Errorf("Expected current progress to be 60, got %d", mockProgress1.getCurrent()) + } + + // Test Write operations + data1 := []byte("hello") + data2 := []byte("world") + progress2.Write(data1) + progress2.Write(data2) + mockProgress2 := progress2.(*mockProgress) + writeCalls := mockProgress2.getWriteCalls() + if len(writeCalls) != 2 { + t.Errorf("Expected 2 write calls, got %d", len(writeCalls)) + } + if !bytes.Equal(writeCalls[0], data1) || !bytes.Equal(writeCalls[1], data2) { + t.Errorf("Unexpected write calls: %v", writeCalls) + } + if mockProgress2.getCurrent() != 10 { + t.Errorf("Expected current progress to be 10, got %d", mockProgress2.getCurrent()) + } + + // Test Close + progress1.Close() + progress2.Close() + if !mockProgress1.isClosed() || !mockProgress2.isClosed() { + t.Error("Expected both progress instances to be closed") + } +} + +func TestConcurrentAccess(t *testing.T) { + defer resetProgressWriter() + + mock := &mockProgressWriter{} + + var wg sync.WaitGroup + numRoutines := 10 + + // Test concurrent SetProgressWriter calls + for i := 0; i < numRoutines; i++ { + wg.Add(1) + go func() { + defer wg.Done() + SetProgressWriter(mock) + _ = getProgressWriter() + }() + } + + wg.Wait() + + // Test concurrent progress creation + for i := 0; i < numRoutines; i++ { + wg.Add(1) + go func(idx int) { + defer wg.Done() + writer := getProgressWriter() + progress := writer.NewCountProgress(int64(idx*10), "concurrent test") + progress.Add(1) + progress.Close() + }(i) + } + + wg.Wait() + + calls := mock.getCountProgressCalls() + if len(calls) != numRoutines { + t.Errorf("Expected %d concurrent progress calls, got %d", numRoutines, len(calls)) + } +} + +func TestSetQuietModeAndCustomWriter(t *testing.T) { + defer resetProgressWriter() + + mock := &mockProgressWriter{} + + // Set custom writer first + SetProgressWriter(mock) + writer1 := getProgressWriter() + if writer1 != mock { + t.Error("Expected custom writer to be set") + } + + // SetQuietMode should override custom writer + SetQuietMode(true) + writer2 := getProgressWriter() + if _, ok := writer2.(*quietProgressWriter); !ok { + t.Errorf("Expected quiet writer after SetQuietMode(true), got %T", writer2) + } + if !IsQuietMode() { + t.Error("Expected IsQuietMode() to return true") + } + + // SetQuietMode(false) should restore default writer, not custom + SetQuietMode(false) + writer3 := getProgressWriter() + if _, ok := writer3.(*defaultProgressWriter); !ok { + t.Errorf("Expected default writer after SetQuietMode(false), got %T", writer3) + } + if IsQuietMode() { + t.Error("Expected IsQuietMode() to return false") + } + + // Setting custom writer again should work + SetProgressWriter(mock) + writer4 := getProgressWriter() + if writer4 != mock { + t.Error("Expected custom writer to be set again") + } +} \ No newline at end of file diff --git a/pmtiles/sync.go b/pmtiles/sync.go index 9516735..fe3b6ed 100644 --- a/pmtiles/sync.go +++ b/pmtiles/sync.go @@ -7,7 +7,6 @@ import ( "fmt" "github.com/cespare/xxhash/v2" "github.com/dustin/go-humanize" - "github.com/schollz/progressbar/v3" "golang.org/x/sync/errgroup" "io" "log" @@ -85,11 +84,14 @@ func Sync(logger *log.Logger, oldVersion string, newVersion string, dryRun bool) if err != nil { return err } - bar := progressbar.DefaultBytes( - resp.ContentLength, - "downloading syncfile", - ) - bufferedReader = bufio.NewReader(io.TeeReader(resp.Body, bar)) + var progress Progress + progressWriter := getProgressWriter() + if progressWriter != nil { + progress = progressWriter.NewBytesProgress(resp.ContentLength, "downloading syncfile") + bufferedReader = bufio.NewReader(io.TeeReader(resp.Body, progress)) + } else { + bufferedReader = bufio.NewReader(resp.Body) + } var syncHeader syncHeader jsonBytes, _ := bufferedReader.ReadSlice('\n') @@ -97,7 +99,9 @@ func Sync(logger *log.Logger, oldVersion string, newVersion string, dryRun bool) json.Unmarshal(jsonBytes, &syncHeader) blocks := deserializeSyncBlocks(syncHeader.NumBlocks, bufferedReader) - bar.Close() + if progress != nil { + progress.Close() + } ctx := context.Background() @@ -122,10 +126,10 @@ func Sync(logger *log.Logger, oldVersion string, newVersion string, dryRun bool) return fmt.Errorf("archive must be clustered for sync") } - bar = progressbar.Default( - int64(len(blocks)), - "calculating diff", - ) + var diffProgress Progress + if progressWriter != nil { + diffProgress = progressWriter.NewCountProgress(int64(len(blocks)), "calculating diff") + } wanted := make([]syncBlock, 0) have := make([]srcDstRange, 0) @@ -172,13 +176,17 @@ func Sync(logger *log.Logger, oldVersion string, newVersion string, dryRun bool) mu.Lock() wanted = append(wanted, blocks[idx]) mu.Unlock() - bar.Add(1) + if diffProgress != nil { + diffProgress.Add(1) + } idx = idx + 1 } if e.TileID == blocks[idx].Start { tasks <- syncTask{NewBlock: blocks[idx], OldOffset: e.Offset} - bar.Add(1) + if diffProgress != nil { + diffProgress.Add(1) + } idx = idx + 1 } } @@ -193,7 +201,9 @@ func Sync(logger *log.Logger, oldVersion string, newVersion string, dryRun bool) mu.Lock() wanted = append(wanted, blocks[idx]) mu.Unlock() - bar.Add(1) + if diffProgress != nil { + diffProgress.Add(1) + } idx = idx + 1 } @@ -282,36 +292,45 @@ func Sync(logger *log.Logger, oldVersion string, newVersion string, dryRun bool) req.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", newHeader.LeafDirectoryOffset, newHeader.LeafDirectoryOffset+newHeader.LeafDirectoryLength-1)) resp, err = client.Do(req) - leafBar := progressbar.DefaultBytes( - int64(newHeader.LeafDirectoryLength), - "downloading leaf directories", - ) - io.Copy(leafWriter, io.TeeReader(resp.Body, leafBar)) - leafBar.Close() + var leafProgress Progress + if progressWriter != nil { + leafProgress = progressWriter.NewBytesProgress(int64(newHeader.LeafDirectoryLength), "downloading leaf directories") + io.Copy(leafWriter, io.TeeReader(resp.Body, leafProgress)) + leafProgress.Close() + } else { + io.Copy(leafWriter, resp.Body) + } - fmt.Println(len(have), "local chunks") - bar := progressbar.DefaultBytes( - int64(totalRemoteBytes-toTransfer), - "copying local chunks", - ) + var copyProgress Progress + if !quietMode { + fmt.Println(len(have), "local chunks") + } + if progressWriter != nil { + copyProgress = progressWriter.NewBytesProgress(int64(totalRemoteBytes-toTransfer), "copying local chunks") + } // write the tile data (from local) for _, h := range haveRanges { chunkWriter := io.NewOffsetWriter(outfile, int64(newHeader.TileDataOffset+h.DstOffset)) r := io.NewSectionReader(oldFile, int64(oldHeader.TileDataOffset+h.SrcOffset), int64(h.Length)) - io.Copy(io.MultiWriter(chunkWriter, bar), r) + if copyProgress != nil { + io.Copy(io.MultiWriter(chunkWriter, copyProgress), r) + } else { + io.Copy(chunkWriter, r) + } } oldFile.Close() // write the tile data (from remote) multiRanges := makeMultiRanges(ranges, int64(newHeader.TileDataOffset), 1048576-200) - fmt.Println("Batched into http requests", len(multiRanges)) - - bar = progressbar.DefaultBytes( - int64(toTransfer), - "fetching remote chunks", - ) + var fetchProgress Progress + if !quietMode { + fmt.Println("Batched into http requests", len(multiRanges)) + } + if progressWriter != nil { + fetchProgress = progressWriter.NewBytesProgress(int64(toTransfer), "fetching remote chunks") + } downloadPart := func(task multiRange) error { req, err := http.NewRequest("GET", newVersion, nil) @@ -332,7 +351,11 @@ func Sync(logger *log.Logger, oldVersion string, newVersion string, dryRun bool) part, _ := mr.NextPart() _ = part.Header.Get("Content-Range") chunkWriter := io.NewOffsetWriter(outfile, int64(newHeader.TileDataOffset+r.DstOffset)) - io.Copy(io.MultiWriter(chunkWriter, bar), part) + if fetchProgress != nil { + io.Copy(io.MultiWriter(chunkWriter, fetchProgress), part) + } else { + io.Copy(chunkWriter, part) + } } return nil } diff --git a/pmtiles/upload.go b/pmtiles/upload.go index b1a3fe6..3c5cde0 100644 --- a/pmtiles/upload.go +++ b/pmtiles/upload.go @@ -3,7 +3,6 @@ package pmtiles import ( "context" "fmt" - "github.com/schollz/progressbar/v3" "gocloud.dev/blob" "io" "log" @@ -49,8 +48,14 @@ func Upload(_ *log.Logger, InputPMTiles string, bucket string, RemotePMTiles str return fmt.Errorf("Failed to obtain writer: %w", err) } - bar := progressbar.Default(filestat.Size()) - io.Copy(io.MultiWriter(w, bar), f) + var progress Progress + progressWriter := getProgressWriter() + if progressWriter != nil { + progress = progressWriter.NewBytesProgress(filestat.Size(), "") + io.Copy(io.MultiWriter(w, progress), f) + } else { + io.Copy(w, f) + } if err := w.Close(); err != nil { return fmt.Errorf("Failed to complete upload: %w", err) From 028626a196a5292b80b67fdced4b83c95a5d437a Mon Sep 17 00:00:00 2001 From: Angus Peart Date: Thu, 25 Sep 2025 16:15:21 -0700 Subject: [PATCH 2/2] Add example for custom progress writer --- .gitignore | 1 + examples/custom_progress.go | 143 ++++++++++++++++++++++++++++++++++++ 2 files changed, 144 insertions(+) create mode 100644 examples/custom_progress.go diff --git a/.gitignore b/.gitignore index a593b7a..7c6c346 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,4 @@ go-pmtiles *.json *.geojson *.tsv.gz +examples/*.pmtiles \ No newline at end of file diff --git a/examples/custom_progress.go b/examples/custom_progress.go new file mode 100644 index 0000000..074c1aa --- /dev/null +++ b/examples/custom_progress.go @@ -0,0 +1,143 @@ +package main + +import ( + "fmt" + "strings" + "time" + + "github.com/protomaps/go-pmtiles/pmtiles" +) + +// CloudProgress represents a custom progress writer +type CloudProgress struct { + serviceName string +} + +// CustomProgress tracks progress for individual operations +type CustomProgress struct { + serviceName string + operation string + total int64 + current int64 + startTime time.Time +} + +// NewCountProgress creates a progress tracker for count-based operations +func (c *CloudProgress) NewCountProgress(total int64, description string) pmtiles.Progress { + fmt.Printf("[%s] Starting operation: %s (total items: %d)\n", c.serviceName, description, total) + return &CustomProgress{ + serviceName: c.serviceName, + operation: description, + total: total, + current: 0, + startTime: time.Now(), + } +} + +// NewBytesProgress creates a progress tracker for byte-based operations +func (c *CloudProgress) NewBytesProgress(total int64, description string) pmtiles.Progress { + fmt.Printf("[%s] Starting operation: %s (total bytes: %s)\n", + c.serviceName, description, formatBytes(total)) + return &CustomProgress{ + serviceName: c.serviceName, + operation: description, + total: total, + current: 0, + startTime: time.Now(), + } +} + +// Write implements io.Writer for byte-based progress tracking +func (p *CustomProgress) Write(data []byte) (int, error) { + p.current += int64(len(data)) + p.reportProgress() + + // Here you would typically send the progress update to some external service + + return len(data), nil +} + +// Add increments the progress counter for count-based operations +func (p *CustomProgress) Add(num int) { + p.current += int64(num) + p.reportProgress() + + // Send progress update to outside service here if needed +} + +// Close finalizes the progress tracking +func (p *CustomProgress) Close() error { + duration := time.Since(p.startTime) + fmt.Printf("[%s] Completed: %s in %v\n", p.serviceName, p.operation, duration.Round(time.Millisecond)) + + // Send completion notification to outside service here if needed + return nil +} + +// reportProgress displays current progress and could send updates to cloud services +func (p *CustomProgress) reportProgress() { + if p.total <= 0 { + return + } + + percentage := float64(p.current) / float64(p.total) * 100 + elapsed := time.Since(p.startTime) + + // Create a simple progress bar + barWidth := 30 + filled := int(percentage / 100 * float64(barWidth)) + bar := strings.Repeat("█", filled) + strings.Repeat("░", barWidth-filled) + + if strings.Contains(p.operation, "bytes") || p.total > 1000 { + // Byte-based progress + fmt.Printf("[%s] %s: %.1f%% [%s] %s/%s (%.2fs)\n", + p.serviceName, p.operation, percentage, bar, + formatBytes(p.current), formatBytes(p.total), elapsed.Seconds()) + } else { + // Count-based progress + fmt.Printf("[%s] %s: %.1f%% [%s] %d/%d items (%.2fs)\n", + p.serviceName, p.operation, percentage, bar, + p.current, p.total, elapsed.Seconds()) + } +} + +// formatBytes converts bytes to human-readable format +func formatBytes(bytes int64) string { + const unit = 1024 + if bytes < unit { + return fmt.Sprintf("%d B", bytes) + } + div, exp := int64(unit), 0 + for n := bytes / unit; n >= unit; n /= unit { + div *= unit + exp++ + } + return fmt.Sprintf("%.1f %cB", float64(bytes)/float64(div), "KMGTPE"[exp]) +} + +func main() { + // Custom progress writer for external services + fmt.Println("\nCustom Cloud Progress Writer:") + cloudProgress := &CloudProgress{serviceName: "Service"} + pmtiles.SetProgressWriter(cloudProgress) + + // Simulate some progress operations + fmt.Println("\nSimulating PMTiles operations with custom progress reporting:") + + // Simulate a count-based operation + progressCount := cloudProgress.NewCountProgress(100, "Processing tiles") + for i := 0; i < 100; i += 10 { + progressCount.Add(10) + time.Sleep(50 * time.Millisecond) // Simulate work + } + progressCount.Close() + + // Simulate a bytes-based operation + progressBytes := cloudProgress.NewBytesProgress(1024*1024, "Uploading tiles") + data := make([]byte, 64*1024) // 64KB chunks + for i := 0; i < 16; i++ { + progressBytes.Write(data) + time.Sleep(30 * time.Millisecond) // Simulate upload time + } + progressBytes.Close() +}