sq/libsq/source/files.go

package source

import (
	"context"
	"io"
	"mime"
	"net/url"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/neilotoole/sq/libsq/core/lg/lga"
	"golang.org/x/exp/slog"

	"github.com/neilotoole/sq/libsq/core/lg/lgm"

	"github.com/neilotoole/sq/libsq/core/lg"

	"github.com/djherbis/fscache"
	"github.com/h2non/filetype"
	"github.com/h2non/filetype/matchers"
	"golang.org/x/sync/errgroup"

	"github.com/neilotoole/sq/libsq/core/cleanup"
	"github.com/neilotoole/sq/libsq/core/errz"
	"github.com/neilotoole/sq/libsq/source/fetcher"
)

// Files is the centralized API for interacting with files.
//
// Why does Files exist? There's a need for functionality to
// transparently get a Reader for remote or local files, and most importantly,
// an ability for multiple goroutines to read/sample a file while
// it's being read (mainly to "sample" the file type, e.g. to determine
// if it's an XLSX file etc.). Currently we use fscache under the hood
// for this, but our implementation is not satisfactory: in particular,
// the implementation currently requires that we read the entire source
// file into fscache before it's available to be read (which is awful
// if we're reading long-running pipe from stdin). This entire thing
// needs to be revisited. Maybe Files even becomes a fs.FS.
type Files struct {
	log       *slog.Logger
	mu        sync.Mutex
	clnup     *cleanup.Cleanup
	fcache    *fscache.FSCache
	detectFns []DriverDetectFunc
}

// NewFiles returns a new Files instance.
func NewFiles(ctx context.Context) (*Files, error) {
	fs := &Files{clnup: cleanup.New(), log: lg.FromContext(ctx)}

	tmpdir, err := os.MkdirTemp("", "sq_files_fscache_*")
	if err != nil {
		return nil, errz.Err(err)
	}

	fcache, err := fscache.New(tmpdir, os.ModePerm, time.Hour)
	if err != nil {
		return nil, errz.Err(err)
	}

	fs.clnup.AddE(fcache.Clean)
	fs.fcache = fcache
	return fs, nil
}

// AddDriverDetectors adds driver type detectors.
func (fs *Files) AddDriverDetectors(detectFns ...DriverDetectFunc) {
	fs.detectFns = append(fs.detectFns, detectFns...)
}

// Size returns the file size of src.Location. This exists
// as a convenience function and something of a replacement
// for using os.Stat to get the file size.
func (fs *Files) Size(src *Source) (size int64, err error) {
	r, err := fs.Open(src)
	if err != nil {
		return 0, err
	}

	defer lg.WarnIfCloseError(fs.log, lgm.CloseFileReader, r)

	size, err = io.Copy(io.Discard, r)
	if err != nil {
		return 0, errz.Err(err)
	}

	return size, nil
}

// AddStdin copies f to fs's cache: the stdin data in f
// is later accessible via fs.Open(src) where src.Handle
// is StdinHandle; f's type can be detected via TypeStdin.
// Note that f is closed by this method.
//
// REVISIT: it's possible we'll ditch AddStdin and TypeStdin
// in some future version; this mechanism is a stopgap.
func (fs *Files) AddStdin(f *os.File) error {
	fs.mu.Lock()
	defer fs.mu.Unlock()

	// We don't need r, but we're responsible for closing it.
	r, err := fs.addFile(f, StdinHandle) // f is closed by addFile
	if err != nil {
		return err
	}

	return r.Close()
}

// TypeStdin detects the type of stdin as previously added
// by AddStdin. An error is returned if AddStdin was not
// first invoked. If the type cannot be detected, TypeNone and
// nil are returned.
func (fs *Files) TypeStdin(ctx context.Context) (DriverType, error) {
	if !fs.fcache.Exists(StdinHandle) {
		return TypeNone, errz.New("must invoke AddStdin before invoking TypeStdin")
	}

	typ, ok, err := fs.detectType(ctx, StdinHandle)
	if err != nil {
		return TypeNone, err
	}

	if !ok {
		return TypeNone, nil
	}

	return typ, nil
}

