Correct message type comparisons.

James William Pye · James William Pye · commit 1d856f718857 · 2009-12-08T19:24:00.000-07:00
Using the 'is' operator only works so long as the message_types data is
consistent. It is possible for that consistency to fail in cases where modules
are reloaded.
diff --git a/postgresql/documentation/changes.txt b/postgresql/documentation/changes.txt
@@ -7,6 +7,19 @@ Changes
  * Correct types.Array binary serialization (8.5 added extra checks)
  * Add more encoding name mapping entries.
  * Documentation improvements.
+ * Correct protocol message type comparison method.
+   In 0.9.1 and before, types were compared using 'is' because all
+   message type objects were loaded from a tuple created by the
+   postgresql.protocol.message_types module. In some environments,
+   sys.modules is cleared. For pure-Python installations, this is unlikely
+   to cause a problem, but when the optimized module is available, it will hold
+   a reference to the "old" version of the message_types data. Subsequently, the
+   use of the 'is' operator is then broken resulting in superfluous protocol
+   errors. The fix for 0.9.2 is to use the actual comparison operator.
+   This does not add any significant overhead, but use of 'is' may return if
+   some permanent storage method becomes available.
+   [Reported by Radomir Stevanovic]
+
 
 0.9.1 released on 2009-08-12
 ----------------------------
diff --git a/postgresql/driver/pq3.py b/postgresql/driver/pq3.py
@@ -34,8 +34,6 @@
 
 from .. import types as pg_types
 
-is_showoption = lambda x: getattr(x, 'type', None) is element.ShowOption.type
-
 IDNS = 'py:%s'
 def ID(s, title = None):
 	'generate an id for a client statement or cursor'
@@ -321,7 +319,10 @@ def __init__(self, statement, parameters):
 		self.database = statement.database
 		Output.__init__(self, '')
 
