summit/backend/venv/lib/python3.12/site-packages/redis/commands/policies.py

from abc import ABC, abstractmethod
from typing import Optional

from redis._parsers.commands import (
    CommandPolicies,
    CommandsParser,
    PolicyRecords,
    RequestPolicy,
    ResponsePolicy,
)

STATIC_POLICIES: PolicyRecords = {
    "ft": {
        "explaincli": CommandPolicies(
            request_policy=RequestPolicy.DEFAULT_KEYLESS,
            response_policy=ResponsePolicy.DEFAULT_KEYLESS,
        ),
        "suglen": CommandPolicies(
            request_policy=RequestPolicy.DEFAULT_KEYED,
            response_policy=ResponsePolicy.DEFAULT_KEYED,
        ),
        "profile": CommandPolicies(
            request_policy=RequestPolicy.DEFAULT_KEYLESS,
            response_policy=ResponsePolicy.DEFAULT_KEYLESS,
        ),
        "dropindex": CommandPolicies(
            request_policy=RequestPolicy.DEFAULT_KEYLESS,
            response_policy=ResponsePolicy.DEFAULT_KEYLESS,
        ),
        "aliasupdate": CommandPolicies(
            request_policy=RequestPolicy.DEFAULT_KEYLESS,
            response_policy=ResponsePolicy.DEFAULT_KEYLESS,
        ),
        "alter": CommandPolicies(
            request_policy=RequestPolicy.DEFAULT_KEYLESS,
            response_policy=ResponsePolicy.DEFAULT_KEYLESS,
        ),
        "aggregate": CommandPolicies(
            request_policy=RequestPolicy.DEFAULT_KEYLESS,
            response_policy=ResponsePolicy.DEFAULT_KEYLESS,
        ),
        "syndump": CommandPolicies(
            request_policy=RequestPolicy.DEFAULT_KEYLESS,
            response_policy=ResponsePolicy.DEFAULT_KEYLESS,
        ),
        "create": CommandPolicies(
            request_policy=RequestPolicy.DEFAULT_KEYLESS,
            response_policy=ResponsePolicy.DEFAULT_KEYLESS,
        ),
        "explain": CommandPolicies(
            request_policy=RequestPolicy.DEFAULT_KEYLESS,
            response_policy=ResponsePolicy.DEFAULT_KEYLESS,
        ),
        "sugget": CommandPolicies(
            request_policy=RequestPolicy.DEFAULT_KEYED,
            response_policy=ResponsePolicy.DEFAULT_KEYED,
        ),
        "dictdel": CommandPolicies(
            request_policy=RequestPolicy.DEFAULT_KEYLESS,
            response_policy=ResponsePolicy.DEFAULT_KEYLESS,
        ),
        "aliasadd": CommandPolicies(
            request_policy=RequestPolicy.DEFAULT_KEYLESS,
            response_policy=ResponsePolicy.DEFAULT_KEYLESS,
        ),
        "dictadd": CommandPolicies(
            request_policy=RequestPolicy.DEFAULT_KEYLESS,
            response_policy=ResponsePolicy.DEFAULT_KEYLESS,
        ),
        "synupdate": CommandPolicies(
            request_policy=RequestPolicy.DEFAULT_KEYLESS,
            response_policy=ResponsePolicy.DEFAULT_KEYLESS,
        ),
        "drop": CommandPolicies(
            request_policy=RequestPolicy.DEFAULT_KEYLESS,
            response_policy=ResponsePolicy.DEFAULT_KEYLESS,
        ),
        "info": CommandPolicies(
            request_policy=RequestPolicy.DEFAULT_KEYLESS,
            response_policy=ResponsePolicy.DEFAULT_KEYLESS,
        ),
        "sugadd": CommandPolicies(
            request_policy=RequestPolicy.DEFAULT_KEYED,
            response_policy=ResponsePolicy.DEFAULT_KEYED,
        ),
        "dictdump": CommandPolicies(
            request_policy=RequestPolicy.DEFAULT_KEYLESS,
            response_policy=ResponsePolicy.DEFAULT_KEYLESS,
        ),
        "cursor": CommandPolicies(
            request_policy=RequestPolicy.SPECIAL,
            response_policy=ResponsePolicy.DEFAULT_KEYLESS,
        ),
        "search": CommandPolicies(
            request_policy=RequestPolicy.DEFAULT_KEYLESS,
            response_policy=ResponsePolicy.DEFAULT_KEYLESS,
        ),
        "tagvals": CommandPolicies(
            request_policy=RequestPolicy.DEFAULT_KEYLESS,
            response_policy=ResponsePolicy.DEFAULT_KEYLESS,
        ),
        "aliasdel": CommandPolicies(
            request_policy=RequestPolicy.DEFAULT_KEYLESS,
            response_policy=ResponsePolicy.DEFAULT_KEYLESS,
        ),
        "sugdel": CommandPolicies(
            request_policy=RequestPolicy.DEFAULT_KEYED,
            response_policy=ResponsePolicy.DEFAULT_KEYED,
        ),
        "spellcheck": CommandPolicies(
            request_policy=RequestPolicy.DEFAULT_KEYLESS,
            response_policy=ResponsePolicy.DEFAULT_KEYLESS,
        ),
    },
    "core": {
        "command": CommandPolicies(
            request_policy=RequestPolicy.DEFAULT_KEYLESS,
            response_policy=ResponsePolicy.DEFAULT_KEYLESS,
        ),
    },
}


