From 8d3cf737bca5f5374fd363873240d2534c1e7f1c Mon Sep 17 00:00:00 2001
From: "Giau. Tran Minh" <hello@giautm.dev>
Date: Wed, 15 Jan 2025 18:17:54 +0700
Subject: [PATCH] cmd/atlas/cmdapi: fallback `env://src` to `env://schema.src`
 if exists

---
 cmd/atlas/internal/cmdapi/project.go      |  4 ++++
 cmd/atlas/internal/cmdapi/project_test.go | 18 ++++++++++++++++++
 2 files changed, 22 insertions(+)

diff --git a/cmd/atlas/internal/cmdapi/project.go b/cmd/atlas/internal/cmdapi/project.go
index b1434486b42..a9528915952 100644
--- a/cmd/atlas/internal/cmdapi/project.go
+++ b/cmd/atlas/internal/cmdapi/project.go
@@ -259,6 +259,10 @@ func (e *Env) VarFromURL(s string) (string, error) {
 		switch {
 		case u.Host == "src":
 			attr, ok = e.Attr("src")
+			if !ok && e.Schema != nil {
+				// Fallback to schema.src
+				attr, ok = e.Schema.Attr("src")
+			}
 		case e.Schema != nil:
 			attr, ok = e.Schema.Attr("src")
 		}
diff --git a/cmd/atlas/internal/cmdapi/project_test.go b/cmd/atlas/internal/cmdapi/project_test.go
index 1b5457d0a70..545f390a9bd 100644
--- a/cmd/atlas/internal/cmdapi/project_test.go
+++ b/cmd/atlas/internal/cmdapi/project_test.go
@@ -129,6 +129,12 @@ env "multi" {
     }
   }
 }
+env "schema" {
+  url = "mysql://root:pass@localhost:3306/"
+  schema {
+    src = "./a.hcl"
+  }
+}
 `
 	path := filepath.Join(d, "atlas.hcl")
 	err := os.WriteFile(path, []byte(h), 0600)
@@ -274,6 +280,18 @@ env "multi" {
 		_, _, err = EnvByName(&cobra.Command{}, "home", nil)
 		require.ErrorContains(t, err, `no such file or directory`)
 	})
+	t.Run("fallback src to schema", func(t *testing.T) {
+		GlobalFlags.ConfigURL = "file://" + path
+		_, envs, err := EnvByName(&cobra.Command{}, "schema", nil)
+		require.NoError(t, err)
+		require.Len(t, envs, 1)
+		srcs, err := envs[0].Sources()
+		require.NoError(t, err)
+		require.EqualValues(t, []string{"./a.hcl"}, srcs)
+		src, err := envs[0].VarFromURL("env://src")
+		require.NoError(t, err)
+		require.Equal(t, "./a.hcl", src)
+	})
 }
 
 func TestUnnamedEnv(t *testing.T) {