update: set custom factory based on config, add context flag

pkg/amigo/amigo.go

@@ -20,14 +20,13 @@ type Amigo struct {
 	CustomSchemaFactory schema.Factory[schema.Schema]
 }
 
-type OptionFn func(a *Amigo)
+type OptionFn func(*Amigo)
 
-// WithCustomSchemaFactory allows to set a custom schema factory
-// A schema factory is a function that returns a schema : func(ctx *MigratorContext, tx DB, db DB) T
-// This is used for custom drivers implementation
-func WithCustomSchemaFactory(factory schema.Factory[schema.Schema]) OptionFn {
+// WithCustomSchemaFactory returns an option function that sets a custom schema factory
+// based on the config
+func WithCustomSchemaFactory(factoryFn func(cfg amigoconfig.Config) schema.Factory[schema.Schema]) OptionFn {
 	return func(a *Amigo) {
-		a.CustomSchemaFactory = factory
+		a.CustomSchemaFactory = factoryFn(*a.Config)
 	}
 }
 

pkg/amigoconfig/cli_context.go

@@ -82,16 +82,20 @@ func (c *Config) OverrideWithYamlConfig(yaml *YamlConfig) {
 	if yaml.ShellPath != "" {
 		c.RootConfig.ShellPath = yaml.ShellPath
 	}
+
 	if yaml.Debug {
 		c.RootConfig.Debug = yaml.Debug
 	}
+
 	if yaml.ShowSQL {
 		c.RootConfig.ShowSQL = yaml.ShowSQL
 	}
+
 	if yaml.ShowSQLSyntaxHighlighting {
 		c.RootConfig.ShowSQLSyntaxHighlighting = yaml.ShowSQLSyntaxHighlighting
 	}
-	if yaml.CurrentContext != "" {
+
+	if yaml.CurrentContext != "" && c.RootConfig.CurrentContext == "" {
 		c.RootConfig.CurrentContext = yaml.CurrentContext
 	}
 

pkg/entrypoint/root.go

@@ -102,7 +102,9 @@ func init() {
 
 	rootCmd.PersistentFlags().BoolVar(&config.Debug, "debug", false, "Print debug information")
 
-	initConfig()
+	rootCmd.PersistentFlags().StringVarP(&config.RootConfig.CurrentContext, "context", "c", "", "Context to use")
+
+	initConfig(cfg)
 }
 
 func initConfig() {