From 18f05c399a137cfe7667461c4a81ef715d876ecb Mon Sep 17 00:00:00 2001
From: Fabian Jakobs <fabian.jakobs@databricks.com>
Date: Thu, 9 Mar 2023 13:19:38 +0100
Subject: [PATCH] Properly detect azure China and Azure GovCloud

---
 config/config.go | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/config/config.go b/config/config.go
index 6dae1d355..ce65345f3 100644
--- a/config/config.go
+++ b/config/config.go
@@ -130,7 +130,11 @@ func (c *Config) Authenticate(r *http.Request) error {
 
 // IsAzure returns true if client is configured for Azure Databricks
 func (c *Config) IsAzure() bool {
-	return strings.Contains(c.Host, ".azuredatabricks.net") || c.AzureResourceID != ""
+	isAzureHost := strings.Contains(c.Host, ".azuredatabricks.net") ||
+		strings.Contains(c.Host, "databricks.azure.cn") ||
+		strings.Contains(c.Host, ".databricks.azure.us")
+
+	return isAzureHost || c.AzureResourceID != ""
 }
 
 // IsGcp returns true if client is configured for GCP