fe/postgresql/protocol/client3.py at v1.3 · python-postgres/fe

558 lines (505 loc) · 15.3 KB
# .protocol.client3
Protocol version 3.0 client and tools.
import weakref
from .buffer import pq_message_stream
from . import element3 as element
from . import xact3 as xact
__all__ = ('Connection',)
client_detected_protocol_error = element.ClientError((
	(b'S', 'FATAL'),
	(b'C', '08P01'),
	(b'M', "wire-data caused exception in protocol transaction"),
	(b'H', "Protocol error detected."),
client_connect_timeout = element.ClientError((
	(b'S', 'FATAL'),
	(b'C', '--TOE'),
	(b'M', "connect timed out"),
not_pq_error = element.ClientError((
	# ProtocolError
	(b'S', 'FATAL'),
	(b'C', '08P01'),
	(b'M', 'server did not support SSL negotiation'),
	(b'H', 'The server is probably not PostgreSQL.'),
no_ssl_error = element.ClientError((
	(b'S', 'FATAL'),
	# InsecurityError
	(b'C', '--SEC'),
	(b'M', 'SSL was required, and the server could not accommodate'),
# Details in __context__
ssl_failed_error = element.ClientError((
	(b'S', 'FATAL'),
	# InsecurityError
	(b'C', '--SEC'),
	(b'M', 'SSL negotiation caused exception'),
# failed to complete the connection, but no error set.
# indicates a programmer error.
partial_connection_error = element.ClientError((
	(b'S', 'FATAL'),
	(b'C', '--XXX'),
	(b'M', "failed to complete negotiation"),
	(b'H',	"Negotiation failed to completed, but no " \
			"error was attributed on the connection."),
eof_error = element.ClientError((
	(b'S', 'FATAL'),
	(b'C', '08006'),
	(b'M', 'unexpected EOF from server'),
	(b'D',	"Zero-length read from the connection's socket."),
class Connection(object):
	A PQv3 connection.
	Operations are designed to not raise exceptions. The user of the
	connection must check for failures. This is done to encourage users
	to use their own Exception hierarchy.
	_tracer = None
	def tracer():
		def fget(self):
			return self._tracer
		def fset(self, value):
			self._tracer = value
			self.write_messages = self.traced_write_messages
			self.read_messages = self.traced_read_messages
		def fdel(self):
			del self._tracer
			self.write_messages = self.standard_write_messages
			self.read_messages = self.standard_read_messages
		doc = 'Callable object to pass protocol trace strings to. '\
			'(Normally a write method.)'
		return locals()
	tracer = property(**tracer())
	def synchronize(self):
		Explicitly send a Synchronize message to the backend.
		Useful for forcing the completion of lazily processed transactions.
		NOTE: This will not cause trash to be taken out.
		if self.xact is not None:
			self.complete()
		x = xact.Instruction((element.SynchronizeMessage,))
		self.xact = x
		self.complete()
	def interrupt(self, timeout = None):
		cq = element.CancelRequest(self.backend_id, self.key).bytes()
		s = self.socket_factory(timeout = timeout)
			s.sendall(cq)
			s.close()
	def connect(self, ssl = None, timeout = None):
		Establish the connection to the server.
		If `ssl` is None, the socket will not be secured.
		If `ssl` is True, the socket will be secured, but it will
		close the connection and return if SSL is not available.
		If `ssl` is False, the socket will attempt to be secured, but
		will continue even in the event of a server that does not
		support SSL.
		`timeout` will be passed directly to the configured `socket_factory`.
		if hasattr(self, 'socket'):
			# If there's a socket attribute it normally means
			# that the connection has already been connected.
			# Successfully or not; doesn't matter.
		# The existence of the socket attribute indicates an attempt was made.
		self.socket = None
			self.socket = self.socket_factory(timeout = timeout)
			self.socket_factory.timeout_exception,
			self.socket_factory.fatal_exception
		) as err:
			self.xact.state = xact.Complete
			self.xact.fatal = True
			self.xact.exception = err
			if self.socket_factory.timed_out(err):
				self.xact.error_message = client_connect_timeout
				errmsg = self.socket_factory.fatal_exception_message(err)
				# It's an error that occurred during socket creation/connection.
				# Even if there isn't a known fatal message,
				# identify it as fatal and set an ambiguous message.
				self.xact.error_message = element.ClientError((
					(b'S', 'FATAL'),
					# ConnectionRejectionError
					(b'C', '08004'),
					(b'M', errmsg or "could not connect"),
		if ssl is not None:
			# if ssl is True, ssl is *required*
			# if ssl is False, ssl will be tried, but not required
			# if ssl is None, no SSL negotiation will happen
			self.ssl_negotiation = supported = self.negotiate_ssl()
			# b'S' or b'N' was *not* received.
			if supported is None:
				# probably not PQv3..
				self.xact.fatal = True
				self.xact.error_message = not_pq_error
				self.xact.state = xact.Complete
			# b'N' was received, but ssl is required.
			if not supported and ssl is True:
				# ssl is required..
				self.xact.fatal = True
				self.xact.error_message = no_ssl_error
				self.xact.state = xact.Complete
			if supported:
				# Make an SSL connection.
					self.socket = self.socket_factory.secure(self.socket)
				except Exception as err:
					# Any exception marks a failure.
					self.xact.exception = err
					self.xact.fatal = True
					self.xact.state = xact.Complete
					self.xact.error_message = ssl_failed_error
					return
		# time to negotiate
		negxact = self.xact
		self.complete()
		if negxact.state is xact.Complete and negxact.fatal is None:
			self.key = negxact.killinfo.key
			self.backend_id = negxact.killinfo.pid
		elif not hasattr(self.xact, 'error_message'):
			# if it's not complete, something strange happened.
			# make sure to clean up...
			self.xact.fatal = True
			self.xact.state = xact.Complete
			self.xact.error_message = partial_connection_error
	def negotiate_ssl(self) -> (bool, None):
		Negotiate SSL
		If SSL is available--received b'S'--return True.
		If SSL is unavailable--received b'N'--return False.
		Otherwise, return None. Indicates non-PQv3 endpoint.
		r = element.NegotiateSSLMessage.bytes()
			r = r[self.socket.send(r):]
		status = self.socket.recv(1)
		if status == b'S':
			return True
		elif status == b'N':
			return False
		# probably not postgresql.
		return None
	def read_into(self, Complete = xact.Complete):
		read data from the wire and write it into the message buffer.
		BUFFER_HAS_MSG = self.message_buffer.has_message
		BUFFER_WRITE_MSG = self.message_buffer.write
		RECV_DATA = self.socket.recv
		RECV_BYTES = self.recvsize
		XACT = self.xact
		while not BUFFER_HAS_MSG():
			if self.read_data is not None:
				BUFFER_WRITE_MSG(self.read_data)
				self.read_data = None
				# If the read_data satisfied a message,
				# no more data should be read.
				continue
				self.read_data = RECV_DATA(RECV_BYTES)
			except self.socket_factory.fatal_exception as e:
				msg = self.socket_factory.fatal_exception_message(e)
				if msg is not None:
					XACT.state = Complete
					XACT.fatal = True
					XACT.exception = e
					XACT.error_message = element.ClientError((
						(b'S', 'FATAL'),
						(b'C', '08006'),
						(b'M', msg),
					return False
					# It's probably a non-fatal error,
					# timeout or try again..
			# nothing read from a blocking socket? it's over.
			if self.read_data == b'':
				XACT.state = Complete
				XACT.fatal = True
				XACT.error_message = eof_error
				return False
			# Got data. Put it in the buffer and clear read_data.
			self.read_data = BUFFER_WRITE_MSG(self.read_data)
		return True
	def standard_read_messages(self):
		Read more messages into self.read when self.read is empty.
		if not self.read:
			# get more data from the wire and
			# write it into the message buffer.
			r = self.read_into()
			self.read = self.message_buffer.read()
	read_messages = standard_read_messages
	def send_message_data(self):
		send all `message_data`.
		If an exception occurs, it will check if the exception
		is fatal or not.
		SEND_DATA = self.socket.send
			while self.message_data:
				# Send data while there is data to send.
				self.message_data = self.message_data[
					SEND_DATA(self.message_data):
		except self.socket_factory.fatal_exception as e:
			msg = self.socket_factory.fatal_exception_message(e)
			if msg is not None:
				# it's fatal.
				self.xact.state = xact.Complete
				self.xact.fatal = True
				self.xact.exception = e
				self.xact.error_message = element.ClientError((
					(b'S', 'FATAL'),
					(b'C', '08006'),
					(b'M', msg),
				return False
				# It wasn't fatal, so just raise
		return True
	def standard_write_messages(self, messages,
		cat_messages = element.cat_messages
		Protocol message writer.
		if self.writing is not self.written:
			self.message_data += cat_messages(self.writing)
			self.written = self.writing
		if messages is not self.writing:
			self.writing = messages
			self.message_data += cat_messages(self.writing)
			self.written = self.writing
		return self.send_message_data()
	write_messages = standard_write_messages
	def traced_write_messages(self, messages):
		`message_writer` used when tracing.
		for msg in messages:
			t = getattr(msg, 'type', None)
			if t is not None:
				data_out = msg.bytes()
				self._tracer('↑ {type}({lend}): {data}{nl}'.format(
					type = repr(t)[2:-1],
					lend = len(data_out),
					data = repr(data_out),
					nl = os.linesep
				# It's not a message instance, so assume raw data.
				self._tracer('↑__(%d): %r%s' %(
					len(msg), msg, os.linesep
		return self.standard_write_messages(messages)
	def traced_read_messages(self):
		`message_reader` used when tracing.
		r = self.standard_read_messages()
		for msg in self.read:
			self._tracer('↓ %r(%d): %r%s' %(
				msg[0], len(msg[1]), msg[1], os.linesep)
	def take_out_trash(self):
		close cursors and statements slated for closure.
		cursors = 0
		for x in self.garbage_cursors:
			xm.append(element.ClosePortal(x))
			cursors += 1
		statements = 0
		for x in self.garbage_statements:
			xm.append(element.CloseStatement(x))
			statements += 1
		xm.append(element.SynchronizeMessage)
		x = xact.Instruction(xm)
		self.xact = x
		del self.garbage_cursors[:cursors]
		del self.garbage_statements[:statements]
		self.complete()
	def push(self, x):
		setup the given transaction to be processed.
		# Push any queued closures onto the transaction or a new transaction.
		if x.state is xact.Complete:
			# It's already complete.
		if self.xact is not None:
			self.complete()
		if self.xact is None:
			if self.garbage_statements or self.garbage_cursors:
				# This *has* to be done before a new transaction
				# is pushed.
				self.take_out_trash()
			if self.xact is None:
				# set it as the current transaction and begin
				self.xact = x
				# start it up
				self.step()
	def step(self):
		Make a single transition on the transaction.
		This should be used during COPY TO STDOUT or large result sets
		to stream information out.
		x = self.xact
			dir, op = x.state
			if dir is xact.Sending:
				self.write_messages(x.messages)
				# The "op" callable will either switch the state, or
				# set the 'messages' attribute with a new sequence
				# of message objects for more writing.
			elif dir is xact.Receiving:
				self.read_messages()
				self.read = self.read[op(self.read):]
				self.state = getattr(x, 'last_ready', self.state)
				raise RuntimeError(
					"unexpected PQ transaction state: " + repr(dir)
		except self.socket_factory.try_again_exception as e:
			# Unlike _complete, this catches at the outermost level
			# as there is no loop here for more transitioning.
			if self.socket_factory.try_again(e):
				# Can't read or write, ATM? Consider it a transition. :(
		if x.state is xact.Complete and \
		getattr(self.xact, 'fatal', None) is not True:
			# only remove the transaction if it's *not* fatal
			self.xact = None
	def complete(self):
		Complete the current transaction.
		# Continue to transition until all transactions have been
		# completed, or an exception occurs that does not signal retry.
		x = self.xact
		R = xact.Receiving
		S = xact.Sending
		C = xact.Complete
		READ_MORE = self.read_messages
		WRITE_MESSAGES = self.write_messages
		while x.state is not C:
				while x.state[0] is R:
					if READ_MORE():
						self.read = self.read[x.state[1](self.read):]
				# push() always takes one step, so it is likely that
				# the transaction is done sending out data by the time
				# complete() is called.
				while x.state[0] is S:
					if WRITE_MESSAGES(x.messages):
						x.state[1]()
					# Multiple calls to get() without signaling
					# completion *should* yield the same set over
					# and over again.
			except self.socket_factory.try_again_exception as e:
				if not self.socket_factory.try_again(e):
			except Exception as proto_exc:
				# If an exception is raised here, it's a protocol or a programming error.
				# XXX: It may be useful to have this closer to the actual
				# message so that a more informative message can be given.
				x.fatal = True
				x.state = xact.Complete
				x.exception = proto_exc
				x.error_message = client_detected_protocol_error
				self.state = b''
		self.state = getattr(x, 'last_ready', self.state)
		if getattr(x, 'fatal', None) is not True:
			# only remove the transaction if it's *not* fatal
			self.xact = None
	def register_cursor(self, cursor, pq_cursor_id):
		trash = self.trash_cursor
		self.cursors[pq_cursor_id] = weakref.ref(cursor, lambda ref: trash(pq_cursor_id))
	def trash_cursor(self, pq_cursor_id):
			del self.cursors[pq_cursor_id]
		except KeyError:
		self.garbage_cursors.append(pq_cursor_id)
	def register_statement(self, statement, pq_statement_id):
		trash = self.trash_statement
		self.statements[pq_statement_id] = weakref.ref(statement, lambda ref: trash(pq_statement_id))
	def trash_statement(self, pq_statement_id):
			del self.statements[pq_statement_id]
		except KeyError:
		self.garbage_statements.append(pq_statement_id)
	def __str__(self):
		if hasattr(self, 'ssl_negotiation'):
			if self.ssl_negotiation is True:
				ssl = 'SSL'
			elif self.ssl_negotiation is False:
				ssl = 'NOSSL after SSL'
			ssl = 'NOSSL'
		excstr = ''.join(self.exception_string(type(self.exception), self.exception))
		return str(self.socket_factory) \
			+ ' -> (' + ssl + ')' \
			+ os.linesep + excstr.strip()
	def __init__(self, socket_factory, startup, password = b'',):
		Create a connection.
		This does not establish the connection, it only initializes it.
		self.key = None
		self.backend_id = None
		self.socket_factory = socket_factory
		self.xact = xact.Negotiation(
			element.Startup(startup), password
		self.cursors = {}
		self.statements = {}
		self.garbage_statements = []
		self.garbage_cursors = []
		self.message_buffer = pq_message_stream()
		self.recvsize = 8192
		self.read = ()
		# bytes received.
		self.read_data = None
		# serialized message data to be written
		self.message_data = b''
		# messages to be written.
		self.writing = None
		# messages that have already been transformed into bytes.
		# (used to detect whether messages have already been written)
		self.written = None
		self.state = 'INITIALIZED'
Provide feedback

Saved searches

Use saved searches to filter your results more quickly

FilesExpand file tree

client3.py

Latest commit

History

client3.py

File metadata and controls
