Skip to content

Commit

Permalink
PYTHON-3879 Rename SocketInfo to Connection (#1329)
Browse files Browse the repository at this point in the history
  • Loading branch information
NoahStapp authored Jul 28, 2023
1 parent c945ec6 commit c88ae79
Show file tree
Hide file tree
Showing 35 changed files with 731 additions and 735 deletions.
2 changes: 1 addition & 1 deletion doc/common-issues.rst
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,7 @@ will receive the following error::
File "/Library/Python/2.7/site-packages/pymongo/collection.py", line 1560, in count
return self._count(cmd, collation, session)
File "/Library/Python/2.7/site-packages/pymongo/collection.py", line 1504, in _count
with self._socket_for_reads() as (sock_info, slave_ok):
with self._socket_for_reads() as (connection, slave_ok):
File "/System/Library/Frameworks/Python.framework/Versions/2.7/lib/python2.7/contextlib.py", line 17, in __enter__
return self.gen.next()
File "/Library/Python/2.7/site-packages/pymongo/mongo_client.py", line 982, in _socket_for_reads
Expand Down
16 changes: 8 additions & 8 deletions pymongo/aggregation.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@
from pymongo.collection import Collection
from pymongo.command_cursor import CommandCursor
from pymongo.database import Database
from pymongo.pool import SocketInfo
from pymongo.pool import Connection
from pymongo.read_preferences import _ServerMode
from pymongo.server import Server
from pymongo.typings import _Pipeline
Expand All @@ -52,7 +52,7 @@ def __init__(
explicit_session: bool,
let: Optional[Mapping[str, Any]] = None,
user_fields: Optional[MutableMapping[str, Any]] = None,
result_processor: Optional[Callable[[Mapping[str, Any], SocketInfo], None]] = None,
result_processor: Optional[Callable[[Mapping[str, Any], Connection], None]] = None,
comment: Any = None,
) -> None:
if "explain" in options:
Expand Down Expand Up @@ -134,7 +134,7 @@ def get_cursor(
self,
session: ClientSession,
server: Server,
sock_info: SocketInfo,
conn: Connection,
read_preference: _ServerMode,
) -> CommandCursor:
# Serialize command.
Expand All @@ -146,7 +146,7 @@ def get_cursor(
# - server version is >= 4.2 or
# - server version is >= 3.2 and pipeline doesn't use $out
if ("readConcern" not in cmd) and (
not self._performs_write or (sock_info.max_wire_version >= 8)
not self._performs_write or (conn.max_wire_version >= 8)
):
read_concern = self._target.read_concern
else:
Expand All @@ -161,7 +161,7 @@ def get_cursor(
write_concern = None

# Run command.
result = sock_info.command(
result = conn.command(
self._database.name,
cmd,
read_preference,
Expand All @@ -176,7 +176,7 @@ def get_cursor(
)

if self._result_processor:
self._result_processor(result, sock_info)
self._result_processor(result, conn)

# Extract cursor from result or mock/fake one if necessary.
if "cursor" in result:
Expand All @@ -193,14 +193,14 @@ def get_cursor(
cmd_cursor = self._cursor_class(
self._cursor_collection(cursor),
cursor,
sock_info.address,
conn.address,
batch_size=self._batch_size or 0,
max_await_time_ms=self._max_await_time_ms,
session=session,
explicit_session=self._explicit_session,
comment=self._options.get("comment"),
)
cmd_cursor._maybe_pin_connection(sock_info)
cmd_cursor._maybe_pin_connection(conn)
return cmd_cursor


Expand Down
70 changes: 33 additions & 37 deletions pymongo/auth.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@

if TYPE_CHECKING:
from pymongo.hello import Hello
from pymongo.pool import SocketInfo
from pymongo.pool import Connection

HAVE_KERBEROS = True
_USE_PRINCIPAL = False
Expand Down Expand Up @@ -220,9 +220,7 @@ def _authenticate_scram_start(
return nonce, first_bare, cmd


def _authenticate_scram(
credentials: MongoCredential, sock_info: SocketInfo, mechanism: str
) -> None:
def _authenticate_scram(credentials: MongoCredential, conn: Connection, mechanism: str) -> None:
"""Authenticate using SCRAM."""
username = credentials.username
if mechanism == "SCRAM-SHA-256":
Expand All @@ -239,13 +237,13 @@ def _authenticate_scram(
# Make local
_hmac = hmac.HMAC

ctx = sock_info.auth_ctx
ctx = conn.auth_ctx
if ctx and ctx.speculate_succeeded():
nonce, first_bare = ctx.scram_data
res = ctx.speculative_authenticate
else:
nonce, first_bare, cmd = _authenticate_scram_start(credentials, mechanism)
res = sock_info.command(source, cmd)
res = conn.command(source, cmd)

server_first = res["payload"]
parsed = _parse_scram_response(server_first)
Expand Down Expand Up @@ -285,7 +283,7 @@ def _authenticate_scram(
("payload", Binary(client_final)),
]
)
res = sock_info.command(source, cmd)
res = conn.command(source, cmd)

parsed = _parse_scram_response(res["payload"])
if not hmac.compare_digest(parsed[b"v"], server_sig):
Expand All @@ -301,7 +299,7 @@ def _authenticate_scram(
("payload", Binary(b"")),
]
)
res = sock_info.command(source, cmd)
res = conn.command(source, cmd)
if not res["done"]:
raise OperationFailure("SASL conversation failed to complete.")

Expand Down Expand Up @@ -345,7 +343,7 @@ def _canonicalize_hostname(hostname: str) -> str:
return name[0].lower()


def _authenticate_gssapi(credentials: MongoCredential, sock_info: SocketInfo) -> None:
def _authenticate_gssapi(credentials: MongoCredential, conn: Connection) -> None:
"""Authenticate using GSSAPI."""
if not HAVE_KERBEROS:
raise ConfigurationError(
Expand All @@ -358,7 +356,7 @@ def _authenticate_gssapi(credentials: MongoCredential, sock_info: SocketInfo) ->
props = credentials.mechanism_properties
# Starting here and continuing through the while loop below - establish
# the security context. See RFC 4752, Section 3.1, first paragraph.
host = sock_info.address[0]
host = conn.address[0]
if props.canonicalize_host_name:
host = _canonicalize_hostname(host)
service = props.service_name + "@" + host
Expand Down Expand Up @@ -413,7 +411,7 @@ def _authenticate_gssapi(credentials: MongoCredential, sock_info: SocketInfo) ->
("autoAuthorize", 1),
]
)
response = sock_info.command("$external", cmd)
response = conn.command("$external", cmd)

# Limit how many times we loop to catch protocol / library issues
for _ in range(10):
Expand All @@ -430,7 +428,7 @@ def _authenticate_gssapi(credentials: MongoCredential, sock_info: SocketInfo) ->
("payload", payload),
]
)
response = sock_info.command("$external", cmd)
response = conn.command("$external", cmd)

if result == kerberos.AUTH_GSS_COMPLETE:
break
Expand All @@ -453,7 +451,7 @@ def _authenticate_gssapi(credentials: MongoCredential, sock_info: SocketInfo) ->
("payload", payload),
]
)
sock_info.command("$external", cmd)
conn.command("$external", cmd)

finally:
kerberos.authGSSClientClean(ctx)
Expand All @@ -462,7 +460,7 @@ def _authenticate_gssapi(credentials: MongoCredential, sock_info: SocketInfo) ->
raise OperationFailure(str(exc))


def _authenticate_plain(credentials: MongoCredential, sock_info: SocketInfo) -> None:
def _authenticate_plain(credentials: MongoCredential, conn: Connection) -> None:
"""Authenticate using SASL PLAIN (RFC 4616)"""
source = credentials.source
username = credentials.username
Expand All @@ -476,52 +474,50 @@ def _authenticate_plain(credentials: MongoCredential, sock_info: SocketInfo) ->
("autoAuthorize", 1),
]
)
sock_info.command(source, cmd)
conn.command(source, cmd)


