mirror of
				https://github.com/zulip/zulip.git
				synced 2025-11-04 05:53:43 +00:00 
			
		
		
		
	embedded bots: Add database config storage.
Storage limititations are only set on the value of a config entry, since this is the only user-accessible part of the schema. Keys are statically set by each embedded bot.
This commit is contained in:
		
							
								
								
									
										44
									
								
								zerver/lib/bot_config.py
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										44
									
								
								zerver/lib/bot_config.py
									
									
									
									
									
										Normal file
									
								
							@@ -0,0 +1,44 @@
 | 
			
		||||
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):
 | 
			
		||||
    # type: (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, key=None):
 | 
			
		||||
    # type: (UserProfile, Optional[Text]) -> 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, key, value):
 | 
			
		||||
    # type: (UserProfile, Text, 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()
 | 
			
		||||
@@ -10,6 +10,7 @@ from zerver.lib.actions import internal_send_message
 | 
			
		||||
from zerver.models import UserProfile
 | 
			
		||||
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_config import get_bot_config
 | 
			
		||||
from zerver.lib.integrations import EMBEDDED_BOTS
 | 
			
		||||
 | 
			
		||||
import configparser
 | 
			
		||||
@@ -99,11 +100,6 @@ class EmbeddedBotHandler:
 | 
			
		||||
                sender_email=message['sender_email'],
 | 
			
		||||
            ))
 | 
			
		||||
 | 
			
		||||
    def get_config_info(self, bot_name, section=None):
 | 
			
		||||
        # type: (str, Optional[str]) -> Dict[str, Any]
 | 
			
		||||
        conf_file_path = os.path.realpath(os.path.join(
 | 
			
		||||
            our_dir, '..', 'bots', bot_name, bot_name + '.conf'))
 | 
			
		||||
        section = section or bot_name
 | 
			
		||||
        config = configparser.ConfigParser()
 | 
			
		||||
        config.readfp(open(conf_file_path))  # type: ignore # likely typeshed issue
 | 
			
		||||
        return dict(config.items(section))
 | 
			
		||||
    def get_config_info(self):
 | 
			
		||||
        # type: () -> Dict[Text, Text]
 | 
			
		||||
        return get_bot_config(self.user_profile)
 | 
			
		||||
 
 | 
			
		||||
							
								
								
									
										30
									
								
								zerver/migrations/0118_botuserconfigdata.py
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										30
									
								
								zerver/migrations/0118_botuserconfigdata.py
									
									
									
									
									
										Normal file
									
								
							@@ -0,0 +1,30 @@
 | 
			
		||||
# -*- coding: utf-8 -*-
 | 
			
		||||
# Generated by Django 1.11.6 on 2017-11-01 19:12
 | 
			
		||||
from __future__ import unicode_literals
 | 
			
		||||
 | 
			
		||||
from django.conf import settings
 | 
			
		||||
from django.db import migrations, models
 | 
			
		||||
import django.db.models.deletion
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
class Migration(migrations.Migration):
 | 
			
		||||
 | 
			
		||||
    dependencies = [
 | 
			
		||||
        ('zerver', '0117_add_desc_to_user_group'),
 | 
			
		||||
    ]
 | 
			
		||||
 | 
			
		||||
    operations = [
 | 
			
		||||
        migrations.CreateModel(
 | 
			
		||||
            name='BotUserConfigData',
 | 
			
		||||
            fields=[
 | 
			
		||||
                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
 | 
			
		||||
                ('key', models.TextField(db_index=True)),
 | 
			
		||||
                ('value', models.TextField()),
 | 
			
		||||
                ('bot_profile', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to=settings.AUTH_USER_MODEL)),
 | 
			
		||||
            ],
 | 
			
		||||
        ),
 | 
			
		||||
        migrations.AlterUniqueTogether(
 | 
			
		||||
            name='botuserconfigdata',
 | 
			
		||||
            unique_together=set([('bot_profile', 'key')]),
 | 
			
		||||
        ),
 | 
			
		||||
    ]
 | 
			
		||||
@@ -2004,3 +2004,11 @@ class BotUserStateData(models.Model):
 | 
			
		||||
 | 
			
		||||
    class Meta:
 | 
			
		||||
        unique_together = ("bot_profile", "key")
 | 
			
		||||
 | 
			
		||||
