Fix 60 seconds delay in gov cloud connections

Author: samikshya-db

src/databricks/sql/auth/retry.py

@@ -36,6 +36,7 @@ class CommandType(Enum):
     CLOSE_SESSION = "CloseSession"
     CLOSE_OPERATION = "CloseOperation"
     GET_OPERATION_STATUS = "GetOperationStatus"
+    FEATURE_FLAGS = "FeatureFlags"
     OTHER = "Other"
 
     @classmethod
@@ -408,6 +409,14 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]:
                 "CloseOperation received 404 code from Databricks. Cursor is already closed."
             )
 
+        # Request failed with 404 for feature flags endpoint (not available in Gov Cloud)
+        # Feature flags are optional, so gracefully degrade without retrying
+        if status_code == 404 and self.command_type == CommandType.FEATURE_FLAGS:
+            return (
+                False,
+                "FeatureFlags endpoint returned 404. Feature flags are optional and will be disabled."
+            )
+
         # Request failed, was an ExecuteStatement and the command may have reached the server
         if (
             self.command_type == CommandType.EXECUTE_STATEMENT

src/databricks/sql/common/unified_http_client.py

@@ -251,11 +251,15 @@ def _prepare_headers(
 
         return request_headers
 
-    def _prepare_retry_policy(self):
+    def _prepare_retry_policy(self, url: str = ""):
         """Set up the retry policy for the current request."""
         if isinstance(self._retry_policy, DatabricksRetryPolicy):
-            # Set command type for HTTP requests to OTHER (not database commands)
-            self._retry_policy.command_type = CommandType.OTHER
+            # Detect feature flags endpoint by URL pattern
+            if "/connector-service/feature-flags/" in url:
+                self._retry_policy.command_type = CommandType.FEATURE_FLAGS
+            else:
+                # Set command type for HTTP requests to OTHER (not database commands)
+                self._retry_policy.command_type = CommandType.OTHER
             # Start the retry timer for duration-based retry limits
             self._retry_policy.start_retry_timer()
 
@@ -285,8 +289,8 @@ def request_context(
 
         request_headers = self._prepare_headers(headers)
 
-        # Prepare retry policy for this request
-        self._prepare_retry_policy()
+        # Prepare retry policy for this request (pass url for detection)
+        self._prepare_retry_policy(url)
 
         # Select appropriate pool manager based on target URL
         pool_manager = self._get_pool_manager_for_url(url)

tests/unit/test_retry.py

@@ -83,3 +83,23 @@ def test_excessive_retry_attempts_error(self, t_mock, retry_policy):
                 retry_policy.sleep(HTTPResponse(status=503))
                 # Internally urllib3 calls the increment function generating a new instance for every retry
                 retry_policy = retry_policy.increment()
+
+    def test_feature_flags_404_does_not_retry(self, retry_policy):
+        """Test that FEATURE_FLAGS CommandType with 404 should not retry"""
+        retry_policy._retry_start_time = time.time()
+        retry_policy.command_type = CommandType.FEATURE_FLAGS
+
+        should_retry, msg = retry_policy.should_retry("POST", 404)
+
+        assert should_retry is False
+        assert "FeatureFlags endpoint returned 404" in msg
+        assert "optional" in msg
+
+    def test_feature_flags_503_does_retry(self, retry_policy):
+        """Test that FEATURE_FLAGS CommandType with 503 should retry normally"""
+        retry_policy._retry_start_time = time.time()
+        retry_policy.command_type = CommandType.FEATURE_FLAGS
+
+        should_retry, msg = retry_policy.should_retry("POST", 503)
+
+        assert should_retry is True

tests/unit/test_unified_http_client.py

@@ -12,6 +12,7 @@
 from databricks.sql.exc import RequestError
 from databricks.sql.auth.common import ClientContext
 from databricks.sql.types import SSLOptions
+from databricks.sql.auth.retry import CommandType
 
 
 class TestUnifiedHttpClientMaxRetryError:
@@ -134,3 +135,66 @@ def test_generic_exception_no_crash(self, http_client):
 
             error = exc_info.value
             assert "HTTP request error" in str(error)
+
+
+class TestUnifiedHttpClientFeatureFlagsDetection:
+    """Test feature flags URL detection and CommandType setting."""
+
+    @pytest.fixture
+    def client_context(self):
+        """Create a minimal ClientContext for testing."""
+        context = Mock(spec=ClientContext)
+        context.hostname = "https://test.databricks.com"
+        context.ssl_options = SSLOptions(
+            tls_verify=True,
+            tls_verify_hostname=True,
+            tls_trusted_ca_file=None,
+            tls_client_cert_file=None,
+            tls_client_cert_key_file=None,
+            tls_client_cert_key_password=None,
+        )
+        context.socket_timeout = 30
+        context.retry_stop_after_attempts_count = 3
+        context.retry_delay_min = 1.0
+        context.retry_delay_max = 10.0
+        context.retry_stop_after_attempts_duration = 300.0
+        context.retry_delay_default = 5.0
+        context.retry_dangerous_codes = []
+        context.proxy_auth_method = None
+        context.pool_connections = 10
+        context.pool_maxsize = 20
+        context.user_agent = "test-agent"
+        return context
+
+    @pytest.fixture
+    def http_client(self, client_context):
+        """Create UnifiedHttpClient instance."""
+        return UnifiedHttpClient(client_context)
+
+    def test_feature_flags_url_sets_feature_flags_command_type(self, http_client):
+        """Test that feature flags URL sets CommandType.FEATURE_FLAGS"""
+        feature_flags_url = "https://test.databricks.com/connector-service/feature-flags/v1"
+
+        # Call _prepare_retry_policy with feature flags URL
+        http_client._prepare_retry_policy(feature_flags_url)
+
+        # Verify CommandType is set to FEATURE_FLAGS
+        assert http_client._retry_policy.command_type == CommandType.FEATURE_FLAGS
+
+    def test_non_feature_flags_url_sets_other_command_type(self, http_client):
+        """Test that non-feature-flags URLs set CommandType.OTHER"""
+        normal_url = "https://test.databricks.com/api/2.0/sql/statements"
+
+        # Call _prepare_retry_policy with normal URL
+        http_client._prepare_retry_policy(normal_url)
+
+        # Verify CommandType is set to OTHER
+        assert http_client._retry_policy.command_type == CommandType.OTHER
+
+    def test_empty_url_sets_other_command_type(self, http_client):
+        """Test that empty URL defaults to CommandType.OTHER"""
+        # Call _prepare_retry_policy with empty URL
+        http_client._prepare_retry_policy("")
+
+        # Verify CommandType is set to OTHER
+        assert http_client._retry_policy.command_type == CommandType.OTHER