def write(self, message): is_stream = hasattr(message, "read_bytes") if (hasattr(message, 'to_bytes') is False or isinstance(getattr(message, 'to_bytes'), collections.Callable) is False) and not is_stream: raise TypeError("invalid message: ({0})".format(message)) logger.debug('=> %s', message) try: if not is_stream: self._socket().sendall(message.to_bytes()) else: # read to end in chunks while True: data = message.read_bytes() if len(data) == 0: break self._socket().sendall(data) except Exception as e: self.close_socket() if str(e) == 'unsupported authentication method: 9': raise errors.ConnectionError( 'Error during authentication. Your password might be expired.' ) else: raise errors.ConnectionError(str(e))
def _socket(self): if self.socket is not None: return self.socket host = self.options.get('host') port = self.options.get('port') connection_timeout = self.options.get('connection_timeout') raw_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) raw_socket.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) if connection_timeout is not None: raw_socket.settimeout(connection_timeout) raw_socket.connect((host, port)) ssl_options = self.options.get('ssl') if ssl_options is not None and ssl_options is not False: from ssl import CertificateError, SSLError raw_socket.sendall(messages.SslRequest().to_bytes()) response = raw_socket.recv(1) if response in ('S', b'S'): try: if isinstance(ssl_options, ssl.SSLContext): raw_socket = ssl_options.wrap_socket( raw_socket, server_hostname=host) else: raw_socket = ssl.wrap_socket(raw_socket) except CertificateError as e: raise errors.ConnectionError('SSL: ' + e.message) except SSLError as e: raise errors.ConnectionError('SSL: ' + e.reason) else: raise SSLNotSupported( "SSL requested but not supported by server") self.socket = raw_socket return self.socket
def write(self, message): is_stream = hasattr(message, "read_bytes") if (hasattr(message, 'to_bytes') is False or callable( getattr(message, 'to_bytes')) is False) and not is_stream: raise TypeError("invalid message: ({0})".format(message)) if getattr(self, 'debug', False): print "=> {0}".format(message) try: if not is_stream: self._socket().sendall(message.to_bytes()) else: # read to end in chunks while True: data = message.read_bytes() if len(data) == 0: break self._socket().sendall(data) except Exception, e: self.close_socket() raise errors.ConnectionError(e.message)
def process_message(self, message): if isinstance(message, messages.ErrorResponse): raise errors.ConnectionError(message.error_message()) elif isinstance(message, messages.NoticeResponse): if getattr(self, 'notice_handler', None) is not None: self.notice_handler(message) elif isinstance(message, messages.BackendKeyData): self.backend_pid = message.pid self.backend_key = message.key elif isinstance(message, messages.ParameterStatus): self.parameters[message.name] = message.value elif isinstance(message, messages.ReadyForQuery): self.transaction_status = message.transaction_status elif isinstance(message, messages.CommandComplete): # TODO: im not ever seeing this actually returned by vertica... # if vertica returns a row count, set the rowcount attribute in cursor #if hasattr(message, 'rows'): # self.cursor.rowcount = message.rows pass elif isinstance(message, messages.EmptyQueryResponse): pass elif isinstance(message, messages.CopyInResponse): pass else: raise errors.MessageError("Unhandled message: {0}".format(message)) # set last message self._cursor._message = message
def cursor(self, cursor_type=None): if self.closed(): raise errors.ConnectionError('Connection is closed') # let user change type if they want? self._cursor.cursor_type = cursor_type return self._cursor
def read_bytes(self, n): results = '' while len(results) < n: bytes = self._socket().recv(n - len(results)) if bytes is None or len(bytes) == 0: raise errors.ConnectionError("Connection closed by Vertica") results = results + bytes return results
def read_bytes(self, n): results = bytes() while len(results) < n: bytes_ = self._socket().recv(n - len(results)) if not bytes_: raise errors.ConnectionError("Connection closed by Vertica") results = results + bytes_ return results
def write(self, message): if hasattr(message, 'to_bytes') is False or callable( getattr(message, 'to_bytes')) is False: raise TypeError("invalid message: ({0})".format(message)) logger.debug('=> %s', message) try: self._socket().sendall(message.to_bytes()) except Exception, e: self.close_socket() raise errors.ConnectionError(e.message)
def read_message(self): try: type_ = self.read_bytes(1) size = unpack('!I', self.read_bytes(4))[0] if size < 4: raise errors.MessageError("Bad message size: {0}".format(size)) message = BackendMessage.factory(type_, self.read_bytes(size - 4)) logger.debug('<= %s', message) return message except (SystemError, IOError) as e: self.close_socket() raise errors.ConnectionError(str(e))
def process_message(self, message): if isinstance(message, messages.ErrorResponse): raise errors.ConnectionError(message.error_message()) elif isinstance(message, messages.NoticeResponse): if getattr(self, 'notice_handler', None) is not None: self.notice_handler(message) elif isinstance(message, messages.BackendKeyData): self.backend_pid = message.pid self.backend_key = message.key elif isinstance(message, messages.ParameterStatus): self.parameters[message.name] = message.value elif isinstance(message, messages.ReadyForQuery): self.transaction_status = message.transaction_status else: raise errors.MessageError("Unhandled message: {0}".format(message))
def read_message(self): try: ready = select.select([self._socket()], [], [], self.options['read_timeout']) if len(ready[0]) > 0: type = self.read_bytes(1) size = unpack('!I', self.read_bytes(4))[0] if size < 4: raise errors.MessageError( "Bad message size: {0}".format(size)) message = BackendMessage.factory(type, self.read_bytes(size - 4)) if getattr(self, 'debug', False): print "<= {0}".format(message) return message else: self.close() raise errors.TimedOutError("Connection timed out") except Exception as e: self.close_socket() raise errors.ConnectionError(e.message)
def rollback(self): if self.closed(): raise errors.ConnectionError('Connection is closed') cur = self.cursor() cur.execute('ROLLBACK;')
def commit(self): if self.closed(): raise errors.ConnectionError('Connection is closed') cur = self.cursor() cur.execute('COMMIT;')