{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "aaf1d224",
   "metadata": {},
   "source": [
    "*Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved*  \n",
    "\n",
    "*SPDX-License-Identifier: MIT-0*    \n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6cde1dbc",
   "metadata": {},
   "source": [
    "## Connecting to Amazon Keyspaces using SigV4 authentication plugin for temporary credentials.  \n",
    "\n",
    "This plugin enables IAM users, roles, and federated identities to add authentication information to Amazon Keyspaces API requests using the [AWS Signature Version 4 process (SigV4)](https://docs.aws.amazon.com/general/latest/gr/signature-version-4.html) .\n",
    "\n",
    "In this example, we do NOT need to generate Keyspaces service-specific credentials."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04a96df1",
   "metadata": {},
   "source": [
    "### Requrements \n",
    "\n",
    "The Notebook execution role must include permissions to access Amazon Keyspaces and Assume the role. \n",
    "\n",
    "*  To access Amazon Keyspaces database - use AmazonKeyspacesReadOnlyAccess or AmazonKeyspacesFullAccess managed policies. Use the _least privileged approach_ for your production application.  \n",
    "See more at\n",
    "[AWS Identity and Access Management for Amazon Keyspaces](https://docs.aws.amazon.com/keyspaces/latest/devguide/security-iam.html).\n",
    "\n",
    "* To assume the role, you need to have [sts:AssumeRole action](https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRole.html) permissions.\n",
    "    ```\n",
    "    {\n",
    "      \"Version\": \"2012-10-17\",  \n",
    "      \"Statement\": [  \n",
    "        {  \n",
    "           \"Action\": [  \n",
    "           \"sts:AssumeRole\"  \n",
    "          ],  \n",
    "          \"Effect\": \"Allow\",  \n",
    "          \"Resource\": \"*\"  \n",
    "        }\n",
    "      ]\n",
    "    }\n",
    "    ```\n",
    "\n",
    "#### Note:\n",
    "Amazon Keyspaces is available in the following [AWS Regions](https://docs.aws.amazon.com/keyspaces/latest/devguide/programmatic.endpoints.html).\n",
    "\n",
    "This notebook was tested with conda_python3 kernel and should work with Python 3.x."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f4e9f23",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Installing Cassanda SigV4 \n",
    "%pip install  cassandra-sigv4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e75d5af3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get Security certificate \n",
    "!curl https://certs.secureserver.net/repository/sf-class2-root.crt -O \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50da3133",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import \n",
    "from sagemaker import get_execution_role\n",
    "from cassandra.cluster import Cluster\n",
    "from ssl import SSLContext, PROTOCOL_TLSv1_2 , CERT_REQUIRED\n",
    "from cassandra_sigv4.auth import SigV4AuthProvider\n",
    "import boto3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32467082",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Getting credentials from the role\n",
    "client = boto3.client('sts')\n",
    "\n",
    "#Get notebook Role\n",
    "role = get_execution_role()\n",
    "role_info = {\n",
    "    'RoleArn':role,\n",
    "    'RoleSessionName': 'session1'\n",
    "}\n",
    "print(role_info)\n",
    "\n",
    "credentials = client.assume_role(**role_info)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "253da9a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Connecting to Cassandra using temporarly credentials. \n",
    "session = boto3.session.Session(\n",
    "    aws_access_key_id=credentials['Credentials']['AccessKeyId'],\n",
    "    aws_secret_access_key=credentials['Credentials']['SecretAccessKey'],\n",
    "    aws_session_token=credentials['Credentials']['SessionToken']\n",
    ")\n",
    "\n",
    "region_name = session.region_name\n",
    "\n",
    "#print(region_name)\n",
    "#print(credentials['Credentials']['AccessKeyId'])\n",
    "#print(credentials['Credentials']['SecretAccessKey'])\n",
    "#print(credentials['Credentials']['SessionToken'])\n",
    "\n",
    "#Set Context\n",
    "ssl_context = SSLContext(PROTOCOL_TLSv1_2)\n",
    "ssl_context.load_verify_locations('sf-class2-root.crt')\n",
    "ssl_context.verify_mode = CERT_REQUIRED\n",
    "\n",
    "auth_provider = SigV4AuthProvider(session)\n",
    "\n",
    "keyspaces_host = 'cassandra.' + region_name + '.amazonaws.com'\n",
    "\n",
    "cluster = Cluster([keyspaces_host], ssl_context=ssl_context, auth_provider=auth_provider,   port=9142)\n",
    "session = cluster.connect()\n",
    "\n",
    "\n",
    "# Read data from Keyspaces system table.  This table   \n",
    "r = session.execute('select * from system_schema.keyspaces')\n",
    "print(r.current_rows)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51a71bcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read Keyspaces row into Panda DataFrame \n",
    "from pandas import DataFrame\n",
    "df = DataFrame(r)\n",
    "print(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05700848",
   "metadata": {},
   "outputs": [],
   "source": [
    "####\n",
    "## Read your own tables. \n",
    "####\n",
    "\n",
    "#keyspaces = 'your_schema'\n",
    "#table ='your_table'\n",
    "\n",
    "#rows = session.execute('select * from '+ keyspaces + '.' + table )\n",
    "#df = DataFrame(rows)\n",
    "#print(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd7ca2f5",
   "metadata": {},
   "source": [
    "## The end."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_python3",
   "language": "python",
   "name": "conda_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.6.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}