Skip to content

Commit 4803185

Browse files
blink1073zaif-yuval
authored andcommitted
PYTHON-3096 Finish implementation and tests for GSSAPI options (mongodb#1985)
1 parent 8d27699 commit 4803185

File tree

10 files changed

+69
-3
lines changed

10 files changed

+69
-3
lines changed

doc/contributors.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,3 +103,4 @@ The following is a list of people who have contributed to
103103
- Terry Patterson
104104
- Romain Morotti
105105
- Navjot Singh (navjots18)
106+
- Yuval Zaif (zaif-yuval)

pymongo/asynchronous/pool.py

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,12 @@ def _set_non_inheritable_non_atomic(fd: int) -> None: # noqa: ARG001
121121
"""Dummy function for platforms that don't provide fcntl."""
122122

123123

124+
try:
125+
from python_socks import ProxyType
126+
from python_socks.sync import Proxy
127+
except ImportError:
128+
Proxy = ProxyType = None
129+
124130
_IS_SYNC = False
125131

126132
_MAX_TCP_KEEPIDLE = 120
@@ -838,7 +844,21 @@ def _create_connection(address: _Address, options: PoolOptions) -> socket.socket
838844
sock.settimeout(timeout)
839845
sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, True)
840846
_set_keepalive_times(sock)
841-
sock.connect(sa)
847+
if proxy := options.proxy:
848+
if Proxy is None:
849+
raise RuntimeError(
850+
"In order to use SOCKS5 proxy, python_socks must be installed. "
851+
"This can be done by re-installing pymongo with `pip install pymongo[socks]`"
852+
)
853+
proxy_host = proxy["host"]
854+
proxy_port = proxy["port"] or 1080
855+
sock.connect((proxy_host, proxy_port))
856+
proxy = Proxy(
857+
ProxyType.SOCKS5, proxy_host, proxy_port, proxy["username"], proxy["password"]
858+
)
859+
proxy.connect(sa[0], dest_port=sa[1], _socket=sock)
860+
else:
861+
sock.connect(sa)
842862
return sock
843863
except OSError as e:
844864
err = e

pymongo/asynchronous/topology.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -967,6 +967,7 @@ def _create_pool_for_monitor(self, address: _Address) -> Pool:
967967
driver=options.driver,
968968
pause_enabled=False,
969969
server_api=options.server_api,
970+
proxy=options.proxy,
970971
)
971972

972973
return self._settings.pool_class(

pymongo/client_options.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -170,6 +170,15 @@ def _parse_pool_options(
170170
ssl_context, tls_allow_invalid_hostnames = _parse_ssl_options(options)
171171
load_balanced = options.get("loadbalanced")
172172
max_connecting = options.get("maxconnecting", common.MAX_CONNECTING)
173+
if proxy_host := options.get("proxyHost"):
174+
proxy = {
175+
"host": proxy_host,
176+
"port": options.get("proxyPort"),
177+
"username": options.get("proxyUserName"),
178+
"password": options.get("proxyPassword"),
179+
}
180+
else:
181+
proxy = None
173182
return PoolOptions(
174183
max_pool_size,
175184
min_pool_size,
@@ -188,6 +197,7 @@ def _parse_pool_options(
188197
load_balanced=load_balanced,
189198
credentials=credentials,
190199
is_sync=is_sync,
200+
proxy=proxy,
191201
)
192202

193203

pymongo/common.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -729,6 +729,10 @@ def validate_server_monitoring_mode(option: str, value: str) -> str:
729729
"srvmaxhosts": validate_non_negative_integer,
730730
"timeoutms": validate_timeoutms,
731731
"servermonitoringmode": validate_server_monitoring_mode,
732+
"proxyhost": validate_string,
733+
"proxyport": validate_positive_integer_or_none,
734+
"proxyusername": validate_string_or_none,
735+
"proxypassword": validate_string_or_none,
732736
}
733737

734738
# Dictionary where keys are the names of URI options specific to pymongo,

pymongo/pool_options.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -312,6 +312,7 @@ class PoolOptions:
312312
"__server_api",
313313
"__load_balanced",
314314
"__credentials",
315+
"__proxy",
315316
)
316317

317318
def __init__(
@@ -334,6 +335,7 @@ def __init__(
334335
load_balanced: Optional[bool] = None,
335336
credentials: Optional[MongoCredential] = None,
336337
is_sync: Optional[bool] = True,
338+
proxy: Optional[dict] = None,
337339
):
338340
self.__max_pool_size = max_pool_size
339341
self.__min_pool_size = min_pool_size
@@ -353,7 +355,7 @@ def __init__(
353355
self.__load_balanced = load_balanced
354356
self.__credentials = credentials
355357
self.__metadata = copy.deepcopy(_METADATA)
356-
358+
self.__proxy = copy.deepcopy(proxy)
357359
if appname:
358360
self.__metadata["application"] = {"name": appname}
359361

@@ -522,3 +524,8 @@ def server_api(self) -> Optional[ServerApi]:
522524
def load_balanced(self) -> Optional[bool]:
523525
"""True if this Pool is configured in load balanced mode."""
524526
return self.__load_balanced
527+
528+
@property
529+
def proxy(self) -> Optional[dict]:
530+
"""Proxy settings, if configured"""
531+
return self.__proxy

pymongo/synchronous/pool.py

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,12 @@ def _set_non_inheritable_non_atomic(fd: int) -> None: # noqa: ARG001
121121
"""Dummy function for platforms that don't provide fcntl."""
122122

123123

124+
try:
125+
from python_socks import ProxyType
126+
from python_socks.sync import Proxy
127+
except ImportError:
128+
Proxy = ProxyType = None
129+
124130
_IS_SYNC = True
125131

126132
_MAX_TCP_KEEPIDLE = 120
@@ -836,7 +842,21 @@ def _create_connection(address: _Address, options: PoolOptions) -> socket.socket
836842
sock.settimeout(timeout)
837843
sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, True)
838844
_set_keepalive_times(sock)
839-
sock.connect(sa)
845+
if proxy := options.proxy:
846+
if Proxy is None:
847+
raise RuntimeError(
848+
"In order to use SOCKS5 proxy, python_socks must be installed. "
849+
"This can be done by re-installing pymongo with `pip install pymongo[socks]`"
850+
)
851+
proxy_host = proxy["host"]
852+
proxy_port = proxy["port"] or 1080
853+
sock.connect((proxy_host, proxy_port))
854+
proxy = Proxy(
855+
ProxyType.SOCKS5, proxy_host, proxy_port, proxy["username"], proxy["password"]
856+
)
857+
proxy.connect(sa[0], dest_port=sa[1], _socket=sock)
858+
else:
859+
sock.connect(sa)
840860
return sock
841861
except OSError as e:
842862
err = e

pymongo/synchronous/topology.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -965,6 +965,7 @@ def _create_pool_for_monitor(self, address: _Address) -> Pool:
965965
driver=options.driver,
966966
pause_enabled=False,
967967
server_api=options.server_api,
968+
proxy=options.proxy,
968969
)
969970

970971
return self._settings.pool_class(

pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@ ocsp = ["requirements/ocsp.txt"]
6767
snappy = ["requirements/snappy.txt"]
6868
test = ["requirements/test.txt"]
6969
zstd = ["requirements/zstd.txt"]
70+
socks = ["requirements/socks.txt"]
7071

7172
[tool.pytest.ini_options]
7273
minversion = "7"

requirements/socks.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
python-socks[asyncio]

0 commit comments

Comments
 (0)