1from abc import ABC, abstractmethod
2from typing import Optional
3
4from redis._parsers.commands import (
5 CommandPolicies,
6 CommandsParser,
7 PolicyRecords,
8 RequestPolicy,
9 ResponsePolicy,
10)
11
12STATIC_POLICIES: PolicyRecords = {
13 "ft": {
14 "explaincli": CommandPolicies(
15 request_policy=RequestPolicy.DEFAULT_KEYLESS,
16 response_policy=ResponsePolicy.DEFAULT_KEYLESS,
17 ),
18 "suglen": CommandPolicies(
19 request_policy=RequestPolicy.DEFAULT_KEYED,
20 response_policy=ResponsePolicy.DEFAULT_KEYED,
21 ),
22 "profile": CommandPolicies(
23 request_policy=RequestPolicy.DEFAULT_KEYLESS,
24 response_policy=ResponsePolicy.DEFAULT_KEYLESS,
25 ),
26 "dropindex": CommandPolicies(
27 request_policy=RequestPolicy.DEFAULT_KEYLESS,
28 response_policy=ResponsePolicy.DEFAULT_KEYLESS,
29 ),
30 "aliasupdate": CommandPolicies(
31 request_policy=RequestPolicy.DEFAULT_KEYLESS,
32 response_policy=ResponsePolicy.DEFAULT_KEYLESS,
33 ),
34 "alter": CommandPolicies(
35 request_policy=RequestPolicy.DEFAULT_KEYLESS,
36 response_policy=ResponsePolicy.DEFAULT_KEYLESS,
37 ),
38 "aggregate": CommandPolicies(
39 request_policy=RequestPolicy.DEFAULT_KEYLESS,
40 response_policy=ResponsePolicy.DEFAULT_KEYLESS,
41 ),
42 "syndump": CommandPolicies(
43 request_policy=RequestPolicy.DEFAULT_KEYLESS,
44 response_policy=ResponsePolicy.DEFAULT_KEYLESS,
45 ),
46 "create": CommandPolicies(
47 request_policy=RequestPolicy.DEFAULT_KEYLESS,
48 response_policy=ResponsePolicy.DEFAULT_KEYLESS,
49 ),
50 "explain": CommandPolicies(
51 request_policy=RequestPolicy.DEFAULT_KEYLESS,
52 response_policy=ResponsePolicy.DEFAULT_KEYLESS,
53 ),
54 "sugget": CommandPolicies(
55 request_policy=RequestPolicy.DEFAULT_KEYED,
56 response_policy=ResponsePolicy.DEFAULT_KEYED,
57 ),
58 "dictdel": CommandPolicies(
59 request_policy=RequestPolicy.DEFAULT_KEYLESS,
60 response_policy=ResponsePolicy.DEFAULT_KEYLESS,
61 ),
62 "aliasadd": CommandPolicies(
63 request_policy=RequestPolicy.DEFAULT_KEYLESS,
64 response_policy=ResponsePolicy.DEFAULT_KEYLESS,
65 ),
66 "dictadd": CommandPolicies(
67 request_policy=RequestPolicy.DEFAULT_KEYLESS,
68 response_policy=ResponsePolicy.DEFAULT_KEYLESS,
69 ),
70 "synupdate": CommandPolicies(
71 request_policy=RequestPolicy.DEFAULT_KEYLESS,
72 response_policy=ResponsePolicy.DEFAULT_KEYLESS,
73 ),
74 "drop": CommandPolicies(
75 request_policy=RequestPolicy.DEFAULT_KEYLESS,
76 response_policy=ResponsePolicy.DEFAULT_KEYLESS,
77 ),
78 "info": CommandPolicies(
79 request_policy=RequestPolicy.DEFAULT_KEYLESS,
80 response_policy=ResponsePolicy.DEFAULT_KEYLESS,
81 ),
82 "sugadd": CommandPolicies(
83 request_policy=RequestPolicy.DEFAULT_KEYED,
84 response_policy=ResponsePolicy.DEFAULT_KEYED,
85 ),
86 "dictdump": CommandPolicies(
87 request_policy=RequestPolicy.DEFAULT_KEYLESS,
88 response_policy=ResponsePolicy.DEFAULT_KEYLESS,
89 ),
90 "cursor": CommandPolicies(
91 request_policy=RequestPolicy.SPECIAL,
92 response_policy=ResponsePolicy.DEFAULT_KEYLESS,
93 ),
94 "search": CommandPolicies(
95 request_policy=RequestPolicy.DEFAULT_KEYLESS,
96 response_policy=ResponsePolicy.DEFAULT_KEYLESS,
97 ),
98 "tagvals": CommandPolicies(
99 request_policy=RequestPolicy.DEFAULT_KEYLESS,
100 response_policy=ResponsePolicy.DEFAULT_KEYLESS,
101 ),
102 "aliasdel": CommandPolicies(
103 request_policy=RequestPolicy.DEFAULT_KEYLESS,
104 response_policy=ResponsePolicy.DEFAULT_KEYLESS,
105 ),
106 "aliaslist": CommandPolicies(
107 request_policy=RequestPolicy.DEFAULT_KEYLESS,
108 response_policy=ResponsePolicy.DEFAULT_KEYLESS,
109 ),
110 "sugdel": CommandPolicies(
111 request_policy=RequestPolicy.DEFAULT_KEYED,
112 response_policy=ResponsePolicy.DEFAULT_KEYED,
113 ),
114 "spellcheck": CommandPolicies(
115 request_policy=RequestPolicy.DEFAULT_KEYLESS,
116 response_policy=ResponsePolicy.DEFAULT_KEYLESS,
117 ),
118 },
119 "core": {
120 "command": CommandPolicies(
121 request_policy=RequestPolicy.DEFAULT_KEYLESS,
122 response_policy=ResponsePolicy.DEFAULT_KEYLESS,
123 ),
124 },
125}
126
127
128class PolicyResolver(ABC):
129 @abstractmethod
130 def resolve(self, command_name: str) -> Optional[CommandPolicies]:
131 """
132 Resolves the command name and determines the associated command policies.
133
134 Args:
135 command_name: The name of the command to resolve.
136
137 Returns:
138 CommandPolicies: The policies associated with the specified command.
139 """
140 pass
141
142 @abstractmethod
143 def with_fallback(self, fallback: "PolicyResolver") -> "PolicyResolver":
144 """
145 Factory method to instantiate a policy resolver with a fallback resolver.
146
147 Args:
148 fallback: Fallback resolver
149
150 Returns:
151 PolicyResolver: Returns a new policy resolver with the specified fallback resolver.
152 """
153 pass
154
155
156class AsyncPolicyResolver(ABC):
157 @abstractmethod
158 async def resolve(self, command_name: str) -> Optional[CommandPolicies]:
159 """
160 Resolves the command name and determines the associated command policies.
161
162 Args:
163 command_name: The name of the command to resolve.
164
165 Returns:
166 CommandPolicies: The policies associated with the specified command.
167 """
168 pass
169
170 @abstractmethod
171 def with_fallback(self, fallback: "AsyncPolicyResolver") -> "AsyncPolicyResolver":
172 """
173 Factory method to instantiate an async policy resolver with a fallback resolver.
174
175 Args:
176 fallback: Fallback resolver
177
178 Returns:
179 AsyncPolicyResolver: Returns a new policy resolver with the specified fallback resolver.
180 """
181 pass
182
183
184class BasePolicyResolver(PolicyResolver):
185 """
186 Base class for policy resolvers.
187 """
188
189 def __init__(
190 self, policies: PolicyRecords, fallback: Optional[PolicyResolver] = None
191 ) -> None:
192 self._policies = policies
193 self._fallback = fallback
194
195 def resolve(self, command_name: str) -> Optional[CommandPolicies]:
196 parts = command_name.split(".")
197
198 if len(parts) > 2:
199 raise ValueError(f"Wrong command or module name: {command_name}")
200
201 module, command = parts if len(parts) == 2 else ("core", parts[0])
202
203 if self._policies.get(module, None) is None:
204 if self._fallback is not None:
205 return self._fallback.resolve(command_name)
206 else:
207 return None
208
209 if self._policies.get(module).get(command, None) is None:
210 if self._fallback is not None:
211 return self._fallback.resolve(command_name)
212 else:
213 return None
214
215 return self._policies.get(module).get(command)
216
217 @abstractmethod
218 def with_fallback(self, fallback: "PolicyResolver") -> "PolicyResolver":
219 pass
220
221
222class AsyncBasePolicyResolver(AsyncPolicyResolver):
223 """
224 Async base class for policy resolvers.
225 """
226
227 def __init__(
228 self, policies: PolicyRecords, fallback: Optional[AsyncPolicyResolver] = None
229 ) -> None:
230 self._policies = policies
231 self._fallback = fallback
232
233 async def resolve(self, command_name: str) -> Optional[CommandPolicies]:
234 parts = command_name.split(".")
235
236 if len(parts) > 2:
237 raise ValueError(f"Wrong command or module name: {command_name}")
238
239 module, command = parts if len(parts) == 2 else ("core", parts[0])
240
241 if self._policies.get(module, None) is None:
242 if self._fallback is not None:
243 return await self._fallback.resolve(command_name)
244 else:
245 return None
246
247 if self._policies.get(module).get(command, None) is None:
248 if self._fallback is not None:
249 return await self._fallback.resolve(command_name)
250 else:
251 return None
252
253 return self._policies.get(module).get(command)
254
255 @abstractmethod
256 def with_fallback(self, fallback: "AsyncPolicyResolver") -> "AsyncPolicyResolver":
257 pass
258
259
260class DynamicPolicyResolver(BasePolicyResolver):
261 """
262 Resolves policy dynamically based on the COMMAND output.
263 """
264
265 def __init__(
266 self, commands_parser: CommandsParser, fallback: Optional[PolicyResolver] = None
267 ) -> None:
268 """
269 Parameters:
270 commands_parser (CommandsParser): COMMAND output parser.
271 fallback (Optional[PolicyResolver]): An optional resolver to be used when the
272 primary policies cannot handle a specific request.
273 """
274 self._commands_parser = commands_parser
275 super().__init__(commands_parser.get_command_policies(), fallback)
276
277 def with_fallback(self, fallback: "PolicyResolver") -> "PolicyResolver":
278 return DynamicPolicyResolver(self._commands_parser, fallback)
279
280
281class StaticPolicyResolver(BasePolicyResolver):
282 """
283 Resolves policy from a static list of policy records.
284 """
285
286 def __init__(self, fallback: Optional[PolicyResolver] = None) -> None:
287 """
288 Parameters:
289 fallback (Optional[PolicyResolver]): An optional fallback policy resolver
290 used for resolving policies if static policies are inadequate.
291 """
292 super().__init__(STATIC_POLICIES, fallback)
293
294 def with_fallback(self, fallback: "PolicyResolver") -> "PolicyResolver":
295 return StaticPolicyResolver(fallback)
296
297
298class AsyncDynamicPolicyResolver(AsyncBasePolicyResolver):
299 """
300 Async version of DynamicPolicyResolver.
301 """
302
303 def __init__(
304 self,
305 policy_records: PolicyRecords,
306 fallback: Optional[AsyncPolicyResolver] = None,
307 ) -> None:
308 """
309 Parameters:
310 policy_records (PolicyRecords): Policy records.
311 fallback (Optional[AsyncPolicyResolver]): An optional resolver to be used when the
312 primary policies cannot handle a specific request.
313 """
314 super().__init__(policy_records, fallback)
315
316 def with_fallback(self, fallback: "AsyncPolicyResolver") -> "AsyncPolicyResolver":
317 return AsyncDynamicPolicyResolver(self._policies, fallback)
318
319
320class AsyncStaticPolicyResolver(AsyncBasePolicyResolver):
321 """
322 Async version of StaticPolicyResolver.
323 """
324
325 def __init__(self, fallback: Optional[AsyncPolicyResolver] = None) -> None:
326 """
327 Parameters:
328 fallback (Optional[AsyncPolicyResolver]): An optional fallback policy resolver
329 used for resolving policies if static policies are inadequate.
330 """
331 super().__init__(STATIC_POLICIES, fallback)
332
333 def with_fallback(self, fallback: "AsyncPolicyResolver") -> "AsyncPolicyResolver":
334 return AsyncStaticPolicyResolver(fallback)