package cli import ( "context" "io" "os" "path/filepath" "sync" "github.com/neilotoole/sq/drivers" "github.com/neilotoole/sq/cli/config/yamlstore" v0_34_0 "github.com/neilotoole/sq/cli/config/yamlstore/upgrades/v0.34.0" "github.com/neilotoole/sq/libsq/core/lg/slogbuf" "github.com/neilotoole/sq/libsq/core/options" "github.com/neilotoole/sq/cli/config" "github.com/neilotoole/sq/cli/flag" "github.com/neilotoole/sq/drivers/csv" "github.com/neilotoole/sq/drivers/json" "github.com/neilotoole/sq/drivers/mysql" "github.com/neilotoole/sq/drivers/postgres" "github.com/neilotoole/sq/drivers/sqlite3" "github.com/neilotoole/sq/drivers/sqlserver" "github.com/neilotoole/sq/drivers/userdriver" "github.com/neilotoole/sq/drivers/userdriver/xmlud" "github.com/neilotoole/sq/drivers/xlsx" "github.com/neilotoole/sq/libsq/core/cleanup" "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/driver" "github.com/neilotoole/sq/libsq/source" "github.com/spf13/cobra" "golang.org/x/exp/slog" ) type runContextKey struct{} // WithRunContext returns ctx with rc added as a value. func WithRunContext(ctx context.Context, rc *RunContext) context.Context { if ctx == nil { ctx = context.Background() } return context.WithValue(ctx, runContextKey{}, rc) } // RunContextFrom extracts the RunContext added to ctx via WithRunContext. func RunContextFrom(ctx context.Context) *RunContext { return ctx.Value(runContextKey{}).(*RunContext) } // getRunContext is a convenience function for getting RunContext // from the cmd.Context(). func getRunContext(cmd *cobra.Command) *RunContext { rc := RunContextFrom(cmd.Context()) if rc.Cmd == nil { // rc.Cmd is usually set by the cmd.PreRunE that is added // by addCmd. But some commands (I'm looking at you __complete) don't // interact with that mechanism. So, we set the field here for those // odd cases. rc.Cmd = cmd } return rc } // RunContext is a container for injectable resources passed // to all execX funcs. The Close method should be invoked when // the RunContext is no longer needed. type RunContext struct { // Stdin typically is os.Stdin, but can be changed for testing. Stdin *os.File // Out is the output destination. // If nil, default to stdout. Out io.Writer // ErrOut is the error output destination. // If nil, default to stderr. ErrOut io.Writer // Cmd is the command instance provided by cobra for // the currently executing command. This field will // be set before the command's runFunc is invoked. Cmd *cobra.Command // Args is the arg slice supplied by cobra for // the currently executing command. This field will // be set before the command's runFunc is invoked. Args []string // Config is the run's config. Config *config.Config // ConfigStore is run's config store. ConfigStore config.Store initOnce sync.Once initErr error // writers holds the various writer types that // the CLI uses to print output. writers *writers driverReg *driver.Registry files *source.Files databases *driver.Databases clnup *cleanup.Cleanup OptionsRegistry *options.Registry } // newDefaultRunContext returns a RunContext configured // with standard values for logging, config, etc. This // effectively is the bootstrap mechanism for sq. // // Note: This func always returns a RunContext, even if // an error occurs during bootstrap of the RunContext (for // example if there's a config error). We do this to provide // enough framework so that such an error can be logged or // printed per the normal mechanisms if at all possible. func newDefaultRunContext(ctx context.Context, stdin *os.File, stdout, stderr io.Writer, args []string, ) (*RunContext, *slog.Logger, error) { // logbuf holds log records until defaultLogging is completed. log, logbuf := slogbuf.New() log = log.With(lga.Pid, os.Getpid()) rc := &RunContext{ Stdin: stdin, Out: stdout, ErrOut: stderr, OptionsRegistry: &options.Registry{}, } RegisterDefaultOpts(rc.OptionsRegistry) upgrades := yamlstore.UpgradeRegistry{ v0_34_0.Version: v0_34_0.Upgrade, } ctx = lg.NewContext(ctx, log) var configErr error rc.Config, rc.ConfigStore, configErr = yamlstore.Load(ctx, args, rc.OptionsRegistry, upgrades) log, logHandler, logCloser, logErr := defaultLogging(ctx, args, rc.Config) rc.clnup = cleanup.New().AddE(logCloser) if logErr != nil { stderrLog, h := stderrLogger() _ = logbuf.Flush(ctx, h) return rc, stderrLog, logErr } if logHandler != nil { if err := logbuf.Flush(ctx, logHandler); err != nil { return rc, log, err } } if log == nil { log = lg.Discard() } log = log.With(lga.Pid, os.Getpid()) if rc.Config == nil { rc.Config = config.New() } if configErr != nil { // configErr is more important, return that first return rc, log, configErr } return rc, log, nil } // init is invoked by cobra prior to the command RunE being // invoked. It sets up the driverReg, databases, writers and related // fundamental components. Subsequent invocations of this method // are no-op. func (rc *RunContext) init(ctx context.Context) error { if rc == nil { return errz.New("fatal: RunContext is nil") } rc.initOnce.Do(func() { rc.initErr = rc.doInit(ctx) }) return rc.initErr } // doInit performs the actual work of initializing rc. // It must only be invoked once. func (rc *RunContext) doInit(ctx context.Context) error { rc.clnup = cleanup.New() cfg, log := rc.Config, lg.FromContext(ctx) // If the --output=/some/file flag is set, then we need to // override rc.Out (which is typically stdout) to point it at // the output destination file. if cmdFlagChanged(rc.Cmd, flag.Output) { fpath, _ := rc.Cmd.Flags().GetString(flag.Output) fpath, err := filepath.Abs(fpath) if err != nil { return errz.Wrapf(err, "failed to get absolute path for --%s", flag.Output) } // Ensure the parent dir exists err = os.MkdirAll(filepath.Dir(fpath), os.ModePerm) if err != nil { return errz.Wrapf(err, "failed to make parent dir for --%s", flag.Output) } f, err := os.Create(fpath) if err != nil { return errz.Wrapf(err, "failed to open file specified by flag --%s", flag.Output) } rc.clnup.AddC(f) // Make sure the file gets closed eventually rc.Out = f } cmdOpts, err := getOptionsFromCmd(rc.Cmd) if err != nil { return err } rc.writers, rc.Out, rc.ErrOut = newWriters(rc.Cmd, cmdOpts, rc.Out, rc.ErrOut) var scratchSrcFunc driver.ScratchSrcFunc // scratchSrc could be nil, and that's ok scratchSrc := cfg.Collection.Scratch() if scratchSrc == nil { scratchSrcFunc = sqlite3.NewScratchSource } else { scratchSrcFunc = func(_ context.Context, name string) (src *source.Source, clnup func() error, err error) { return scratchSrc, nil, nil } } rc.files, err = source.NewFiles(ctx) if err != nil { lg.WarnIfFuncError(log, lga.Cleanup, rc.clnup.Run) return err } // Note: it's important that files.Close is invoked // after databases.Close (hence added to clnup first), // because databases could depend upon the existence of // files (such as a sqlite db file). rc.clnup.AddE(rc.files.Close) rc.files.AddDriverDetectors(source.DetectMagicNumber) rc.driverReg = driver.NewRegistry(log) rc.databases = driver.NewDatabases(log, rc.driverReg, scratchSrcFunc) rc.clnup.AddC(rc.databases) rc.driverReg.AddProvider(sqlite3.Type, &sqlite3.Provider{Log: log}) rc.driverReg.AddProvider(postgres.Type, &postgres.Provider{Log: log}) rc.driverReg.AddProvider(sqlserver.Type, &sqlserver.Provider{Log: log}) rc.driverReg.AddProvider(mysql.Type, &mysql.Provider{Log: log}) csvp := &csv.Provider{Log: log, Scratcher: rc.databases, Files: rc.files} rc.driverReg.AddProvider(csv.TypeCSV, csvp) rc.driverReg.AddProvider(csv.TypeTSV, csvp) rc.files.AddDriverDetectors(csv.DetectCSV, csv.DetectTSV) jsonp := &json.Provider{Log: log, Scratcher: rc.databases, Files: rc.files} rc.driverReg.AddProvider(json.TypeJSON, jsonp) rc.driverReg.AddProvider(json.TypeJSONA, jsonp) rc.driverReg.AddProvider(json.TypeJSONL, jsonp) sampleSize := drivers.OptIngestSampleSize.Get(cfg.Options) rc.files.AddDriverDetectors( json.DetectJSON(sampleSize), json.DetectJSONA(sampleSize), json.DetectJSONL(sampleSize), ) rc.driverReg.AddProvider(xlsx.Type, &xlsx.Provider{Log: log, Scratcher: rc.databases, Files: rc.files}) rc.files.AddDriverDetectors(xlsx.DetectXLSX) // One day we may have more supported user driver genres. userDriverImporters := map[string]userdriver.ImportFunc{ xmlud.Genre: xmlud.Import, } for i, userDriverDef := range cfg.Ext.UserDrivers { userDriverDef := userDriverDef errs := userdriver.ValidateDriverDef(userDriverDef) if len(errs) > 0 { err := errz.Combine(errs...) err = errz.Wrapf(err, "failed validation of user driver definition [%d] {%s} from config", i, userDriverDef.Name) return err } importFn, ok := userDriverImporters[userDriverDef.Genre] if !ok { return errz.Errorf("unsupported genre {%s} for user driver {%s} specified via config", userDriverDef.Genre, userDriverDef.Name) } // For each user driver definition, we register a // distinct userdriver.Provider instance. udp := &userdriver.Provider{ Log: log, DriverDef: userDriverDef, ImportFn: importFn, Scratcher: rc.databases, Files: rc.files, } rc.driverReg.AddProvider(source.DriverType(userDriverDef.Name), udp) rc.files.AddDriverDetectors(udp.Detectors()...) } return nil } // Close should be invoked to dispose of any open resources // held by rc. If an error occurs during Close and rc.Log // is not nil, that error is logged at WARN level before // being returned. func (rc *RunContext) Close() error { if rc == nil { return nil } return errz.Wrap(rc.clnup.Run(), "Close RunContext") }