class BotUserConfigData(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
 | 
			
		||||
 | 
			
		||||
    class Meta(object):
 | 
			
		||||
        unique_together = ("bot_profile", "key")
 | 
			
		||||
 
 | 
			
		||||
@@ -11,8 +11,9 @@ from zerver.lib.actions import (
 | 
			
		||||
    do_create_user,
 | 
			
		||||
    get_service_bot_events,
 | 
			
		||||
)
 | 
			
		||||
from zerver.lib.bot_lib import StateHandler
 | 
			
		||||
from zerver.lib.bot_lib import StateHandler, EmbeddedBotHandler
 | 
			
		||||
from zerver.lib.bot_storage import StateError
 | 
			
		||||
from zerver.lib.bot_config import set_bot_config, ConfigError
 | 
			
		||||
from zerver.lib.test_classes import ZulipTestCase
 | 
			
		||||
from zerver.models import (
 | 
			
		||||
    get_realm,
 | 
			
		||||
@@ -176,6 +177,42 @@ class TestServiceBotStateHandler(ZulipTestCase):
 | 
			
		||||
        self.assertTrue(storage.contains('another key'))
 | 
			
		||||
        self.assertRaises(StateError, lambda: storage.remove('some key'))
 | 
			
		||||
 | 
			
		||||
class TestServiceBotConfigHandler(ZulipTestCase):
 | 
			
		||||
    def setUp(self):
 | 
			
		||||
        # type: () -> None
 | 
			
		||||
        self.user_profile = self.example_user("othello")
 | 
			
		||||
        self.bot_profile = self.create_test_bot('embedded-bot@zulip.testserver', self.user_profile, 'Embedded bot',
 | 
			
		||||
                                                'embedded', UserProfile.EMBEDDED_BOT, service_name='helloworld')
 | 
			
		||||
        self.bot_handler = EmbeddedBotHandler(self.bot_profile)
 | 
			
		||||
 | 
			
		||||
    def test_basic_storage_and_retrieval(self):
 | 
			
		||||
        # type: () -> None
 | 
			
		||||
        config_dict = {"entry 1": "value 1", "entry 2": "value 2"}
 | 
			
		||||
        for key, value in config_dict.items():
 | 
			
		||||
            set_bot_config(self.bot_profile, key, value)
 | 
			
		||||
        self.assertEqual(self.bot_handler.get_config_info(), config_dict)
 | 
			
		||||
 | 
			
		||||
        config_update = {"entry 2": "new value", "entry 3": "value 3"}
 | 
			
		||||
        for key, value in config_update.items():
 | 
			
		||||
            set_bot_config(self.bot_profile, key, value)
 | 
			
		||||
        config_dict.update(config_update)
 | 
			
		||||
        self.assertEqual(self.bot_handler.get_config_info(), config_dict)
 | 
			
		||||
 | 
			
		||||
    @override_settings(BOT_CONFIG_SIZE_LIMIT=100)
 | 
			
		||||
    def test_config_entry_limit(self):
 | 
			
		||||
        # type: () -> None
 | 
			
		||||
        set_bot_config(self.bot_profile, "some key", 'x' * (settings.BOT_CONFIG_SIZE_LIMIT-8))
 | 
			
		||||
        self.assertRaisesMessage(ConfigError,
 | 
			
		||||
                                 "Cannot store configuration. Request would require 101 characters. "
 | 
			
		||||
                                 "The current configuration size limit is 100 characters.",
 | 
			
		||||
                                 lambda: set_bot_config(self.bot_profile, "some key", 'x' * (settings.BOT_CONFIG_SIZE_LIMIT-8+1)))
 | 
			
		||||
        set_bot_config(self.bot_profile, "some key", 'x' * (settings.BOT_CONFIG_SIZE_LIMIT-20))
 | 
			
		||||
        set_bot_config(self.bot_profile, "another key", 'x')
 | 
			
		||||
        self.assertRaisesMessage(ConfigError,
 | 
			
		||||
                                 "Cannot store configuration. Request would require 116 characters. "
 | 
			
		||||
                                 "The current configuration size limit is 100 characters.",
 | 
			
		||||
                                 lambda: set_bot_config(self.bot_profile, "yet another key", 'x'))
 | 
			
		||||
 | 
			
		||||
class TestServiceBotEventTriggers(ZulipTestCase):
 | 
			
		||||
 | 
			
		||||
    def setUp(self) -> None:
 | 
			
		||||
 
 | 
			
		||||
@@ -162,6 +162,8 @@ DEFAULT_SETTINGS = {
 | 
			
		||||
    # Max state storage per user
 | 
			
		||||
    # TODO: Add this to zproject/prod_settings_template.py once stateful bots are fully functional.
 | 
			
		||||
    'USER_STATE_SIZE_LIMIT': 10000000,
 | 
			
		||||
    # Max size of a single configuration entry of an embedded bot.
 | 
			
		||||
    'BOT_CONFIG_SIZE_LIMIT': 10000,
 | 
			
		||||
 | 
			
		||||
    # External service configuration
 | 
			
		||||
    'CAMO_URI': '',
 | 
			
		||||
 
 | 
			
		||||
		Reference in New Issue
	
	Block a user