diff --git a/zerver/lib/bot_lib.py b/zerver/lib/bot_lib.py index fb2d311817..dda46a5300 100644 --- a/zerver/lib/bot_lib.py +++ b/zerver/lib/bot_lib.py @@ -7,14 +7,15 @@ import time import re import importlib from zerver.lib.actions import internal_send_message -from zerver.models import UserProfile +from zerver.models import UserProfile, \ + get_bot_state, set_bot_state, get_bot_state_size, is_key_in_bot_state from zerver.lib.integrations import EMBEDDED_BOTS from six.moves import configparser if False: from mypy_extensions import NoReturn -from typing import Any, Optional, List, Dict +from typing import Any, Optional, List, Dict, Text from types import ModuleType our_dir = os.path.dirname(os.path.abspath(__file__)) @@ -32,6 +33,36 @@ def get_bot_handler(service_name): bot_module = importlib.import_module(bot_module_name) # type: Any return bot_module.handler_class() +class StateHandlerError(Exception): + pass + +class StateHandler(object): + state_size_limit = 10000000 # type: int # TODO: Store this in the server configuration model. + + def __init__(self, user_profile): + # type: (UserProfile) -> None + self.user_profile = user_profile + + def __getitem__(self, key): + # type: (Text) -> Text + return get_bot_state(self.user_profile, key) + + def __setitem__(self, key, value): + # type: (Text, Text) -> None + old_entry_size = get_bot_state_size(self.user_profile, key) + new_entry_size = len(key) + len(value) + old_state_size = get_bot_state_size(self.user_profile) + new_state_size = old_state_size + (new_entry_size - old_entry_size) + if new_state_size > self.state_size_limit: + raise StateHandlerError("Cannot set state. Request would require {} bytes storage. " + "The current storage limit is {}.".format(new_state_size, self.state_size_limit)) + else: + set_bot_state(self.user_profile, key, value) + + def __contains__(self, key): + # type: (Text) -> bool + return is_key_in_bot_state(self.user_profile, key) + class EmbeddedBotHandler(object): def __init__(self, user_profile): # type: (UserProfile) -> None @@ -40,6 +71,7 @@ class EmbeddedBotHandler(object): self._rate_limit = RateLimit(20, 5) self.full_name = user_profile.full_name self.email = user_profile.email + self.state = StateHandler(user_profile) def send_message(self, message): # type: (Dict[str, Any]) -> None diff --git a/zerver/worker/queue_processors.py b/zerver/worker/queue_processors.py index d921f45be1..2c66d3dd9a 100644 --- a/zerver/worker/queue_processors.py +++ b/zerver/worker/queue_processors.py @@ -9,7 +9,6 @@ from functools import wraps import smtplib import socket -from zulip_bots.lib import ExternalBotHandler, StateHandler from django.conf import settings from django.db import connection from django.core.handlers.wsgi import WSGIRequest @@ -506,11 +505,6 @@ class EmbeddedBotWorker(QueueProcessingWorker): # type: (UserProfile) -> EmbeddedBotHandler return EmbeddedBotHandler(user_profile) - # TODO: Handle stateful bots properly - def get_state_handler(self): - # type: () -> StateHandler - return StateHandler() - def consume(self, event): # type: (Mapping[str, Any]) -> None user_profile_id = event['user_profile_id'] @@ -528,4 +522,4 @@ class EmbeddedBotWorker(QueueProcessingWorker): bot_handler.handle_message( message=message, bot_handler=self.get_bot_api_client(user_profile), - state_handler=self.get_state_handler()) + state_handler=None)