def _authenticate_x509(credentials: MongoCredential, sock_info: SocketInfo) -> None:
def _authenticate_x509(credentials: MongoCredential, conn: Connection) -> None:
"""Authenticate using MONGODB-X509."""
ctx = sock_info.auth_ctx
ctx = conn.auth_ctx
if ctx and ctx.speculate_succeeded():
# MONGODB-X509 is done after the speculative auth step.
return

cmd = _X509Context(credentials, sock_info.address).speculate_command()
sock_info.command("$external", cmd)
cmd = _X509Context(credentials, conn.address).speculate_command()
conn.command("$external", cmd)


def _authenticate_mongo_cr(credentials: MongoCredential, sock_info: SocketInfo) -> None:
def _authenticate_mongo_cr(credentials: MongoCredential, conn: Connection) -> None:
"""Authenticate using MONGODB-CR."""
source = credentials.source
username = credentials.username
password = credentials.password
# Get a nonce
response = sock_info.command(source, {"getnonce": 1})
response = conn.command(source, {"getnonce": 1})
nonce = response["nonce"]
key = _auth_key(nonce, username, password)

# Actually authenticate
query = SON([("authenticate", 1), ("user", username), ("nonce", nonce), ("key", key)])
sock_info.command(source, query)
conn.command(source, query)


