mirror of
https://github.com/zulip/zulip.git
synced 2025-11-06 06:53:25 +00:00
60 lines
2.6 KiB
Python
60 lines
2.6 KiB
Python
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 typing import Text, Optional, List, Tuple
|
|
|
|
class StateError(Exception):
|
|
pass
|
|
|
|
def get_bot_state(bot_profile, key):
|
|
# type: (UserProfile, Text) -> Text
|
|
try:
|
|
return BotUserStateData.objects.get(bot_profile=bot_profile, key=key).value
|
|
except BotUserStateData.DoesNotExist:
|
|
raise StateError("Key does not exist.")
|
|
|
|
def get_bot_state_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
|
|
else:
|
|
try:
|
|
return len(key) + len(BotUserStateData.objects.get(bot_profile=bot_profile, key=key).value)
|
|
except BotUserStateData.DoesNotExist:
|
|
return 0
|
|
|
|
def set_bot_state(bot_profile, entries):
|
|
# type: (UserProfile, List[Tuple[str, str]]) -> None
|
|
state_size_limit = settings.USER_STATE_SIZE_LIMIT
|
|
state_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:
|
|
raise StateError("Request exceeds storage limit by {} characters. The limit is {} characters."
|
|
.format(new_state_size - state_size_limit, state_size_limit))
|
|
else:
|
|
for key, value in entries:
|
|
BotUserStateData.objects.update_or_create(bot_profile=bot_profile, key=key,
|
|
defaults={'value': value})
|
|
|
|
def remove_bot_state(bot_profile, keys):
|
|
# type: (UserProfile, List[Text]) -> None
|
|
queryset = BotUserStateData.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):
|
|
# type: (UserProfile, Text) -> bool
|
|
return BotUserStateData.objects.filter(bot_profile=bot_profile, key=key).exists()
|