fix e2e

Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

Author: vikrantpuppala

src/databricks/sql/auth/common.py

@@ -69,12 +69,10 @@ def __init__(
         # HTTP client configuration
         self.ssl_options = ssl_options
         self.socket_timeout = socket_timeout
-        self.retry_stop_after_attempts_count = retry_stop_after_attempts_count or 30
+        self.retry_stop_after_attempts_count = retry_stop_after_attempts_count or 5
         self.retry_delay_min = retry_delay_min or 1.0
         self.retry_delay_max = retry_delay_max or 60.0
-        self.retry_stop_after_attempts_duration = (
-            retry_stop_after_attempts_duration or 900.0
-        )
+        self.retry_stop_after_attempts_duration = retry_stop_after_attempts_duration or 900.0
         self.retry_delay_default = retry_delay_default or 5.0
         self.retry_dangerous_codes = retry_dangerous_codes or []
         self.http_proxy = http_proxy

src/databricks/sql/auth/retry.py

@@ -294,7 +294,7 @@ def sleep_for_retry(self, response: BaseHTTPResponse) -> bool:
         else:
             proposed_wait = self.get_backoff_time()
 
-        proposed_wait = max(proposed_wait, self.delay_max)
+        proposed_wait = min(proposed_wait, self.delay_max)
         self.check_proposed_wait(proposed_wait)
         logger.debug(f"Retrying after {proposed_wait} seconds")
         time.sleep(proposed_wait)
@@ -355,8 +355,8 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]:
         logger.info(f"Received status code {status_code} for {method} request")
 
         # Request succeeded. Don't retry.