// add file copies f to fs's cache, returning a reader which the
// caller is responsible for closing. f is closed by this method.
func (fs *Files) addFile(f *os.File, key string) (fscache.ReadAtCloser, error) {
	fs.log.Debug("Adding file", lga.Key, key, lga.Path, f.Name())
	r, w, err := fs.fcache.Get(key)
	if err != nil {
		return nil, errz.Err(err)
	}

	if w == nil {
		lg.WarnIfCloseError(fs.log, lgm.CloseFileReader, r)
		return nil, errz.Errorf("failed to add to fscache (possibly previously added): %s", key)
	}

	// TODO: Problematically, we copy the entire contents of f into fscache.
	// If f is a large file (e.g. piped over stdin), this means that
	// everything is held up until f is fully copied. Hopefully we can
	// do something with fscache so that the readers returned from
	// fscache can lazily read from f.
	_, err = io.Copy(w, f)
	if err != nil {
		lg.WarnIfCloseError(fs.log, lgm.CloseFileReader, r)
		return nil, errz.Err(err)
	}

	err = errz.Combine(w.Close(), f.Close())
	if err != nil {
		lg.WarnIfCloseError(fs.log, lgm.CloseFileReader, r)
		return nil, err
	}

	return r, nil
}

// Open returns a new io.ReadCloser for src.Location.
// If src.Handle is StdinHandle, AddStdin must first have
// been invoked. The caller must close the reader.
func (fs *Files) Open(src *Source) (io.ReadCloser, error) {
	fs.mu.Lock()
	defer fs.mu.Unlock()

	return fs.newReader(src.Location)
}

// OpenFunc returns a func that invokes fs.Open for src.Location.
func (fs *Files) OpenFunc(src *Source) func() (io.ReadCloser, error) {
	return func() (io.ReadCloser, error) {
		return fs.Open(src)
	}
}

// ReadAll is a convenience method to read the bytes of a source.
func (fs *Files) ReadAll(src *Source) ([]byte, error) {
	r, err := fs.newReader(src.Location)
	if err != nil {
		return nil, err
	}

	var data []byte
	data, err = io.ReadAll(r)
	closeErr := r.Close()
	if err != nil {
		return nil, err
	}
	if closeErr != nil {
		return nil, closeErr
	}

	return data, nil
}

func (fs *Files) newReader(loc string) (io.ReadCloser, error) {
	if loc == StdinHandle {
		r, w, err := fs.fcache.Get(StdinHandle)
		if err != nil {
			return nil, errz.Err(err)
		}
		if w != nil {
			return nil, errz.New("@stdin not cached: has AddStdin been invoked yet?")
		}

		return r, nil
	}

	if !fs.fcache.Exists(loc) {
		// cache miss
		f, err := fs.openLocation(loc)
		if err != nil {
			return nil, err
		}

		// Note that addFile closes f
		r, err := fs.addFile(f, loc)
		if err != nil {
			return nil, err
		}
		return r, nil
	}

	r, _, err := fs.fcache.Get(loc)
	if err != nil {
		return nil, err
	}

	return r, nil
}

// openLocation returns a file for loc. It is the caller's
// responsibility to close the returned file.
func (fs *Files) openLocation(loc string) (*os.File, error) {
	var fpath string
	var ok bool
	var err error

	fpath, ok = isFpath(loc)
	if !ok {
		// It's not a local file path, maybe it's remote (http)
		var u *url.URL
		u, ok = httpURL(loc)
		if !ok {
			// We're out of luck, it's not a valid file location
			return nil, errz.Errorf("invalid src location: %s", loc)
		}

		// It's a remote file
		fpath, err = fs.fetch(u.String())
		if err != nil {
			return nil, err
		}
	}

	// we have a legitimate fpath
	return fs.openFile(fpath)
}

