Fixed CredentialsProvider examples (redis#2587)

Co-authored-by: dvora-h <[email protected]>
aksinha334 · Mar 15, 2023 · 7b09b3d · 7b09b3d
1 parent a4feb3c
commit 7b09b3d
Showing 1 changed file with 30 additions and 20 deletions.
diff --git a/docs/examples/connection_examples.ipynb b/docs/examples/connection_examples.ipynb
@@ -222,18 +222,23 @@
     "import json\n",
     "import cachetools.func\n",
     "\n",
-    "sm_client = boto3.client('secretsmanager')\n",
-    "  \n",
-    "def sm_auth_provider(self, secret_id, version_id=None, version_stage='AWSCURRENT'):\n",
-    "    @cachetools.func.ttl_cache(maxsize=128, ttl=24 * 60 * 60) #24h\n",
-    "    def get_sm_user_credentials(secret_id, version_id, version_stage):\n",
-    "        secret = sm_client.get_secret_value(secret_id, version_id)\n",
-    "        return json.loads(secret['SecretString'])\n",
-    "    creds = get_sm_user_credentials(secret_id, version_id, version_stage)\n",
-    "    return creds['username'], creds['password']\n",
+    "class SecretsManagerProvider(redis.CredentialProvider):\n",
+    "    def __init__(self, secret_id, version_id=None, version_stage='AWSCURRENT'):\n",
+    "        self.sm_client = boto3.client('secretsmanager')\n",
+    "        self.secret_id = secret_id\n",
+    "        self.version_id = version_id\n",
+    "        self.version_stage = version_stage\n",
     "\n",
-    "secret_id = \"EXAMPLE1-90ab-cdef-fedc-ba987SECRET1\"\n",
-    "creds_provider = redis.CredentialProvider(supplier=sm_auth_provider, secret_id=secret_id)\n",
+    "    def get_credentials(self) -> Union[Tuple[str], Tuple[str, str]]:\n",
+    "        @cachetools.func.ttl_cache(maxsize=128, ttl=24 * 60 * 60) #24h\n",
+    "        def get_sm_user_credentials(secret_id, version_id, version_stage):\n",
+    "            secret = self.sm_client.get_secret_value(secret_id, version_id)\n",
+    "            return json.loads(secret['SecretString'])\n",
+    "        creds = get_sm_user_credentials(self.secret_id, self.version_id, self.version_stage)\n",
+    "        return creds['username'], creds['password']\n",
+    "\n",
+    "my_secret_id = \"EXAMPLE1-90ab-cdef-fedc-ba987SECRET1\"\n",
+    "creds_provider = SecretsManagerProvider(secret_id=my_secret_id)\n",
     "user_connection = redis.Redis(host=\"localhost\", port=6379, credential_provider=creds_provider)\n",
     "user_connection.ping()"
    ]
@@ -266,19 +271,24 @@
     "import boto3\n",
     "import cachetools.func\n",
     "\n",
-    "ec_client = boto3.client('elasticache')\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",
+    "        self.user = user\n",
+    "        self.endpoint = endpoint\n",
+    "        self.port = port\n",
+    "        self.region = region\n",
     "\n",
-    "def iam_auth_provider(self, user, endpoint, port=6379, region=\"us-east-1\"):\n",
-    "    @cachetools.func.ttl_cache(maxsize=128, ttl=15 * 60) # 15m\n",
-    "    def get_iam_auth_token(user, endpoint, port, region):\n",
-    "        return ec_client.generate_iam_auth_token(user, endpoint, port, region)\n",
-    "    iam_auth_token = get_iam_auth_token(endpoint, port, user, region)\n",
-    "    return iam_auth_token\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",
     "\n",
     "username = \"barshaul\"\n",
     "endpoint = \"test-001.use1.cache.amazonaws.com\"\n",
-    "creds_provider = redis.CredentialProvider(supplier=iam_auth_provider, user=username,\n",
-    "                                           endpoint=endpoint)\n",
+    "creds_provider = ElastiCacheIAMProvider(user=username, endpoint=endpoint)\n",
     "user_connection = redis.Redis(host=endpoint, port=6379, credential_provider=creds_provider)\n",
     "user_connection.ping()"
    ]