-        if status_code == 200:
-            return False, "200 codes are not retried"
+        if status_code // 100 == 2:
+            return False, "2xx codes are not retried"
 
         if status_code == 401:
             return (

src/databricks/sql/backend/thrift_backend.py

@@ -194,7 +194,7 @@ def __init__(
 
         if _max_redirects:
             if _max_redirects > self._retry_stop_after_attempts_count:
-                logger.warn(
+                logger.warning(
                     "_retry_max_redirects > _retry_stop_after_attempts_count so it will have no affect!"
                 )
             urllib3_kwargs = {"redirect": _max_redirects}

src/databricks/sql/client.py

@@ -376,21 +376,17 @@ def _build_client_context(self, server_hostname: str, **kwargs):
             hostname=server_hostname,
             ssl_options=ssl_options,
             socket_timeout=kwargs.get("_socket_timeout"),
-            retry_stop_after_attempts_count=kwargs.get(
-                "_retry_stop_after_attempts_count", 30
-            ),
-            retry_delay_min=kwargs.get("_retry_delay_min", 1.0),
-            retry_delay_max=kwargs.get("_retry_delay_max", 60.0),
-            retry_stop_after_attempts_duration=kwargs.get(
-                "_retry_stop_after_attempts_duration", 900.0
-            ),
-            retry_delay_default=kwargs.get("_retry_delay_default", 1.0),
-            retry_dangerous_codes=kwargs.get("_retry_dangerous_codes", []),
+            retry_stop_after_attempts_count=kwargs.get("_retry_stop_after_attempts_count"),
+            retry_delay_min=kwargs.get("_retry_delay_min"),
+            retry_delay_max=kwargs.get("_retry_delay_max"),
+            retry_stop_after_attempts_duration=kwargs.get("_retry_stop_after_attempts_duration"),
+            retry_delay_default=kwargs.get("_retry_delay_default"),
+            retry_dangerous_codes=kwargs.get("_retry_dangerous_codes"),
             http_proxy=kwargs.get("_http_proxy"),
             proxy_username=kwargs.get("_proxy_username"),
             proxy_password=kwargs.get("_proxy_password"),
-            pool_connections=kwargs.get("_pool_connections", 1),
-            pool_maxsize=kwargs.get("_pool_maxsize", 1),
+            pool_connections=kwargs.get("_pool_connections"),
+            pool_maxsize=kwargs.get("_pool_maxsize"),
             user_agent=user_agent,
         )
 

src/databricks/sql/common/unified_http_client.py

@@ -9,7 +9,7 @@
 from urllib3.util import make_headers
 from urllib3.exceptions import MaxRetryError
 
-from databricks.sql.auth.retry import DatabricksRetryPolicy
+from databricks.sql.auth.retry import DatabricksRetryPolicy, CommandType
 from databricks.sql.exc import RequestError
 
 logger = logging.getLogger(__name__)
@@ -33,6 +33,7 @@ def __init__(self, client_context):
         """
         self.config = client_context
         self._pool_manager = None
+        self._retry_policy = None
         self._setup_pool_manager()
 
     def _setup_pool_manager(self):
@@ -69,20 +70,25 @@ def _setup_pool_manager(self):
                 )
 
         # Create retry policy
-        retry_policy = DatabricksRetryPolicy(
+        self._retry_policy = DatabricksRetryPolicy(
             delay_min=self.config.retry_delay_min,
             delay_max=self.config.retry_delay_max,
             stop_after_attempts_count=self.config.retry_stop_after_attempts_count,
             stop_after_attempts_duration=self.config.retry_stop_after_attempts_duration,
             delay_default=self.config.retry_delay_default,
             force_dangerous_codes=self.config.retry_dangerous_codes,
         )
+        
+        # Initialize the required attributes that DatabricksRetryPolicy expects
+        # but doesn't initialize in its constructor
+        self._retry_policy._command_type = None
+        self._retry_policy._retry_start_time = None
 
         # Common pool manager kwargs
         pool_kwargs = {
             "num_pools": self.config.pool_connections,
             "maxsize": self.config.pool_maxsize,
-            "retries": retry_policy,
+            "retries": self._retry_policy,
             "timeout": urllib3.Timeout(
                 connect=self.config.socket_timeout, read=self.config.socket_timeout
             )
@@ -119,6 +125,14 @@ def _prepare_headers(
 
         return request_headers
 
+    def _prepare_retry_policy(self):
+        """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
+            # Start the retry timer for duration-based retry limits
+            self._retry_policy.start_retry_timer()
+
     @contextmanager
     def request_context(
         self, method: str, url: str, headers: Optional[Dict[str, str]] = None, **kwargs
@@ -138,6 +152,10 @@ def request_context(
         logger.debug("Making %s request to %s", method, url)
 
         request_headers = self._prepare_headers(headers)
+        
+        # Prepare retry policy for this request
+        self._prepare_retry_policy()
+        
         response = None
 
         try:

tests/e2e/common/retry_test_mixins.py

@@ -247,6 +247,7 @@ def test_retry_exponential_backoff(self, mock_send_telemetry, extra_params):
         """
         retry_policy = self._retry_policy.copy()
         retry_policy["_retry_delay_min"] = 1
+        retry_policy["_retry_delay_max"] = 10
 
         time_start = time.time()
         with mocked_server_response(
@@ -282,9 +283,11 @@ def test_retry_max_duration_not_exceeded(self, extra_params):
         WHEN the server sends a Retry-After header of 60 seconds
         THEN the connector raises a MaxRetryDurationError
         """
+        retry_policy = self._retry_policy.copy()
+        retry_policy["_retry_delay_max"] = 60
         with mocked_server_response(status=429, headers={"Retry-After": "60"}):
             with pytest.raises(RequestError) as cm:
-                extra_params = {**extra_params, **self._retry_policy}
+                extra_params = {**extra_params, **retry_policy}
                 with self.connection(extra_params=extra_params) as conn:
                     pass
             assert isinstance(cm.value.args[1], MaxRetryDurationError)

tests/e2e/common/staging_ingestion_tests.py

@@ -68,15 +68,19 @@ def test_staging_ingestion_life_cycle(self, ingestion_user):
         # REMOVE should succeed
 
         remove_query = f"REMOVE 'stage://tmp/{ingestion_user}/tmp/11/16/file1.csv'"
-
-        with self.connection(extra_params={"staging_allowed_local_path": "/"}) as conn:
+        # Use minimal retry settings to fail fast for staging operations
+        extra_params = {
+            "staging_allowed_local_path": "/",
+            "_retry_stop_after_attempts_count": 1,
+        }
+        with self.connection(extra_params=extra_params) as conn:
             cursor = conn.cursor()
             cursor.execute(remove_query)
 
             # GET after REMOVE should fail
 
             with pytest.raises(
-                Error, match="Staging operation over HTTP was unsuccessful: 404"
+                Error, match="too many 404 error responses"
             ):
                 cursor = conn.cursor()
                 query = f"GET 'stage://tmp/{ingestion_user}/tmp/11/16/file1.csv' TO '{new_temp_path}'"

tests/e2e/common/uc_volume_tests.py

@@ -68,14 +68,19 @@ def test_uc_volume_life_cycle(self, catalog, schema):
 
         remove_query = f"REMOVE '/Volumes/{catalog}/{schema}/e2etests/file1.csv'"
 
-        with self.connection(extra_params={"staging_allowed_local_path": "/"}) as conn:
+        # Use minimal retry settings to fail fast
+        extra_params = {
+            "staging_allowed_local_path": "/",
+            "_retry_stop_after_attempts_count": 1,
+        }
+        with self.connection(extra_params=extra_params) as conn:
             cursor = conn.cursor()
             cursor.execute(remove_query)
 
             # GET after REMOVE should fail
 
             with pytest.raises(
-                Error, match="Staging operation over HTTP was unsuccessful: 404"
+                Error, match="too many 404 error responses"
             ):
                 cursor = conn.cursor()
                 query = f"GET '/Volumes/{catalog}/{schema}/e2etests/file1.csv' TO '{new_temp_path}'"

tests/e2e/test_driver.py

@@ -60,12 +60,14 @@
 unsafe_logger.addHandler(logging.FileHandler("./tests-unsafe.log"))
 
 # manually decorate DecimalTestsMixin to need arrow support
-for name in loader.getTestCaseNames(DecimalTestsMixin, "test_"):
-    fn = getattr(DecimalTestsMixin, name)
-    decorated = skipUnless(pysql_supports_arrow(), "Decimal tests need arrow support")(
-        fn
-    )
-    setattr(DecimalTestsMixin, name, decorated)
+test_loader = loader.TestLoader()
+for name in test_loader.getTestCaseNames(DecimalTestsMixin):
+    if name.startswith("test_"):
+        fn = getattr(DecimalTestsMixin, name)
+        decorated = skipUnless(pysql_supports_arrow(), "Decimal tests need arrow support")(
+            fn
+        )
+        setattr(DecimalTestsMixin, name, decorated)
 
 
 class PySQLPytestTestCase: