Ensure that the socket is closed on connection close. (3.2 ResourceWarning fixes)

James William Pye · James William Pye · commit 1e7b7558bdb1 · 2011-09-10T10:52:29.000-05:00
diff --git a/postgresql/driver/pq3.py b/postgresql/driver/pq3.py
@@ -2293,26 +2293,27 @@ def cursor_from_id(self, cursor_id : str) -> Cursor:
 
 	@property
 	def closed(self) -> bool:
-		if not hasattr(self, 'pq'):
+		if getattr(self, 'pq', None) is None:
 			return True
 		if hasattr(self.pq, 'socket') and self.pq.xact is not None:
 			return self.pq.xact.fatal is True
 		return False
 
-	def close(self):
+	def close(self, getattr = getattr):
 		if self.closed:
 			return
 		# 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.
-		x = self.pq.xact
+		x = getattr(self.pq, 'xact', None)
 		if x:
 			# don't raise?
 			self.pq.complete()
 			if self.closed:
 				return
 		self.pq.push(xact.Closing())
 		self.pq.complete()
+		self.pq.socket.close()
 
 	@property
 	def state(self) -> str:
@@ -2347,13 +2348,22 @@ def connect(self):
 		if self.closed is False:
 			# already connected? just return.
 			return
-		# It's closed.
 
 		if hasattr(self, 'pq'):
 			# It's closed, *but* there's a PQ connection..
 			x = self.pq.xact
 			self.typio.raise_error(x.error_message, cause = getattr(x, 'exception', None), creator = self)
 
+		# It's closed.
+		try:
+			self._establish()
+		except Exception:
+			# Close it up on failure.
+			self.close()
+			raise
+
+	def _establish(self):
+		# guts of connect()
 		self.pq = None
 		# if any exception occurs past this point, the connection
 		# will not be usable.
