Check for the producer_fault during CopyFail.

James William Pye · James William Pye · commit 06f3589f5b71 · 2010-03-25T18:47:04.000-07:00
diff --git a/postgresql/copyman.py b/postgresql/copyman.py
@@ -520,13 +520,12 @@ def __exit__(self, typ, val, tb):
 		if typ is None or issubclass(typ, Exception):
 			db = self.statement.database
 			if not db.closed and self._chunks._xact is not None:
+				# The COPY transaction is still happening,
+				# force an interrupt if the connection still exists.
 				db.interrupt()
 				if db.pq.xact:
-					try:
-						db._pq_complete()
-					except Exception:
-						# Let the copy manager indicate the failure.
-						pass
+					# Raise, CopyManager should trap.
+					db._pq_complete()
 		super().__exit__(typ, val, tb)
 
 class NullReceiver(Receiver):
@@ -724,9 +723,9 @@ def __exit__(self, typ, val, tb):
 				pass
 
 			self.producer.__exit__(typ, val, tb)
-		except Exception as profail:
+		except Exception as x:
 			# reference profail later.
-			pass
+			profail = x
 
 		# No receivers? It wasn't a success.
 		if not self.receivers:
diff --git a/postgresql/test/test_copyman.py b/postgresql/test/test_copyman.py
@@ -8,6 +8,7 @@
 from ..temporal import pg_tmp
 # The asyncs, and alternative termination.
 from ..protocol.element3 import Notice, Notify, Error, cat_messages
+from .. import exceptions as pg_exc
 
 # state manager can handle empty data messages, right? =)
 emptysource = """
@@ -530,6 +531,7 @@ def testProducerFailure(self):
 		except copyman.CopyFail as cf:
 			# Expecting to see CopyFail
 			self.failUnless(True)
+			self.failUnless(isinstance(cf.producer_fault, pg_exc.ConnectionFailureError))
 		self.failUnless(done)
 		self.failUnlessRaises(Exception, sqlexec, 'select 1')
 		self.failUnlessEqual(dst.prepare(dstcount).first(), 0)
