Add support for SOCKS proxy

Set the SOCKS proxy with the environment variable SOCKS_PROXY and a string with the format {host}:{port}.
prestodb · Jul 20, 2017 · 307a6a4 · 307a6a4
1 parent 80643d2
commit 307a6a4
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/integration_tests/fixtures.py b/integration_tests/fixtures.py
@@ -41,7 +41,7 @@ def get_latest_release():
         raise Exception('not release found')
 
 
-PRESTO_VERSION = os.environ.get('PRESTO_VERSION', get_latest_release())
+PRESTO_VERSION = os.environ.get('PRESTO_VERSION') or get_latest_release()
 PRESTO_HOST = '127.0.0.1'
 PRESTO_PORT = 8080
 

diff --git a/prestodb/client.py b/prestodb/client.py
@@ -35,6 +35,7 @@
 
 from urllib.parse import urlunparse
 import logging
+import os
 from typing import Any, Dict, List, Optional, Text, Tuple, Union  # NOQA for mypy types
 
 import requests
@@ -50,7 +51,14 @@
 
 
 MAX_ATTEMPTS = constants.DEFAULT_MAX_ATTEMPTS
-
+SOCKS_PROXY = os.environ.get('SOCKS_PROXY')
+if SOCKS_PROXY:
+    PROXIES = {
+        'http': 'socks5://' + SOCKS_PROXY,
+        'https': 'socks5://' + SOCKS_PROXY,
+    }
+else:
+    PROXIES = None
 
 class ClientSession(object):
     def __init__(
@@ -294,19 +302,22 @@ def post(self, sql):
             data=sql.encode('utf-8'),
             headers=self.http_headers,
             timeout=self._request_timeout,
+            proxies=PROXIES,
         )
 
     def get(self, url):
         return self._get(
             url,
             headers=self.http_headers,
             timeout=self._request_timeout,
+            proxies=PROXIES,
         )
 
     def delete(self, url):
         return self._delete(
             url,
             timeout=self._request_timeout,
+            proxies=PROXIES,
         )
 
     def _process_error(self, error):