Add the ProducerFault back.

James William Pye · James William Pye · commit e7a1da503c77 · 2010-03-23T12:22:06.000-07:00
The user needs to be able to identify what caused an exception. Producer faults
are usually fatal, but at least the failure will be identifiable.
diff --git a/postgresql/copyman.py b/postgresql/copyman.py
@@ -20,9 +20,25 @@
 default_buffer_size = 1024 * 10
 
 class Fault(Exception):
+	pass
+
+class ProducerFault(Fault):
+	"""
+	Exception raised when the Producer caused an exception.
+
+	Normally, Producer faults are fatal.
+	"""
+	def __init__(self, manager):
+		self.manager = manager
+
+	def __str__(self):
+		return "producer raised exception"
+
+class ReceiverFaults(Fault):
 	"""
-	Receivers raised exceptions. This happens in cases where a receiver raises
-	an exception. Faults should be trapped if recovery from an exception is
+	Exception raised when Receivers cause an exception.
+
+	Faults should be trapped if recovery from an exception is
 	possible, or if the failed receiver is optional to the succes of the
 	operation.
 
@@ -46,13 +62,20 @@ class CopyFail(Exception):
 
 	The 'reason' attribute is a string indicating why it failed.
 
-	The 'faults' attribute is a mapping of receivers to exceptions that were
+	The 'receiver_faults' attribute is a mapping of receivers to exceptions that were
 	raised on exit.
+
+	The 'producer_fault' attribute specifies if the producer raise an exception
+	on exit.
 	"""
-	def __init__(self, manager, reason = None, faults = None):
+	def __init__(self, manager, reason = None,
+		receiver_faults = None,
+		producer_fault = None,
+	):
 		self.manager = manager
 		self.reason = reason
-		self.faults = faults or {}
+		self.receiver_faults = receiver_faults or {}
+		self.producer_fault = producer_fault
 
 	def __str__(self):
 		return self.reason or 'copy aborted'
@@ -683,27 +706,29 @@ def __exit__(self, typ, val, tb):
 			# Don't recover on interrupts.
 			return
 
-		# Does nothing if the COPY was successful.
-		self.producer.realign()
+		profail = None
 		try:
-			##
-			# If the producer is not aligned to a message boundary,
-			# it can emit completion data that will put the receivers
-			# back on track.
-			# This last service call will move that data onto the receivers.
-			self._service_producer()
-			##
-			# The receivers need to handle any new data in their __exit__.
-		except StopIteration:
-			# No re-alignment needed.
-			pass
+			# Does nothing if the COPY was successful.
+			self.producer.realign()
+			try:
+				##
+				# If the producer is not aligned to a message boundary,
+				# it can emit completion data that will put the receivers
+				# back on track.
+				# This last service call will move that data onto the receivers.
+				self._service_producer()
+				##
+				# The receivers need to handle any new data in their __exit__.
+			except StopIteration:
+				# No re-alignment needed.
+				pass
 
-		self.producer.__exit__(typ, val, tb)
+			self.producer.__exit__(typ, val, tb)
+		except Exception as profail:
+			pass
 
 		# No receivers? It wasn't a success.
 		if not self.receivers:
-			if typ is CopyFail:
-				raise
 			raise CopyFail(self, "no receivers")
 
 		exit_faults = {}
@@ -736,13 +761,15 @@ def _service_producer(self):
 		# Setup current data.
 		if not self.receivers:
 			# No receivers to take the data.
-			raise CopyFail(self, "no receivers")
+			raise StopIteration
 
 		try:
 			nextdata = next(self.producer)
 		except StopIteration:
 			# Should be over.
 			raise
+		except Exception:
+			raise ProducerFault(self)
 
 		self.transformer(nextdata)
 
@@ -762,7 +789,7 @@ def _service_receivers(self):
 			# The CopyManager is eager to continue the operation.
 			for x in faults:
 				self.receivers.discard(x)
-			raise Fault(self, faults)
+			raise ReceiverFaults(self, faults)
 
 	# Run the COPY to completion.
 	def run(self):
diff --git a/postgresql/test/test_copyman.py b/postgresql/test/test_copyman.py
@@ -328,7 +328,7 @@ def testCopyWithFailure(self):
 					else:
 						# Done with copy.
 						break
-				except copyman.Fault as cf:
+				except copyman.ReceiverFaults as cf:
 					if sr2 not in cf.faults:
 						raise
 		self.failUnless(done)
@@ -388,17 +388,17 @@ def testNoReceivers(self):
 		done = False
 		try:
 			with copyman.CopyManager(sp, sr1) as copy:
-				while True:
+				while not done:
 					try:
 						for x in copy:
 							if not done:
 								done = True
 								dst.pq.socket.close()
 							else:
 								self.fail("failed to detect dead socket")
-					except copyman.Fault as cf:
+					except copyman.ReceiverFaults as cf:
 						self.failUnless(sr1 in cf.faults)
-						# Don't reconcile.
+						# Don't reconcile. Let the manager drop the receiver.
 		except copyman.CopyFail:
 			self.failUnless(not bool(copy.receivers))
 			# Success.
@@ -436,7 +436,7 @@ def failed_write(*args):
 					else:
 						# Done with COPY, break out of while copy.receivers.
 						break
-				except copyman.Fault as cf:
+				except copyman.ReceiverFaults as cf:
 					if isinstance(cf.faults[sr], RecoverableError):
 						if done is True:
 							self.fail("failed_write was called twice?")
@@ -485,7 +485,7 @@ def failed_write(*args):
 					else:
 						# Done with COPY, break out of while copy.receivers.
 						break
-				except copyman.Fault as cf:
+				except copyman.ReceiverFaults as cf:
 					self.failUnless(isinstance(cf.faults[sr2], TheCause))
 					if done is True:
 						self.fail("failed_write was called twice?")
@@ -509,7 +509,32 @@ def failed_write(*args):
 		self.failUnlessEqual(sp.count(), stdrowcount)
 		self.failUnlessEqual(sp.command(), "COPY")
 
+	@pg_tmp
+	def testProducerFailure(self):
+		sqlexec(stdsource)
+		dst = new()
+		dst.execute(stddst)
+		sp = copyman.StatementProducer(prepare(srcsql))
+		sr = copyman.StatementReceiver(dst.prepare(dstsql))
+		done = False
+		try:
+			with copyman.CopyManager(sp, sr) as copy:
+				try:
+					for x in copy:
+						if not done:
+							done = True
+							db.pq.socket.close()
+				except copyman.ProducerFault as pf:
+					self.failUnless(pf.__context__ is not None)
+			self.fail('expected CopyManager to raise CopyFail')
+		except copyman.CopyFail as cf:
+			pass
+		self.failUnless(done)
+		self.failUnlessRaises(Exception, sqlexec, 'select 1')
+		self.failUnlessEqual(dst.prepare(dstcount).first(), 0)
+
 from ..copyman import WireState
+
 class test_WireState(unittest.TestCase):
 	def testNormal(self):
 		WS=WireState()
