zulip/zerver/lib/bot_config.py

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 BotUserConfigData, UserProfile

from typing import Text, Dict, Optional

class ConfigError(Exception):
    pass

def get_bot_config(bot_profile: UserProfile) -> Dict[Text, Text]:
    entries = BotUserConfigData.objects.filter(bot_profile=bot_profile)
    return {entry.key: entry.value for entry in entries}

def get_bot_config_size(bot_profile: UserProfile, key: Optional[Text]=None) -> int:
    if key is None:
        return BotUserConfigData.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(BotUserConfigData.objects.get(bot_profile=bot_profile, key=key).value)
        except BotUserConfigData.DoesNotExist:
            return 0

def set_bot_config(bot_profile: UserProfile, key: Text, value: Text) -> None:
    config_size_limit = settings.BOT_CONFIG_SIZE_LIMIT
    old_entry_size = get_bot_config_size(bot_profile, key)
    new_entry_size = len(key) + len(value)
    old_config_size = get_bot_config_size(bot_profile)
    new_config_size = old_config_size + (new_entry_size - old_entry_size)
    if new_config_size > config_size_limit:
        raise ConfigError("Cannot store configuration. Request would require {} characters. "
                          "The current configuration size limit is {} characters.".format(new_config_size,
                                                                                          config_size_limit))
    obj, created = BotUserConfigData.objects.get_or_create(bot_profile=bot_profile, key=key,
                                                           defaults={'value': value})
    if not created:
        obj.value = value
        obj.save()