Only cache TLS sessions after successful connection

Previously, TLS sessions were stored immediately after wrap_socket()
completed, but before the actual TCP connection was established. This
meant we could cache sessions for connections that subsequently failed.

Now sessions are only stored after the connection is fully established
and validated, ensuring we only cache sessions for successful connections.

The session storage logic has been moved from _wrap_socket_from_context()
to _connect_socket(), after _initiate_connection() and _validate_hostname()
have succeeded.

Co-authored-by: mykaul <[email protected]>

Author: Copilot

cassandra/connection.py

@@ -1046,15 +1046,8 @@ def _wrap_socket_from_context(self):
 
         ssl_socket = self.ssl_context.wrap_socket(self._socket, **opts)
 
-        # Store the session for future reuse
-        if self.tls_session_cache and ssl_socket.session:
-            self.tls_session_cache.set_session(
-                self.endpoint.address, self.endpoint.port, ssl_socket.session)
-            # Track if the session was reused
-            self.session_reused = ssl_socket.session_reused
-            if self.session_reused:
-                log.debug("TLS session was reused for %s:%s", 
-                         self.endpoint.address, self.endpoint.port)
+        # Note: Session is NOT stored here - it will be stored after successful connection
+        # in _connect_socket() to ensure we only cache sessions for successful connections
 
         return ssl_socket
 
@@ -1111,6 +1104,19 @@ def _connect_socket(self):
                 # run that here.
                 if self._check_hostname:
                     self._validate_hostname()
+                
+                # Store the TLS session after successful connection
+                # This ensures we only cache sessions for connections that actually succeeded
+                if self.tls_session_cache and self.ssl_context and hasattr(self._socket, 'session'):
+                    if self._socket.session:
+                        self.tls_session_cache.set_session(
+                            self.endpoint.address, self.endpoint.port, self._socket.session)
+                        # Track if the session was reused
+                        self.session_reused = self._socket.session_reused
+                        if self.session_reused:
+                            log.debug("TLS session was reused for %s:%s", 
+                                     self.endpoint.address, self.endpoint.port)
+                
                 sockerr = None
                 break
             except socket.error as err: