sq/libsq/core/ioz/download/handler.go
Neil O'Toole db55986980
#307: Ingest cache (#354)
- Support for ingest cache, download cache, and progress bars.
2024-01-14 18:45:34 -07:00

95 lines
2.6 KiB
Go

package download
import (
"bytes"
"log/slog"
"sync"
"github.com/neilotoole/sq/libsq/core/ioz"
"github.com/neilotoole/sq/libsq/core/lg/lga"
)
// Handler is a callback invoked by Download.Get. Exactly one of the
// handler functions will be invoked, exactly one time.
type Handler struct {
// Cached is invoked when the download is already cached on disk. The
// fp arg is the path to the downloaded file.
Cached func(fp string)
// Uncached is invoked when the download is not cached. The handler should
// return an ioz.WriteErrorCloser, which the download contents will be written
// to (as well as being written to the disk cache). On success, the dest
// writer is closed. If an error occurs during download or writing,
// WriteErrorCloser.Error is invoked (but Close is not invoked). If the
// handler returns a nil dest, the Download will log a warning and return.
Uncached func() (dest ioz.WriteErrorCloser)
// Error is invoked on any error, other than writing to the destination
// io.WriteCloser returned by Handler.Uncached, which has its own error
// handling mechanism.
Error func(err error)
}
// SinkHandler is a download.Handler that records the results of the callbacks
// it receives. This is useful for testing.
type SinkHandler struct {
Handler
mu sync.Mutex
log *slog.Logger
// Errors records the errors received via Handler.Error.
Errors []error
// CachedFiles records the cached files received via Handler.Cached.
CachedFiles []string
// UncachedBufs records in bytes.Buffer instances the data written
// via Handler.Uncached.
UncachedBufs []*bytes.Buffer
// WriteErrors records the write errors received via Handler.Uncached.
WriteErrors []error
}
// Reset resets the handler sinks.
func (sh *SinkHandler) Reset() {
sh.mu.Lock()
defer sh.mu.Unlock()
sh.Errors = nil
sh.CachedFiles = nil
sh.UncachedBufs = nil
sh.WriteErrors = nil
}
// NewSinkHandler returns a new SinkHandler.
func NewSinkHandler(log *slog.Logger) *SinkHandler {
h := &SinkHandler{log: log}
h.Cached = func(fp string) {
log.Info("Cached", lga.File, fp)
h.mu.Lock()
defer h.mu.Unlock()
h.CachedFiles = append(h.CachedFiles, fp)
}
h.Uncached = func() ioz.WriteErrorCloser {
log.Info("Uncached")
h.mu.Lock()
defer h.mu.Unlock()
buf := &bytes.Buffer{}
h.UncachedBufs = append(h.UncachedBufs, buf)
return ioz.NewFuncWriteErrorCloser(ioz.WriteCloser(buf), func(err error) {
h.mu.Lock()
defer h.mu.Unlock()
h.WriteErrors = append(h.WriteErrors, err)
})
}
h.Error = func(err error) {
log.Info("Error", lga.Err, err)
h.mu.Lock()
defer h.mu.Unlock()
h.Errors = append(h.Errors, err)
}
return h
}