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__()