From c3887d64d841851347c62b4fe4a6d81122b95c55 Mon Sep 17 00:00:00 2001
From: Alexis Viscogliosi <alexis.viscogliosi@outlook.fr>
Date: Thu, 12 Dec 2024 09:43:22 +0100
Subject: [PATCH] update: set custom factory based on config, add context flag

---
 pkg/amigo/amigo.go             | 11 +++++------
 pkg/amigoconfig/cli_context.go |  6 +++++-
 pkg/entrypoint/root.go         |  4 +++-
 3 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/pkg/amigo/amigo.go b/pkg/amigo/amigo.go
index 247be13..64cd56e 100644
--- a/pkg/amigo/amigo.go
+++ b/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)
 	}
 }
 
diff --git a/pkg/amigoconfig/cli_context.go b/pkg/amigoconfig/cli_context.go
index 981f9ec..1c1dd77 100644
--- a/pkg/amigoconfig/cli_context.go
+++ b/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
 	}
 
diff --git a/pkg/entrypoint/root.go b/pkg/entrypoint/root.go
index 4e0697a..c2abc70 100644
--- a/pkg/entrypoint/root.go
+++ b/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() {