package files import ( "bytes" "context" "fmt" "os" "path/filepath" "strings" "time" "github.com/neilotoole/sq/libsq/core/errz" "github.com/neilotoole/sq/libsq/core/ioz" "github.com/neilotoole/sq/libsq/core/ioz/checksum" "github.com/neilotoole/sq/libsq/core/ioz/lockfile" "github.com/neilotoole/sq/libsq/core/lg" "github.com/neilotoole/sq/libsq/core/lg/lga" "github.com/neilotoole/sq/libsq/core/options" "github.com/neilotoole/sq/libsq/core/progress" "github.com/neilotoole/sq/libsq/core/stringz" "github.com/neilotoole/sq/libsq/source" "github.com/neilotoole/sq/libsq/source/drivertype" "github.com/neilotoole/sq/libsq/source/location" ) // OptCacheLockTimeout is the time allowed to acquire a cache lock. // // See also: driver.OptIngestCache. var OptCacheLockTimeout = options.NewDuration( "cache.lock.timeout", "", 0, time.Second*5, "Wait timeout to acquire cache lock", `Wait timeout to acquire cache lock. During this period, retry will occur if the lock is already held by another process. If zero, no retry occurs.`, ) // CacheDir returns the cache dir. It is not guaranteed that the // returned dir exists. func (fs *Files) CacheDir() string { return fs.cacheDir } // TempDir returns the temp dir. It is not guaranteed that the // returned dir exists. func (fs *Files) TempDir() string { return fs.tempDir } // CacheDirFor gets the cache dir for handle. It is not guaranteed // that the returned dir exists or is accessible. func (fs *Files) CacheDirFor(src *source.Source) (dir string, err error) { handle := src.Handle if err = source.ValidHandle(handle); err != nil { return "", errz.Wrapf(err, "cache dir: invalid handle: %s", handle) } if handle == source.StdinHandle { // stdin is different input every time, so we need a unique // cache dir. In practice, stdin probably isn't using this function. handle += "_" + stringz.UniqN(32) } dir = filepath.Join( fs.cacheDir, "sources", filepath.Join(strings.Split(strings.TrimPrefix(handle, "@"), "/")...), fs.sourceHash(src), ) return dir, nil } // WriteIngestChecksum is invoked (after successful ingestion) to write the // checksum for the ingest cache db. func (fs *Files) WriteIngestChecksum(ctx context.Context, src, backingSrc *source.Source) (err error) { log := lg.FromContext(ctx) ingestFilePath, err := fs.filepath(src) if err != nil { return err } // Write the checksums file. var sum checksum.Checksum if sum, err = checksum.ForFile(ingestFilePath); err != nil { log.Warn("Failed to compute checksum for source file; caching not in effect", lga.Src, src, lga.Dest, backingSrc, lga.Path, ingestFilePath, lga.Err, err) return err } var checksumsPath string if _, _, checksumsPath, err = fs.CachePaths(src); err != nil { return err } if err = checksum.WriteFile(checksumsPath, sum, ingestFilePath); err != nil { log.Warn("Failed to write checksum; file caching not in effect", lga.Src, src, lga.Dest, backingSrc, lga.Path, ingestFilePath, lga.Err, err) } return err } // CachedBackingSourceFor returns the underlying backing source for src, if // it exists. If it does not exist, ok returns false. func (fs *Files) CachedBackingSourceFor(ctx context.Context, src *source.Source) (backingSrc *source.Source, ok bool, err error, ) { fs.mu.Lock() defer fs.mu.Unlock() switch location.TypeOf(src.Location) { case location.TypeLocalFile: return fs.cachedBackingSourceForFile(ctx, src) case location.TypeRemoteFile: return fs.cachedBackingSourceForRemoteFile(ctx, src) default: return nil, false, errz.Errorf("caching not applicable for source: %s", src.Handle) } } // cachedBackingSourceForFile returns the underlying cached backing // source for src, if it exists. func (fs *Files) cachedBackingSourceForFile(ctx context.Context, src *source.Source) (*source.Source, bool, error) { if _, dlFileExists := fs.downloadedFiles[src.Handle]; !dlFileExists { // It's NOT the case that the file has just been downloaded. // So, we check if there's an active download happening... if _, dlStreamExists := fs.streams[src.Handle]; dlStreamExists { lg.FromContext(ctx).Debug("Source has download stream, so backing cache DB not valid", lga.Src, src) return nil, false, nil } } _, cacheDBPath, checksumsPath, err := fs.CachePaths(src) if err != nil { return nil, false, err } if !ioz.FileAccessible(checksumsPath) { return nil, false, nil } mChecksums, err := checksum.ReadFile(checksumsPath) if err != nil { return nil, false, err } srcFilepath, err := fs.filepath(src) if err != nil { return nil, false, err } cachedChecksum, ok := mChecksums[srcFilepath] if !ok { return nil, false, nil } srcChecksum, err := checksum.ForFile(srcFilepath) if err != nil { return nil, false, err } if srcChecksum != cachedChecksum { return nil, false, nil } // The checksums match, so we can use the cached DB, // if it exists. if !ioz.FileAccessible(cacheDBPath) { return nil, false, nil } backingSrc := &source.Source{ Handle: src.Handle + "_cached", Location: "sqlite3://" + cacheDBPath, Type: drivertype.SQLite, } lg.FromContext(ctx).Debug("Found cached backing source DB", lga.Src, src, "backing_src", backingSrc) return backingSrc, true, nil } // cachedBackingSourceForRemoteFile returns the underlying cached backing // source for src, if it exists. func (fs *Files) cachedBackingSourceForRemoteFile(ctx context.Context, src *source.Source) (*source.Source, bool, error, ) { log := lg.FromContext(ctx) downloadedFile, _, err := fs.maybeStartDownload(ctx, src, true) if err != nil { return nil, false, err } if downloadedFile == "" { log.Debug("No cached download file", lga.Src, src) return nil, false, nil } log.Debug("Found cached download file", lga.Src, src, lga.Path, downloadedFile) return fs.cachedBackingSourceForFile(ctx, src) } // CachePaths returns the paths to the cache files for src. // There is no guarantee that these files exist, or are accessible. // It's just the paths. func (fs *Files) CachePaths(src *source.Source) (srcCacheDir, cacheDB, checksums string, err error) { if srcCacheDir, err = fs.CacheDirFor(src); err != nil { return "", "", "", err } checksums = filepath.Join(srcCacheDir, "checksums.txt") cacheDB = filepath.Join(srcCacheDir, "cache.sqlite.db") return srcCacheDir, cacheDB, checksums, nil } // sourceHash generates a hash for src. The hash is based on the // member fields of src, with special handling for src.Options. // Only the opts that affect data ingestion (options.TagIngestMutate) // are incorporated in the hash. The generated hash is used to // determine the cache dir for src. Thus, if a source is mutated // (e.g. the remote http location changes), a new cache dir results. func (fs *Files) sourceHash(src *source.Source) string { if src == nil { return "" } buf := bytes.Buffer{} buf.WriteString(src.Handle) buf.WriteString(string(src.Type)) buf.WriteString(src.Location) buf.WriteString(src.Catalog) buf.WriteString(src.Schema) mUsedKeys := make(map[string]any) if src.Options != nil { keys := src.Options.Keys() for _, k := range keys { opt := fs.optRegistry.Get(k) switch { case opt == nil, !opt.IsSet(src.Options), !opt.HasTag(options.TagIngestMutate): continue default: } buf.WriteString(k) v := src.Options[k] buf.WriteString(fmt.Sprintf("%v", v)) mUsedKeys[k] = v } } sum := checksum.Sum(buf.Bytes()) return sum } // cacheLockFor returns the lock file for src's cache. func (fs *Files) cacheLockFor(src *source.Source) (lockfile.Lockfile, error) { cacheDir, err := fs.CacheDirFor(src) if err != nil { return "", errz.Wrapf(err, "cache lock for %s", src.Handle) } lf, err := lockfile.New(filepath.Join(cacheDir, "pid.lock")) if err != nil { return "", errz.Wrapf(err, "cache lock for %s", src.Handle) } return lf, nil } // CacheLockAcquire acquires the cache lock for src. The caller must invoke // the returned unlock func. func (fs *Files) CacheLockAcquire(ctx context.Context, src *source.Source) (unlock func(), err error) { lock, err := fs.cacheLockFor(src) if err != nil { return nil, err } lockTimeout := OptCacheLockTimeout.Get(options.FromContext(ctx)) log := lg.FromContext(ctx).With(lga.Src, src, lga.Timeout, lockTimeout, lga.Lock, lock) log.Debug("Acquiring cache lock for source") bar := progress.FromContext(ctx).NewTimeoutWaiter( src.Handle+": acquire lock", time.Now().Add(lockTimeout), ) err = lock.Lock(ctx, lockTimeout) bar.Stop() if err != nil { return nil, errz.Wrap(err, src.Handle+": acquire cache lock") } return func() { if err = lock.Unlock(); err != nil { log.Warn("Failed to release cache lock", lga.Err, err) } }, nil } // CacheClearAll clears the entire cache dir. // Note that this operation is distinct from Files.doCacheSweep. func (fs *Files) CacheClearAll(ctx context.Context) error { fs.mu.Lock() defer fs.mu.Unlock() return fs.doCacheClearAll(ctx) } // CacheClearSource clears the ingest cache for src. If arg downloads is true, // the source's download dir is also cleared. The caller should typically // first acquire the cache lock for src via Files.cacheLockFor. func (fs *Files) CacheClearSource(ctx context.Context, src *source.Source, clearDownloads bool) error { fs.mu.Lock() defer fs.mu.Unlock() return fs.doCacheClearSource(ctx, src, clearDownloads) } func (fs *Files) doCacheClearSource(ctx context.Context, src *source.Source, clearDownloads bool) error { cacheDir, err := fs.CacheDirFor(src) if err != nil { return err } if !ioz.DirExists(cacheDir) { return nil } entries, err := os.ReadDir(cacheDir) if err != nil { return errz.Wrapf(err, "%s: clear cache", src.Handle) } for _, entry := range entries { switch entry.Name() { case "pid.lock": continue case "download": if !clearDownloads { continue } default: } if err = os.RemoveAll(filepath.Join(cacheDir, entry.Name())); err != nil { return errz.Wrapf(err, "%s: clear cache", src.Handle) } } lg.FromContext(ctx). With("clear_downloads", clearDownloads, lga.Src, src, lga.Dir, cacheDir). Info("Cleared source cache") return nil } func (fs *Files) doCacheClearAll(ctx context.Context) error { log := lg.FromContext(ctx).With(lga.Dir, fs.cacheDir) log.Debug("Clearing cache dir") if !ioz.DirExists(fs.cacheDir) { log.Debug("Cache dir does not exist") return nil } // Instead of directly deleting the existing cache dir, we first // move it to /tmp, and then try to delete it. This should probably // help with the situation where another sq instance has an open pid // lock in the cache dir. tmpDir := DefaultTempDir() if err := ioz.RequireDir(tmpDir); err != nil { return errz.Wrap(err, "cache clear") } relocateDir := filepath.Join(tmpDir, "dead_cache_"+stringz.Uniq8()) if err := os.Rename(fs.cacheDir, relocateDir); err != nil { return errz.Wrap(err, "cache clear: relocate") } if err := os.RemoveAll(relocateDir); err != nil { log.Warn("Could not delete relocated cache dir", lga.Path, relocateDir, lga.Err, err) } // Recreate the cache dir. if err := ioz.RequireDir(fs.cacheDir); err != nil { return errz.Wrap(err, "cache clear") } return nil } // doCacheSweep sweeps the cache dir, making a best-effort attempt // to remove any empty directories. Note that this operation is // distinct from Files.CacheClearAll. // // REVISIT: This doesn't really do as much as desired. It should // also be able to detect orphaned src cache dirs and delete those. func (fs *Files) doCacheSweep() { ctx, cancelFn := context.WithTimeout(context.Background(), time.Millisecond*100) defer cancelFn() log := fs.log.With(lga.Dir, fs.cacheDir) if unlock, err := fs.cfgLockFn(ctx); err != nil { log.Error("Sweep cache dir: failed to acquire config lock", lga.Lock, fs.cfgLockFn, lga.Err, err) return } else { defer unlock() } count, err := ioz.PruneEmptyDirTree(ctx, fs.cacheDir) if err != nil { log.Warn("Problem sweeping cache dir", lga.Err, err, "deleted_dirs", count) return } log.Info("Swept cache dir", "deleted_dirs", count) } // DefaultCacheDir returns the sq cache dir. This is generally // in USER_CACHE_DIR/*/sq, but could also be in TEMP_DIR/*/sq/cache // or similar. It is not guaranteed that the returned dir exists // or is accessible. func DefaultCacheDir() (dir string) { var err error if dir, err = os.UserCacheDir(); err != nil { // Some systems may not have a user cache dir, so we fall back // to the system temp dir. dir = filepath.Join(DefaultTempDir(), "cache") return dir } dir = filepath.Join(dir, "sq") return dir } // DefaultTempDir returns the default sq temp dir. It is not // guaranteed that the returned dir exists or is accessible. func DefaultTempDir() (dir string) { return filepath.Join(os.TempDir(), "sq") }