Fix close() to always close pq.socket if pq.socket is not None.

James William Pye · James William Pye · commit 4ddc8ce27ae6 · 2011-09-16T11:35:17.000-05:00
If there isn't a fatal transaction, be sure push Closing().
diff --git a/postgresql/driver/pq3.py b/postgresql/driver/pq3.py
@@ -2303,17 +2303,26 @@ def close(self, getattr = getattr):
 		# Write out the disconnect message if the socket is around.
 		# If the connection is known to be lost, don't bother. It will
 		# generate an extra exception.
+		if getattr(self, 'pq', None) is None or getattr(self.pq, 'socket', None) is None:
+			# No action to take.
+			return
+
 		x = getattr(self.pq, 'xact', None)
-		if x:
-			# complete the existing transaction, if possible.
-			# XXX: don't raise?
+		if x is not None and x.fatal is not True:
+			# finish the existing pq transaction iff it's not Closing.
 			self.pq.complete()
-			if self.closed:
-				return
-		if getattr(self, 'pq', None):
+
+		if self.pq.xact is None:
+			# It completed the existing transaction.
 			self.pq.push(xact.Closing())
 			self.pq.complete()
+			if self.pq.socket:
+				self.pq.complete()
+
+		# Close the socket if there is one.
+		if self.pq.socket:
 			self.pq.socket.close()
+			self.pq.socket = None
 
 	@property
 	def state(self) -> str:
@@ -2446,6 +2455,8 @@ def _establish(self):
 				# Close out the sockets ourselves.
 				pq.socket.close()
 
+			# Identify whether or not we can skip the attempt.
+			# Whether or not we can skip depends entirely on the SSL parameter.
 			if sslmode == 'prefer' and ssl is False and didssl is False:
 				# In this case, the server doesn't support SSL or it's
 				# turned off. Therefore, the "without_ssl" attempt need
diff --git a/postgresql/test/test_driver.py b/postgresql/test/test_driver.py
@@ -1720,6 +1720,7 @@ def testNamedTuples(self):
 
 	@pg_tmp
 	def testBadFD(self):
+		db.pq.socket.close()
 		# bad fd now.
 		self.assertRaises(
 			pg_exc.ConnectionFailureError,
