From 6d32503a58d841a41556cebcf188add324c6ce86 Mon Sep 17 00:00:00 2001 From: Nick Gerow Date: Wed, 3 May 2023 09:27:32 -0400 Subject: Updated AWS Elasticache IAM Connection Example (#2702) Co-authored-by: Nick Gerow --- docs/examples/connection_examples.ipynb | 56 ++++++++++++++++++++++++++------- 1 file changed, 44 insertions(+), 12 deletions(-) diff --git a/docs/examples/connection_examples.ipynb b/docs/examples/connection_examples.ipynb index 7f5ac53..d15d964 100644 --- a/docs/examples/connection_examples.ipynb +++ b/docs/examples/connection_examples.ipynb @@ -267,28 +267,60 @@ } ], "source": [ + "from typing import Tuple, Union\n", + "from urllib.parse import ParseResult, urlencode, urlunparse\n", + "\n", + "import botocore.session\n", "import redis\n", - "import boto3\n", - "import cachetools.func\n", + "from botocore.model import ServiceId\n", + "from botocore.signers import RequestSigner\n", + "from cachetools import TTLCache, cached\n", "\n", "class ElastiCacheIAMProvider(redis.CredentialProvider):\n", - " def __init__(self, user, endpoint, port=6379, region=\"us-east-1\"):\n", - " self.ec_client = boto3.client('elasticache')\n", + " def __init__(self, user, cluster_name, region=\"us-east-1\"):\n", " self.user = user\n", - " self.endpoint = endpoint\n", - " self.port = port\n", + " self.cluster_name = cluster_name\n", " self.region = region\n", "\n", + " session = botocore.session.get_session()\n", + " self.request_signer = RequestSigner(\n", + " ServiceId(\"elasticache\"),\n", + " self.region,\n", + " \"elasticache\",\n", + " \"v4\",\n", + " session.get_credentials(),\n", + " session.get_component(\"event_emitter\"),\n", + " )\n", + "\n", + " # Generated IAM tokens are valid for 15 minutes\n", + " @cached(cache=TTLCache(maxsize=128, ttl=900))\n", " def get_credentials(self) -> Union[Tuple[str], Tuple[str, str]]:\n", - " @cachetools.func.ttl_cache(maxsize=128, ttl=15 * 60) # 15m\n", - " def get_iam_auth_token(user, endpoint, port, region):\n", - " return self.ec_client.generate_iam_auth_token(user, endpoint, port, region)\n", - " iam_auth_token = get_iam_auth_token(self.endpoint, self.port, self.user, self.region)\n", - " return iam_auth_token\n", + " query_params = {\"Action\": \"connect\", \"User\": self.user}\n", + " url = urlunparse(\n", + " ParseResult(\n", + " scheme=\"https\",\n", + " netloc=self.cluster_name,\n", + " path=\"/\",\n", + " query=urlencode(query_params),\n", + " params=\"\",\n", + " fragment=\"\",\n", + " )\n", + " )\n", + " signed_url = self.request_signer.generate_presigned_url(\n", + " {\"method\": \"GET\", \"url\": url, \"body\": {}, \"headers\": {}, \"context\": {}},\n", + " operation_name=\"connect\",\n", + " expires_in=900,\n", + " region_name=self.region,\n", + " )\n", + " # RequestSigner only seems to work if the URL has a protocol, but\n", + " # Elasticache only accepts the URL without a protocol\n", + " # So strip it off the signed URL before returning\n", + " return (self.user, signed_url.removeprefix(\"https://\"))\n", "\n", "username = \"barshaul\"\n", + "cluster_name = \"test-001\"\n", "endpoint = \"test-001.use1.cache.amazonaws.com\"\n", - "creds_provider = ElastiCacheIAMProvider(user=username, endpoint=endpoint)\n", + "creds_provider = ElastiCacheIAMProvider(user=username, cluster_name=cluster_name)\n", "user_connection = redis.Redis(host=endpoint, port=6379, credential_provider=creds_provider)\n", "user_connection.ping()" ] -- cgit v1.2.1