|
267 | 267 | }
|
268 | 268 | ],
|
269 | 269 | "source": [
|
| 270 | + "from typing import Tuple, Union\n", |
| 271 | + "from urllib.parse import ParseResult, urlencode, urlunparse\n", |
| 272 | + "\n", |
| 273 | + "import botocore.session\n", |
270 | 274 | "import redis\n",
|
271 |
| - "import boto3\n", |
272 |
| - "import cachetools.func\n", |
| 275 | + "from botocore.model import ServiceId\n", |
| 276 | + "from botocore.signers import RequestSigner\n", |
| 277 | + "from cachetools import TTLCache, cached\n", |
273 | 278 | "\n",
|
274 | 279 | "class ElastiCacheIAMProvider(redis.CredentialProvider):\n",
|
275 |
| - " def __init__(self, user, endpoint, port=6379, region=\"us-east-1\"):\n", |
276 |
| - " self.ec_client = boto3.client('elasticache')\n", |
| 280 | + " def __init__(self, user, cluster_name, region=\"us-east-1\"):\n", |
277 | 281 | " self.user = user\n",
|
278 |
| - " self.endpoint = endpoint\n", |
279 |
| - " self.port = port\n", |
| 282 | + " self.cluster_name = cluster_name\n", |
280 | 283 | " self.region = region\n",
|
281 | 284 | "\n",
|
| 285 | + " session = botocore.session.get_session()\n", |
| 286 | + " self.request_signer = RequestSigner(\n", |
| 287 | + " ServiceId(\"elasticache\"),\n", |
| 288 | + " self.region,\n", |
| 289 | + " \"elasticache\",\n", |
| 290 | + " \"v4\",\n", |
| 291 | + " session.get_credentials(),\n", |
| 292 | + " session.get_component(\"event_emitter\"),\n", |
| 293 | + " )\n", |
| 294 | + "\n", |
| 295 | + " # Generated IAM tokens are valid for 15 minutes\n", |
| 296 | + " @cached(cache=TTLCache(maxsize=128, ttl=900))\n", |
282 | 297 | " def get_credentials(self) -> Union[Tuple[str], Tuple[str, str]]:\n",
|
283 |
| - " @cachetools.func.ttl_cache(maxsize=128, ttl=15 * 60) # 15m\n", |
284 |
| - " def get_iam_auth_token(user, endpoint, port, region):\n", |
285 |
| - " return self.ec_client.generate_iam_auth_token(user, endpoint, port, region)\n", |
286 |
| - " iam_auth_token = get_iam_auth_token(self.endpoint, self.port, self.user, self.region)\n", |
287 |
| - " return iam_auth_token\n", |
| 298 | + " query_params = {\"Action\": \"connect\", \"User\": self.user}\n", |
| 299 | + " url = urlunparse(\n", |
| 300 | + " ParseResult(\n", |
| 301 | + " scheme=\"https\",\n", |
| 302 | + " netloc=self.cluster_name,\n", |
| 303 | + " path=\"/\",\n", |
| 304 | + " query=urlencode(query_params),\n", |
| 305 | + " params=\"\",\n", |
| 306 | + " fragment=\"\",\n", |
| 307 | + " )\n", |
| 308 | + " )\n", |
| 309 | + " signed_url = self.request_signer.generate_presigned_url(\n", |
| 310 | + " {\"method\": \"GET\", \"url\": url, \"body\": {}, \"headers\": {}, \"context\": {}},\n", |
| 311 | + " operation_name=\"connect\",\n", |
| 312 | + " expires_in=900,\n", |
| 313 | + " region_name=self.region,\n", |
| 314 | + " )\n", |
| 315 | + " # RequestSigner only seems to work if the URL has a protocol, but\n", |
| 316 | + " # Elasticache only accepts the URL without a protocol\n", |
| 317 | + " # So strip it off the signed URL before returning\n", |
| 318 | + " return (self.user, signed_url.removeprefix(\"https://\"))\n", |
288 | 319 | "\n",
|
289 | 320 | "username = \"barshaul\"\n",
|
| 321 | + "cluster_name = \"test-001\"\n", |
290 | 322 | "endpoint = \"test-001.use1.cache.amazonaws.com\"\n",
|
291 |
| - "creds_provider = ElastiCacheIAMProvider(user=username, endpoint=endpoint)\n", |
| 323 | + "creds_provider = ElastiCacheIAMProvider(user=username, cluster_name=cluster_name)\n", |
292 | 324 | "user_connection = redis.Redis(host=endpoint, port=6379, credential_provider=creds_provider)\n",
|
293 | 325 | "user_connection.ping()"
|
294 | 326 | ]
|
|
0 commit comments