// openFile opens the file at fpath. It is the caller's
// responsibility to close the returned file.
func (fs *Files) openFile(fpath string) (*os.File, error) {
	f, err := os.OpenFile(fpath, os.O_RDWR, 0o666)
	if err != nil {
		return nil, errz.Err(err)
	}

	return f, nil
}

// fetch ensures that loc exists locally as a file. This may
// entail downloading the file via HTTPS etc.
func (fs *Files) fetch(loc string) (fpath string, err error) {
	// This impl is a vestigial abomination from an early
	// experiment.

	var ok bool
	if fpath, ok = isFpath(loc); ok {
		// loc is already a local file path
		return fpath, nil
	}

	var u *url.URL
	if u, ok = httpURL(loc); !ok {
		return "", errz.Errorf("not a valid file location: %s", loc)
	}

	var dlFile *os.File
	dlFile, err = os.CreateTemp("", "")
	if err != nil {
		return "", errz.Err(err)
	}

	fetchr := &fetcher.Fetcher{}
	// TOOD: ultimately should be passing a real context here
	err = fetchr.Fetch(context.Background(), u.String(), dlFile)
	if err != nil {
		return "", errz.Err(err)
	}

	// dlFile is kept open until fs is closed.
	fs.clnup.AddC(dlFile)

	return dlFile.Name(), nil
}

// Close closes any open resources.
func (fs *Files) Close() error {
	fs.log.Debug("Files.Close invoked: executing clean funcs", lga.Count, fs.clnup.Len())

	return fs.clnup.Run()
}

// CleanupE adds fn to the cleanup sequence invoked by fs.Close.
func (fs *Files) CleanupE(fn func() error) {
	fs.clnup.AddE(fn)
}

// DriverType returns the driver type of loc.
func (fs *Files) DriverType(ctx context.Context, loc string) (DriverType, error) {
	ploc, err := parseLoc(loc)
	if err != nil {
		return TypeNone, err
	}

	if ploc.typ != TypeNone {
		return ploc.typ, nil
	}

	if ploc.ext != "" {
		mtype := mime.TypeByExtension(ploc.ext)
		if mtype == "" {
			fs.log.Debug(
				"unknown mime type",
				lga.Type, mtype,
				lga.Loc, loc,
			)
		} else {
			if typ, ok := typeFromMediaType(mtype); ok {
				return typ, nil
			}
			fs.log.Debug(
				"unknown driver type for media type",
				lga.Type, mtype,
				lga.Loc, loc,
			)
		}
	}

	// Fall back to the byte detectors
	typ, ok, err := fs.detectType(ctx, loc)
	if err != nil {
		return TypeNone, err
	}

	if !ok {
		return TypeNone, errz.Errorf("unable to determine driver type: %s", loc)
	}

	return typ, nil
}

func (fs *Files) detectType(ctx context.Context, loc string) (typ DriverType, ok bool, err error) {
	if len(fs.detectFns) == 0 {
		return TypeNone, false, nil
	}

	type result struct {
		typ   DriverType
		score float32
	}

	resultCh := make(chan result, len(fs.detectFns))
	openFn := func() (io.ReadCloser, error) {
		fs.mu.Lock()
		defer fs.mu.Unlock()

		return fs.newReader(loc)
	}

	select {
	case <-ctx.Done():
		return TypeNone, false, ctx.Err()
	default:
	}

	g, gCtx := errgroup.WithContext(ctx)
	gCtx = lg.NewContext(gCtx, fs.log)

	for _, detectFn := range fs.detectFns {
		detectFn := detectFn

		g.Go(func() error {
			select {
			case <-gCtx.Done():
				return gCtx.Err()
			default:
			}

			gTyp, gScore, gErr := detectFn(gCtx, openFn)
			if gErr != nil {
				return gErr
			}

			if gScore > 0 {
				resultCh <- result{typ: gTyp, score: gScore}
			}
			return nil
		})
	}

	err = g.Wait()
	if err != nil {
		fs.log.Error(err.Error())
		return TypeNone, false, errz.Err(err)
	}
	close(resultCh)

	var highestScore float32
	for res := range resultCh {
		if res.score > highestScore {
			highestScore = res.score
			typ = res.typ
		}
	}

	const detectScoreThreshold = 0.5
	if highestScore >= detectScoreThreshold {
		return typ, true, nil
	}

	return TypeNone, false, nil
}

