Repository: ariga/atlas
Pull request #1758 (2023-06-18 18:27:10)

cmd/atlas: defer HCL normalization until ReadState is called
Also fixed https://github.com/ariga/atlas/issues/1732


Patch:
diff --git a/cmd/atlas/internal/cmdext/reader.go b/cmd/atlas/internal/cmdext/reader.go
index 41955fe1..9f61fe6e 100644
--- a/cmd/atlas/internal/cmdext/reader.go
+++ b/cmd/atlas/internal/cmdext/reader.go
@@ -159,17 +159,28 @@ func stateReaderHCL(ctx context.Context, config *StateReaderConfig, paths []stri
 			config.Dev.URL.Schema,
 		)
 	}
-	if nr, ok := client.Driver.(schema.Normalizer); ok && config.Dev != nil { // only normalize on a dev database
-		if config.Dev.URL.Schema != "" {
-			realm.Schemas[0], err = nr.NormalizeSchema(ctx, realm.Schemas[0])
-		} else {
-			realm, err = nr.NormalizeRealm(ctx, realm)
-		}
-		if err != nil {
-			return nil, err
-		}
-	}
-	return &StateReadCloser{StateReader: migrate.Realm(realm), HCL: true}, nil
+	var normalized bool
+	return &StateReadCloser{
+		HCL: true,
+		// Defer normalization until the first call to ReadState. This is required because the same
+		// dev-database is used for both migration replaying and schema normalization. As a result,
+		// objects created by the migrations, which are not yet supported by Atlas, such as functions,
+		// won't be cleaned and can be referenced by the HCL schema.
+		StateReader: migrate.StateReaderFunc(func(ctx context.Context) (*schema.Realm, error) {
+			// Normalize once, only on dev database connection.
+			if nr, ok := client.Driver.(schema.Normalizer); ok && !normalized && config.Dev != nil {
+				if config.Dev.URL.Schema != "" {
+					realm.Schemas[0], err = nr.NormalizeSchema(ctx, realm.Schemas[0])
+				} else {
+					realm, err = nr.NormalizeRealm(ctx, realm)
+				}
+				if err != nil {
+					return nil, err
+				}
+			}
+			return realm, nil
+		}),
+	}, nil
 }
 
 // parseHCLPaths parses the HCL files in the given paths. If a path represents a directory,
diff --git a/sql/schema/dsl.go b/sql/schema/dsl.go
index d8485a3e..e9ae9e52 100644
--- a/sql/schema/dsl.go
+++ b/sql/schema/dsl.go
@@ -469,6 +469,13 @@ func TimePrecision(precision int) TimeOption {
 	}
 }
 
+// TimeScale configures the scale of the time type.
+func TimeScale(scale int) TimeOption {
+	return func(b *TimeType) {
+		b.Scale = &scale
+	}
+}
+
 // NewTimeColumn creates a new TimeType column.
 func NewTimeColumn(name, typ string, opts ...TimeOption) *Column {
 	t := &TimeType{T: typ}
diff --git a/sql/schema/schema.go b/sql/schema/schema.go
index 7edb34ef..9e6d70c7 100644
--- a/sql/schema/schema.go
+++ b/sql/schema/schema.go
@@ -265,6 +265,7 @@ type (
 	TimeType struct {
 		T         string
 		Precision *int
+		Scale     *int
 	}
 
 	// JSONType represents a JSON type.


Test patch:
diff --git a/internal/integration/testdata/postgres/cli-migrate-diff-unsupported.txt b/internal/integration/testdata/postgres/cli-migrate-diff-unsupported.txt
new file mode 100644
index 00000000..1037bb08
--- /dev/null
+++ b/internal/integration/testdata/postgres/cli-migrate-diff-unsupported.txt
@@ -0,0 +1,37 @@
+atlas migrate hash
+atlas migrate diff --dev-url URL --to file://./schema.hcl
+cmpmig 1 2.expected.sql
+
+# Test no diff.
+atlas migrate diff --dev-url URL --to file://./schema.hcl
+stdout 'The migration directory is synced with the desired state, no changes to be made'
+
+-- migrations/1_baseline.sql --
+CREATE OR REPLACE FUNCTION random_id(n INTEGER) RETURNS TEXT AS $$
+DECLARE
+  chars TEXT := '0123456789';
+  result TEXT := '';
+  i INTEGER := 0;
+BEGIN
+  FOR i IN 1..n LOOP
+    result := result || substr(chars, trunc(random()*10)::integer + 1, 1);
+  END LOOP;
+  RETURN result;
+END;
+$$ LANGUAGE plpgsql;
+
+-- schema.hcl --
+schema "script_cli_migrate_diff_unsupported" {}
+
+table "users" {
+  schema = schema.script_cli_migrate_diff_unsupported
+  column "id" {
+    null = false
+    type = text
+    default = sql("random_id(10)")
+  }
+}
+
+-- 2.expected.sql --
+-- Create "users" table
+CREATE TABLE "users" ("id" text NOT NULL DEFAULT random_id(10));
\ No newline at end of file
diff --git a/sql/schema/dsl_test.go b/sql/schema/dsl_test.go
index 8a74ae48..e4b52ff9 100644
--- a/sql/schema/dsl_test.go
+++ b/sql/schema/dsl_test.go
@@ -43,6 +43,7 @@ func TestSchema_AddTables(t *testing.T) {
 		schema.NewBoolColumn("active", "boolean"),
 		schema.NewNullStringColumn("name", "varchar", schema.StringSize(255)),
 		schema.NewTimeColumn("registered_at", "timestamp", schema.TimePrecision(6)),
+		schema.NewTimeColumn("open_time", "time", schema.TimeScale(2)),
 	}
 	users := schema.NewTable("users").
 		AddColumns(userColumns...).
@@ -74,6 +75,7 @@ func TestSchema_AddTables(t *testing.T) {
 		t,
 		func() *schema.Schema {
 			p := 6
+			scale := 2
 			s := &schema.Schema{Name: "public"}
 			users := &schema.Table{
 				Name:   "users",
@@ -86,6 +88,7 @@ func TestSchema_AddTables(t *testing.T) {
 					{Name: "active", Type: &schema.ColumnType{Type: &schema.BoolType{T: "boolean"}}},
 					{Name: "name", Type: &schema.ColumnType{Null: true, Type: &schema.StringType{T: "varchar", Size: 255}}},
 					{Name: "registered_at", Type: &schema.ColumnType{Null: false, Type: &schema.TimeType{T: "timestamp", Precision: &p}}},
+					{Name: "open_time", Type: &schema.ColumnType{Null: false, Type: &schema.TimeType{T: "time", Precision: nil, Scale: &scale}}},
 				},
 			}
 			s.Tables = append(s.Tables, users)