upstream-zulip/zerver/lib/bot_config.py

import configparser
import importlib
import os
from collections import defaultdict

from django.conf import settings
from django.db.models import F, Sum
from django.db.models.functions import Length

from zerver.models import BotConfigData, UserProfile


class ConfigError(Exception):
    pass


def get_bot_config(bot_profile: UserProfile) -> dict[str, str]:
    entries = BotConfigData.objects.filter(bot_profile=bot_profile)
    if not entries:
        raise ConfigError("No config data available.")
    return {entry.key: entry.value for entry in entries}


def get_bot_configs(bot_profile_ids: list[int]) -> dict[int, dict[str, str]]:
    if not bot_profile_ids:
        return {}
    entries = BotConfigData.objects.filter(bot_profile_id__in=bot_profile_ids)
    entries_by_uid: dict[int, dict[str, str]] = defaultdict(dict)
    for entry in entries:
        entries_by_uid[entry.bot_profile_id].update({entry.key: entry.value})
    return entries_by_uid


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


def set_bot_config(bot_profile: UserProfile, key: str, value: str) -> 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(
            f"Cannot store configuration. Request would require {new_config_size} characters. "
            f"The current configuration size limit is {config_size_limit} characters."
        )
    obj, created = BotConfigData.objects.get_or_create(
        bot_profile=bot_profile, key=key, defaults={"value": value}
    )
    if not created:
        obj.value = value
        obj.save()


def load_bot_config_template(bot: str) -> dict[str, str]:
    bot_module_name = f"zulip_bots.bots.{bot}"
    bot_module = importlib.import_module(bot_module_name)
    assert bot_module.__file__ is not None
    bot_module_path = os.path.dirname(bot_module.__file__)
    config_path = os.path.join(bot_module_path, f"{bot}.conf")
    if os.path.isfile(config_path):
        config = configparser.ConfigParser()
        with open(config_path) as conf:
            config.read_file(conf)
        return dict(config.items(bot))
    else:
        return {}