zulip/zerver/views/storage.py

from typing import Dict, List, Optional

from django.http import HttpRequest, HttpResponse

from zerver.decorator import REQ, has_request_variables
from zerver.lib.bot_storage import (
    StateError,
    get_bot_storage,
    get_keys_in_bot_storage,
    remove_bot_storage,
    set_bot_storage,
)
from zerver.lib.response import json_error, json_success
from zerver.lib.validator import check_dict, check_list, check_string
from zerver.models import UserProfile


@has_request_variables
def update_storage(request: HttpRequest, user_profile: UserProfile,
                   storage: Dict[str, str]=REQ(validator=check_dict([]))) -> HttpResponse:
    try:
        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_storage(
        request: HttpRequest,
        user_profile: UserProfile,
        keys: Optional[List[str]]=REQ(validator=check_list(check_string), default=None),
) -> HttpResponse:
    keys = keys or get_keys_in_bot_storage(user_profile)
    try:
        storage = {key: get_bot_storage(user_profile, key) for key in keys}
    except StateError as e:
        return json_error(str(e))
    return json_success({'storage': storage})

@has_request_variables
def remove_storage(
        request: HttpRequest,
        user_profile: UserProfile,
        keys: Optional[List[str]]=REQ(validator=check_list(check_string), default=None),
) -> HttpResponse:
    keys = keys or get_keys_in_bot_storage(user_profile)
    try:
        remove_bot_storage(user_profile, keys)
    except StateError as e:
        return json_error(str(e))
    return json_success()