feat(fxorm): Prevent connection auto close on test mode (#233)

Author: ekkinox

fxorm/README.md

@@ -59,12 +59,12 @@ type Model struct {
 
 func main() {
 	fx.New(
-		fxconfig.FxConfigModule,      // load the module dependencies
+		fxconfig.FxConfigModule,          // load the module dependencies
 		fxlog.FxLogModule,
 		fxtrace.FxTraceModule,
-		fxorm.FxOrmModule,            // load the module
-		fx.Invoke(func(db *gorm.DB) { // invoke the orm
-			db.Create(&Model{Name: "some name"})
+		fxorm.FxOrmModule,                // load the module
+		fx.Invoke(func(gormDB *gorm.DB) { // invoke the orm
+			gormDB.Create(&Model{Name: "some name"})
 		}),
 	).Run()
 }
@@ -91,7 +91,7 @@ app:
 modules:
   orm:
     driver: mysql                                               # driver to use
-    dsn: "user:pass@tcp(dbhost:3306)/dbname?parseTime=True"     # database DSN to connect to
+    dsn: "user:password@tcp(localhost:3306)/db?parseTime=True"  # database DSN to use
     config:
       dry_run: false                                            # disabled by default
       skip_default_transaction: false                           # disabled by default
@@ -113,7 +113,17 @@ modules:
       values: true   # by adding or not clear SQL queries parameters values in trace spans, disabled by default
 ```
 
-See [GORM Config](https://github.com/go-gorm/gorm/blob/master/gorm.go) for more details about the configuration.
+See [GORM Config](https://github.com/go-gorm/gorm/blob/master/gorm.go) for more details about the `modules.orm.config` configuration keys.
+
+For security reasons, you should avoid to hardcode DSN sensible parts (like the password) in your config files, you can use the [env vars placeholders](https://github.com/ankorstore/yokai/tree/main/fxconfig#configuration-env-var-placeholders) instead:
+
+```yaml
+# ./configs/config.yaml
+modules:
+  orm:
+    driver: mysql
+    dsn: "${MYSQL_USER}:${MYSQL_PASSWORD}@tcp(${MYSQL_HOST}:${MYSQL_PORT})/${MYSQL_DATABASE}?parseTime=True"
+```
 
 ### Auto migrations
 
@@ -143,8 +153,8 @@ func main() {
 		fxtrace.FxTraceModule,
 		fxorm.FxOrmModule,                   // load the module
 		fxorm.RunFxOrmAutoMigrate(&Model{}), // run auto migration for Model
-		fx.Invoke(func(db *gorm.DB) { // invoke the orm
-			db.Create(&Model{Name: "some name"})
+		fx.Invoke(func(gormDB *gorm.DB) {        // invoke the orm
+			gormDB.Create(&Model{Name: "some name"})
 		}),
 	).Run()
 }
@@ -163,17 +173,11 @@ You can disable it in the configuration:
 
 ```yaml
 # ./configs/config.yaml
-app:
-  name: app
-  env: dev
-  version: 0.1.0
-  debug: false
+
 modules:
   orm:
-    driver: mysql                                               # driver to use
-    dsn: user:pass@tcp(127.0.0.1:3306)/dbname?parseTime=True"   # database DSN to connect to
     config:
-      skip_default_transaction: true                            # disable default transaction
+      skip_default_transaction: true # disable default transaction
 ```
 
 #### Cache Prepared Statement
@@ -182,17 +186,10 @@ To create a prepared statement when executing any SQL (and cache them to speed u
 
 ```yaml
 # ./configs/config.yaml
-app:
-  name: app
-  env: dev
-  version: 0.1.0
-  debug: false
 modules:
   orm:
-    driver: mysql                                               # driver to use
-    dsn: user:pass@tcp(127.0.0.1:3306)/dbname?parseTime=True"   # database DSN to connect to
     config:
-      prepare_stmt: true                                        # enable prepared statements
+      prepare_stmt: true # enable prepared statements
 ```
 
 ### Override

fxorm/module.go

@@ -36,24 +36,6 @@ type FxOrmParam struct {
 	TracerProvider trace.TracerProvider
 }
 
-// RunFxOrmAutoMigrate performs auto migrations for a provided list of models.
-func RunFxOrmAutoMigrate(models ...any) fx.Option {
-	return fx.Invoke(func(logger *log.Logger, db *gorm.DB) error {
-		logger.Info().Msg("starting ORM auto migration")
-
-		err := db.AutoMigrate(models...)
-		if err != nil {
-			logger.Error().Err(err).Msg("error during ORM auto migration")
-
-			return err
-		}
-
-		logger.Info().Msg("ORM auto migration success")
-
-		return nil
-	})
-}
-
 // NewFxOrm returns a [gorm.DB].
 func NewFxOrm(p FxOrmParam) (*gorm.DB, error) {
 	ormConfig := gorm.Config{
@@ -80,37 +62,54 @@ func NewFxOrm(p FxOrmParam) (*gorm.DB, error) {
 
 	driver := orm.FetchDriver(p.Config.GetString("modules.orm.driver"))
 
-	db, err := p.Factory.Create(
+	gormDB, err := p.Factory.Create(
 		orm.WithDsn(p.Config.GetString("modules.orm.dsn")),
 		orm.WithDriver(driver),
 		orm.WithConfig(ormConfig),
 	)
-
 	if err != nil {
 		return nil, err
 	}
 
 	if p.Config.GetBool("modules.orm.trace.enabled") {
-		err = db.Use(plugin.NewOrmTracerPlugin(p.TracerProvider, p.Config.GetBool("modules.orm.trace.values")))
+		err = gormDB.Use(plugin.NewOrmTracerPlugin(p.TracerProvider, p.Config.GetBool("modules.orm.trace.values")))
 		if err != nil {
 			return nil, err
 		}
 	}
 
 	p.LifeCycle.Append(fx.Hook{
 		OnStop: func(ctx context.Context) error {
-			if driver != orm.Sqlite {
-				ormDb, err := db.DB()
+			if !p.Config.IsTestEnv() {
+				db, err := gormDB.DB()
 				if err != nil {
 					return err
 				}
 
-				return ormDb.Close()
+				return db.Close()
 			}
 
 			return nil
 		},
 	})
 
-	return db, nil
+	return gormDB, nil
+}
+
+// RunFxOrmAutoMigrate performs auto migrations for a provided list of models.
+func RunFxOrmAutoMigrate(models ...any) fx.Option {
+	return fx.Invoke(func(logger *log.Logger, gormDB *gorm.DB) error {
+		logger.Info().Msg("starting ORM auto migration")
+
+		err := gormDB.AutoMigrate(models...)
+		if err != nil {
+			logger.Error().Err(err).Msg("error during ORM auto migration")
+
+			return err
+		}
+
+		logger.Info().Msg("ORM auto migration success")
+
+		return nil
+	})
 }

fxorm/module_test.go

@@ -39,6 +39,7 @@ func (c *fakeT) FailNow() {
 }
 
 func TestModuleWithSqliteAndWithLogEnabledWithValuesAndWithTracesEnabledWithValues(t *testing.T) {
+	t.Setenv("APP_ENV", "test")
 	t.Setenv("APP_CONFIG_PATH", "testdata/config")
 	t.Setenv("ORM_DRIVER", "sqlite")
 	t.Setenv("ORM_DSN", ":memory:")
@@ -50,6 +51,7 @@ func TestModuleWithSqliteAndWithLogEnabledWithValuesAndWithTracesEnabledWithValu
 
 	ctx := context.Background()
 
+	var gormDB *gorm.DB
 	var repository *model.TestModelRepository
 	var logBuffer logtest.TestLogBuffer
 	var traceExporter tracetest.TestTraceExporter
@@ -68,7 +70,7 @@ func TestModuleWithSqliteAndWithLogEnabledWithValuesAndWithTracesEnabledWithValu
 				Name: "test",
 			})
 		}),
-		fx.Populate(&repository, &logBuffer, &traceExporter),
+		fx.Populate(&gormDB, &repository, &logBuffer, &traceExporter),
 	).RequireStart().RequireStop()
 
 	// assert on DB insertion
@@ -102,9 +104,17 @@ func TestModuleWithSqliteAndWithLogEnabledWithValuesAndWithTracesEnabledWithValu
 		semconv.DBSystemKey.String("sqlite"),
 		semconv.DBStatementKey.String("INSERT INTO `test_models` (`name`) VALUES (\"test\")"),
 	)
+
+	// close
+	db, err := gormDB.DB()
+	assert.NoError(t, err)
+
+	err = db.Close()
+	assert.NoError(t, err)
 }
 
 func TestModuleWithSqliteAndWithLogEnabledWithoutValuesAndWithTracesEnabledWithoutValues(t *testing.T) {
+	t.Setenv("APP_ENV", "test")
 	t.Setenv("APP_CONFIG_PATH", "testdata/config")
 	t.Setenv("ORM_DRIVER", "sqlite")
 	t.Setenv("ORM_DSN", ":memory:")
@@ -116,6 +126,7 @@ func TestModuleWithSqliteAndWithLogEnabledWithoutValuesAndWithTracesEnabledWitho
 
 	ctx := context.Background()
 
+	var gormDB *gorm.DB
 	var repository *model.TestModelRepository
 	var logBuffer logtest.TestLogBuffer
 	var traceExporter tracetest.TestTraceExporter
@@ -134,7 +145,7 @@ func TestModuleWithSqliteAndWithLogEnabledWithoutValuesAndWithTracesEnabledWitho
 				Name: "test",
 			})
 		}),
-		fx.Populate(&repository, &logBuffer, &traceExporter),
+		fx.Populate(&gormDB, &repository, &logBuffer, &traceExporter),
 	).RequireStart().RequireStop()
 
 	// assert on DB insertion
@@ -168,9 +179,17 @@ func TestModuleWithSqliteAndWithLogEnabledWithoutValuesAndWithTracesEnabledWitho
 		semconv.DBSystemKey.String("sqlite"),
 		semconv.DBStatementKey.String("INSERT INTO `test_models` (`name`) VALUES (\"?\")"),
 	)
+
+	// close
+	db, err := gormDB.DB()
+	assert.NoError(t, err)
+
+	err = db.Close()
+	assert.NoError(t, err)
 }
 
 func TestModuleWithSqliteAndWithLogDisabledAndWithTracesDisabled(t *testing.T) {
+	t.Setenv("APP_ENV", "test")
 	t.Setenv("APP_CONFIG_PATH", "testdata/config")
 	t.Setenv("ORM_DRIVER", "sqlite")
 	t.Setenv("ORM_DSN", ":memory:")
@@ -179,6 +198,7 @@ func TestModuleWithSqliteAndWithLogDisabledAndWithTracesDisabled(t *testing.T) {
 
 	ctx := context.Background()
 
+	var gormDB *gorm.DB
 	var repository *model.TestModelRepository
 	var logBuffer logtest.TestLogBuffer
 	var traceExporter tracetest.TestTraceExporter
@@ -197,7 +217,7 @@ func TestModuleWithSqliteAndWithLogDisabledAndWithTracesDisabled(t *testing.T) {
 				Name: "test",
 			})
 		}),
-		fx.Populate(&repository, &logBuffer, &traceExporter),
+		fx.Populate(&gormDB, &repository, &logBuffer, &traceExporter),
 	).RequireStart().RequireStop()
 
 	// assert on DB insertion
@@ -226,9 +246,17 @@ func TestModuleWithSqliteAndWithLogDisabledAndWithTracesDisabled(t *testing.T) {
 
 	// assert on SQL traces
 	assert.False(t, traceExporter.HasSpan("orm.Create"))
+
+	// close
+	db, err := gormDB.DB()
+	assert.NoError(t, err)
+
+	err = db.Close()
+	assert.NoError(t, err)
 }
 
 func TestModuleWithAutoMigrationError(t *testing.T) {
+	t.Setenv("APP_ENV", "test")
 	t.Setenv("APP_CONFIG_PATH", "testdata/config")
 	t.Setenv("ORM_DRIVER", "sqlite")
 	t.Setenv("ORM_DSN", ":memory:")
@@ -253,7 +281,7 @@ func TestModuleWithAutoMigrationError(t *testing.T) {
 func TestModuleDecoration(t *testing.T) {
 	t.Setenv("APP_CONFIG_PATH", "testdata/config")
 
-	var db *gorm.DB
+	var gormDB *gorm.DB
 
 	fxtest.New(
 		t,
@@ -262,8 +290,15 @@ func TestModuleDecoration(t *testing.T) {
 		fxtrace.FxTraceModule,
 		fxorm.FxOrmModule,
 		fx.Decorate(factory.NewTestOrmFactory),
-		fx.Populate(&db),
+		fx.Populate(&gormDB),
 	).RequireStart().RequireStop()
 
-	assert.True(t, db.DryRun)
+	assert.True(t, gormDB.DryRun)
+
+	// close
+	db, err := gormDB.DB()
+	assert.NoError(t, err)
+
+	err = db.Close()
+	assert.NoError(t, err)
 }

fxorm/testdata/config/config.test.yaml

@@ -0,0 +1,8 @@
+modules:
+  log:
+    level: debug
+    output: test
+  trace:
+    processor:
+      type: test
+

fxorm/testdata/config/config.yaml

@@ -1,12 +1,6 @@
 app:
   name: test
 modules:
-  log:
-    level: debug
-    output: test
-  trace:
-    processor:
-      type: test
   orm:
     driver: ${ORM_DRIVER}
     dsn: ${ORM_DSN}