Updated AWS Elasticache IAM Connection Example (#2702)

Co-authored-by: Nick Gerow <nick.gerow@enlightedinc.com>

1 files 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()"
    ]