class PolicyResolver(ABC):
    @abstractmethod
    def resolve(self, command_name: str) -> Optional[CommandPolicies]:
        """
        Resolves the command name and determines the associated command policies.

        Args:
            command_name: The name of the command to resolve.

        Returns:
            CommandPolicies: The policies associated with the specified command.
        """
        pass

    @abstractmethod
    def with_fallback(self, fallback: "PolicyResolver") -> "PolicyResolver":
        """
        Factory method to instantiate a policy resolver with a fallback resolver.

        Args:
            fallback: Fallback resolver

        Returns:
            PolicyResolver: Returns a new policy resolver with the specified fallback resolver.
        """
        pass


class AsyncPolicyResolver(ABC):
    @abstractmethod
    async def resolve(self, command_name: str) -> Optional[CommandPolicies]:
        """
        Resolves the command name and determines the associated command policies.

        Args:
            command_name: The name of the command to resolve.

        Returns:
            CommandPolicies: The policies associated with the specified command.
        """
        pass

    @abstractmethod
    def with_fallback(self, fallback: "AsyncPolicyResolver") -> "AsyncPolicyResolver":
        """
        Factory method to instantiate an async policy resolver with a fallback resolver.

        Args:
            fallback: Fallback resolver

        Returns:
            AsyncPolicyResolver: Returns a new policy resolver with the specified fallback resolver.
        """
        pass


class BasePolicyResolver(PolicyResolver):
    """
    Base class for policy resolvers.
    """

    def __init__(
        self, policies: PolicyRecords, fallback: Optional[PolicyResolver] = None
    ) -> None:
        self._policies = policies
        self._fallback = fallback

    def resolve(self, command_name: str) -> Optional[CommandPolicies]:
        parts = command_name.split(".")

        if len(parts) > 2:
            raise ValueError(f"Wrong command or module name: {command_name}")

        module, command = parts if len(parts) == 2 else ("core", parts[0])

        if self._policies.get(module, None) is None:
            if self._fallback is not None:
                return self._fallback.resolve(command_name)
            else:
                return None

        if self._policies.get(module).get(command, None) is None:
            if self._fallback is not None:
                return self._fallback.resolve(command_name)
            else:
                return None

        return self._policies.get(module).get(command)

    @abstractmethod
    def with_fallback(self, fallback: "PolicyResolver") -> "PolicyResolver":
        pass


class AsyncBasePolicyResolver(AsyncPolicyResolver):
    """
    Async base class for policy resolvers.
    """

    def __init__(
        self, policies: PolicyRecords, fallback: Optional[AsyncPolicyResolver] = None
    ) -> None:
        self._policies = policies
        self._fallback = fallback

    async def resolve(self, command_name: str) -> Optional[CommandPolicies]:
        parts = command_name.split(".")

        if len(parts) > 2:
            raise ValueError(f"Wrong command or module name: {command_name}")

        module, command = parts if len(parts) == 2 else ("core", parts[0])

        if self._policies.get(module, None) is None:
            if self._fallback is not None:
                return await self._fallback.resolve(command_name)
            else:
                return None

        if self._policies.get(module).get(command, None) is None:
            if self._fallback is not None:
                return await self._fallback.resolve(command_name)
            else:
                return None

        return self._policies.get(module).get(command)

    @abstractmethod
    def with_fallback(self, fallback: "AsyncPolicyResolver") -> "AsyncPolicyResolver":
        pass


class DynamicPolicyResolver(BasePolicyResolver):
    """
    Resolves policy dynamically based on the COMMAND output.
    """

    def __init__(
        self, commands_parser: CommandsParser, fallback: Optional[PolicyResolver] = None
    ) -> None:
        """
        Parameters:
            commands_parser (CommandsParser): COMMAND output parser.
            fallback (Optional[PolicyResolver]): An optional resolver to be used when the
                primary policies cannot handle a specific request.
        """
        self._commands_parser = commands_parser
        super().__init__(commands_parser.get_command_policies(), fallback)

    def with_fallback(self, fallback: "PolicyResolver") -> "PolicyResolver":
        return DynamicPolicyResolver(self._commands_parser, fallback)


class StaticPolicyResolver(BasePolicyResolver):
    """
    Resolves policy from a static list of policy records.
    """

    def __init__(self, fallback: Optional[PolicyResolver] = None) -> None:
        """
        Parameters:
            fallback (Optional[PolicyResolver]): An optional fallback policy resolver
            used for resolving policies if static policies are inadequate.
        """
        super().__init__(STATIC_POLICIES, fallback)

    def with_fallback(self, fallback: "PolicyResolver") -> "PolicyResolver":
        return StaticPolicyResolver(fallback)


class AsyncDynamicPolicyResolver(AsyncBasePolicyResolver):
    """
    Async version of DynamicPolicyResolver.
    """

    def __init__(
        self,
        policy_records: PolicyRecords,
        fallback: Optional[AsyncPolicyResolver] = None,
    ) -> None:
        """
        Parameters:
            policy_records (PolicyRecords): Policy records.
            fallback (Optional[AsyncPolicyResolver]): An optional resolver to be used when the
                primary policies cannot handle a specific request.
        """
        super().__init__(policy_records, fallback)

    def with_fallback(self, fallback: "AsyncPolicyResolver") -> "AsyncPolicyResolver":
        return AsyncDynamicPolicyResolver(self._policies, fallback)


class AsyncStaticPolicyResolver(AsyncBasePolicyResolver):
    """
    Async version of StaticPolicyResolver.
    """

    def __init__(self, fallback: Optional[AsyncPolicyResolver] = None) -> None:
        """
        Parameters:
            fallback (Optional[AsyncPolicyResolver]): An optional fallback policy resolver
            used for resolving policies if static policies are inadequate.
        """
        super().__init__(STATIC_POLICIES, fallback)

    def with_fallback(self, fallback: "AsyncPolicyResolver") -> "AsyncPolicyResolver":
        return AsyncStaticPolicyResolver(fallback)