File tree Expand file tree Collapse file tree 2 files changed +41
-6
lines changed Expand file tree Collapse file tree 2 files changed +41
-6
lines changed Original file line number Diff line number Diff line change @@ -161,6 +161,7 @@ class Connection:
161
161
"""
162
162
163
163
_sock = None
164
+ _rfile = None
164
165
_auth_plugin_name = ""
165
166
_closed = False
166
167
_secure = False
@@ -430,6 +431,8 @@ def open(self):
430
431
431
432
def _force_close (self ):
432
433
"""Close connection without QUIT message."""
434
+ if self ._rfile :
435
+ self ._rfile .close ()
433
436
if self ._sock :
434
437
try :
435
438
self ._sock .close ()
@@ -696,12 +699,7 @@ def connect(self, sock=None):
696
699
if self .autocommit_mode is not None :
697
700
self .autocommit (self .autocommit_mode )
698
701
except BaseException as e :
699
- self ._rfile = None
700
- if sock is not None :
701
- try :
702
- sock .close ()
703
- except : # noqa
704
- pass
702
+ self ._force_close ()
705
703
706
704
if isinstance (e , (OSError , IOError )):
707
705
exc = err .OperationalError (
Original file line number Diff line number Diff line change @@ -883,3 +883,40 @@ def test_commit_during_multi_result(self):
883
883
con .commit ()
884
884
cur .execute ("SELECT 3" )
885
885
self .assertEqual (cur .fetchone ()[0 ], 3 )
886
+
887
+ def test_force_close_closes_socketio (self ):
888
+ con = self .connect ()
889
+ sock = con ._sock
890
+ fileno = sock .fileno ()
891
+ rfile = con ._rfile
892
+
893
+ con ._force_close ()
894
+ assert rfile .closed
895
+ assert sock ._closed
896
+ assert sock .fileno () != fileno # should be set to -1
897
+
898
+ def test_socket_closed_on_exception_in_connect (self ):
899
+ con = self .connect (defer_connect = True )
900
+ sock = None
901
+ rfile = None
902
+ fileno = - 1
903
+
904
+ def _request_authentication ():
905
+ nonlocal sock , rfile , fileno
906
+ sock = con ._sock
907
+ assert sock is not None
908
+ fileno = sock .fileno ()
909
+ rfile = con ._rfile
910
+ assert rfile is not None
911
+ raise TypeError
912
+
913
+ con ._request_authentication = _request_authentication
914
+
915
+ with pytest .raises (TypeError ):
916
+ con .connect ()
917
+ assert not con .open
918
+ assert con ._rfile is None
919
+ assert con ._sock is None
920
+ assert rfile .cl
3DA9
osed
921
+ assert sock ._closed
922
+ assert sock .fileno () != fileno # should be set to -1
You can’t perform that action at this time.
0 commit comments