// FileOpenFunc returns a func that opens a ReadCloser. The caller
// is responsible for closing the returned ReadCloser.
type FileOpenFunc func() (io.ReadCloser, error)

// DriverDetectFunc interrogates a byte stream to determine
// the source driver type. A score is returned indicating
// the confidence that the driver type has been detected.
// A score <= 0 is failure, a score >= 1 is success; intermediate
// values indicate some level of confidence.
// An error is returned only if an IO problem occurred.
// The implementation gets access to the byte stream by invoking openFn,
// and is responsible for closing any reader it opens.
type DriverDetectFunc func(ctx context.Context, openFn FileOpenFunc) (detected DriverType, score float32, err error)

var _ DriverDetectFunc = DetectMagicNumber

// DetectMagicNumber is a DriverDetectFunc that uses an external
// pkg (h2non/filetype) to detect the "magic number" from
// the start of files.
func DetectMagicNumber(ctx context.Context, openFn FileOpenFunc,
) (detected DriverType, score float32, err error) {
	log := lg.FromContext(ctx)
	var r io.ReadCloser
	r, err = openFn()
	if err != nil {
		return TypeNone, 0, errz.Err(err)
	}
	defer lg.WarnIfCloseError(log, lgm.CloseFileReader, r)

	// We only have to pass the file header = first 261 bytes
	head := make([]byte, 261)
	_, err = r.Read(head)
	if err != nil {
		return TypeNone, 0, errz.Wrapf(err, "failed to read header")
	}

	ftype, err := filetype.Match(head)
	if err != nil {
		if err != nil {
			return TypeNone, 0, errz.Err(err)
		}
	}

	switch ftype {
	default:
		return TypeNone, 0, nil
	case matchers.TypeXlsx:
		// This doesn't seem to work, because .xlsx files are
		// zipped, so the type returns as a zip. Perhaps there's
		// something we can do about it, such as first extracting
		// the zip, and then reading the inner magic number, but
		// the xlsx.DetectXLSX func should catch the type anyway.
		return typeXLSX, 1.0, nil
	case matchers.TypeXls:
		// TODO: our xlsx driver doesn't yet support XLS
		return typeXLSX, 1.0, errz.Errorf("Microsoft XLS (%s) not currently supported", ftype)
	case matchers.TypeSqlite:
		return typeSL3, 1.0, nil
	}
}

// httpURL tests if s is a well-structured HTTP or HTTPS url, and
// if so, returns the url and true.
func httpURL(s string) (u *url.URL, ok bool) {
	var err error
	u, err = url.Parse(s)
	if err != nil || u.Host == "" || !(u.Scheme == "http" || u.Scheme == "https") {
		return nil, false
	}

	return u, true
}

// TempDirFile creates a new temporary file in a new temp dir,
// opens the file for reading and writing, and returns the resulting *os.File,
// as well as the parent dir.
// It is the caller's responsibility to close the file and remove the temp
// dir, which the returned cleanFn encapsulates.
func TempDirFile(filename string) (dir string, f *os.File, cleanFn func() error, err error) {
	dir, err = os.MkdirTemp("", "sq_")
	if err != nil {
		return "", nil, nil, errz.Err(err)
	}

	name := filepath.Join(dir, filename)
	f, err = os.OpenFile(name, os.O_RDWR|os.O_CREATE|os.O_EXCL, 0o600)
	if err != nil {
		// Silently delete the temp dir
		_ = os.RemoveAll(dir)

		return "", nil, nil, errz.Err(err)
	}

	cleanFn = func() error {
		return errz.Append(f.Close(), os.RemoveAll(dir))
	}

	return dir, f, cleanFn, nil
}