SNOW-1524206: introduce a private parameter to enable raising raw http error not handled by connector logic (snowflakedb#2041)

Author: sfc-gh-aling

src/snowflake/connector/network.py

@@ -892,6 +892,14 @@ def _request_exec_wrapper(
         full_url = retry_ctx.add_retry_params(full_url)
         full_url = SnowflakeRestful.add_request_guid(full_url)
         is_fetch_query_status = kwargs.pop("is_fetch_query_status", False)
+        # raise_raw_http_failure is not a public parameter and may change in the future
+        # it enables raising raw http errors that are not handled by
+        # connector, connector handles the following http error:
+        #  1. FORBIDDEN error when trying to login
+        #  2. retryable http code defined in method is_retryable_http_code
+        #  3. UNAUTHORIZED error when using okta authentication
+        # raise_raw_http_failure doesn't work for the 3 mentioned cases.
+        raise_raw_http_failure = kwargs.pop("raise_raw_http_failure", False)
         try:
             return_object = self._request_exec(
                 session=session,
@@ -900,6 +908,7 @@ def _request_exec_wrapper(
                 headers=headers,
                 data=data,
                 token=token,
+                raise_raw_http_failure=raise_raw_http_failure,
                 **kwargs,
             )
             if return_object is not None:
@@ -967,7 +976,9 @@ def _request_exec_wrapper(
                     )
             return None  # retry
         except Exception as e:
-            if not no_retry:
+            if (
+                raise_raw_http_failure and isinstance(e, requests.exceptions.HTTPError)
+            ) or not no_retry:
                 raise e
             logger.debug("Ignored error", exc_info=True)
             return {}
@@ -1036,6 +1047,7 @@ def _request_exec(
         binary_data_handler=None,
         socket_timeout: int | None = None,
         is_okta_authentication: bool = False,
+        raise_raw_http_failure: bool = False,
     ):
         if socket_timeout is None:
             if self._connection.socket_timeout is not None:
@@ -1113,6 +1125,8 @@ def _request_exec(
                     raise_okta_unauthorized_error(self._connection, raw_ret)
                     return None  # required for tests
                 else:
+                    if raise_raw_http_failure:
+                        raw_ret.raise_for_status()
                     raise_failed_request_error(
                         self._connection, full_url, method, raw_ret
                     )

test/unit/test_network.py

@@ -0,0 +1,69 @@
+#!/usr/bin/env python
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+
+import io
+import unittest.mock
+from test.unit.mock_utils import mock_connection
+
+import pytest
+
+try:
+    from snowflake.connector import Error, InterfaceError
+    from snowflake.connector.network import SnowflakeRestful
+    from snowflake.connector.vendored.requests import HTTPError, Response
+except ImportError:
+    # skipping old driver test
+    pass
+
+pytestmark = pytest.mark.skipolddriver
+
+
+def test_fetch():
+    connection = mock_connection()
+    connection.errorhandler = Error.default_errorhandler
+    rest = SnowflakeRestful(
+        host="testaccount.snowflakecomputing.com",
+        port=443,
+        connection=connection,
+    )
+
+    default_parameters = {
+        "method": "POST",
+        "full_url": "https://testaccount.snowflakecomputing.com/",
+        "headers": {},
+        "data": '{"code": 12345}',
+        "token": None,
+    }
+
+    failed_response = Response()
+    failed_response.status_code = 409
+    failed_response.raw = io.StringIO("error")
+    failed_response.url = default_parameters["full_url"]
+    failed_response.reason = "conflict"
+    with unittest.mock.patch(
+        "snowflake.connector.vendored.requests.sessions.Session.request",
+        return_value=failed_response,
+    ):
+        with pytest.raises(HTTPError) as exc:
+            rest.fetch(**default_parameters, raise_raw_http_failure=True)
+        assert exc.value.response.status_code == failed_response.status_code
+        assert exc.value.response.reason == failed_response.reason
+
+        with pytest.raises(HTTPError) as exc:
+            rest.fetch(**default_parameters, raise_raw_http_failure=True, no_retry=True)
+        assert exc.value.response.status_code == failed_response.status_code
+        assert exc.value.response.reason == failed_response.reason
+
+        # if not setting the flag, the function returns an empty dictionary
+        assert (
+            rest.fetch(
+                **default_parameters, raise_raw_http_failure=False, no_retry=True
+            )
+            == {}
+        )
+        assert rest.fetch(**default_parameters, no_retry=True) == {}
+        # if no retry is set to False, the function raises an InterfaceError
+        with pytest.raises(InterfaceError) as exc:
+            assert rest.fetch(**default_parameters, no_retry=False)