SNOW-810810: Fix 'SOCKSProxyManager' object has no attribute 'proxy_headers' error (#1551)

Author: stan-janssen

DESCRIPTION.md

@@ -12,6 +12,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
 
   - Fixed a bug in retry logic for okta authentication to refresh token.
   - Support `RSAPublicKey` when constructing `AuthByKeyPair` in addition to raw bytes.
+  - Fixed a bug when connecting through SOCKS5 proxy, the attribute `proxy_header` is missing on `SOCKSProxyManager`.
 
 - v3.1.0(July 31,2023)
 

src/snowflake/connector/network.py

@@ -111,6 +111,7 @@
 )
 from .vendored.requests.utils import prepend_scheme_if_needed, select_proxy
 from .vendored.urllib3.exceptions import ProtocolError
+from .vendored.urllib3.poolmanager import ProxyManager
 from .vendored.urllib3.util.url import parse_url
 
 if TYPE_CHECKING:
@@ -268,8 +269,14 @@ def get_connection(
                 )
             proxy_manager = self.proxy_manager_for(proxy)
 
-            # Add Host to proxy header SNOW-232777
-            proxy_manager.proxy_headers["Host"] = parsed_url.hostname
+            if isinstance(proxy_manager, ProxyManager):
+                # Add Host to proxy header SNOW-232777
+                proxy_manager.proxy_headers["Host"] = parsed_url.hostname
+            else:
+                logger.debug(
+                    f"Unable to set 'Host' to proxy manager of type {type(proxy_manager)} as"
+                    f" it does not have attribute 'proxy_headers'."
+                )
             conn = proxy_manager.connection_from_url(url)
         else:
             # Only scheme should be lower case

test/unit/test_proxies.py

@@ -5,7 +5,14 @@
 
 from __future__ import annotations
 
+import logging
 import os
+import unittest.mock
+
+import pytest
+
+import snowflake.connector
+from snowflake.connector.errors import OperationalError
 
 
 def test_set_proxies():
@@ -31,3 +38,60 @@ def test_set_proxies():
     # NOTE environment variable is set if the proxy parameter is specified.
     del os.environ["HTTP_PROXY"]
     del os.environ["HTTPS_PROXY"]
+
+
+@pytest.mark.skipolddriver
+def test_socks_5_proxy_missing_proxy_header_attribute(caplog):
+    from snowflake.connector.vendored.urllib3.poolmanager import ProxyManager
+
+    os.environ["HTTPS_PROXY"] = "socks5://localhost:8080"
+
+    class MockSOCKSProxyManager:
+        def __init__(self):
+            pass
+
+        def connection_from_url(self, url):
+            pass
+
+    def mock_proxy_manager_for_url_no_header(*args, **kwargs):
+        return MockSOCKSProxyManager()
+
+    def mock_proxy_manager_for_url_wiht_header(*args, **kwargs):
+        return ProxyManager("testurl")
+
+    # connection
+    caplog.set_level(logging.DEBUG, "snowflake.connector")
+
+    # bad path
+    with unittest.mock.patch(
+        "snowflake.connector.network.ProxySupportAdapter.proxy_manager_for",
+        mock_proxy_manager_for_url_no_header,
+    ):
+        with pytest.raises(OperationalError):
+            snowflake.connector.connect(
+                account="testaccount",
+                user="testuser",
+                password="testpassword",
+                database="TESTDB",
+                warehouse="TESTWH",
+            )
+    assert "Unable to set 'Host' to proxy manager of type" in caplog.text
+
+    caplog.clear()
+
+    # happy path
+    with unittest.mock.patch(
+        "snowflake.connector.network.ProxySupportAdapter.proxy_manager_for",
+        mock_proxy_manager_for_url_wiht_header,
+    ):
+        with pytest.raises(OperationalError):
+            snowflake.connector.connect(
+                account="testaccount",
+                user="testuser",
+                password="testpassword",
+                database="TESTDB",
+                warehouse="TESTWH",
+            )
+    assert "Unable to set 'Host' to proxy manager of type" not in caplog.text
+
+    del os.environ["HTTPS_PROXY"]