embedded bots: Add functional state handler.

This replaces the former non-functional StateHandler
stub with a dictionary-like state object. Accessing it will
will read and store strings in the BotUserStateData model.

Each bot has a limited state size. To enforce this limit while
keeping data updates efficient, StateHandler caches the expensive
query for getting a bot's total state size. Assignments to a key
then only need to fetch that entry's previous size, if any, and
compare it to the new entry's size.

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

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)