embedded bots: Consistently use 'storage' instead of 'state.'

This commit is contained in:
derAnfaenger 2017-11-24 10:18:29 +01:00 committed by Tim Abbott
parent 9645c8b31f
commit c8a5ae753c
7 changed files with 98 additions and 79 deletions

View File

@ -9,8 +9,8 @@ import importlib
from zerver.lib.actions import internal_send_private_message, \ from zerver.lib.actions import internal_send_private_message, \
internal_send_stream_message, internal_send_huddle_message internal_send_stream_message, internal_send_huddle_message
from zerver.models import UserProfile, get_user from zerver.models import UserProfile, get_user
from zerver.lib.bot_storage import get_bot_state, set_bot_state, \ from zerver.lib.bot_storage import get_bot_storage, set_bot_storage, \
is_key_in_bot_state, get_bot_state_size, remove_bot_state is_key_in_bot_storage, get_bot_storage_size, remove_bot_storage
from zerver.lib.bot_config import get_bot_config from zerver.lib.bot_config import get_bot_config
from zerver.lib.integrations import EMBEDDED_BOTS from zerver.lib.integrations import EMBEDDED_BOTS
@ -38,7 +38,7 @@ def get_bot_handler(service_name: str) -> Any:
class StateHandler: class StateHandler:
state_size_limit = 10000000 # type: int # TODO: Store this in the server configuration model. storage_size_limit = 10000000 # type: int # TODO: Store this in the server configuration model.
def __init__(self, user_profile: UserProfile) -> None: def __init__(self, user_profile: UserProfile) -> None:
self.user_profile = user_profile self.user_profile = user_profile
@ -46,16 +46,16 @@ class StateHandler:
self.demarshal = lambda obj: json.loads(obj) self.demarshal = lambda obj: json.loads(obj)
def get(self, key: Text) -> Text: def get(self, key: Text) -> Text:
return self.demarshal(get_bot_state(self.user_profile, key)) return self.demarshal(get_bot_storage(self.user_profile, key))
def put(self, key: Text, value: Text) -> None: def put(self, key: Text, value: Text) -> None:
set_bot_state(self.user_profile, [(key, self.marshal(value))]) set_bot_storage(self.user_profile, [(key, self.marshal(value))])
def remove(self, key: Text) -> None: def remove(self, key: Text) -> None:
remove_bot_state(self.user_profile, [key]) remove_bot_storage(self.user_profile, [key])
def contains(self, key: Text) -> bool: def contains(self, key: Text) -> bool:
return is_key_in_bot_state(self.user_profile, key) return is_key_in_bot_storage(self.user_profile, key)
class EmbeddedBotHandler: class EmbeddedBotHandler:
def __init__(self, user_profile: UserProfile) -> None: def __init__(self, user_profile: UserProfile) -> None:

View File

