{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Connection Examples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Connecting to a default Valkey instance, running locally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import valkey\n",
    "\n",
    "connection = valkey.Valkey()\n",
    "connection.ping()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### By default Valkey return binary responses, to decode them use decode_responses=True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import valkey\n",
    "\n",
    "decoded_connection = valkey.Valkey(decode_responses=True)\n",
    "decoded_connection.ping()"
   ]
  },
  {
   "cell_type": "markdown",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### By default this library uses the RESP 2 protocol. To enable RESP3, set protocol=3."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "import valkey\n",
    "\n",
    "r = valkey.Valkey(protocol=3)\n",
    "rcon.ping()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Connecting to a valkey instance, specifying a host and port with credentials."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import valkey\n",
    "\n",
    "user_connection = valkey.Valkey(host='localhost', port=6380, username='dvora', password='valkey', decode_responses=True)\n",
    "user_connection.ping()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Connecting to a valkey instance with username and password credential provider"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import valkey\n",
    "\n",
    "creds_provider = valkey.UsernamePasswordCredentialProvider(\"username\", \"password\")\n",
    "user_connection = valkey.Valkey(host=\"localhost\", port=6379, credential_provider=creds_provider)\n",
    "user_connection.ping()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Connecting to a valkey instance with standard credential provider"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Tuple\n",
    "import valkey\n",
    "\n",
    "creds_map = {\"user_1\": \"pass_1\",\n",
    "             \"user_2\": \"pass_2\"}\n",
    "\n",
    "class UserMapCredentialProvider(valkey.CredentialProvider):\n",
    "    def __init__(self, username: str):\n",
    "        self.username = username\n",
    "\n",
    "    def get_credentials(self) -> Tuple[str, str]:\n",
    "        return self.username, creds_map.get(self.username)\n",
    "\n",
    "# Create a default connection to set the ACL user\n",
    "default_connection = valkey.Valkey(host=\"localhost\", port=6379)\n",
    "default_connection.acl_setuser(\n",
    "    \"user_1\",\n",
    "    enabled=True,\n",
    "    passwords=[\"+\" + \"pass_1\"],\n",
    "    keys=\"~*\",\n",
    "    commands=[\"+ping\", \"+command\", \"+info\", \"+select\", \"+flushdb\"],\n",
    ")\n",
    "\n",
    "# Create a UserMapCredentialProvider instance for user_1\n",
    "creds_provider = UserMapCredentialProvider(\"user_1\")\n",
    "# Initiate user connection with the credential provider\n",
    "user_connection = valkey.Valkey(host=\"localhost\", port=6379,\n",
    "                              credential_provider=creds_provider)\n",
    "user_connection.ping()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Connecting to a valkey instance first with an initial credential set and then calling the credential provider"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Union\n",
    "import valkey\n",
    "\n",
    "class InitCredsSetCredentialProvider(valkey.CredentialProvider):\n",
    "    def __init__(self, username, password):\n",
    "        self.username = username\n",
    "        self.password = password\n",
    "        self.call_supplier = False\n",
    "\n",
    "    def call_external_supplier(self) -> Union[Tuple[str], Tuple[str, str]]:\n",
    "        # Call to an external credential supplier\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def get_credentials(self) -> Union[Tuple[str], Tuple[str, str]]:\n",
    "        if self.call_supplier:\n",
    "            return self.call_external_supplier()\n",
    "        # Use the init set only for the first time\n",
    "        self.call_supplier = True\n",
    "        return self.username, self.password\n",
    "\n",
    "cred_provider = InitCredsSetCredentialProvider(username=\"init_user\", password=\"init_pass\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Connecting to a valkey instance with AWS Secrets Manager credential provider."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import valkey\n",
    "import boto3\n",
    "import json\n",
    "import cachetools.func\n",
    "\n",
    "class SecretsManagerProvider(valkey.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",
    "    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 = valkey.Valkey(host=\"localhost\", port=6379, credential_provider=creds_provider)\n",
    "user_connection.ping()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Connecting to a valkey instance with ElastiCache IAM credential provider."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from typing import Tuple, Union\n",
    "from urllib.parse import ParseResult, urlencode, urlunparse\n",
    "\n",
    "import botocore.session\n",
    "import valkey\n",
    "from botocore.model import ServiceId\n",
    "from botocore.signers import RequestSigner\n",
    "from cachetools import TTLCache, cached\n",
    "\n",
    "class ElastiCacheIAMProvider(valkey.CredentialProvider):\n",
    "    def __init__(self, user, cluster_name, region=\"us-east-1\"):\n",
    "        self.user = user\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",
    "        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, cluster_name=cluster_name)\n",
    "user_connection = valkey.Valkey(host=endpoint, port=6379, credential_provider=creds_provider)\n",
    "user_connection.ping()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Connecting to Valkey instances by specifying a URL scheme.\n",
    "Parameters are passed to the following schems, as parameters to the url scheme.\n",
    "\n",
    "Three URL schemes are supported:\n",
    "\n",
    "- `valkey://` creates a TCP socket connection. <https://www.iana.org/assignments/uri-schemes/prov/valkey>\n",
    "- `valkeys://` creates a SSL wrapped TCP socket connection. <https://www.iana.org/assignments/uri-schemes/prov/valkeys>\n",
    "- ``unix://``: creates a Unix Domain Socket connection.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "url_connection = valkey.from_url(\"valkey://localhost:6379?decode_responses=True&health_check_interval=2\")\n",
    "url_connection.ping()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Connecting to Valkey instances by specifying a URL scheme and the RESP3 protocol.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "url_connection = valkey.from_url(\"valkey://localhost:6379?decode_responses=True&health_check_interval=2&protocol=3\")\n",
    "url_connection.ping()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Connecting to a Sentinel instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from valkey.sentinel import Sentinel\n",
    "sentinel = Sentinel([('localhost', 26379)], socket_timeout=0.1)\n",
    "sentinel.discover_master(\"valkey-py-test\")"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "d45c99ba0feda92868abafa8257cbb4709c97f1a0b5dc62bbeebdf89d4fad7fe"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
