Make Cursor.seek() return the number of rows it navigated over

elprans · James William Pye · commit bb780ab24f51 · 2010-08-18T08:23:59.000+08:00
Cursor.seek() now returns the number of rows seek navigated over to
reach the specified offset.  For this to become useful, this commit
also adds support for relative 'FORWARD'/'BACKWARD' modes.

While currently implemented "RELATIVE" mode makes relative seeks possible,
it does not support exhaustive seek mode, i. e., "FORWARD ALL" and
"BACKWARD ALL".  It is important to distinguish those modes from ABSOLUTE
seeks since ABSOLUTE seek always returns a single row, while relative
seeks return total number of rows.  Thus, seek('ALL', 'FORWARD') is the
only way to obtain the total number of rows returned by the query
without actually iterating over it.
diff --git a/postgresql/api.py b/postgresql/api.py
@@ -228,6 +228,8 @@ class Cursor(
 		0 : 'ABSOLUTE',
 		1 : 'RELATIVE',
 		2 : 'FROM_END',
+		3 : 'FORWARD',
+		4 : 'BACKWARD'
 	}
 	_direction_map = {
 		True : 'FORWARD',
@@ -293,6 +295,10 @@ def seek(self, offset, whence = 'ABSOLUTE'):
 		  Relative.
 		 ``2`` or ``"FROM_END"``
 		  Absolute from end.
+		 ``3`` or ``"FORWARD"``
+		  Relative forward.
+		 ``4`` or ``"BACKWARD"``
+		  Relative backward.
 
 		Direction effects whence. If direction is BACKWARD, ABSOLUTE positioning
 		will effectively be FROM_END, RELATIVE's position will be negated, and
diff --git a/postgresql/driver/pq3.py b/postgresql/driver/pq3.py
@@ -1092,23 +1092,40 @@ def seek(self, offset, whence = 'ABSOLUTE'):
 				"unknown whence parameter, %r" %(whence,)
 			)
 		rwhence = rwhence.upper()
-		if self.direction is False:
-			if rwhence == 'RELATIVE':
+
+		if offset == 'ALL':
+			if rwhence not in ('BACKWARD', 'FORWARD'):
+				rwhence = 'BACKWARD' if self.direction is False else 'FORWARD'
+		else:
+			if offset < 0 and rwhence == 'BACKWARD':
 				offset = -offset
-			elif rwhence == 'ABSOLUTE':
-				rwhence = 'FROM_END'
-			else:
-				rwhence = 'ABSOLUTE'
+				rwhence = 'FORWARD'
 
-		if rwhence == 'RELATIVE':
-			if offset < 0:
+			if self.direction is False:
+				if offset == 'ALL' and rwhence != 'FORWARD':
+					rwhence = 'BACKWARD'
+				else:
+					if rwhence == 'RELATIVE':
+						offset = -offset
+					elif rwhence == 'ABSOLUTE':
+						rwhence = 'FROM_END'
+					else:
+						rwhence = 'ABSOLUTE'
+
+		if rwhence in ('RELATIVE', 'BACKWARD', 'FORWARD'):
+			if offset == 'ALL':
 				cmd = self._pq_xp_move(
-					str(-offset).encode('ascii'), b'BACKWARD'
+					str(offset).encode('ascii'), str(rwhence).encode('ascii')
 				)
 			else:
-				cmd = self._pq_xp_move(
-					str(offset).encode('ascii'), b'RELATIVE'
-				)
+				if offset < 0:
+					cmd = self._pq_xp_move(
+						str(-offset).encode('ascii'), b'BACKWARD'
+					)
+				else:
+					cmd = self._pq_xp_move(
+						str(offset).encode('ascii'), str(rwhence).encode('ascii')
+					)
 		elif rwhence == 'ABSOLUTE':
 			cmd = self._pq_xp_move(str(offset).encode('ascii'), b'ABSOLUTE')
 		else:
@@ -1122,6 +1139,16 @@ def seek(self, offset, whence = 'ABSOLUTE'):
 		self.database._pq_push(x, self)
 		self.database._pq_complete()
 
+		count = None
+		complete = element.Complete.type
+		for cm in x.messages_received():
+			if getattr(cm, 'type', None) == complete:
+				count = cm.extract_count()
+				break
+
+		# XXX: Raise if count is None?
+		return count
+
 class Statement(pg_api.Statement):
 	string = None
 	database = None
diff --git a/postgresql/test/test_driver.py b/postgresql/test/test_driver.py
@@ -879,6 +879,22 @@ def testScroll(self, direction = True):
 		c.seek(10, 2)
 		self.failUnlessEqual(r1, c.read(10))
 
+	@pg_tmp
+	def testSeek(self):
+		ps = db.prepare("SELECT i FROM generate_series(0, (2^6)::int - 1) AS g(i)")
+		c = ps.declare()
+
+		self.failUnlessEqual(c.seek(4, 'FORWARD'), 4)
+		self.failUnlessEqual([x for x, in c.read(10)], list(range(4, 14)))
+
+		self.failUnlessEqual(c.seek(2, 'BACKWARD'), 2)
+		self.failUnlessEqual([x for x, in c.read(10)], list(range(12, 22)))
+
+		self.failUnlessEqual(c.seek(-5, 'BACKWARD'), 5)
+		self.failUnlessEqual([x for x, in c.read(10)], list(range(27, 37)))
+
+		self.failUnlessEqual(c.seek('ALL'), 27)
+
 	def testScrollBackwards(self):
 		# testScroll again, but backwards this time.
 		self.testScroll(direction = False)
