From d6abc61b013fdd0b67a07da5e7903892eabe9492 Mon Sep 17 00:00:00 2001
From: Ariel Mashraki <ariel@mashraki.co.il>
Date: Fri, 4 Oct 2024 23:25:03 +0300
Subject: [PATCH] sql/internal/sqlx: pass diff options to drivers

---
 cmd/atlas/internal/cmdapi/cmdapi.go | 17 +++++++++++++----
 sql/internal/sqlx/diff.go           |  4 ++--
 sql/mysql/diff.go                   |  2 +-
 sql/postgres/driver_oss.go          |  2 +-
 sql/sqlite/diff.go                  |  2 +-
 5 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/cmd/atlas/internal/cmdapi/cmdapi.go b/cmd/atlas/internal/cmdapi/cmdapi.go
index c182b33b95d..0d22ab33624 100644
--- a/cmd/atlas/internal/cmdapi/cmdapi.go
+++ b/cmd/atlas/internal/cmdapi/cmdapi.go
@@ -105,6 +105,11 @@ type (
 	AbortError struct {
 		Err error
 	}
+	// Aborter allows errors to signal if the error is an abort error.
+	Aborter interface {
+		error
+		IsAbort()
+	}
 )
 
 func (e *FormattedError) Error() string { return e.Err.Error() }
@@ -129,10 +134,14 @@ func (e *AbortError) FormatError(cmd *cobra.Command) {
 
 // RunE wraps the command cobra.Command.RunE function with additional postrun logic.
 func RunE(f func(*cobra.Command, []string) error) func(*cobra.Command, []string) error {
-	return func(cmd *cobra.Command, args []string) error {
-		err := f(cmd, args)
-		if ef, ok := err.(ErrorFormatter); ok {
-			ef.FormatError(cmd)
+	return func(cmd *cobra.Command, args []string) (err error) {
+		if err = f(cmd, args); err != nil {
+			if err1 := (Aborter)(nil); errors.As(err, &err1) {
+				err = &AbortError{Err: err}
+			}
+			if ef, ok := err.(ErrorFormatter); ok {
+				ef.FormatError(cmd)
+			}
 		}
 		return err
 	}
diff --git a/sql/internal/sqlx/diff.go b/sql/internal/sqlx/diff.go
index 40bf4e3d11e..536894fc0fe 100644
--- a/sql/internal/sqlx/diff.go
+++ b/sql/internal/sqlx/diff.go
@@ -37,7 +37,7 @@ type (
 
 		// SchemaObjectDiff returns a changeset for migrating schema objects from
 		// one state to the other. For example, changing schema custom types.
-		SchemaObjectDiff(from, to *schema.Schema) ([]schema.Change, error)
+		SchemaObjectDiff(from, to *schema.Schema, _ *schema.DiffOptions) ([]schema.Change, error)
 
 		// TableAttrDiff returns a changeset for migrating table attributes from
 		// one state to the other. For example, dropping or adding a `CHECK` constraint.
@@ -196,7 +196,7 @@ func (d *Diff) schemaDiff(from, to *schema.Schema, opts *schema.DiffOptions) ([]
 		})
 	}
 	// Add, drop or modify objects.
-	change, err := d.SchemaObjectDiff(from, to)
+	change, err := d.SchemaObjectDiff(from, to, opts)
 	if err != nil {
 		return nil, err
 	}
diff --git a/sql/mysql/diff.go b/sql/mysql/diff.go
index 899b9791e60..9804fb08d89 100644
--- a/sql/mysql/diff.go
+++ b/sql/mysql/diff.go
@@ -61,7 +61,7 @@ func (*diff) RealmObjectDiff(_, _ *schema.Realm) ([]schema.Change, error) {
 
 // SchemaObjectDiff returns a changeset for migrating schema objects from
 // one state to the other.
-func (*diff) SchemaObjectDiff(_, _ *schema.Schema) ([]schema.Change, error) {
+func (*diff) SchemaObjectDiff(_, _ *schema.Schema, _ *schema.DiffOptions) ([]schema.Change, error) {
 	return nil, nil
 }
 
diff --git a/sql/postgres/driver_oss.go b/sql/postgres/driver_oss.go
index 0d06de2a0bf..2eb7110736e 100644
--- a/sql/postgres/driver_oss.go
+++ b/sql/postgres/driver_oss.go
@@ -204,7 +204,7 @@ func (*diff) RealmObjectDiff(_, _ *schema.Realm) ([]schema.Change, error) {
 
 // SchemaObjectDiff returns a changeset for migrating schema objects from
 // one state to the other.
-func (*diff) SchemaObjectDiff(from, to *schema.Schema) ([]schema.Change, error) {
+func (*diff) SchemaObjectDiff(from, to *schema.Schema, _ *schema.DiffOptions) ([]schema.Change, error) {
 	var changes []schema.Change
 	// Drop or modify enums.
 	for _, o1 := range from.Objects {
diff --git a/sql/sqlite/diff.go b/sql/sqlite/diff.go
index c728b4bbba9..bbf5205630c 100644
--- a/sql/sqlite/diff.go
+++ b/sql/sqlite/diff.go
@@ -36,7 +36,7 @@ func (*diff) RealmObjectDiff(_, _ *schema.Realm) ([]schema.Change, error) {
 
 // SchemaObjectDiff returns a changeset for migrating schema objects from
 // one state to the other.
-func (*diff) SchemaObjectDiff(_, _ *schema.Schema) ([]schema.Change, error) {
+func (*diff) SchemaObjectDiff(_, _ *schema.Schema, _ *schema.DiffOptions) ([]schema.Change, error) {
 	return nil, nil
 }