Fix leaking proxy env vars

Author: ozturkberkay

src/snowflake/connector/connection.py

@@ -32,7 +32,7 @@
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey
 
-from . import errors, proxy
+from . import errors
 from ._query_context_cache import QueryContextCache
 from .auth import (
     FIRST_PARTY_AUTHENTICATORS,
@@ -927,16 +927,16 @@ def __open_connection(self):
             use_numpy=self._numpy, support_negative_year=self._support_negative_year
         )
 
-        proxy.set_proxies(
-            self.proxy_host, self.proxy_port, self.proxy_user, self.proxy_password
-        )
-
         self._rest = SnowflakeRestful(
             host=self.host,
             port=self.port,
             protocol=self._protocol,
             inject_client_pause=self._inject_client_pause,
             connection=self,
+            proxy_host=self.proxy_host,
+            proxy_port=self.proxy_port,
+            proxy_user=self.proxy_user,
+            proxy_password=self.proxy_password,
         )
         logger.debug("REST API object was created: %s:%s", self.host, self.port)
 

src/snowflake/connector/network.py

@@ -87,6 +87,7 @@
     ServiceUnavailableError,
     TooManyRequests,
 )
+from .proxy import get_proxy_url
 from .sqlstate import (
     SQLSTATE_CONNECTION_NOT_EXISTS,
     SQLSTATE_CONNECTION_REJECTED,
@@ -363,9 +364,14 @@ def __init__(
         protocol: str = "http",
         inject_client_pause: int = 0,
         connection: SnowflakeConnection | None = None,
+        proxy_host: str | None = None,
+        proxy_port: str | None = None,
+        proxy_user: str | None = None,
+        proxy_password: str | None = None,
     ) -> None:
         self._host = host
         self._port = port
+        self._proxy = get_proxy_url(proxy_host, proxy_port, proxy_user, proxy_password)
         self._protocol = protocol
         self._inject_client_pause = inject_client_pause
         self._connection = connection
@@ -1160,6 +1166,7 @@ def make_requests_session(self) -> Session:
         s.mount("http://", ProxySupportAdapter(max_retries=REQUESTS_RETRY))
         s.mount("https://", ProxySupportAdapter(max_retries=REQUESTS_RETRY))
         s._reuse_count = itertools.count()
+        s.proxies = {"http": self._proxy, "https": self._proxy}
         return s
 
     @contextlib.contextmanager

src/snowflake/connector/proxy.py

@@ -5,43 +5,26 @@
 
 from __future__ import annotations
 
-import os
 
-
-def set_proxies(
+def get_proxy_url(
     proxy_host: str | None,
     proxy_port: str | None,
     proxy_user: str | None = None,
     proxy_password: str | None = None,
-) -> dict[str, str] | None:
-    """Sets proxy dict for requests."""
-    PREFIX_HTTP = "http://"
-    PREFIX_HTTPS = "https://"
-    proxies = None
+) -> str | None:
+    http_prefix = "http://"
+    https_prefix = "https://"
+
     if proxy_host and proxy_port:
-        if proxy_host.startswith(PREFIX_HTTP):
-            proxy_host = proxy_host[len(PREFIX_HTTP) :]
-        elif proxy_host.startswith(PREFIX_HTTPS):
-            proxy_host = proxy_host[len(PREFIX_HTTPS) :]
-        if proxy_user or proxy_password:
-            proxy_auth = "{proxy_user}:{proxy_password}@".format(
-                proxy_user=proxy_user if proxy_user is not None else "",
-                proxy_password=proxy_password if proxy_password is not None else "",
-            )
+        if proxy_host.startswith(http_prefix):
+            host = proxy_host[len(http_prefix) :]
+        elif proxy_host.startswith(https_prefix):
+            host = proxy_host[len(https_prefix) :]
         else:
-            proxy_auth = ""
-        proxies = {
-            "http": "http://{proxy_auth}{proxy_host}:{proxy_port}".format(
-                proxy_host=proxy_host,
-                proxy_port=str(proxy_port),
-                proxy_auth=proxy_auth,
-            ),
-            "https": "http://{proxy_auth}{proxy_host}:{proxy_port}".format(
-                proxy_host=proxy_host,
-                proxy_port=str(proxy_port),
-                proxy_auth=proxy_auth,
-            ),
-        }
-        os.environ["HTTP_PROXY"] = proxies["http"]
-        os.environ["HTTPS_PROXY"] = proxies["https"]
-    return proxies
+            host = proxy_host
+        auth = (
+            f"{proxy_user or ''}:{proxy_password or ''}@"
+            if proxy_user or proxy_password
+            else ""
+        )
+        return f"{http_prefix}{auth}{host}:{proxy_port}"

test/integ/test_connection.py

@@ -524,6 +524,8 @@ def test_invalid_account_timeout():
 
 @pytest.mark.timeout(15)
 def test_invalid_proxy(db_parameters):
+    http_proxy = os.environ.get("HTTP_PROXY")
+    https_proxy = os.environ.get("HTTPS_PROXY")
     with pytest.raises(OperationalError):
         snowflake.connector.connect(
             protocol="http",
@@ -532,13 +534,13 @@ def test_invalid_proxy(db_parameters):
             password=db_parameters["password"],
             host=db_parameters["host"],
             port=db_parameters["port"],
-            login_timeout=5,
+            login_timeout=0,
             proxy_host="localhost",
             proxy_port="3333",
         )
-    # NOTE environment variable is set if the proxy parameter is specified.
-    del os.environ["HTTP_PROXY"]
-    del os.environ["HTTPS_PROXY"]
+    # Proxy environment variables should not change
+    assert os.environ.get("HTTP_PROXY") == http_proxy
+    assert os.environ.get("HTTPS_PROXY") == https_proxy
 
 
 @pytest.mark.timeout(15)

test/unit/test_proxies.py

@@ -14,30 +14,18 @@
 import snowflake.connector
 from snowflake.connector.errors import OperationalError
 
-
-def test_set_proxies():
-    from snowflake.connector.proxy import set_proxies
-
-    assert set_proxies("proxyhost", "8080") == {
-        "http": "http://proxyhost:8080",
-        "https": "http://proxyhost:8080",
-    }
-    assert set_proxies("http://proxyhost", "8080") == {
-        "http": "http://proxyhost:8080",
-        "https": "http://proxyhost:8080",
-    }
-    assert set_proxies("http://proxyhost", "8080", "testuser", "testpass") == {
-        "http": "http://testuser:testpass@proxyhost:8080",
-        "https": "http://testuser:testpass@proxyhost:8080",
-    }
-    assert set_proxies("proxyhost", "8080", "testuser", "testpass") == {
-        "http": "http://testuser:testpass@proxyhost:8080",
-        "https": "http://testuser:testpass@proxyhost:8080",
-    }
-
-    # NOTE environment variable is set if the proxy parameter is specified.
-    del os.environ["HTTP_PROXY"]
-    del os.environ["HTTPS_PROXY"]
+def test_get_proxy_url():
+    from snowflake.connector.proxy import get_proxy_url
+
+    assert get_proxy_url("host", "port", "user", "password") == (
+        "http://user:password@host:port"
+    )
+    assert get_proxy_url("host", "port") == "http://host:port"
+
+    assert get_proxy_url("http://host", "port") == "http://host:port"
+    assert get_proxy_url("https://host", "port", "user", "password") == (
+        "http://user:password@host:port"
+    )
 
 
 @pytest.mark.skipolddriver