def _extract_rpc_error(self, exec_method, error):
"""Raises a VitessError for a proto3 vtrpc.RPCError structure, if set.
Args:
exec_method: name of the method to use in VitessError.
error: vtrpc.RPCError structure.
Raises:
vtgate_utils.VitessError: if an error was set.
"""
if error.code:
raise vtgate_utils.VitessError(exec_method, error.code,
error.message)
elif error.legacy_code:
raise vtgate_utils.VitessError(
exec_method, legacy_code_to_code_map[error.legacy_code],
error.message)
def _stream_execute(self,
sql,
bind_variables,
keyspace_name,
tablet_type,
keyspace_ids=None,
keyranges=None,
not_in_transaction=False,
effective_caller_id=None):
"""Return a generator and the fields for the response.
This method takes ownership of self.client, since multiple
stream_executes can be active at once.
Args:
sql: Str sql.
bind_variables: A (str: value) dict.
keyspace_name: Str keyspace name.
tablet_type: Str tablet_type.
keyspace_ids: List of uint64 or bytes keyspace_ids.
keyranges: KeyRange objects.
not_in_transaction: bool.
effective_caller_id: CallerID.
Returns:
Generator, fields pair.
Raises:
dbexceptions.ProgrammingError: On bad input.
"""
exec_method = None
req = None
if keyspace_ids is not None:
req = _create_req_with_keyspace_ids(sql, bind_variables,
keyspace_name, tablet_type,
keyspace_ids,
not_in_transaction)
exec_method = 'VTGate.StreamExecuteKeyspaceIds2'
elif keyranges is not None:
req = _create_req_with_keyranges(sql, bind_variables,
keyspace_name, tablet_type,
keyranges, not_in_transaction)
exec_method = 'VTGate.StreamExecuteKeyRanges2'
else:
raise dbexceptions.ProgrammingError(
'_stream_execute called without specifying keyspace_ids or keyranges'
)
self._add_caller_id(req, effective_caller_id)
stream_fields = []
stream_conversions = []
rpc_client = self._create_client()
try:
rpc_client.dial()
except gorpc.GoRpcError as e:
raise self._convert_exception(e)
def drain_conn_after_streaming_app_error():
"""Drain connection of all incoming streaming packets (ignoring them).
This is necessary for streaming calls which return application
errors inside the RPC response (instead of through the usual GoRPC
error return). This is because GoRPC always expects the last
packet to be an error; either the usual GoRPC application error
return, or a special "end-of-stream" error.
If an application error is returned with the RPC response, there
will still be at least one more packet coming, as GoRPC has not
seen anything that it considers to be an error. If the connection
is not drained of this last packet, future reads from the wire
will be off by one and will return errors.
"""
next_result = rpc_client.stream_next()
if next_result is not None:
rpc_client.close()
raise gorpc.GoRpcError(
'Connection should only have one packet remaining'
' after streaming app error in RPC response.')
try:
rpc_client.stream_call(exec_method, req)
first_response = rpc_client.stream_next()
if first_response:
if first_response.reply.get('Err'):
drain_conn_after_streaming_app_error()
raise vtgate_utils.VitessError(exec_method,
first_response.reply['Err'])
reply = first_response.reply['Result']
for field in reply['Fields']:
stream_fields.append((field['Name'], field['Type']))
stream_conversions.append(
field_types.conversions.get(field['Type']))
except (gorpc.GoRpcError, vtgate_utils.VitessError) as e:
self.logger_object.log_private_data(bind_variables)
raise self._convert_exception(e,
sql,
keyspace_ids,
keyranges,
keyspace=keyspace_name,
tablet_type=tablet_type)
except Exception:
logging.exception('gorpc low-level error')
raise
def row_generator():
try:
while True:
stream_result = rpc_client.stream_next()
if stream_result is None:
break
if stream_result.reply.get('Result'):
for result_item in stream_result.reply['Result'][
'Rows']:
yield tuple(
_make_row(result_item, stream_conversions))
except gorpc.GoRpcError as e:
raise self._convert_exception(e)
except Exception:
logging.exception('gorpc low-level error')
raise
finally:
rpc_client.close()
return row_generator(), stream_fields
def _stream_execute(self,
sql,
bind_variables,
tablet_type,
not_in_transaction=False):
"""Start a streaming query via VTGate.StreamExecute2."""
req = _create_req(sql, bind_variables, tablet_type, not_in_transaction)
self._add_session(req)
rpc_client = self._get_client()
stream_fields = []
stream_conversions = []
def drain_conn_after_streaming_app_error():
"""Drain connection of all incoming streaming packets (ignoring them).
This is necessary for streaming calls which return application
errors inside the RPC response (instead of through the usual GoRPC
error return). This is because GoRPC always expects the last
packet to be an error; either the usual GoRPC application error
return, or a special "end-of-stream" error.
If an application error is returned with the RPC response, there
will still be at least one more packet coming, as GoRPC has not
seen anything that it considers to be an error. If the connection
is not drained of this last packet, future reads from the wire
will be off by one and will return errors.
"""
next_result = rpc_client.stream_next()
if next_result is not None:
rpc_client.close()
raise gorpc.GoRpcError(
'Connection should only have one packet remaining'
' after streaming app error in RPC response.')
try:
method_name = 'VTGate.StreamExecute2'
rpc_client.stream_call(method_name, req)
first_response = rpc_client.stream_next()
if first_response.reply.get('Err'):
drain_conn_after_streaming_app_error()
raise vtgate_utils.VitessError(method_name,
first_response.reply['Err'])
reply = first_response.reply['Result']
for field in reply['Fields']:
stream_fields.append((field['Name'], field['Type']))
stream_conversions.append(
field_types.conversions.get(field['Type']))
except (gorpc.GoRpcError, vtgate_utils.VitessError) as e:
self.logger_object.log_private_data(bind_variables)
raise convert_exception(e, str(self), sql)
except Exception:
logging.exception('gorpc low-level error')
raise
# Take the BsonRpcClient from VTGateConnection. The row_generator
# will manage the BsonRpcClient. This VTGateConnection will connect
# to a new client if needed.
self.client = None
def row_generator():
try:
while True:
try:
stream_result = rpc_client.stream_next()
if stream_result is None:
break
# A session message, if any comes separately with no rows.
# I am not sure if we can ignore this.
if stream_result.reply.get('Session'):
self.session = stream_result.reply['Session']
else:
for result_item in stream_result.reply['Result'][
'Rows']:
yield tuple(
_make_row(result_item, stream_conversions))
except gorpc.GoRpcError as e:
raise convert_exception(e, str(self))
except Exception:
logging.exception('gorpc low-level error')
raise
finally:
rpc_client.close()
return row_generator(), stream_fields
def _extract_rpc_error(exec_method, error):
if error.code:
raise vtgate_utils.VitessError(exec_method, {
'Code': error.code,
'Message': error.message,
})