From 85e5b3e2f8a7f0e8de21e731302cfa9f9c447c4b Mon Sep 17 00:00:00 2001
From: Karolis <karolis.arbaciauskas@gmail.com>
Date: Thu, 5 Oct 2023 17:24:26 +0300
Subject: [PATCH] gormschema: support custom gorm.Config option (#14)

---
 .github/workflows/ci.yaml |  2 +-
 README.md                 | 17 ++++++++++++++++-
 go.mod                    |  2 +-
 go.sum                    |  4 ++++
 gormschema/gorm.go        | 29 +++++++++++++++++++++++------
 gormschema/gorm_test.go   | 22 ++++++++++++++++++++++
 6 files changed, 67 insertions(+), 9 deletions(-)
 create mode 100644 gormschema/gorm_test.go

diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
index c0f1b62..11a271e 100644
--- a/.github/workflows/ci.yaml
+++ b/.github/workflows/ci.yaml
@@ -15,7 +15,7 @@ jobs:
       - name: Run Go linters
         uses: golangci/golangci-lint-action@v3
         with:
-          args: --verbose
+          args: --verbose --timeout=5m
           skip-pkg-cache: true
   unit-tests:
     runs-on: ubuntu-latest
diff --git a/README.md b/README.md
index d387e66..b921015 100644
--- a/README.md
+++ b/README.md
@@ -108,6 +108,21 @@ env "gorm" {
 }
 ```
 
+### Additional Configuration
+
+To supply custom `gorm.Config{}` object to the provider use the [Go Program Mode](#as-go-file) with
+the `WithConfig` option. For example, to disable foreign keys:
+
+```go
+loader := New("sqlite", WithConfig(
+    &gorm.Config{
+        DisableForeignKeyConstraintWhenMigrating: true,
+    },
+))
+```
+
+For a full list of options, see the [GORM documentation](https://gorm.io/docs/gorm_config.html).
+
 ### Usage
 
 Once you have the provider installed, you can use it to apply your GORM schema to the database:
@@ -127,7 +142,7 @@ target database.
 
 #### Diff
 
-Atlas supports a [version migration](https://atlasgo.io/concepts/declarative-vs-versioned#versioned-migrations) 
+Atlas supports a [versioned migration](https://atlasgo.io/concepts/declarative-vs-versioned#versioned-migrations) 
 workflow, where each change to the database is versioned and recorded in a migration file. You can use the
 `atlas migrate diff` command to automatically generate a migration file that will migrate the database
 from its latest revision to the current GORM schema.
diff --git a/go.mod b/go.mod
index 1fb3ea6..c2371c9 100644
--- a/go.mod
+++ b/go.mod
@@ -3,7 +3,7 @@ module ariga.io/atlas-provider-gorm
 go 1.20
 
 require (
-	ariga.io/atlas-go-sdk v0.0.0-20230709063453-1058d6508503
+	ariga.io/atlas-go-sdk v0.1.1-0.20231001054405-7edfcfc14f1c
 	github.com/alecthomas/kong v0.7.1
 	github.com/stretchr/testify v1.8.4
 	golang.org/x/tools v0.10.0
diff --git a/go.sum b/go.sum
index 68cb8e7..ba18070 100644
--- a/go.sum
+++ b/go.sum
@@ -1,5 +1,9 @@
 ariga.io/atlas-go-sdk v0.0.0-20230709063453-1058d6508503 h1:D12EzAAjhL7xEJk5jFnkKYUXVrDQ4mh0IjB8vqTmo8I=
 ariga.io/atlas-go-sdk v0.0.0-20230709063453-1058d6508503/go.mod h1:fwi5nIOFLedo6CqZ0a172dhykLWBnoD25bqmZhvW948=
+ariga.io/atlas-go-sdk v0.1.0 h1:sSPV26Lv2DIzbc+oZxgnAqYgKEwBFU0FsFDUUnUNHHs=
+ariga.io/atlas-go-sdk v0.1.0/go.mod h1:738TvNdlvLnkRhB+euXvrhKPJkeV+LPoVa4xarUGCaQ=
+ariga.io/atlas-go-sdk v0.1.1-0.20231001054405-7edfcfc14f1c h1:jvi4KB/7DmYYT+Wy2TFImccaBU0+dw7V8Un67NDGuio=
+ariga.io/atlas-go-sdk v0.1.1-0.20231001054405-7edfcfc14f1c/go.mod h1:MLvZ9QwZx1KhI6+8XguxHPUPm0/PTTUr46S5GQAe9WI=
 github.com/alecthomas/assert/v2 v2.1.0 h1:tbredtNcQnoSd3QBhQWI7QZ3XHOVkw1Moklp2ojoH/0=
 github.com/alecthomas/kong v0.7.1 h1:azoTh0IOfwlAX3qN9sHWTxACE2oV8Bg2gAwBsMwDQY4=
 github.com/alecthomas/kong v0.7.1/go.mod h1:n1iCIO2xS46oE8ZfYCNDqdR0b0wZNrXAIAqro/2132U=
diff --git a/gormschema/gorm.go b/gormschema/gorm.go
index 3448bee..f306b6d 100644
--- a/gormschema/gorm.go
+++ b/gormschema/gorm.go
@@ -13,15 +13,32 @@ import (
 )
 
 // New returns a new Loader.
-func New(dialect string) *Loader {
-	return &Loader{dialect: dialect}
+func New(dialect string, opts ...Option) *Loader {
+	l := &Loader{dialect: dialect, config: &gorm.Config{}}
+	for _, opt := range opts {
+		opt(l)
+	}
+	return l
 }
 
-// Loader is a Loader for gorm schema.
-type Loader struct {
-	dialect string
+type (
+	// Loader is a Loader for gorm schema.
+	Loader struct {
+		dialect string
+		config  *gorm.Config
+	}
+	// Option configures the Loader.
+	Option func(*Loader)
+)
+
+// WithConfig sets the gorm config.
+func WithConfig(cfg *gorm.Config) Option {
+	return func(l *Loader) {
+		l.config = cfg
+	}
 }
 
+// Load loads the models and returns the DDL statements representing the schema.
 func (l *Loader) Load(models ...any) (string, error) {
 	var di gorm.Dialector
 	switch l.dialect {
@@ -52,7 +69,7 @@ func (l *Loader) Load(models ...any) (string, error) {
 	default:
 		return "", fmt.Errorf("unsupported engine: %s", l.dialect)
 	}
-	db, err := gorm.Open(di, &gorm.Config{})
+	db, err := gorm.Open(di, l.config)
 	if err != nil {
 		return "", err
 	}
diff --git a/gormschema/gorm_test.go b/gormschema/gorm_test.go
new file mode 100644
index 0000000..a929156
--- /dev/null
+++ b/gormschema/gorm_test.go
@@ -0,0 +1,22 @@
+package gormschema
+
+import (
+	"testing"
+
+	"ariga.io/atlas-provider-gorm/internal/testdata/models"
+	"github.com/stretchr/testify/require"
+	"gorm.io/gorm"
+)
+
+func TestConfig(t *testing.T) {
+	l := New("sqlite", WithConfig(
+		&gorm.Config{
+			DisableForeignKeyConstraintWhenMigrating: true,
+		},
+	))
+	sql, err := l.Load(models.Pet{}, models.User{})
+	require.NoError(t, err)
+	require.Contains(t, sql, "CREATE TABLE `pets`")
+	require.Contains(t, sql, "CREATE TABLE `users`")
+	require.NotContains(t, sql, "FOREIGN KEY")
+}