1313# See the License for the specific language governing permissions and
1414# limitations under the License.
1515
16+ import logging
1617from collections import OrderedDict
1718import socket
1819from typing import Union
2829
2930CLIENT_STATUS_AUTH_FAILURE = 2000
3031
32+ logger = logging .getLogger ('.' .join (__name__ .split ('.' )[:- 1 ]))
33+
3134
3235class BaseConnection :
3336 def __init__ (self , client , host : str = None , port : int = None , username : str = None , password : str = None ,
@@ -78,21 +81,53 @@ def protocol_context(self):
7881 return self .client .protocol_context
7982
8083 def _process_handshake_error (self , response ):
81- error_text = f'Handshake error: { response .message } '
8284 # if handshake fails for any reason other than protocol mismatch
8385 # (i.e. authentication error), server version is 0.0.0
86+ if response .client_status == CLIENT_STATUS_AUTH_FAILURE :
87+ raise AuthenticationError (response .message )
88+
8489 protocol_version = self .client .protocol_context .version
8590 server_version = (response .version_major , response .version_minor , response .version_patch )
86-
91+ error_text = f'Handshake error: { response . message } '
8792 if any (server_version ):
8893 error_text += f' Server expects binary protocol version ' \
8994 f'{ server_version [0 ]} .{ server_version [1 ]} .{ server_version [2 ]} . ' \
9095 f'Client provides ' \
9196 f'{ protocol_version [0 ]} .{ protocol_version [1 ]} .{ protocol_version [2 ]} .'
92- elif response .client_status == CLIENT_STATUS_AUTH_FAILURE :
93- raise AuthenticationError (error_text )
9497 raise HandshakeError (server_version , error_text )
9598
99+ def _on_handshake_start (self ):
100+ if logger .isEnabledFor (logging .DEBUG ):
101+ logger .debug ("Connecting to node(address=%s, port=%d) with protocol context %s" ,
102+ self .host , self .port , self .client .protocol_context )
103+
104+ def _on_handshake_success (self , result ):
105+ features = BitmaskFeature .from_array (result .get ('features' , None ))
106+ self .client .protocol_context .features = features
107+ self .uuid = result .get ('node_uuid' , None ) # version-specific (1.4+)
108+ self .failed = False
109+
110+ if logger .isEnabledFor (logging .DEBUG ):
111+ logger .debug ("Connected to node(address=%s, port=%d, node_uuid=%s) with protocol context %s" ,
112+ self .host , self .port , self .uuid , self .client .protocol_context )
113+
114+ def _on_handshake_fail (self , err ):
115+ if isinstance (err , AuthenticationError ):
116+ logger .error ("Authentication failed while connecting to node(address=%s, port=%d): %s" ,
117+ self .host , self .port , err )
118+ else :
119+ logger .error ("Failed to perform handshake, connection to node(address=%s, port=%d) "
120+ "with protocol context %s failed: %s" ,
121+ self .host , self .port , self .client .protocol_context , err , exc_info = True )
122+
123+ def _on_connection_lost (self , err = None , expected = False ):
124+ if expected and logger .isEnabledFor (logging .DEBUG ):
125+ logger .debug ("Connection closed to node(address=%s, port=%d, node_uuid=%s)" ,
126+ self .host , self .port , self .uuid )
127+ else :
128+ logger .info ("Connection lost to node(address=%s, port=%d, node_uuid=%s): %s" ,
129+ self .host , self .port , self .uuid , err )
130+
96131
97132class Connection (BaseConnection ):
98133 """
@@ -168,24 +203,26 @@ def connect(self):
168203 self .client .protocol_context = ProtocolContext (max (PROTOCOLS ), BitmaskFeature .all_supported ())
169204
170205 try :
206+ self ._on_handshake_start ()
171207 result = self ._connect_version ()
172208 except HandshakeError as e :
173209 if e .expected_version in PROTOCOLS :
174210 self .client .protocol_context .version = e .expected_version
175211 result = self ._connect_version ()
176212 else :
213+ self ._on_handshake_fail (e )
177214 raise e
178- except connection_errors :
215+ except AuthenticationError as e :
216+ self ._on_handshake_fail (e )
217+ raise e
218+ except Exception as e :
179219 # restore undefined protocol version
180220 if detecting_protocol :
181221 self .client .protocol_context = None
182- raise
222+ self ._on_handshake_fail (e )
223+ raise e
183224
184- # connection is ready for end user
185- features = BitmaskFeature .from_array (result .get ('features' , None ))
186- self .client .protocol_context .features = features
187- self .uuid = result .get ('node_uuid' , None ) # version-specific (1.4+)
188- self .failed = False
225+ self ._on_handshake_success (result )
189226
190227 def _connect_version (self ) -> Union [dict , OrderedDict ]:
191228 """
@@ -258,11 +295,12 @@ def send(self, data: Union[bytes, bytearray], flags=None, reconnect=True):
258295
259296 try :
260297 self ._socket .sendall (data , ** kwargs )
261- except connection_errors :
298+ except connection_errors as e :
262299 self .failed = True
263300 if reconnect :
301+ self ._on_connection_lost (e )
264302 self .reconnect ()
265- raise
303+ raise e
266304
267305 def recv (self , flags = None , reconnect = True ) -> bytearray :
268306 """
@@ -287,11 +325,12 @@ def recv(self, flags=None, reconnect=True) -> bytearray:
287325 if bytes_received == 0 :
288326 raise SocketError ('Connection broken.' )
289327 bytes_total_received += bytes_received
290- except connection_errors :
328+ except connection_errors as e :
291329 self .failed = True
292330 if reconnect :
331+ self ._on_connection_lost (e )
293332 self .reconnect ()
294- raise
333+ raise e
295334
296335 if bytes_total_received < 4 :
297336 continue
@@ -325,5 +364,5 @@ def close(self):
325364 self ._socket .close ()
326365 except connection_errors :
327366 pass
328-
367+ self . _on_connection_lost ( expected = True )
329368 self ._socket = None
0 commit comments