Fix command extraction.

James William Pye · James William Pye · commit fdf9a3f0f083 · 2010-02-19T16:39:46.000-07:00
Prior, only the first field was extracted. There may be multiple words in the command tag. Fixes #22
diff --git a/postgresql/protocol/element3.py b/postgresql/protocol/element3.py
@@ -244,18 +244,19 @@ def extract_count(self):
 		"""
 		Extract the last set of digits as an integer.
 		"""
-		rms = self.data.strip().split()
-		if rms[0].lower() == b'copy':
-			if len(rms) > 1:
-				return int(rms[-1])
-		elif rms[-1].isdigit():
-			return int(rms[-1])
+		# Find the last sequence of digits.
+		# If there are no fields consisting only of digits, there is no count.
+		for x in reversed(self.data.split()):
+			if x.isdigit():
+				return int(x)
+		return None
 
 	def extract_command(self):
-		t = self.data.strip().split()
-		if t:
-			return t[0]
-		return None
+		"""
+		Strip all the *surrounding* digits and spaces from the command tag,
+		and return that string.
+		"""
+		return self.data.strip(b'\c\n\t 0123456789') or None
 
 class Null(EmptyMessage):
 	'Null command'
diff --git a/postgresql/test/test_lib.py b/postgresql/test/test_lib.py
@@ -78,7 +78,7 @@ def _testILF(self, lib):
 		self.failUnlessEqual(b.sym_const_column, [1])
 		self.failUnlessEqual(b.sym_const_rows, [(1,)])
 		self.failUnlessEqual(b.sym_const_chunks, [[(1,)]])
-		self.failUnlessEqual(b.sym_const_ddl, ('CREATE', None))
+		self.failUnlessEqual(b.sym_const_ddl, ('CREATE TABLE', None))
 		self.failUnlessEqual(b.sym_preload(), 1)
 		# now stored procs
 		self.failUnlessEqual(b.sym_proc(2,), 2)
diff --git a/postgresql/test/test_protocol.py b/postgresql/test/test_protocol.py
@@ -250,8 +250,26 @@ def testUnknownNoticeFields(self):
 		self.failUnlessEqual(E[b'P'], b'foobar')
 		self.failUnlessEqual(len(N), 4)
 		self.failUnlessEqual(len(E), 4)
+
+	def testCompleteExtracts(self):
+		x = e3.Complete(b'FOO BAR 1321')
+		self.failUnlessEqual(x.extract_command(), b'FOO BAR')
+		self.failUnlessEqual(x.extract_count(), 1321)
+		x = e3.Complete(b' CREATE  	TABLE 13210  ')
+		self.failUnlessEqual(x.extract_command(), b'CREATE  	TABLE')
+		self.failUnlessEqual(x.extract_count(), 13210)
+		x = e3.Complete(b'  CREATE  	TABLE  \t713210  ')
+		self.failUnlessEqual(x.extract_command(), b'CREATE  	TABLE')
+		self.failUnlessEqual(x.extract_count(), 713210)
+		x = e3.Complete(b'  CREATE  	TABLE  0 \t13210  ')
+		self.failUnlessEqual(x.extract_command(), b'CREATE  	TABLE')
+		self.failUnlessEqual(x.extract_count(), 13210)
+		x = e3.Complete(b' 0 \t13210 ')
+		self.failUnlessEqual(x.extract_command(), None)
+		self.failUnlessEqual(x.extract_count(), 13210)
+
 ##
-# xact3 tests
+# .protocol.xact3 tests
 ##
 
 xact_samples = [
