def DeleteItem(self, request: api_pb2.DeleteItemRequest,
context: grpc.ServicerContext) -> api_pb2.DeleteItemResponse:
try:
db.models.Item.objects.get(id=request.id).delete()
return api_pb2.DeleteItemResponse()
except db.models.Item.DoesNotExist:
context.set_code(grpc.StatusCode.NOT_FOUND)
context.set_details('Item does not exist')
return api_pb2.DeleteItemResponse()
def AddItem(self, request: api_pb2.Item, context: grpc.ServicerContext) -> api_pb2.AddItemResponse:
try:
item = db.models.Item(**_json_format.MessageToDict(request))
item.full_clean()
item.save()
return api_pb2.AddItemResponse(item=convert(item))
except ValidationError as e: # TODO: Return correct status code by exception types
context.set_code(grpc.StatusCode.INTERNAL)
context.set_details(json.dumps(e.messages))
return api_pb2.AddItemResponse()
def GetItem(self, request: api_pb2.GetItemRequest, context: grpc.ServicerContext) -> api_pb2.Item:
try:
item = db.models.Item.objects.get(id=request.id)
item.pv = F('pv') + 1
item.save()
item.refresh_from_db()
return convert(item)
except db.models.Item.DoesNotExist:
context.set_code(grpc.StatusCode.NOT_FOUND)
context.set_details('Item does not exist')
return api_pb2.Item()
def set_grpc_err(
context: grpc.ServicerContext,
code: grpc.StatusCode,
details: str,
):
"""
Sets status code and details for a gRPC context. Removes commas from
the details message (see https://github.com/grpc/grpc-node/issues/769)
"""
context.set_code(code)
context.set_details(details.replace(',', ''))
def GetBeer(self, get_beer_request: bartender_pb2.GetBeerRequest,
context: grpc.ServicerContext) -> bartender_pb2.Beer:
types = self._cache.get(get_beer_request.brand, [])
if not types:
context.set_code(grpc.StatusCode.NOT_FOUND)
return bartender_pb2.Beer()
beer = [b for b in types if b.name == get_beer_request.name]
if not beer:
context.set_code(grpc.StatusCode.NOT_FOUND)
return bartender_pb2.Beer()
return beer[0]
def GetByToken(self, request: app_pb2.GetByTokenReq,
context: grpc.ServicerContext) -> app_pb2.GetByTokenResp:
# Handle Error
# https://github.com/grpc/grpc/blob/master/doc/statuscodes.md
if request.token not in self.maxCacheKeysByToken:
context.set_code(grpc.StatusCode.NOT_FOUND)
context.set_details('Token not found {}'.format(request.token))
return app_pb2.GetByTokenResp()
max_cache_keys = self.maxCacheKeysByToken[request.token]
account = app_pb2.Account(max_cache_keys=max_cache_keys)
return app_pb2.GetByTokenResp(account=account)
def __init__(self,
message: str = None,
context: ServicerContext = None,
status_code: tuple = None):
self.status_code = self.status_code if status_code is None\
else status_code
self.message = self.message if message is None\
else message
if context:
context.set_code(self.status_code)
context.set_details(self.message)
super().__init__(message)
def wrapped_catch_them_all(self, request, context: grpc.ServicerContext):
try:
return func(self, request, context)
except Exception as e:
start_time = getattr(context, "start_time", None)
if start_time is not None:
delta = time.perf_counter() - start_time
self._log.exception("FAIL %.3f", delta)
else:
self._log.exception("FAIL ?")
context.set_code(grpc.StatusCode.INTERNAL)
context.set_details("%s: %s" % (type(e), e))
context.error = True
submit_event("error", 1)
return EventResponse()
def ListItem(self, request: api_pb2.ListItemRequest, context: grpc.ServicerContext) -> api_pb2.ListItemResponse:
page = request.page
limit = request.limit
items = Paginator(db.models.Item.objects.all(), limit)
try:
return api_pb2.ListItemResponse(
items=map(convert, items.page(page).object_list),
total=db.models.Item.objects.count(),
prevPage=max(1, page - 1),
nextPage=min(page + 1, items.num_pages),
)
except InvalidPage:
context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
context.set_details('Invalid page number')
return api_pb2.ListItemResponse()
def dispatch(self, request_object: Message,
context_object: grpc.ServicerContext, func: Callable):
ctx = contextvars.copy_context()
context.set(context_object)
bound_logger = self.logger.bind(peer=context_object.peer())
with bound_logger.catch():
try:
return ctx.run(func, request_object)
except exceptions.GRPCError as e:
context_object.set_code(e.code)
if e.details:
context_object.set_details(e.details)
except NotImplementedError:
context_object.set_code(grpc.StatusCode.UNIMPLEMENTED)
def UpdateItem(self, request: api_pb2.UpdateItemRequest,
context: grpc.ServicerContext) -> api_pb2.UpdateItemResponse:
try:
item = db.models.Item.objects.get(id=request.item.id)
item.__dict__.update(_json_format.MessageToDict(request)['item'])
item.full_clean()
item.save()
item.refresh_from_db()
return api_pb2.UpdateItemResponse(item=convert(item))
except db.models.Item.DoesNotExist:
context.set_code(grpc.StatusCode.NOT_FOUND)
context.set_details('Item does not exist')
return api_pb2.UpdateItemResponse()
except ValidationError as e: # TODO: Return correct status code by exception types
context.set_code(grpc.StatusCode.INTERNAL)
context.set_details(json.dumps(e.messages))
return api_pb2.UpdateItemResponse()
def wrapper(instance, req: Message, ctx: Context) -> Message:
try:
res = func(instance, req, ctx)
return res
except ExitGRPCCallWithCode:
return Message()
except Exception as e:
if logger is not None:
logger.error(e)
ctx.set_code(grpc.StatusCode.INTERNAL)
ctx.set_details(
"Unknown error happened during processing request.")
return Message()
def CreateBeer(self, create_beer_request: bartender_pb2.CreateBeerRequest,
context: grpc.ServicerContext) -> bartender_pb2.Beer:
beer = create_beer_request.beer
if not beer or not beer.name or not create_beer_request.brand:
context.set_details("Well ain't that cute? But it's WRONG!")
context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
return bartender_pb2.Beer()
types = self._cache.get(create_beer_request.brand, [])
if beer.name in types:
context.set_details("Plagiarism!")
context.set_code(grpc.StatusCode.ALREADY_EXISTS)
return bartender_pb2.Beer()
types.append(beer)
self._cache[create_beer_request.brand] = types
return beer
def Store(self, request: app_pb2.StoreReq, context: grpc.ServicerContext) -> app_pb2.StoreResp:
accountsCtx = grpc.ServicerContext()
getByTokenResp = self.accounts.GetByToken(app_pb2.GetByTokenReq(token=request.account_token))
if request.account_token not in self.keysByAccount:
self.keysByAccount[request.account_token] = 0
if self.keysByAccount[request.account_token] >= getByTokenResp.account.max_cache_keys:
context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
context.set_details('Account {} exceeds max key limit {}'.format(request.account_token, getByTokenResp.account.max_cache_keys))
return app_pb2.GetResp()
if request.key not in self.store:
self.keysByAccount[request.account_token] += 1
self.store[request.key] = request.val
return app_pb2.StoreResp()
def intercept(
self,
method: Callable,
request: Any,
context: grpc.ServicerContext,
method_name: str,
) -> Any:
"""Do not call this directly; use the interceptor kwarg on grpc.server()."""
try:
return method(request, context)
except GrpcException as e:
context.set_code(e.status_code)
context.set_details(e.details)
raise
except Exception as e:
if self._status_on_unknown_exception is not None:
context.set_code(self._status_on_unknown_exception)
context.set_details(repr(e))
raise
def intercept(
self,
method: Callable,
request: Any,
context: grpc.ServicerContext,
method_name: str,
) -> Any:
"""Override this method to implement a custom interceptor.
You should call method(request, context) to invoke the
next handler (either the RPC method implementation, or the
next interceptor in the list).
Args:
method: The next interceptor, or method implementation.
request: The RPC request, as a protobuf message.
context: The ServicerContext pass by gRPC to the service.
method_name: A string of the form
"/protobuf.package.Service/Method"
Returns:
This should generally return the result of
method(request, context), which is typically the RPC
method response, as a protobuf message. The interceptor
is free to modify this in some way, however.
"""
try:
return method(request, context)
except GrpcException as e:
context.set_code(e.status_code)
context.set_details(e.details)
logger.error(e.details)
return any_pb2.Any()
except marshmallow.ValidationError as e:
context.set_code(InvalidArgument.status_code)
msg = ' '.join([
"%s: %s" % (key, str(value))
for key, value in e.messages.items()
])
context.set_details(msg)
logger.error(e)
return any_pb2.Any()
except mongoengine.errors.DoesNotExist as e:
context.set_code(NotFound.status_code)
context.set_details(str(e))
logger.error(str(e))
return any_pb2.Any()
except Exception as e:
context.set_code(Unknown.status_code)
context.set_details(str(e))
logger.error(str(e))
return any_pb2.Any()
def UpdateExecutionInfo(
self, request: execution_watcher_pb2.UpdateExecutionInfoRequest,
context: grpc.ServicerContext
) -> execution_watcher_pb2.UpdateExecutionInfoResponse:
"""Updates the `custom_properties` field of Execution object in MLMD."""
logging.info(
'Received request to update execution info: updates %s, '
'execution_id %s', request.updates, request.execution_id)
if request.execution_id != self._execution.id:
context.set_code(grpc.StatusCode.NOT_FOUND)
context.set_details(
'Execution with given execution_id not tracked by server: '
f'{request.execution_id}')
return execution_watcher_pb2.UpdateExecutionInfoResponse()
for key, value in request.updates.items():
self._execution.custom_properties[key].CopyFrom(value)
# Only the execution is needed
with self._mlmd_connection as m:
m.store.put_executions((self._execution, ))
return execution_watcher_pb2.UpdateExecutionInfoResponse()
def Delete(self, request: crudata.DeleteRequest, context: grpc.ServicerContext) -> crudata.DeleteResponse:
if not request.originator or not request.originator.id:
context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
context.set_details("missing originator.id")
return crudata.DeleteResponse()
validate_originator(request.originator)
if not request.entity_type:
context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
context.set_details("missing entity_type")
return crudata.DeleteResponse()
delete_originator = self.crud_store.delete(
entity_type=request.entity_type,
originator=request.originator,
)
return crudata.DeleteResponse(
originator=delete_originator,
)
def GetSalesCount(
self, request: product_insight_api_pb2.GetSalesCountRequest,
context: grpc.ServicerContext
) -> product_insight_api_pb2.GetSalesCountResponse:
# FIXME! Observability.
# authorization
client_access_token = ''
for metadata in context.invocation_metadata():
if 'authorization' == metadata.key:
client_access_token = str(metadata.value).replace(
'Bearer ', '', 1)
if client_access_token != self.access_token:
context.set_code(grpc.StatusCode.UNAUTHENTICATED)
return product_insight_api_pb2.GetSalesCountResponse()
start_time = request.start_time # type: timestamp_pb2.Timestamp
end_time = request.end_time # type: timestamp_pb2.Timestamp
product_id = request.product_id # type: int
start_time_epoch = start_time.seconds + start_time.nanos / 1e9 # type: float
end_time_epoch = end_time.seconds + end_time.nanos / 1e9 # type: float
if product_id not in self.products_inverted_index:
return product_insight_api_pb2.GetSalesCountResponse(
product_id=product_id,
sales_count=0,
)
sales_count = 0 # type: int
sales_times = self.products_inverted_index[product_id]
for t in sales_times:
if start_time_epoch <= t <= end_time_epoch:
sales_count += 1
return product_insight_api_pb2.GetSalesCountResponse(
product_id=product_id,
sales_count=sales_count,
)
def Process(self, fs_msg: common_pb2.Message, context: grpc.ServicerContext):
"""Processes a single fleetspeak message."""
try:
validation_info = dict(fs_msg.validation_info.tags)
if fs_msg.message_type == "GrrMessage":
grr_message = rdf_flows.GrrMessage.FromSerializedBytes(
fs_msg.data.value)
self._ProcessGRRMessages(fs_msg.source.client_id, [grr_message],
validation_info)
elif fs_msg.message_type == "MessageList":
packed_messages = rdf_flows.PackedMessageList.FromSerializedBytes(
fs_msg.data.value)
message_list = communicator.Communicator.DecompressMessageList(
packed_messages)
self._ProcessGRRMessages(fs_msg.source.client_id, message_list.job,
validation_info)
else:
logging.error("Received message with unrecognized message_type: %s",
fs_msg.message_type)
context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
except Exception:
logging.exception("Exception processing message: %s", fs_msg)
raise
def _set_unknown_err(context: grpc.ServicerContext, err):
logging.error(traceback.format_exc())
context.set_code(grpc.StatusCode.UNKNOWN)
context.set_details(str(err))
def _set_invalid_arg_err(context: grpc.ServicerContext, err):
logging.error(traceback.format_exc())
context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
context.set_details(str(err))
def _propagate_grpc_code_err(context: grpc.ServicerContext,
err: grpc.RpcError):
logging.error(traceback.format_exc())
context.set_code(err.code())
context.set_details(str(err))
def _set_unavailable_error(context: grpc.ServicerContext, err):
logging.error(traceback.format_exc())
context.set_code(grpc.StatusCode.UNAVAILABLE)
context.set_details(str(err))
def set_code_unary_unary(request: bytes,
context: grpc.ServicerContext):
context.set_code(grpc.StatusCode.INTERNAL)
def __init__(self, ctx: Context, status_code, details: str = ""):
ctx.set_code(status_code)
ctx.set_details(details)
super().__init__()
