Fix copy.

James William Pye · James William Pye · commit ab5ed571a495 · 2009-02-15T14:25:12.000-07:00
load wasn't detecting non-parameterized queries properly and copy wasn't even
working. :P
diff --git a/postgresql/driver/pq3.py b/postgresql/driver/pq3.py
@@ -1075,9 +1075,9 @@ def first(self, *parameters):
 				return None
 			return cm.extract_count() or cm.extract_command()
 
-	def copy(self,
+	def _copy_data_in(self,
 		iterable,
-		tps : "tuples per *set*" = 1000,
+		tps : "tuples per *set*" = None,
 	):
 		"""
 		Given an iterable, execute the COPY ... FROM STDIN statement and
@@ -1086,39 +1086,45 @@ def copy(self,
 		`tps` is the number of tuples to buffer prior to giving the data
 		to the socket's send.
 		"""
+		tps = tps or 500
 		iterable = iter(iterable)
 		x = pq.Transaction((
 			pq.element.Bind(
 				b'',
 				self._pq_statement_id,
-				(), () (),
+				(), (), (),
 			),
 			pq.element.Execute(b'', 1),
-
+			pq.element.SynchronizeMessage,
 		))
 		self.ife_descend(x)
 		self.database._pq_push(x)
+
+		# Get the COPY started.
 		while x.state is not pq.Complete:
+			self.database._pq_step()
+			if hasattr(x, 'CopyFailSequence') and x.messages is x.CopyFailSequence:
+				break
+		else:
+			# Oh, it's not a COPY at all.
+			e = pg_exc.OperationError(
+				"load() used on a non-COPY FROM STDIN query",
+			)
+			x.ife_descend(e)
+			e.raise_exception()
+
+		while x.messages:
 			# Process any messages setup for sending.
 			while x.messages is not x.CopyFailSequence:
 				self.database._pq_step()
-				if x.state is pq.Complete:
-					e = pg_exc.OperationError(
-						"load() used on a non-COPY FROM STDIN query"
-					)
-					x.ife_descend(e)
-					e.raise_exception()
-
-				copyseq = [
-					pq.element.CopyData(l) for l in islice(iterable, tps)
-				]
-			x.messages = copyseq
-		x.messages = self._pq_xact.CopyDoneSequence
+			x.messages = [
+				pq.element.CopyData(l) for l in islice(iterable, tps)
+			]
+		x.messages = x.CopyDoneSequence
 		self.database._pq_complete()
 
-	def load(self, tupleseq, tps = 40):
-		if self.closed:
-			self.prepare()
+	def _load_bulk_tuples(self, tupleseq, tps = None):
+		tps = tps or 64
 		last = pq.element.FlushMessage
 		tupleseqiter = iter(tupleseq)
 		try:
@@ -1164,6 +1170,19 @@ def load(self, tupleseq, tps = 40):
 			self.database.synchronize()
 			raise
 
+	def load(self, iterable, tps = None):
+		"""
+		Execute the query for each parameter set in `iterable`.
+
+		In cases of ``COPY ... FROM STDIN``, iterable must be an iterable `bytes`.
+		"""
+		if self.closed:
+			self.prepare()
+		if not self._input:
+			return self._copy_data_in(iterable, tps = tps)
+		else:
+			return self._load_bulk_tuples(iterable, tps = tps)
+
 class StoredProcedure(pg_api.StoredProcedure):
 	ife_ancestor = None
 	procedure_id = None
diff --git a/postgresql/test/test_driver.py b/postgresql/test/test_driver.py
@@ -95,6 +95,29 @@ def raise_exc(l):
 				e, v, tb = rl[0]
 				raise v
 		self.failUnlessRaises(pg_exc.QueryCanceledError, raise_exc, rl)
+	
+	def testCopyToSTDOUT(self):
+		with self.db.xact:
+			self.db.execute("CREATE TABLE foo (i int)")
+			foo = self.db.prepare('insert into foo values ($1)')
+			foo.load(((x,) for x in range(500)))
+
+			copy_foo = self.db.prepare('copy foo to stdout')
+			foo_content = set(copy_foo)
+			expected = set((str(i).encode('ascii') + b'\n' for i in range(500)))
+			self.failUnlessEqual(expected, foo_content)
+			self.db.execute("DROP TABLE foo")
+
+	def testCopyFromSTDIN(self):
+		with self.db.xact:
+			self.db.execute("CREATE TABLE foo (i int)")
+			foo = self.db.prepare('copy foo from stdin')
+			foo.load((str(i).encode('ascii') + b'\n' for i in range(1000)))
+			foo_content = list((
+				x for (x,) in self.db.prepare('select * from foo order by 1 ASC')
+			))
+			self.failUnlessEqual(foo_content, list(range(1000)))
+			self.db.execute("DROP TABLE foo")
 
 	def testLookupProcByName(self):
 		self.db.execute(
