SNOW-847497: recreate session connection if reset by server (#1639)

DESCRIPTION.md

@@ -38,6 +38,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
   - Fixed a bug where `write_pandas` fails when user does not have the privilege to create stage or file format in the target schema, but has the right privilege for the current schema.
   - Remove Python 3.7 support.
   - Worked around a segfault which sometimes occurred during cache serialization in multi-threaded scenarios.
+  - Improved error handling of connection reset error.
 
 - v3.0.4(May 23,2023)
   - Fixed a bug in which `cursor.execute()` could modify the argument statement_params dictionary object when executing a multistatement query.

src/snowflake/connector/network.py

@@ -958,6 +958,21 @@ def _request_exec_wrapper(
             retry_ctx.increment_retry(reason)
             if retry_ctx.timeout is not None:
                 retry_ctx.timeout -= sleeping_time
+            if "Connection aborted" in repr(e) and "ECONNRESET" in repr(e):
+                # connection is reset by the server, the underlying connection is broken and can not be reused
+                # we need a new urllib3 http(s) connection in this case.
+                # We need to first close the old one so that urllib3 pool manager can create a new connection
+                # for new requests
+                try:
+                    logger.debug(
+                        "shutting down requests session adapter due to connection aborted"
+                    )
+                    session.get_adapter(full_url).close()
+                except Exception as close_adapter_exc:
+                    logger.debug(
+                        "Ignored error caused by closing https connection failure: %s",
+                        close_adapter_exc,
+                    )
             return None  # retry
         except Exception as e:
             if not no_retry:

test/unit/test_retry_network.py

@@ -43,6 +43,7 @@
 
 # We need these for our OldDriver tests. We run most up to date tests with the oldest supported driver version
 try:
+    import snowflake.connector.vendored.urllib3.contrib.pyopenssl
     from snowflake.connector.vendored import requests, urllib3
 except ImportError:  # pragma: no cover
     import requests
@@ -369,3 +370,45 @@ def fake_request_exec(**kwargs):
     assert '"TOKEN": "****' in caplog.text
     assert '"PASSWORD": "****' in caplog.text
     assert ret == {}
+
+
+def test_retry_connection_reset_error(caplog):
+    connection = MagicMock()
+    connection.errorhandler = Mock(return_value=None)
+
+    rest = SnowflakeRestful(
+        host="testaccount.snowflakecomputing.com", port=443, connection=connection
+    )
+
+    data = (
+        '{"code": 12345,'
+        ' "data": {"TOKEN": "_Y1ZNETTn5/qfUWj3Jedb", "PASSWORD": "dummy_pass"}'
+        "}"
+    )
+    default_parameters = {
+        "method": "POST",
+        "full_url": "https://testaccount.snowflakecomputing.com/",
+        "headers": {},
+        "data": data,
+    }
+
+    def error_recv_into(*args, **kwargs):
+        raise OSError(104, "ECONNRESET")
+
+    with patch.object(
+        snowflake.connector.vendored.urllib3.contrib.pyopenssl.WrappedSocket,
+        "recv_into",
+        new=error_recv_into,
+    ):
+        with caplog.at_level(logging.DEBUG):
+            rest.fetch(timeout=10, **default_parameters)
+
+        assert (
+            "shutting down requests session adapter due to connection aborted"
+            in caplog.text
+        )
+        assert (
+            "Ignored error caused by closing https connection failure"
+            not in caplog.text
+        )
+        assert caplog.text.count("Starting new HTTPS connection") > 1