Implement load_rows and load_chunks to replace load().

James William Pye · James William Pye · commit 516e54b1f07b · 2009-05-17T10:39:33.000-07:00
Added: itertools.chunk for supporting load_rows() via load_chunks().

Notably, this change cleaned up the implmentation of COPY TO STDIN and
bulk statement execution.

closes #1010602
diff --git a/postgresql/api.py b/postgresql/api.py
@@ -534,7 +534,7 @@ def first(self, *parameters) -> "'First' object that is returned by the query":
 		"""
 
 	@abstractmethod
-	def load(self,
+	def load_rows(self,
 		iterable : "A iterable of tuples to execute the statement with"
 	):
 		"""
@@ -548,8 +548,29 @@ def load(self,
 			...  q(*i)
 
 		Its purpose is to allow the implementation to take advantage of the
-		knowledge that a series of parameters are to be loaded and subsequently
-		optimize the operation.
+		knowledge that a series of parameters are to be loaded so that the
+		operation can be optimized.
+		"""
+
+	@abstractmethod
+	def load_chunks(self,
+		iterable : "A iterable of chunks of tuples to execute the statement with"
+	):
+		"""
+		Given an iterable, `iterable`, feed the produced parameters of the chunks
+		produced by the iterable to the query. This is a bulk-loading interface
+		for parameterized queries.
+
+		Effectively, it is equivalent to:
+
+			>>> ps = db.prepare(...)
+			>>> for c in iterable:
+			...  for i in c:
+			...   q(*i)
+
+		Its purpose is to allow the implementation to take advantage of the
+		knowledge that a series of chunks of parameters are to be loaded so
+		that the operation can be optimized.
 		"""
 
 	@abstractmethod
diff --git a/postgresql/documentation/driver.py b/postgresql/documentation/driver.py
@@ -435,11 +435,11 @@
 Prepared statement objects have a few execution methods:
 
  ``ps(*parameters)``
-  As shown before, statement objects can be simply invoked like a function to get
+  As shown before, statement objects can be invoked like a function to get
   the statement's results.
 
  ``ps.rows(*parameters)``
-  Return a simple iterator to all the rows produced by the statement. This
+  Return a iterator to all the rows produced by the statement. This
   method will stream rows on demand, so it is ideal for situations where
   each individual row in a large result-set must be processed.
 
@@ -496,6 +496,20 @@
  ``ps.close()``
   Close the statement inhibiting further use.
 
+ ``ps.load_rows(collections.Iterable(parameters))``
+  Given an iterable producing parameters, execute the statement for each
+  iteration. Always returns `None`.
+
+ ``ps.load_chunks(collections.Iterable(collections.Iterable(parameters)))``
+  Given an iterable of iterables producing parameters, execute the statement
+  for each parameter produced. However, send the all execution commands with
+  the corresponding parameters of each chunk before reading any results.
+  Always returns `None`. This access point is designed to be used in conjunction
+  with ``ps.chunks()`` for transferring rows from one connection to another with
+  great efficiency::
+
+   >>> dst.prepare(...).load_chunks(src.prepare(...).chunks())
+
 
 Statement Metadata
 ------------------
@@ -649,23 +663,24 @@
 
 Using the call interface is fine for making a single insert, but when multiple
 records need to be inserted, it's not the most efficient means to load data. For
-multiple records, the ``ps.load([...])`` provides an efficient way to load large
-quantities of structured data::
+multiple records, the ``ps.load_rows([...])`` provides an efficient way to load
+large quantities of structured data::
 
 	>>> from datetime import date
-	>>> mkemp.load([
+	>>> mkemp.load_rows([
 	...  ("Jack Johnson", "85000", date(1962, 11, 23), date(1990, 3, 5)),
 	...  ("Debra McGuffer", "52000", date(1973, 3, 4), date(2002, 1, 14)),
 	...  ("Barbara Smith", "86000", date(1965, 2, 24), date(2005, 7, 19)),
 	... ])
 
-While small, the above illustrates the ``ps.load()`` method taking an iterable of
-tuples that provides parameters for the each execution of the statement.
+While small, the above illustrates the ``ps.load_rows()`` method taking an
+iterable of tuples that provides parameters for the each execution of the
+statement.
 
-Load is also used to support ``COPY ... FROM STDIN`` statements::
+``load_rows`` is also used to support ``COPY ... FROM STDIN`` statements::
 
 	>>> copy_emps_in = db.prepare("COPY employee FROM STDIN")
-	>>> copy_emps_in.load([
+	>>> copy_emps_in.load_rows([
 	...  b'Emp Name1\t72000\t1970-2-01\t1980-10-22\n',
 	...  b'Emp Name2\t62000\t1968-9-11\t1985-11-1\n',
 	...  b'Emp Name3\t62000\t1968-9-11\t1985-11-1\n',
@@ -694,8 +709,8 @@
 In situations where other actions are invoked during a ``COPY FROM STDIN``, a
 COPY failure error will occur. The driver manages the connection state in such
 a way that will purposefully cause the error as the COPY was inappropriately
-interrupted. This not usually a problem as the ``load(...)`` method must
-complete the COPY command before returning.
+interrupted. This not usually a problem as ``load_rows(...)`` and
+``load_chunks(...)`` methods must complete the COPY command before returning.
 
 Copy data is always transferred using ``bytes`` objects. Even in cases where the
 COPY is not in ``BINARY`` mode. Any needed encoding transformations *must* be
@@ -722,9 +737,9 @@
 	<lots of data>
 
 ``COPY FROM STDIN`` commands are supported via
-`postgresql.api.PreparedStatement.load`. Each invocation to ``load``
-is a single invocation of COPY. ``load`` takes an iterable of COPY lines
-to send to the server::
+`postgresql.api.PreparedStatement.load_rows`. Each invocation to
+``load_rows`` is a single invocation of COPY. ``load_rows`` takes an iterable of
+COPY lines to send to the server::
 
 	>>> db.execute("""
 	... CREATE TABLE sample_copy (
@@ -733,21 +748,20 @@
 	... );
 	... """)
 	>>> copyin = db.prepare('COPY sample_copy FROM STDIN')
-	>>> copyin.load([
+	>>> copyin.load_rows([
 	... 	b'123\tone twenty three\n',
 	... 	b'350\ttree fitty\n',
 	... ])
 
-The ``load()`` method is trained to identify chunk iterators so that direct
-transfers from a source database to a destination database could be
-made in a streaming fashion::
+For direct connection-to-connection COPY, use of ``load_chunks(...)`` is
+recommended as it will provide the most efficient transfer method::
 
 	>>> copyout = src.prepare('COPY atable TO STDOUT')
 	>>> copyin = dst.prepare('COPY atable FROM STDIN')
-	>>> copyin.load(copyout.chunks())
+	>>> copyin.load_chunks(copyout.chunks())
 
 Specifically, each chunk of row data produced by ``chunks()`` will be written in
-full by ``load()`` before getting another chunk to write.
+full by ``load_chunks()`` before getting another chunk to write.
 
 
 Cursors
diff --git a/postgresql/driver/pq3.py b/postgresql/driver/pq3.py
@@ -34,7 +34,7 @@
 from .. import api as pg_api
 from ..encodings import aliases as pg_enc_aliases
 
-from ..python.itertools import interlace
+from ..python.itertools import interlace, chunk
 from ..python.socket import SocketFactory
 
 from ..protocol import xact3 as xact
@@ -1025,26 +1025,20 @@ def first(self, *parameters):
 				return None
 			return cm.extract_count() or cm.extract_command()
 
-	def _copy_data_in(self,
-		iterable,
-		tps : "tuples per *set*" = None,
-	):
+	def _load_copy_chunks(self, chunks):
 		"""
-		Given an iterable, execute the COPY ... FROM STDIN statement and
-		send the copy lines produced by the iterable to the remote end.
-
-		`tps` is the number of tuples to buffer prior to giving the data
-		to the socket's send.
+		Given an chunks of COPY lines, execute the COPY ... FROM STDIN
+		statement and send the copy lines produced by the iterable to
+		the remote end.
 		"""
-		tps = tps or 500
 		x = xact.Instruction((
 				element.Bind(
 					b'',
 					self._pq_statement_id,
 					(), (), (),
 				),
 				element.Execute(b'', 1),
-				element.SynchronizeMessage,
+				element.FlushMessage,
 			),
 			asynchook = self.database._receive_async
 		)
@@ -1058,70 +1052,45 @@ def _copy_data_in(self,
 		else:
 			# Oh, it's not a COPY at all.
 			e = pg_exc.OperationError(
-				"_copy_data_in() used on a non-COPY FROM STDIN query",
+				"_load_copy_chunks() used on a non-COPY FROM STDIN query",
 				creator = self
 			)
 			raise e
 
-		if isinstance(iterable, Chunks):
-			# optimized, each iteration == row sequence
-			while x.messages:
-				while x.messages is not x.CopyFailSequence:
-					self.database._pq_step()
-				x.messages = []
-				for rows in iterable:
-					x.messages.extend([
-						element.CopyData(l) for l in rows
-					])
-					if len(x.messages) > tps:
-						break
-		else:
-			# each iteration == one row
-			iterable = iter(iterable)
-			while x.messages:
-				# Process any messages setup for sending.
-				while x.messages is not x.CopyFailSequence:
-					self.database._pq_step()
-				x.messages = [
-					element.CopyData(l) for l in islice(iterable, tps)
-				]
+		for chunk in chunks:
+			x.messages = [element.CopyData(l) for l in chunk]
+			while x.messages is not x.CopyFailSequence:
+				self.database._pq_step()
 		x.messages = x.CopyDoneSequence
 		self.database._pq_complete()
+		self.database.pq.synchronize()
 
-	def _load_bulk_tuples(self, tupleseq, tps = None):
-		tps = tps or 64
-		if isinstance(tupleseq, Chunks):
-			tupleseqiter = chain.from_iterable(tupleseq)
-		else:
-			tupleseqiter = iter(tupleseq)
+	def _load_tuple_chunks(self, chunks):
 		pte = self._raise_parameter_tuple_error
-		last = element.FlushMessage
+		last = (element.SynchronizeMessage,)
 		try:
-			while last is element.FlushMessage:
-				c = 0
-				xm = []
-				for t in tupleseqiter:
-					params = pg_typio.process_tuple(
-						self._input_io, tuple(t), pte
-					)
-					xm.extend((
+			for chunk in chunks:
+				bindings = [
+					(
 						element.Bind(
 							b'',
 							self._pq_statement_id,
 							self._input_formats,
-							params,
+							pg_typio.process_tuple(
+								self._input_io, tuple(t), pte
+							),
 							(),
 						),
 						element.Execute(b'', 1),
-					))
-					if c == tps:
-						break
-					c += 1
-				else:
-					last = element.SynchronizeMessage
-				xm.append(last)
+					)
+					for t in chunk
+				]
+				bindings.append(last)
 				self.database._pq_push(
-					xact.Instruction(xm, asynchook = self.database._receive_async),
+					xact.Instruction(
+						chain.from_iterable(bindings),
+						asynchook = self.database._receive_async
+					),
 					self
 				)
 			self.database._pq_complete()
@@ -1141,16 +1110,32 @@ def _load_bulk_tuples(self, tupleseq, tps = None):
 
 	def load(self, iterable, tps = None):
 		"""
-		Execute the query for each parameter set in `iterable`.
+		WARNING: Deprecated, use load_chunks and load_rows instead.
+		"""
+		if self.closed is None:
+			self._fini()
+		if isinstance(iterable, Chunks):
+			l = self.load_chunks
+		else:
+			l = self.load_rows
+		return l(iterable)
+
+	def load_chunks(self, chunks):
+		"""
+		Execute the query for each row-parameter set in `iterable`.
 
-		In cases of ``COPY ... FROM STDIN``, iterable must be an iterable `bytes`.
+		In cases of ``COPY ... FROM STDIN``, iterable must be an iterable of
+		sequences of `bytes`.
 		"""
 		if self.closed is None:
 			self._fini()
 		if not self._input:
-			return self._copy_data_in(iterable, tps = tps)
+			return self._load_copy_chunks(chunks)
 		else:
-			return self._load_bulk_tuples(iterable, tps = tps)
+			return self._load_tuple_chunks(chunks)
+
+	def load_rows(self, rows, chunksize = 256):
+		return self.load_chunks(chunk(rows, chunksize))
 
 class StoredProcedure(pg_api.StoredProcedure):
 	_e_factors = ('database', 'procedure_id')
diff --git a/postgresql/protocol/typio.py b/postgresql/protocol/typio.py
@@ -434,7 +434,9 @@ def process_tuple(procs, tup, exception_handler):
 	i = len(procs)
 	if len(tup) != i:
 		raise ValueError(
-			"inconsistent items, %d processors and %d objects"
+			"inconsistent items, %d processors and %d objects" %(
+				i, len(tup)
+			)
 		)
 	r = [None] * i
 	try:
diff --git a/postgresql/python/itertools.py b/postgresql/python/itertools.py
@@ -6,7 +6,7 @@
 itertools extensions
 """
 import collections
-from itertools import cycle
+from itertools import cycle, islice
 
 def interlace(*iters) -> collections.Iterable:
 	"""
@@ -20,3 +20,21 @@ def interlace(*iters) -> collections.Iterable:
 	)
 	"""
 	return map(next, cycle([iter(x) for x in iters]))
+
+def chunk(iterable, chunksize = 256):
+	"""
+	Given an iterable, return an iterable producing chunks of the objects
+	produced by the given iterable.
+
+	chunks([o1,o2,o3,o4], chunksize = 2) -> [
+		[o1,o2],
+		[o3,o4],
+	]
+	"""
+	iterable = iter(iterable)
+	last = ()
+	lastsize = chunksize
+	while lastsize == chunksize:
+		last = list(islice(iterable, chunksize))
+		lastsize = len(last)
+		yield last
diff --git a/postgresql/test/test_driver.py b/postgresql/test/test_driver.py
