SNOW-469340 add _no_retry option when submitting queries (#879)

Author: asafsarid

src/snowflake/connector/connection.py

@@ -913,6 +913,7 @@ def cmd_query(
         describe_only: bool = False,
         _no_results: bool = False,
         _update_current_object: bool = True,
+        _no_retry: bool = False,
     ):
         """Executes a query with a sequence counter."""
         logger.debug("_cmd_query")
@@ -953,6 +954,7 @@ def cmd_query(
             client=client,
             _no_results=_no_results,
             _include_retry_params=True,
+            _no_retry=_no_retry,
         )
 
         if ret is None:

src/snowflake/connector/cursor.py

@@ -408,6 +408,7 @@ def _execute_helper(
         describe_only: bool = False,
         _no_results: bool = False,
         _is_put_get=None,
+        _no_retry: bool = False,
     ):
         del self.messages[:]
 
@@ -512,6 +513,7 @@ def interrupt_handler(*_):  # pragma: no cover
                 is_internal=is_internal,
                 describe_only=describe_only,
                 _no_results=_no_results,
+                _no_retry=_no_retry,
             )
         finally:
             try:
@@ -561,6 +563,7 @@ def execute(
         _bind_stage: str | None = None,
         timeout: int | None = None,
         _exec_async: bool = False,
+        _no_retry: bool = False,
         _do_reset: bool = True,
         _put_callback: SnowflakeProgressPercentage = None,
         _put_azure_callback: SnowflakeProgressPercentage = None,
@@ -587,6 +590,7 @@ def execute(
             _bind_stage: Path in temporary stage where binding parameters are uploaded as CSV files.
             timeout: Number of seconds after which to abort the query.
             _exec_async: Whether to execute this query asynchronously.
+            _no_retry: Whether or not to retry on known errors.
             _do_reset: Whether or not the result set needs to be reset before executing query.
             _put_callback: Function to which GET command should call back to.
             _put_azure_callback: Function to which an Azure GET command should call back to.
@@ -636,6 +640,7 @@ def execute(
             "describe_only": _describe_only,
             "_no_results": _no_results,
             "_is_put_get": _is_put_get,
+            "_no_retry": _no_retry,
         }
 
         if self._connection.is_pyformat:

src/snowflake/connector/network.py

@@ -431,6 +431,7 @@ def request(
         _no_results=False,
         timeout=None,
         _include_retry_params=False,
+        _no_retry=False,
     ):
         if body is None:
             body = {}
@@ -470,6 +471,7 @@ def request(
                 _no_results=_no_results,
                 timeout=timeout,
                 _include_retry_params=_include_retry_params,
+                no_retry=_no_retry,
             )
         else:
             return self._get_request(url, headers, token=self.token, timeout=timeout)
@@ -871,28 +873,16 @@ def _request_exec_wrapper(
                     exception=str(e),
                     stack_trace=traceback.format_exc(),
                 )
-            if no_retry:
-                return {}
             cause = e.args[0]
+            if no_retry:
+                self.log_and_handle_http_error_with_cause(e, full_url, method, retry_ctx.total_timeout, retry_ctx.cnt,
+                                                          conn, timed_out=False)
+                return {}  # required for tests
             if retry_ctx.timeout is not None:
                 retry_ctx.timeout -= int(time.time() - start_request_thread)
                 if retry_ctx.timeout <= 0:
-                    logger.error(cause, exc_info=True)
-                    TelemetryService.get_instance().log_http_request_error(
-                        "HttpRequestRetryTimeout",
-                        full_url,
-                        method,
-                        SQLSTATE_IO_ERROR,
-                        ER_FAILED_TO_REQUEST,
-                        retry_timeout=retry_ctx.total_timeout,
-                        retry_count=retry_ctx.cnt,
-                        exception=str(e),
-                        stack_trace=traceback.format_exc(),
-                    )
-                    if isinstance(cause, Error):
-                        Error.errorhandler_wrapper_from_cause(conn, cause)
-                    else:
-                        self.handle_invalid_certificate_error(conn, full_url, cause)
+                    self.log_and_handle_http_error_with_cause(e, full_url, method, retry_ctx.total_timeout,
+                                                              retry_ctx.cnt, conn)
                     return {}  # required for tests
             sleeping_time = retry_ctx.next_sleep()
             logger.debug(
@@ -916,6 +906,34 @@ def _request_exec_wrapper(
             logger.debug("Ignored error", exc_info=True)
             return {}
 
+    def log_and_handle_http_error_with_cause(
+            self,
+            e: Exception,
+            full_url: str,
+            method: str,
+            retry_timeout: int,
+            retry_count: int,
+            conn: SnowflakeConnection,
+            timed_out: bool = True
+    ) -> None:
+        cause = e.args[0]
+        logger.error(cause, exc_info=True)
+        TelemetryService.get_instance().log_http_request_error(
+            "HttpRequestRetryTimeout" if timed_out else f"HttpRequestError: {cause}",
+            full_url,
+            method,
+            SQLSTATE_IO_ERROR,
+            ER_FAILED_TO_REQUEST,
+            retry_timeout=retry_timeout,
+            retry_count=retry_count,
+            exception=str(e),
+            stack_trace=traceback.format_exc(),
+        )
+        if isinstance(cause, Error):
+            Error.errorhandler_wrapper_from_cause(conn, cause)
+        else:
+            self.handle_invalid_certificate_error(conn, full_url, cause)
+
     def handle_invalid_certificate_error(self, conn, full_url, cause):
         # all other errors raise exception
         Error.errorhandler_wrapper(

test/unit/test_retry_network.py

@@ -209,3 +209,11 @@ def fake_request_exec(**kwargs):
     cnt.set(NOT_RETRYABLE)
     with pytest.raises(NotRetryableException):
         rest.fetch(timeout=7, **default_parameters)
+
+    # first attempt fails and will not retry
+    cnt.reset()
+    default_parameters["no_retry"] = True
+    ret = rest.fetch(timeout=10, **default_parameters)
+    assert ret == {}
+    assert cnt.c == 1  # failed on first call - did not retry
+    assert rest._connection.errorhandler.called  # error