@ -2,62 +2,62 @@ from django.conf import settings
from django.db.models import Sum from django.db.models import Sum
from django.db.models.query import F from django.db.models.query import F
from django.db.models.functions import Length from django.db.models.functions import Length
from zerver.models import BotUserStateData, UserProfile, Length from zerver.models import BotStorageData, UserProfile, Length
from typing import Text, Optional, List, Tuple from typing import Text, Optional, List, Tuple
class StateError(Exception): class StateError(Exception):
pass pass
def get_bot_state(bot_profile, key): def get_bot_storage(bot_profile, key):
# type: (UserProfile, Text) -> Text # type: (UserProfile, Text) -> Text
try: try:
return BotUserStateData.objects.get(bot_profile=bot_profile, key=key).value return BotStorageData.objects.get(bot_profile=bot_profile, key=key).value
except BotUserStateData.DoesNotExist: except BotStorageData.DoesNotExist:
raise StateError("Key does not exist.") raise StateError("Key does not exist.")
def get_bot_state_size(bot_profile, key=None): def get_bot_storage_size(bot_profile, key=None):
# type: (UserProfile, Optional[Text]) -> int # type: (UserProfile, Optional[Text]) -> int
if key is None: if key is None:
return BotUserStateData.objects.filter(bot_profile=bot_profile) \ return BotStorageData.objects.filter(bot_profile=bot_profile) \
.annotate(key_size=Length('key'), value_size=Length('value')) \ .annotate(key_size=Length('key'), value_size=Length('value')) \
.aggregate(sum=Sum(F('key_size')+F('value_size')))['sum'] or 0 .aggregate(sum=Sum(F('key_size')+F('value_size')))['sum'] or 0
else: else:
try: try:
return len(key) + len(BotUserStateData.objects.get(bot_profile=bot_profile, key=key).value) return len(key) + len(BotStorageData.objects.get(bot_profile=bot_profile, key=key).value)
except BotUserStateData.DoesNotExist: except BotStorageData.DoesNotExist:
return 0 return 0
def set_bot_state(bot_profile, entries): def set_bot_storage(bot_profile, entries):
# type: (UserProfile, List[Tuple[str, str]]) -> None # type: (UserProfile, List[Tuple[str, str]]) -> None
state_size_limit = settings.USER_STATE_SIZE_LIMIT storage_size_limit = settings.USER_STATE_SIZE_LIMIT
state_size_difference = 0 storage_size_difference = 0
for key, value in entries: for key, value in entries:
if type(key) is not str: if type(key) is not str:
raise StateError("Key type is {}, but should be str.".format(type(key))) raise StateError("Key type is {}, but should be str.".format(type(key)))
if type(value) is not str: if type(value) is not str:
raise StateError("Value type is {}, but should be str.".format(type(value))) raise StateError("Value type is {}, but should be str.".format(type(value)))
state_size_difference += (len(key) + len(value)) - get_bot_state_size(bot_profile, key) storage_size_difference += (len(key) + len(value)) - get_bot_storage_size(bot_profile, key)
new_state_size = get_bot_state_size(bot_profile) + state_size_difference new_storage_size = get_bot_storage_size(bot_profile) + storage_size_difference
if new_state_size > state_size_limit: if new_storage_size > storage_size_limit:
raise StateError("Request exceeds storage limit by {} characters. The limit is {} characters." raise StateError("Request exceeds storage limit by {} characters. The limit is {} characters."
.format(new_state_size - state_size_limit, state_size_limit)) .format(new_storage_size - storage_size_limit, storage_size_limit))
else: else:
for key, value in entries: for key, value in entries:
BotUserStateData.objects.update_or_create(bot_profile=bot_profile, key=key, BotStorageData.objects.update_or_create(bot_profile=bot_profile, key=key,
defaults={'value': value}) defaults={'value': value})
def remove_bot_state(bot_profile, keys): def remove_bot_storage(bot_profile, keys):
# type: (UserProfile, List[Text]) -> None # type: (UserProfile, List[Text]) -> None
queryset = BotUserStateData.objects.filter(bot_profile=bot_profile, key__in=keys) queryset = BotStorageData.objects.filter(bot_profile=bot_profile, key__in=keys)
if len(queryset) < len(keys): if len(queryset) < len(keys):
raise StateError("Key does not exist.") raise StateError("Key does not exist.")
queryset.delete() queryset.delete()
def is_key_in_bot_state(bot_profile, key): def is_key_in_bot_storage(bot_profile, key):
# type: (UserProfile, Text) -> bool # type: (UserProfile, Text) -> bool
return BotUserStateData.objects.filter(bot_profile=bot_profile, key=key).exists() return BotStorageData.objects.filter(bot_profile=bot_profile, key=key).exists()
def get_keys_in_bot_state(bot_profile): def get_keys_in_bot_storage(bot_profile):
# type: (UserProfile) -> List[Text] # type: (UserProfile) -> List[Text]
return list(BotUserStateData.objects.filter(bot_profile=bot_profile).values_list('key', flat=True)) return list(BotStorageData.objects.filter(bot_profile=bot_profile).values_list('key', flat=True))

View File

@ -0,0 +1,19 @@
# -*- coding: utf-8 -*-
# Generated by Django 1.11.6 on 2017-11-24 09:10
from __future__ import unicode_literals
from django.db import migrations
class Migration(migrations.Migration):
dependencies = [
('zerver', '0121_realm_signup_notifications_stream'),
]
operations = [
migrations.RenameModel(
old_name='BotUserStateData',
new_name='BotStorageData',
),
]

