## # .protocol.client3 ## """ Protocol version 3.0 client and tools. """ import os 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) try: s.sendall(cq) finally: 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. return # The existence of the socket attribute indicates an attempt was made. self.socket = None try: self.socket = self.socket_factory(timeout = timeout) except ( 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 else: 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"), )) return 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 return # 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 return if supported: # Make an SSL connection. try: 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() while r: 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 try: 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 else: # It's probably a non-fatal error, # timeout or try again.. raise ## # 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. """ r = True 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() return r 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 try: 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 else: # It wasn't fatal, so just raise 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 )) else: # 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) ) return r def take_out_trash(self): """ close cursors and statements slated for closure. """ xm = [] 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. return 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 try: 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. op() elif dir is xact.Receiving: self.read_messages() self.read = self.read[op(self.read):] self.state = getattr(x, 'last_ready', self.state) else: 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. :( return else: raise 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: try: 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): raise 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'' return 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): try: del self.cursors[pq_cursor_id] except KeyError: pass 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): try: del self.statements[pq_statement_id] except KeyError: pass 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' else: 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'