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;')