View File

@ -2005,7 +2005,7 @@ def get_service_profile(user_profile_id, service_name):
return Service.objects.get(user_profile__id=user_profile_id, name=service_name) return Service.objects.get(user_profile__id=user_profile_id, name=service_name)
class BotUserStateData(models.Model): class BotStorageData(models.Model):
bot_profile = models.ForeignKey(UserProfile, on_delete=CASCADE) # type: UserProfile bot_profile = models.ForeignKey(UserProfile, on_delete=CASCADE) # type: UserProfile
key = models.TextField(db_index=True) # type: Text key = models.TextField(db_index=True) # type: Text
value = models.TextField() # type: Text value = models.TextField() # type: Text

View File

@ -17,7 +17,7 @@ from zerver.lib.bot_config import set_bot_config, ConfigError
from zerver.lib.test_classes import ZulipTestCase from zerver.lib.test_classes import ZulipTestCase
from zerver.models import ( from zerver.models import (
get_realm, get_realm,
BotUserStateData, BotStorageData,
UserProfile, UserProfile,
Recipient, Recipient,
) )
@ -144,7 +144,7 @@ class TestServiceBotStateHandler(ZulipTestCase):
with self.assertRaisesMessage(StateError, "Key type is <class 'dict'>, but should be str."): with self.assertRaisesMessage(StateError, "Key type is <class 'dict'>, but should be str."):
storage.put(serializable_obj, 'some value') # type: ignore # We intend to test an invalid type. storage.put(serializable_obj, 'some value') # type: ignore # We intend to test an invalid type.
# Reduce maximal state size for faster test string construction. # Reduce maximal storage size for faster test string construction.
@override_settings(USER_STATE_SIZE_LIMIT=100) @override_settings(USER_STATE_SIZE_LIMIT=100)
def test_storage_limit(self) -> None: def test_storage_limit(self) -> None:
storage = StateHandler(self.bot_profile) storage = StateHandler(self.bot_profile)
@ -182,80 +182,80 @@ class TestServiceBotStateHandler(ZulipTestCase):
# Store some data. # Store some data.
initial_dict = {'key 1': 'value 1', 'key 2': 'value 2', 'key 3': 'value 3'} initial_dict = {'key 1': 'value 1', 'key 2': 'value 2', 'key 3': 'value 3'}
params = { params = {
'state': ujson.dumps(initial_dict) 'storage': ujson.dumps(initial_dict)
} }
result = self.client_put('/json/user_state', params) result = self.client_put('/json/bot_storage', params)
self.assert_json_success(result) self.assert_json_success(result)
# Assert the stored data for some keys. # Assert the stored data for some keys.
params = { params = {
'keys': ujson.dumps(['key 1', 'key 3']) 'keys': ujson.dumps(['key 1', 'key 3'])
} }
result = self.client_get('/json/user_state', params) result = self.client_get('/json/bot_storage', params)
self.assert_json_success(result) self.assert_json_success(result)
self.assertEqual(result.json()['state'], {'key 3': 'value 3', 'key 1': 'value 1'}) self.assertEqual(result.json()['storage'], {'key 3': 'value 3', 'key 1': 'value 1'})
# Assert the stored data for all keys. # Assert the stored data for all keys.
result = self.client_get('/json/user_state') result = self.client_get('/json/bot_storage')
self.assert_json_success(result) self.assert_json_success(result)
self.assertEqual(result.json()['state'], initial_dict) self.assertEqual(result.json()['storage'], initial_dict)
# Store some more data; update an entry and store a new entry # Store some more data; update an entry and store a new entry
dict_update = {'key 1': 'new value', 'key 4': 'value 4'} dict_update = {'key 1': 'new value', 'key 4': 'value 4'}
params = { params = {
'state': ujson.dumps(dict_update) 'storage': ujson.dumps(dict_update)
} }
result = self.client_put('/json/user_state', params) result = self.client_put('/json/bot_storage', params)
self.assert_json_success(result) self.assert_json_success(result)
# Assert the data was updated. # Assert the data was updated.
updated_dict = initial_dict.copy() updated_dict = initial_dict.copy()
updated_dict.update(dict_update) updated_dict.update(dict_update)
result = self.client_get('/json/user_state') result = self.client_get('/json/bot_storage')
self.assert_json_success(result) self.assert_json_success(result)
self.assertEqual(result.json()['state'], updated_dict) self.assertEqual(result.json()['storage'], updated_dict)
# Assert errors on invalid requests. # Assert errors on invalid requests.
params = { # type: ignore # Ignore 'incompatible type "str": "List[str]"; expected "str": "str"' for testing params = { # type: ignore # Ignore 'incompatible type "str": "List[str]"; expected "str": "str"' for testing
'keys': ["This is a list, but should be a serialized string."] 'keys': ["This is a list, but should be a serialized string."]
} }
result = self.client_get('/json/user_state', params) result = self.client_get('/json/bot_storage', params)
self.assert_json_error(result, 'Argument "keys" is not valid JSON.') self.assert_json_error(result, 'Argument "keys" is not valid JSON.')
params = { params = {
'keys': ujson.dumps(["key 1", "nonexistent key"]) 'keys': ujson.dumps(["key 1", "nonexistent key"])
} }
result = self.client_get('/json/user_state', params) result = self.client_get('/json/bot_storage', params)
self.assert_json_error(result, "Key does not exist.") self.assert_json_error(result, "Key does not exist.")
params = { params = {
'state': ujson.dumps({'foo': [1, 2, 3]}) 'storage': ujson.dumps({'foo': [1, 2, 3]})
} }
result = self.client_put('/json/user_state', params) result = self.client_put('/json/bot_storage', params)
self.assert_json_error(result, "Value type is <class 'list'>, but should be str.") self.assert_json_error(result, "Value type is <class 'list'>, but should be str.")
# Remove some entries. # Remove some entries.
keys_to_remove = ['key 1', 'key 2'] keys_to_remove = ['key 1', 'key 2']
params = { params = {
'keys': ujson.dumps(keys_to_remove) 'keys': ujson.dumps(keys_to_remove)
} }
result = self.client_delete('/json/user_state', params) result = self.client_delete('/json/bot_storage', params)
self.assert_json_success(result) self.assert_json_success(result)
# Assert the entries were removed. # Assert the entries were removed.
for key in keys_to_remove: for key in keys_to_remove:
updated_dict.pop(key) updated_dict.pop(key)
result = self.client_get('/json/user_state') result = self.client_get('/json/bot_storage')
self.assert_json_success(result) self.assert_json_success(result)
self.assertEqual(result.json()['state'], updated_dict) self.assertEqual(result.json()['storage'], updated_dict)
# Try to remove an existing and a nonexistent key. # Try to remove an existing and a nonexistent key.
params = { params = {
'keys': ujson.dumps(['key 3', 'nonexistent key']) 'keys': ujson.dumps(['key 3', 'nonexistent key'])
} }
result = self.client_delete('/json/user_state', params) result = self.client_delete('/json/bot_storage', params)
self.assert_json_error(result, "Key does not exist.") self.assert_json_error(result, "Key does not exist.")
# Assert an error has been thrown and no entries were removed. # Assert an error has been thrown and no entries were removed.
result = self.client_get('/json/user_state') result = self.client_get('/json/bot_storage')
self.assert_json_success(result) self.assert_json_success(result)
self.assertEqual(result.json()['state'], updated_dict) self.assertEqual(result.json()['storage'], updated_dict)
# Remove the entire state. # Remove the entire storage.
result = self.client_delete('/json/user_state') result = self.client_delete('/json/bot_storage')
self.assert_json_success(result) self.assert_json_success(result)
# Assert the entire state has been removed. # Assert the entire storage has been removed.
result = self.client_get('/json/user_state') result = self.client_get('/json/bot_storage')
self.assert_json_success(result) self.assert_json_success(result)
self.assertEqual(result.json()['state'], {}) self.assertEqual(result.json()['storage'], {})
class TestServiceBotConfigHandler(ZulipTestCase): class TestServiceBotConfigHandler(ZulipTestCase):
def setUp(self) -> None: def setUp(self) -> None:

