package files import ( "context" "errors" "io" "mime" "os" "strings" "time" "github.com/h2non/filetype" "github.com/h2non/filetype/matchers" "golang.org/x/sync/errgroup" "github.com/neilotoole/sq/libsq/core/errz" "github.com/neilotoole/sq/libsq/core/lg" "github.com/neilotoole/sq/libsq/core/lg/lga" "github.com/neilotoole/sq/libsq/core/lg/lgm" "github.com/neilotoole/sq/libsq/source" "github.com/neilotoole/sq/libsq/source/drivertype" "github.com/neilotoole/sq/libsq/source/location" ) // TypeDetectFunc 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 // newRdrFn, and is responsible for closing any reader it opens. type TypeDetectFunc func(ctx context.Context, newRdrFn NewReaderFunc) ( detected drivertype.Type, score float32, err error) // AddDriverDetectors adds driver type detectors. func (fs *Files) AddDriverDetectors(detectFns ...TypeDetectFunc) { fs.mu.Lock() defer fs.mu.Unlock() fs.detectFns = append(fs.detectFns, detectFns...) } // DetectType returns the driver type of loc. // This may result in loading files into the cache. func (fs *Files) DetectType(ctx context.Context, handle, loc string) (drivertype.Type, error) { log := lg.FromContext(ctx).With(lga.Loc, loc) fields, err := location.Parse(loc) if err != nil { return drivertype.None, err } if fields.DriverType != drivertype.None { return fields.DriverType, nil } if fields.Ext != "" { mtype := mime.TypeByExtension(fields.Ext) if mtype == "" { log.Debug("unknown mime type", lga.Type, mtype) } else { if typ, ok := driverFromMediaType(mtype); ok { return typ, nil } log.Debug("unknown driver type for media type", lga.Type, mtype) } } // REVISIT: We really should try to be smarter here, esp with sqlite files. fs.mu.Lock() defer fs.mu.Unlock() // Fall back to the byte detectors typ, ok, err := fs.detectType(ctx, handle, loc) if err != nil { return drivertype.None, err } if !ok { return drivertype.None, errz.Errorf("unable to determine driver type: %s", loc) } return typ, nil } // DetectStdinType 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) DetectStdinType(ctx context.Context) (drivertype.Type, error) { fs.mu.Lock() defer fs.mu.Unlock() if _, ok := fs.streams[source.StdinHandle]; !ok { return drivertype.None, errz.New("must invoke Files.AddStdin before invoking DetectStdinType") } typ, ok, err := fs.detectType(ctx, source.StdinHandle, source.StdinHandle) if err != nil { return drivertype.None, err } if !ok { return drivertype.None, nil } return typ, nil } // detectType detects the type of src's location. The value of Source.Type // is ignored. If the type cannot be detected, drivertype.None and false are // returned. func (fs *Files) detectType(ctx context.Context, handle, loc string) (typ drivertype.Type, ok bool, err error) { if len(fs.detectFns) == 0 { return drivertype.None, false, nil } log := lg.FromContext(ctx).With(lga.Loc, loc) start := time.Now() var newRdrFn NewReaderFunc if location.TypeOf(loc) == location.TypeFile { newRdrFn = func(ctx context.Context) (io.ReadCloser, error) { return errz.Return(os.Open(loc)) } } else { newRdrFn = func(ctx context.Context) (io.ReadCloser, error) { src := &source.Source{Handle: handle, Location: loc} return fs.newReader(ctx, src, false) } } // We do the magic number first, because it's so fast. detected, score, err := DetectMagicNumber(ctx, newRdrFn) if err == nil && score >= 1.0 { return detected, true, nil } type result struct { typ drivertype.Type score float32 } resultCh := make(chan result, len(fs.detectFns)) select { case <-ctx.Done(): return drivertype.None, false, ctx.Err() default: } g, gCtx := errgroup.WithContext(ctx) for _, detectFn := range fs.detectFns { detectFn := detectFn g.Go(func() error { select { case <-gCtx.Done(): return gCtx.Err() default: } gTyp, gScore, gErr := detectFn(gCtx, newRdrFn) if gErr != nil { return gErr } if gScore > 0 { resultCh <- result{typ: gTyp, score: gScore} } return nil }) } // REVISIT: We shouldn't have to wait for all goroutines to complete. // This logic could be refactored to return as soon as a single // goroutine returns a score >= 1.0 (then cancelling the other detector // goroutines). if err = g.Wait(); err != nil { log.Error(err.Error()) return drivertype.None, 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 { log.Debug("Type detected", lga.Type, typ, lga.Elapsed, time.Since(start)) return typ, true, nil } log.Warn("No type detected", lga.Type, typ, lga.Elapsed, time.Since(start)) return drivertype.None, false, nil } var _ TypeDetectFunc = DetectMagicNumber // DetectMagicNumber is a TypeDetectFunc that detects the "magic number" // from the start of files. func DetectMagicNumber(ctx context.Context, newRdrFn NewReaderFunc, ) (detected drivertype.Type, score float32, err error) { log := lg.FromContext(ctx) var r io.ReadCloser r, err = newRdrFn(ctx) if err != nil { return drivertype.None, 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) if _, err = io.ReadFull(r, head); err != nil && !errors.Is(err, io.ErrUnexpectedEOF) { return drivertype.None, 0, errz.Wrapf(err, "failed to read header") } ftype, err := filetype.Match(head) if err != nil { return drivertype.None, 0, errz.Err(err) } switch ftype { default: return drivertype.None, 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 drivertype.XLSX, 1.0, nil case matchers.TypeXls: // TODO: our xlsx driver doesn't yet support XLS return drivertype.XLSX, 1.0, errz.Errorf("Microsoft XLS (%s) not currently supported", ftype) case matchers.TypeSqlite: return drivertype.SQLite, 1.0, nil } } // driverFromMediaType returns the driver type corresponding to mediatype. // For example: // // xlsx application/vnd.openxmlformats-officedocument.spreadsheetml.sheet // csv text/csv // // Note that we don't rely on this function for types such // as application/json, because JSON can map to multiple // driver types (json, jsona, jsonl). func driverFromMediaType(mediatype string) (typ drivertype.Type, ok bool) { switch { case strings.Contains(mediatype, `application/vnd.openxmlformats-officedocument.spreadsheetml.sheet`): return drivertype.XLSX, true case strings.Contains(mediatype, `text/csv`): return drivertype.CSV, true case strings.Contains(mediatype, `text/tab-separated-values`): return drivertype.TSV, true } return drivertype.None, false }