keep it simple :)

Author: samikshya-db

src/databricks/sql/auth/retry.py

@@ -36,7 +36,6 @@ class CommandType(Enum):
     CLOSE_SESSION = "CloseSession"
     CLOSE_OPERATION = "CloseOperation"
     GET_OPERATION_STATUS = "GetOperationStatus"
-    FEATURE_FLAGS = "FeatureFlags"
     OTHER = "Other"
 
     @classmethod
@@ -374,6 +373,13 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]:
         if status_code == 403:
             return False, "403 codes are not retried"
 
+        # Request failed with 404. Don't retry for any command type.
+        if status_code == 404:
+            return (
+                False,
+                "Received 404 - NOT_FOUND. The requested resource does not exist.",
+            )
+
         # Request failed and server said NotImplemented. This isn't recoverable. Don't retry.
         if status_code == 501:
             return False, "Received code 501 from server."
@@ -382,41 +388,6 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]:
         if not self._is_method_retryable(method):
             return False, "Only POST requests are retried"
 
-        # Request failed with 404 and was a GetOperationStatus. This is not recoverable. Don't retry.
-        if status_code == 404 and self.command_type == CommandType.GET_OPERATION_STATUS:
-            return (
-                False,
-                "GetOperationStatus received 404 code from Databricks. Operation was canceled.",
-            )
-
-        # Request failed with 404 because CloseSession returns 404 if you repeat the request.
-        if (
-            status_code == 404
-            and self.command_type == CommandType.CLOSE_SESSION
-            and len(self.history) > 0
-        ):
-            raise SessionAlreadyClosedError(
-                "CloseSession received 404 code from Databricks. Session is already closed."
-            )
-
-        # Request failed with 404 because CloseOperation returns 404 if you repeat the request.
-        if (
-            status_code == 404
-            and self.command_type == CommandType.CLOSE_OPERATION
-            and len(self.history) > 0
-        ):
-            raise CursorAlreadyClosedError(
-                "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,15 +251,11 @@ def _prepare_headers(
 
         return request_headers
 
-    def _prepare_retry_policy(self, url: str = ""):
+    def _prepare_retry_policy(self):
         """Set up the retry policy for the current request."""
         if isinstance(self._retry_policy, DatabricksRetryPolicy):
-            # 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
+            # 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()
 
@@ -289,8 +285,8 @@ def request_context(
 
         request_headers = self._prepare_headers(headers)
 
-        # Prepare retry policy for this request (pass url for detection)
-        self._prepare_retry_policy(url)
+        # Prepare retry policy for this request
+        self._prepare_retry_policy()
 
         # Select appropriate pool manager based on target URL
         pool_manager = self._get_pool_manager_for_url(url)

tests/unit/test_retry.py

@@ -84,22 +84,14 @@ def test_excessive_retry_attempts_error(self, t_mock, retry_policy):
                 # 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"""
+    def test_404_does_not_retry_for_any_command_type(self, retry_policy):
+        """Test that 404 never retries for any CommandType"""
         retry_policy._retry_start_time = time.time()
-        retry_policy.command_type = CommandType.FEATURE_FLAGS
 
-        should_retry, msg = retry_policy.should_retry("POST", 404)
+        # Test for each CommandType
+        for command_type in CommandType:
+            retry_policy.command_type = command_type
+            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
+            assert should_retry is False, f"404 should not retry for {command_type}"
+            assert "404" in msg or "NOT_FOUND" in msg

tests/unit/test_unified_http_client.py

@@ -12,7 +12,6 @@
 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:
@@ -135,66 +134,3 @@ 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