def _authenticate_default(credentials: MongoCredential, sock_info: SocketInfo) -> None:
if sock_info.max_wire_version >= 7:
if sock_info.negotiated_mechs:
mechs = sock_info.negotiated_mechs
def _authenticate_default(credentials: MongoCredential, conn: Connection) -> None:
if conn.max_wire_version >= 7:
if conn.negotiated_mechs:
mechs = conn.negotiated_mechs
else:
source = credentials.source
cmd = sock_info.hello_cmd()
cmd = conn.hello_cmd()
cmd["saslSupportedMechs"] = source + "." + credentials.username
mechs = sock_info.command(source, cmd, publish_events=False).get(
"saslSupportedMechs", []
)
mechs = conn.command(source, cmd, publish_events=False).get("saslSupportedMechs", [])
if "SCRAM-SHA-256" in mechs:
return _authenticate_scram(credentials, sock_info, "SCRAM-SHA-256")
return _authenticate_scram(credentials, conn, "SCRAM-SHA-256")
else:
return _authenticate_scram(credentials, sock_info, "SCRAM-SHA-1")
return _authenticate_scram(credentials, conn, "SCRAM-SHA-1")
else:
return _authenticate_scram(credentials, sock_info, "SCRAM-SHA-1")
return _authenticate_scram(credentials, conn, "SCRAM-SHA-1")


_AUTH_MAP: Mapping[str, Callable] = {
Expand Down Expand Up @@ -606,12 +602,12 @@ def speculate_command(self) -> Optional[MutableMapping[str, Any]]:


def authenticate(
credentials: MongoCredential, sock_info: SocketInfo, reauthenticate: bool = False
credentials: MongoCredential, conn: Connection, reauthenticate: bool = False
) -> None:
"""Authenticate sock_info."""
"""Authenticate connection."""
mechanism = credentials.mechanism
auth_func = _AUTH_MAP[mechanism]
if mechanism == "MONGODB-OIDC":
_authenticate_oidc(credentials, sock_info, reauthenticate)
_authenticate_oidc(credentials, conn, reauthenticate)
else:
auth_func(credentials, sock_info)
auth_func(credentials, conn)
10 changes: 5 additions & 5 deletions pymongo/auth_aws.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ def set_cached_credentials(creds):
if TYPE_CHECKING:
from bson.typings import _ReadableBuffer
from pymongo.auth import MongoCredential
from pymongo.pool import SocketInfo
from pymongo.pool import Connection


class _AwsSaslContext(AwsSaslContext): # type: ignore
Expand All @@ -67,15 +67,15 @@ def bson_decode(self, data: _ReadableBuffer) -> Mapping[str, Any]:
return bson.decode(data)


def _authenticate_aws(credentials: MongoCredential, sock_info: SocketInfo) -> None:
def _authenticate_aws(credentials: MongoCredential, conn: Connection) -> None:
"""Authenticate using MONGODB-AWS."""
if not _HAVE_MONGODB_AWS:
raise ConfigurationError(
"MONGODB-AWS authentication requires pymongo-auth-aws: "
"install with: python -m pip install 'pymongo[aws]'"
)

if sock_info.max_wire_version < 9:
if conn.max_wire_version < 9:
raise ConfigurationError("MONGODB-AWS authentication requires MongoDB version 4.4 or later")

try:
Expand All @@ -90,7 +90,7 @@ def _authenticate_aws(credentials: MongoCredential, sock_info: SocketInfo) -> No
client_first = SON(
[("saslStart", 1), ("mechanism", "MONGODB-AWS"), ("payload", client_payload)]
)
server_first = sock_info.command("$external", client_first)
server_first = conn.command("$external", client_first)
res = server_first
# Limit how many times we loop to catch protocol / library issues
for _ in range(10):
Expand All @@ -102,7 +102,7 @@ def _authenticate_aws(credentials: MongoCredential, sock_info: SocketInfo) -> No
("payload", client_payload),
]
)
res = sock_info.command("$external", cmd)
res = conn.command("$external", cmd)
if res["done"]:
# SASL complete.
break
Expand Down
Loading

0 comments on commit c88ae79

Please sign in to comment.