-	def _init(self):
+	def _init(self,
+		null = element.Null.type,
+		complete = element.Complete.type,
+	):
 		expect = self._expect
 		self._xact = self._ins(
 			self._pq_xp_fetchall() + (element.SynchronizeMessage,)
@@ -330,19 +331,19 @@ def _init(self):
 		while self._xact.state != xact.Complete:
 			self.database._pq_step()
 			for x in self._xact.messages_received():
-				if x.type is element.Null.type:
+				if x.type == null:
 					self.database._pq_complete()
 					self._xact = None
 					return
-				elif x.type is element.Complete.type:
+				elif x.type == complete:
 					self._complete_message = x
 					self.database._pq_complete()
 					# If this was a select/copy cursor,
 					# the data messages would have caused an earlier
 					# return.
 					self._xact = None
 					return
-				elif x.type is expect:
+				elif x.type == expect:
 					# no need to step once this is seen
 					return
 				elif x.type in (
@@ -387,9 +388,11 @@ class SingleXactCopy(FetchAll):
 class SingleXactFetch(FetchAll):
 	_expect = element.Tuple.type
 	_process_chunk_ = FetchAll._process_tuple_chunk_Row
-	def _process_chunk(self, x):
+	def _process_chunk(self, x,
+		tuple_type = element.Tuple.type
+	):
 		return self._process_chunk_((
-			y for y in x if y.type is element.Tuple.type
+			y for y in x if y.type == tuple_type
 		))
 
 class MultiXactStream(Chunks):
@@ -433,7 +436,7 @@ def __next__(self):
 			self.database._pq_complete()
 
 		chunk = [
-			y for y in x.messages_received() if y.type is element.Tuple.type
+			y for y in x.messages_received() if y.type == element.Tuple.type
 		]
 		if len(chunk) == self.chunksize:
 			# there may be more, dispatch the request for the next chunk
@@ -496,7 +499,9 @@ def _which_way(self, direction):
 		else:
 			return self.direction
 
-	def _init(self):
+	def _init(self,
+		tupledesc = element.TupleDescriptor.type,
+	):
 		"""
 		Based on the cursor parameters and the current transaction state,
 		select a cursor strategy for managing the response from the server.
@@ -510,7 +515,7 @@ def _init(self):
 			self.database._pq_push(x, self)
 			self.database._pq_complete()
 			for m in x.messages_received():
-				if m.type is element.TupleDescriptor.type:
+				if m.type == tupledesc:
 					self._output = m
 					self._output_attmap = \
 						self.database.typio.attribute_map(self._output)
@@ -545,7 +550,7 @@ def _fetch(self, direction, quantity):
 		self.database._pq_push(x, self)
 		self.database._pq_complete()
 		return self._process_tuple((
-			y for y in x.messages_received() if y.type is element.Tuple.type
+			y for y in x.messages_received() if y.type == element.Tuple.type
 		))
 
 	def seek(self, offset, whence = 'ABSOLUTE'):
@@ -979,6 +984,7 @@ def first(self, *parameters):
 					params,
 					self._output_formats,
 				),
+				# Get all
 				element.Execute(b'', 0xFFFFFFFF),
 				element.SynchronizeMessage
 			),
@@ -990,8 +996,9 @@ def first(self, *parameters):
 		if self._output_io:
 			##
 			# Look for the first tuple.
+			tuple_type = element.Tuple.type
 			for xt in x.messages_received():
-				if xt.type is element.Tuple.type:
+				if xt.type == tuple_type:
 					break
 			else:
 				return None
@@ -1012,8 +1019,14 @@ def first(self, *parameters):
 		else:
 			##
 			# It doesn't return rows, so return a count.
+			##
+			# This loop searches through the received messages
+			# for the Complete message which contains the count.
+			complete = element.Complete.type
 			for cm in x.messages_received():
-				if getattr(cm, 'type', None) == element.Complete.type:
+				# Use getattr because COPY doesn't produce
+				# element.Message instances.
+				if getattr(cm, 'type', None) == complete:
 					break
 			else:
 				# Probably a Null command.
@@ -1800,7 +1813,7 @@ def connect(self):
 
 		# When ssl is None: SSL negotiation will not occur.
 		# When ssl is True: SSL negotiation will occur *and* it must succeed.
-		# When ssl is False: SSL negotiation will occur but NOSSL is okay.
+		# When ssl is False: SSL negotiation will occur but it may fail(NOSSL).
 		if sslmode == 'allow':
 			# first, without ssl, then with. :)
 			socket_factories = interlace(
@@ -1838,15 +1851,21 @@ def connect(self):
 				sf, self.connector._startup_parameters,
 				password = self.connector._password,
 			)
+			# Grab the negotiation transaction before
+			# connecting as it will be needed later if successful.
 			neg = pq.xact
 			pq.connect(ssl = ssl, timeout = timeout)
 
 			didssl = getattr(pq, 'ssl_negotiation', -1)
+
+			# It successfully connected if pq.xact is None.
 			if pq.xact is None:
 				self.pq = pq
-				for x in filter(is_showoption, neg.asyncs):
-					self._receive_async(x)
 				self.security = 'ssl' if didssl is True else None
+				showoption_type = element.ShowOption.type
+				for x in neg.asyncs:
+					if x.type == showoption_type:
+						self._receive_async(x)
 				# success!
 				break
 
@@ -2037,20 +2056,24 @@ def _error_lookup(self, om : element.Error,) -> pg_exc.Error:
 		err.database = self
 		return err
 
-	def _receive_async(self, msg, controller = None):
+	def _receive_async(self, msg, controller = None,
+		showoption = element.ShowOption.type,
+		notice = element.Notice.type,
+		notify = element.Notify.type,
+	):
 		c = controller or getattr(self, '_controller', self)
-		if msg.type is element.ShowOption.type:
+		if msg.type == showoption:
 			if msg.name == b'client_encoding':
 				self.typio.set_encoding(msg.value.decode('ascii'))
 			self.settings._notify(msg)
-		elif msg.type is element.Notice.type:
+		elif msg.type == notice:
 			src = 'SERVER'
 			if type(msg) is element.ClientNotice:
 				src = 'CLIENT'
 			m = self._convert_pq_message(msg, source = src)
 			m.creator = c
 			m.raise_message()
-		elif msg.type is element.Notify.type:
+		elif msg.type == notify:
 			subs = getattr(self, '_subscriptions', {})
 			for x in subs.get(msg.relation, ()):
 				x(self, msg)
diff --git a/postgresql/protocol/message_types.py b/postgresql/protocol/message_types.py
@@ -6,8 +6,11 @@
 Data module providing a sequence of bytes objects whose value corresponds to its
 index in the sequence.
 
-This provides resource for buffer objects to use common message type objects,
-and the additional privilege of using the `is` operator in all situations
-involving message type comparisons.
+This provides resource for buffer objects to use common message type objects.
+
+WARNING: It's tempting to use the 'is' operator and in some circumstances that
+may be okay. However, it's possible (sys.modules.clear()) for the extension
+modules' copy of this to become inconsistent with what protocol.element3 and
+protocol.xact3 are using, so it's important to **not** use 'is'.
 """
 message_types = tuple([bytes((x,)) for x in range(256)])
diff --git a/postgresql/protocol/pbuffer.py b/postgresql/protocol/pbuffer.py
@@ -69,12 +69,14 @@ def __len__(self):
 		rpos = self._start
 		self._strio.seek(self._start)
 		while True:
+			# get the message metadata
 			header = self._strio.read(5)
 			rpos += 5
 			if len(header) < 5:
+				# not enough data for another message
 				break
+			# unpack the length from the header
 			length, = xl_unpack(header)
-			typ = message_types[header[0]]
 			rpos += length - 4
 
 			if length < 4:
@@ -86,12 +88,12 @@ def __len__(self):
 			count += 1
 		return count
 
-	def _get_message(self):
+	def _get_message(self, mtypes = message_types):
 		header = self._strio.read(5)
 		if len(header) < 5:
 			return
 		length, = xl_unpack(header)
-		typ = message_types[header[0]]
+		typ = mtypes[header[0]]
 
 		if length < 4:
 			raise ValueError("invalid message size '%d'" %(length,))
diff --git a/postgresql/protocol/xact3.py b/postgresql/protocol/xact3.py
@@ -134,7 +134,7 @@ def put_messages(self, messages):
 		try:
 			for x in messages:
 				count += 1
-				if x[0] is element.Error.type:
+				if x[0] == element.Error.type:
 					if self.fatal is None:
 						self.error_message = element.Error.parse(x[1])
 						self.fatal = True
@@ -177,7 +177,7 @@ def state_machine(self):
 		"""
 		x = (yield (self.startup_message,))
 
-		if x[0] is not element.Authentication.type:
+		if x[0] != element.Authentication.type:
 			self.fatal = True
 			self.error_message = element.ClientError(
 				message = \
@@ -231,7 +231,7 @@ def state_machine(self):
 
 		# Done authenticating, pick up the killinfo and the ready message.
 		x = (yield None)
-		if x[0] is not element.KillInformation.type:
+		if x[0] != element.KillInformation.type:
 			self.fatal = True
 			self.error_message = element.ClientError(
 				message = \
@@ -246,7 +246,7 @@ def state_machine(self):
 		self.killinfo = element.KillInformation.parse(x[1])
 
 		x = (yield None)
-		if x[0] is not element.Ready.type:
+		if x[0] != element.Ready.type:
 			self.fatal = True
 			self.error_message = element.ClientError(
 				message = \
@@ -476,7 +476,7 @@ def standard_put(self, messages):
 
 			if path is None:
 				# No path for message type, could be a protocol error.
-				if x[0] is element.Error.type:
+				if x[0] == element.Error.type:
 					em = element.Error.parse(x[1])
 					fatal = em['severity'].upper() in (b'FATAL', b'PANIC')
 					self.error_message = em
@@ -541,7 +541,7 @@ def standard_put(self, messages):
 					current_step = next_step
 				else:
 					current_step = 0
-					if r.type is element.Ready.type:
+					if r.type == element.Ready.type:
 						self.last_ready = r.xact_state
 					# Done with the current command. Increment the offset, and
 					# try to process the new command with the remaining data.
@@ -582,11 +582,11 @@ def standard_put(self, messages):
 			last = processed[-1]
 			if type(last) is bytes:
 				self.state = (Receiving, self.put_copydata)
-			elif last.type is element.CopyToBegin.type:
+			elif last.type == element.CopyToBegin.type:
 				self.state = (Receiving, self.put_copydata)
-			elif last.type is element.Tuple.type:
+			elif last.type == element.Tuple.type:
 				self.state = (Receiving, self.put_tupledata)
-			elif last.type is element.CopyFromBegin.type:
+			elif last.type == element.CopyFromBegin.type:
 				self.CopyFailSequence = (self.CopyFailMessage,) + \
 					self.commands[offset+1:]
 				self.CopyDoneSequence = (element.CopyDoneMessage,) + \
@@ -601,13 +601,13 @@ def put_copydata(self, messages):
 		message is received, it reverts the ``state`` attribute back to
 		`standard_put` to process the message-sequence.
 		"""
+		copydata = element.CopyData.type
 		# "Fail" quickly if the last message is not copy data.
-		if messages[-1][0] is not element.CopyData.type:
+		if messages[-1][0] != copydata:
 			self.state = (Receiving, self.standard_put)
 			return self.standard_put(messages)
 
-		cdt = element.CopyData.type
-		lines = [x[1] for x in messages if x[0] is cdt]
+		lines = [x[1] for x in messages if x[0] == copydata]
 		if len(lines) != len(messages):
 			self.state = (Receiving, self.standard_put)
 			return self.standard_put(messages)
@@ -624,13 +624,13 @@ def put_tupledata(self, messages):
 		"""
 		# Fallback to `standard_put` quickly if the last
 		# message is not tuple data.
-		if messages[-1][0] is not element.Tuple.type:
+		if messages[-1][0] != element.Tuple.type:
 			self.state = (Receiving, self.standard_put)
 			return self.standard_put(messages)
 
 		p = element.Tuple.parse
 		t = element.Tuple.type
-		tuplemessages = [p(x[1]) for x in messages if x[0] is t]
+		tuplemessages = [p(x[1]) for x in messages if x[0] == t]
 		if len(tuplemessages) != len(messages):
 			self.state = (Receiving, self.standard_put)
 			return self.standard_put(messages)
diff --git a/postgresql/test/testall.py b/postgresql/test/testall.py
@@ -8,7 +8,6 @@
 import warnings
 from ..installation import Installation
 
-
 from .test_exceptions import *
 from .test_bytea_codec import *
 from .test_iri import *
