From aa8adb561cc30773ece6ffe98de280d3af1093db Mon Sep 17 00:00:00 2001
From: Mehul Goyal <mehul.goyal@astronomer.io>
Date: Thu, 30 Nov 2023 10:41:39 -0800
Subject: [PATCH] override the domain in more places (#1463)

---
 cmd/cloud/setup.go | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/cmd/cloud/setup.go b/cmd/cloud/setup.go
index 31d1c5ba0..47688954c 100644
--- a/cmd/cloud/setup.go
+++ b/cmd/cloud/setup.go
@@ -254,7 +254,10 @@ func checkAPIKeys(astroClient astro.Client, coreClient astrocore.CoreClient, isD
 	c, err := context.GetCurrentContext() // get current context
 	if err != nil {
 		// set context
-		domain := defaultDomain
+		var domain string
+		if domain = os.Getenv("ASTRO_DOMAIN"); domain == "" {
+			domain = defaultDomain
+		}
 		if !context.Exists(domain) {
 			err := context.SetContext(domain)
 			if err != nil {
@@ -369,7 +372,10 @@ func checkAPIToken(isDeploymentFile bool, coreClient astrocore.CoreClient) (bool
 	c, err := context.GetCurrentContext() // get current context
 	if err != nil {
 		// set context
-		domain := defaultDomain
+		var domain string
+		if domain = os.Getenv("ASTRO_DOMAIN"); domain == "" {
+			domain = defaultDomain
+		}
 		if !context.Exists(domain) {
 			err := context.SetContext(domain)
 			if err != nil {