View File

@ -1,11 +1,11 @@
from django.http import HttpRequest, HttpResponse from django.http import HttpRequest, HttpResponse
from django.utils.translation import ugettext as _ from django.utils.translation import ugettext as _
from zerver.lib.bot_storage import ( from zerver.lib.bot_storage import (
get_bot_state, get_bot_storage,
set_bot_state, set_bot_storage,
remove_bot_state, remove_bot_storage,
get_keys_in_bot_state, get_keys_in_bot_storage,
is_key_in_bot_state, is_key_in_bot_storage,
StateError, StateError,
) )
from zerver.decorator import has_request_variables, REQ from zerver.decorator import has_request_variables, REQ
@ -16,30 +16,30 @@ from zerver.models import UserProfile
from typing import Dict, List, Optional from typing import Dict, List, Optional
@has_request_variables @has_request_variables
def update_state(request, user_profile, state=REQ(validator=check_dict([]))): def update_storage(request, user_profile, storage=REQ(validator=check_dict([]))):
# type: (HttpRequest, UserProfile, Optional[Dict[str, str]]) -> HttpResponse # type: (HttpRequest, UserProfile, Optional[Dict[str, str]]) -> HttpResponse
try: try:
set_bot_state(user_profile, list(state.items())) set_bot_storage(user_profile, list(storage.items()))
except StateError as e: except StateError as e:
return json_error(str(e)) return json_error(str(e))
return json_success() return json_success()
@has_request_variables @has_request_variables
def get_state(request, user_profile, keys=REQ(validator=check_list(check_string), default=None)): def get_storage(request, user_profile, keys=REQ(validator=check_list(check_string), default=None)):
# type: (HttpRequest, UserProfile, Optional[List[str]]) -> HttpResponse # type: (HttpRequest, UserProfile, Optional[List[str]]) -> HttpResponse
keys = keys or get_keys_in_bot_state(user_profile) keys = keys or get_keys_in_bot_storage(user_profile)
try: try:
state = {key: get_bot_state(user_profile, key) for key in keys} storage = {key: get_bot_storage(user_profile, key) for key in keys}
except StateError as e: except StateError as e:
return json_error(str(e)) return json_error(str(e))
return json_success({'state': state}) return json_success({'storage': storage})
@has_request_variables @has_request_variables
def remove_state(request, user_profile, keys=REQ(validator=check_list(check_string), default=None)): def remove_storage(request, user_profile, keys=REQ(validator=check_list(check_string), default=None)):
# type: (HttpRequest, UserProfile, Optional[List[str]]) -> HttpResponse # type: (HttpRequest, UserProfile, Optional[List[str]]) -> HttpResponse
keys = keys or get_keys_in_bot_state(user_profile) keys = keys or get_keys_in_bot_storage(user_profile)
try: try:
remove_bot_state(user_profile, keys) remove_bot_storage(user_profile, keys)
except StateError as e: except StateError as e:
return json_error(str(e)) return json_error(str(e))
return json_success() return json_success()

View File

@ -200,11 +200,11 @@ v1_api_and_json_patterns = [
url(r'^user_uploads$', rest_dispatch, url(r'^user_uploads$', rest_dispatch,
{'POST': 'zerver.views.upload.upload_file_backend'}), {'POST': 'zerver.views.upload.upload_file_backend'}),
# user_state -> zerver.views.state # bot_storage -> zerver.views.storage
url(r'^user_state$', rest_dispatch, url(r'^bot_storage$', rest_dispatch,
{'PUT': 'zerver.views.state.update_state', {'PUT': 'zerver.views.storage.update_storage',
'GET': 'zerver.views.state.get_state', 'GET': 'zerver.views.storage.get_storage',
'DELETE': 'zerver.views.state.remove_state'}), 'DELETE': 'zerver.views.storage.remove_storage'}),
# users/me -> zerver.views # users/me -> zerver.views
url(r'^users/me$', rest_dispatch, url(r'^users/me$', rest_dispatch,