mirror of https://github.com/zulip/zulip.git
embedded bots: Consistently use 'storage' instead of 'state.'
This commit is contained in:
derAnfaenger
Tim Abbott
parent
9645c8b31f
commit
c8a5ae753c
|
|
@ -9,8 +9,8 @@ import importlib
|
|||
from zerver.lib.actions import internal_send_private_message, \
|
||||
internal_send_stream_message, internal_send_huddle_message
|
||||
from zerver.models import UserProfile, get_user
|
||||
from zerver.lib.bot_storage import get_bot_state, set_bot_state, \
|
||||
is_key_in_bot_state, get_bot_state_size, remove_bot_state
|
||||
from zerver.lib.bot_storage import get_bot_storage, set_bot_storage, \
|
||||
is_key_in_bot_storage, get_bot_storage_size, remove_bot_storage
|
||||
from zerver.lib.bot_config import get_bot_config
|
||||
from zerver.lib.integrations import EMBEDDED_BOTS
|
||||
|
||||
|
|
@ -38,7 +38,7 @@ def get_bot_handler(service_name: str) -> Any:
|
|||
|
||||
|
||||
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:
|
||||
self.user_profile = user_profile
|
||||
|
|
@ -46,16 +46,16 @@ class StateHandler:
|
|||
self.demarshal = lambda obj: json.loads(obj)
|
||||
|
||||
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:
|
||||
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:
|
||||
remove_bot_state(self.user_profile, [key])
|
||||
remove_bot_storage(self.user_profile, [key])
|
||||
|
||||
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:
|
||||
def __init__(self, user_profile: UserProfile) -> None:
|
||||
|
|
|
|||
|
|
@ -2,62 +2,62 @@ from django.conf import settings
|
|||
from django.db.models import Sum
|
||||
from django.db.models.query import F
|
||||
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
|
||||
|
||||
class StateError(Exception):
|
||||
pass
|
||||
|
||||
def get_bot_state(bot_profile, key):
|
||||
def get_bot_storage(bot_profile, key):
|
||||
# type: (UserProfile, Text) -> Text
|
||||
try:
|
||||
return BotUserStateData.objects.get(bot_profile=bot_profile, key=key).value
|
||||
except BotUserStateData.DoesNotExist:
|
||||
return BotStorageData.objects.get(bot_profile=bot_profile, key=key).value
|
||||
except BotStorageData.DoesNotExist:
|
||||
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
|
||||
if key is None:
|
||||
return BotUserStateData.objects.filter(bot_profile=bot_profile) \
|
||||
.annotate(key_size=Length('key'), value_size=Length('value')) \
|
||||
.aggregate(sum=Sum(F('key_size')+F('value_size')))['sum'] or 0
|
||||
return BotStorageData.objects.filter(bot_profile=bot_profile) \
|
||||
.annotate(key_size=Length('key'), value_size=Length('value')) \
|
||||
.aggregate(sum=Sum(F('key_size')+F('value_size')))['sum'] or 0
|
||||
else:
|
||||
try:
|
||||
return len(key) + len(BotUserStateData.objects.get(bot_profile=bot_profile, key=key).value)
|
||||
except BotUserStateData.DoesNotExist:
|
||||
return len(key) + len(BotStorageData.objects.get(bot_profile=bot_profile, key=key).value)
|
||||
except BotStorageData.DoesNotExist:
|
||||
return 0
|
||||
|
||||
def set_bot_state(bot_profile, entries):
|
||||
def set_bot_storage(bot_profile, entries):
|
||||
# type: (UserProfile, List[Tuple[str, str]]) -> None
|
||||
state_size_limit = settings.USER_STATE_SIZE_LIMIT
|
||||
state_size_difference = 0
|
||||
storage_size_limit = settings.USER_STATE_SIZE_LIMIT
|
||||
storage_size_difference = 0
|
||||
for key, value in entries:
|
||||
if type(key) is not str:
|
||||
raise StateError("Key type is {}, but should be str.".format(type(key)))
|
||||
if type(value) is not str:
|
||||
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)
|
||||
new_state_size = get_bot_state_size(bot_profile) + state_size_difference
|
||||
if new_state_size > state_size_limit:
|
||||
storage_size_difference += (len(key) + len(value)) - get_bot_storage_size(bot_profile, key)
|
||||
new_storage_size = get_bot_storage_size(bot_profile) + storage_size_difference
|
||||
if new_storage_size > storage_size_limit:
|
||||
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:
|
||||
for key, value in entries:
|
||||
BotUserStateData.objects.update_or_create(bot_profile=bot_profile, key=key,
|
||||
defaults={'value': value})
|
||||
BotStorageData.objects.update_or_create(bot_profile=bot_profile, key=key,
|
||||
defaults={'value': value})
|
||||
|
||||
def remove_bot_state(bot_profile, keys):
|
||||
def remove_bot_storage(bot_profile, keys):
|
||||
# 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):
|
||||
raise StateError("Key does not exist.")
|
||||
queryset.delete()
|
||||
|
||||
def is_key_in_bot_state(bot_profile, key):
|
||||
def is_key_in_bot_storage(bot_profile, key):
|
||||
# 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]
|
||||
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))
|
||||
|
|
|
|||
|
|
@ -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',
|
||||
),
|
||||
]
|
||||
|
|
@ -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)
|
||||
|
||||
|
||||
class BotUserStateData(models.Model):
|
||||
class BotStorageData(models.Model):
|
||||
bot_profile = models.ForeignKey(UserProfile, on_delete=CASCADE) # type: UserProfile
|
||||
key = models.TextField(db_index=True) # type: Text
|
||||
value = models.TextField() # type: Text
|
||||
|
|
|
|||
|
|
@ -17,7 +17,7 @@ from zerver.lib.bot_config import set_bot_config, ConfigError
|
|||
from zerver.lib.test_classes import ZulipTestCase
|
||||
from zerver.models import (
|
||||
get_realm,
|
||||
BotUserStateData,
|
||||
BotStorageData,
|
||||
UserProfile,
|
||||
Recipient,
|
||||
)
|
||||
|
|
@ -144,7 +144,7 @@ class TestServiceBotStateHandler(ZulipTestCase):
|
|||
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.
|
||||
|
||||
# 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)
|
||||
def test_storage_limit(self) -> None:
|
||||
storage = StateHandler(self.bot_profile)
|
||||
|
|
@ -182,80 +182,80 @@ class TestServiceBotStateHandler(ZulipTestCase):
|
|||
# Store some data.
|
||||
initial_dict = {'key 1': 'value 1', 'key 2': 'value 2', 'key 3': 'value 3'}
|
||||
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)
|
||||
# Assert the stored data for some keys.
|
||||
params = {
|
||||
'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.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.
|
||||
result = self.client_get('/json/user_state')
|
||||
result = self.client_get('/json/bot_storage')
|
||||
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
|
||||
dict_update = {'key 1': 'new value', 'key 4': 'value 4'}
|
||||
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)
|
||||
# Assert the data was updated.
|
||||
updated_dict = initial_dict.copy()
|
||||
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.assertEqual(result.json()['state'], updated_dict)
|
||||
self.assertEqual(result.json()['storage'], updated_dict)
|
||||
# Assert errors on invalid requests.
|
||||
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."]
|
||||
}
|
||||
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.')
|
||||
params = {
|
||||
'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.")
|
||||
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.")
|
||||
# Remove some entries.
|
||||
keys_to_remove = ['key 1', 'key 2']
|
||||
params = {
|
||||
'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)
|
||||
# Assert the entries were removed.
|
||||
for key in keys_to_remove:
|
||||
updated_dict.pop(key)
|
||||
result = self.client_get('/json/user_state')
|
||||
result = self.client_get('/json/bot_storage')
|
||||
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.
|
||||
params = {
|
||||
'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.")
|
||||
# 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.assertEqual(result.json()['state'], updated_dict)
|
||||
# Remove the entire state.
|
||||
result = self.client_delete('/json/user_state')
|
||||
self.assertEqual(result.json()['storage'], updated_dict)
|
||||
# Remove the entire storage.
|
||||
result = self.client_delete('/json/bot_storage')
|
||||
self.assert_json_success(result)
|
||||
# Assert the entire state has been removed.
|
||||
result = self.client_get('/json/user_state')
|
||||
# Assert the entire storage has been removed.
|
||||
result = self.client_get('/json/bot_storage')
|
||||
self.assert_json_success(result)
|
||||
self.assertEqual(result.json()['state'], {})
|
||||
self.assertEqual(result.json()['storage'], {})
|
||||
|
||||
class TestServiceBotConfigHandler(ZulipTestCase):
|
||||
def setUp(self) -> None:
|
||||
|
|
|
|||
|
|
@ -1,11 +1,11 @@
|
|||
from django.http import HttpRequest, HttpResponse
|
||||
from django.utils.translation import ugettext as _
|
||||
from zerver.lib.bot_storage import (
|
||||
get_bot_state,
|
||||
set_bot_state,
|
||||
remove_bot_state,
|
||||
get_keys_in_bot_state,
|
||||
is_key_in_bot_state,
|
||||
get_bot_storage,
|
||||
set_bot_storage,
|
||||
remove_bot_storage,
|
||||
get_keys_in_bot_storage,
|
||||
is_key_in_bot_storage,
|
||||
StateError,
|
||||
)
|
||||
from zerver.decorator import has_request_variables, REQ
|
||||
|
|
@ -16,30 +16,30 @@ from zerver.models import UserProfile
|
|||
from typing import Dict, List, Optional
|
||||
|
||||
@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
|
||||
try:
|
||||
set_bot_state(user_profile, list(state.items()))
|
||||
set_bot_storage(user_profile, list(storage.items()))
|
||||
except StateError as e:
|
||||
return json_error(str(e))
|
||||
return json_success()
|
||||
|
||||
@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
|
||||
keys = keys or get_keys_in_bot_state(user_profile)
|
||||
keys = keys or get_keys_in_bot_storage(user_profile)
|
||||
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:
|
||||
return json_error(str(e))
|
||||
return json_success({'state': state})
|
||||
return json_success({'storage': storage})
|
||||
|
||||
@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
|
||||
keys = keys or get_keys_in_bot_state(user_profile)
|
||||
keys = keys or get_keys_in_bot_storage(user_profile)
|
||||
try:
|
||||
remove_bot_state(user_profile, keys)
|
||||
remove_bot_storage(user_profile, keys)
|
||||
except StateError as e:
|
||||
return json_error(str(e))
|
||||
return json_success()
|
||||
|
|
@ -200,11 +200,11 @@ v1_api_and_json_patterns = [
|
|||
url(r'^user_uploads$', rest_dispatch,
|
||||
{'POST': 'zerver.views.upload.upload_file_backend'}),
|
||||
|
||||
# user_state -> zerver.views.state
|
||||
url(r'^user_state$', rest_dispatch,
|
||||
{'PUT': 'zerver.views.state.update_state',
|
||||
'GET': 'zerver.views.state.get_state',
|
||||
'DELETE': 'zerver.views.state.remove_state'}),
|
||||
# bot_storage -> zerver.views.storage
|
||||
url(r'^bot_storage$', rest_dispatch,
|
||||
{'PUT': 'zerver.views.storage.update_storage',
|
||||
'GET': 'zerver.views.storage.get_storage',
|
||||
'DELETE': 'zerver.views.storage.remove_storage'}),
|
||||
|
||||
# users/me -> zerver.views
|
||||
url(r'^users/me$', rest_dispatch,
|
||||
|
|
|
|||
Loading…
Reference in New Issue