examples/rest/notebook_query.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "8549cdb7-3b82-416b-b8d0-900732227212",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import timedelta, timezone, datetime\n",
    "import jwt\n",
    "from cryptography.hazmat.primitives.serialization import load_pem_private_key\n",
    "from cryptography.hazmat.primitives.serialization import Encoding\n",
    "from cryptography.hazmat.primitives.serialization import PublicFormat\n",
    "from cryptography.hazmat.backends import default_backend\n",
    "import base64\n",
    "from getpass import getpass\n",
    "import hashlib\n",
    "import requests\n",
    "import json\n",
    "import os\n",
    "import pandas as pd\n",
    "pd.options.display.max_colwidth = 1000\n",
    "\n",
    "# account parameters\n",
    "SNOWFLAKE_ACCOUNT = \"<ACCOUNT_NAME>\" # must be capitalized\n",
    "SNOWFLAKE_USER = \"<USERNAME>\" # must be capitalized\n",
    "SNOWFLAKE_URL = \"https://org-acc.snowflakecomputing.com\"\n",
    "PRIVATE_KEY_PATH = \"/path/to/your/rsa_key.p8\"\n",
    "\n",
    "# service parameters\n",
    "CORTEX_SEARCH_DATABASE = \"<DATABASE>\"\n",
    "CORTEX_SEARCH_SCHEMA = \"<SCHEMA>\"\n",
    "CORTEX_SEARCH_SERVICE = \"<SERVICE>\"\n",
    "\n",
    "# columns to query in the service\n",
    "COLUMNS = [\n",
    "    \"COL1\",\n",
    "    \"COL2\",\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "4ab6c048-b34c-47f9-a032-aa20fadb91d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_JWT_token():\n",
    "    \"\"\"\n",
    "    https://docs.snowflake.com/en/developer-guide/sql-api/authenticating#generating-a-jwt-in-python\n",
    "    Generate a valid JWT token from snowflake account name, user name, private key and private key passphrase.\n",
    "    \"\"\"\n",
    "    # Prompt for private key passphrase\n",
    "    def get_private_key_passphrase():\n",
    "        return getpass('Private Key Passphrase: ')\n",
    "\n",
    "    # Generate encoded public key\n",
    "    with open(PRIVATE_KEY_PATH, 'rb') as pem_in:\n",
    "        pemlines = pem_in.read()\n",
    "        try:\n",
    "            private_key = load_pem_private_key(pemlines, None, default_backend())\n",
    "        except TypeError:\n",
    "            private_key = load_pem_private_key(pemlines, get_private_key_passphrase().encode(), default_backend())\n",
    "    public_key_raw = private_key.public_key().public_bytes(Encoding.DER, PublicFormat.SubjectPublicKeyInfo)\n",
    "    sha256hash = hashlib.sha256()\n",
    "    sha256hash.update(public_key_raw)\n",
    "    public_key_fp = 'SHA256:' + base64.b64encode(sha256hash.digest()).decode('utf-8')\n",
    "\n",
    "    # Generate JWT payload\n",
    "    qualified_username = SNOWFLAKE_ACCOUNT + \".\" + SNOWFLAKE_USER\n",
    "    now = datetime.now(timezone.utc)\n",
    "    lifetime = timedelta(minutes=60)\n",
    "    payload = {\n",
    "        \"iss\": qualified_username + '.' + public_key_fp,\n",
    "        \"sub\": qualified_username,\n",
    "        \"iat\": now,\n",
    "        \"exp\": now + lifetime\n",
    "    }\n",
    "    return jwt.encode(payload, key=private_key, algorithm=\"RS256\")\n",
    "    \n",
    "jwt_token = generate_JWT_token()\n",
    "\n",
    "headers = {\n",
    "    'X-Snowflake-Authorization-Token-Type': 'KEYPAIR_JWT',\n",
    "    'Content-Type': 'application/json',\n",
    "    'Accept': 'application/json',\n",
    "    'Authorization': f'Bearer {jwt_token}',\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "1ae0e607-a9aa-400e-8bb7-a9b6945cccaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def query_service(query):\n",
    "    \"\"\"\n",
    "    Query the specified Snowflake service with the given query string.\n",
    "    \"\"\"\n",
    "    url = f\"{SNOWFLAKE_URL}/api/v2/databases/{CORTEX_SEARCH_DATABASE}/schemas/{CORTEX_SEARCH_SCHEMA}/cortex-search-services/{CORTEX_SEARCH_SERVICE}:query\"\n",
    "    data = {\n",
    "        \"query\": query,\n",
    "        \"columns\": COLUMNS,\n",
    "        \"filter\": \"\",\n",
    "        \"limit\": 10\n",
    "    }\n",
    "    \n",
    "    jwt_token = generate_JWT_token()\n",
    "    headers = {\n",
    "        'X-Snowflake-Authorization-Token-Type': 'KEYPAIR_JWT',\n",
    "        'Content-Type': 'application/json',\n",
    "        'Accept': 'application/json',\n",
    "        'Authorization': f'Bearer {jwt_token}',\n",
    "    }\n",
    "    \n",
    "    try:\n",
    "        response = requests.post(url, headers=headers, json=data)\n",
    "        response.raise_for_status()\n",
    "    except requests.exceptions.HTTPError as http_err:\n",
    "        print(f\"HTTP error occurred: {http_err} - Status code: {response.status_code}\")\n",
    "    except Exception as err:\n",
    "        print(f\"An error occurred: {err}\")\n",
    "    else:\n",
    "        return response.json()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "00202c33-97e0-4860-b34a-fe599d3948b2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>TALK_ID</th>\n",
       "      <th>SPEAKER_1</th>\n",
       "      <th>TITLE</th>\n",
       "      <th>DESCRIPTION</th>\n",
       "      <th>URL</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>220</td>\n",
       "      <td>Joseph Lekuton</td>\n",
       "      <td>A parable for Kenya</td>\n",
       "      <td>Joseph Lekuton, a member of parliament in Kenya, starts with the story of his remarkable education, then offers a parable of how Africa can grow. His message of hope has never been more relevant.</td>\n",
       "      <td>https://www.ted.com/talks/joseph_lekuton_a_parable_for_kenya/</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2221</td>\n",
       "      <td>Boniface Mwangi</td>\n",
       "      <td>The day I stood up alone</td>\n",
       "      <td>Photographer Boniface Mwangi wanted to protest against corruption in his home country of Kenya. So he made a plan: He and some friends would stand up and heckle during a public mass meeting. But when the moment came ... he stood alone. What happened next, he says, showed him who he truly was. As he says, \"There are two most powerful days in your life. The day you are born, and the day you discover why.\" Graphic images.</td>\n",
       "      <td>https://www.ted.com/talks/boniface_mwangi_the_day_i_stood_up_alone/</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>523</td>\n",
       "      <td>Erik Hersman</td>\n",
       "      <td>Reporting crisis via texting</td>\n",
       "      <td>At TEDU 2009, Erik Hersman presents the remarkable story of Ushahidi, a GoogleMap mashup that allowed Kenyans to report and track violence via cell phone texts following the 2008 elections, and has evolved to continue saving lives in other countries.</td>\n",
       "      <td>https://www.ted.com/talks/erik_hersman_reporting_crisis_via_texting/</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>2653</td>\n",
       "      <td>Charity Wayua</td>\n",
       "      <td>A few ways to fix a government</td>\n",
       "      <td>Charity Wayua put her skills as a cancer researcher to use on an unlikely patient: the government of her native Kenya. She shares how she helped her government drastically improve its process for opening up new businesses, a crucial part of economic health and growth, leading to new investments and a World Bank recognition as a top reformer.</td>\n",
       "      <td>https://www.ted.com/talks/charity_wayua_a_few_ways_to_fix_a_government/</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>21033</td>\n",
       "      <td>Mary Maker</td>\n",
       "      <td>Why I fight for the education of refugee girls (like me)</td>\n",
       "      <td>After fleeing war-torn South Sudan as a child, Mary Maker found security and hope in the school at Kenya's Kakuma Refugee Camp. Now a teacher of young refugees herself, she sees education as an essential tool for rebuilding lives -- and empowering a generation of girls who are too often denied entrance into the classroom. \"For the child of war, an education can turn their tears of loss into a passion for peace,\" Maker says.</td>\n",
       "      <td>https://www.ted.com/talks/mary_maker_why_i_fight_for_the_education_of_refugee_girls_like_me/</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  TALK_ID        SPEAKER_1  \\\n",
       "0     220   Joseph Lekuton   \n",
       "1    2221  Boniface Mwangi   \n",
       "2     523     Erik Hersman   \n",
       "3    2653    Charity Wayua   \n",
       "4   21033       Mary Maker   \n",
       "\n",
       "                                                      TITLE  \\\n",
       "0                                       A parable for Kenya   \n",
       "1                                  The day I stood up alone   \n",
       "2                              Reporting crisis via texting   \n",
       "3                            A few ways to fix a government   \n",
       "4  Why I fight for the education of refugee girls (like me)   \n",
       "\n",
       "                                                                                                                                                                                                                                                                                                                                                                                                                                   DESCRIPTION  \\\n",
       "0                                                                                                                                                                                                                                          Joseph Lekuton, a member of parliament in Kenya, starts with the story of his remarkable education, then offers a parable of how Africa can grow. His message of hope has never been more relevant.   \n",
       "1       Photographer Boniface Mwangi wanted to protest against corruption in his home country of Kenya. So he made a plan: He and some friends would stand up and heckle during a public mass meeting. But when the moment came ... he stood alone. What happened next, he says, showed him who he truly was. As he says, \"There are two most powerful days in your life. The day you are born, and the day you discover why.\" Graphic images.   \n",
       "2                                                                                                                                                                                   At TEDU 2009, Erik Hersman presents the remarkable story of Ushahidi, a GoogleMap mashup that allowed Kenyans to report and track violence via cell phone texts following the 2008 elections, and has evolved to continue saving lives in other countries.   \n",
       "3                                                                                      Charity Wayua put her skills as a cancer researcher to use on an unlikely patient: the government of her native Kenya. She shares how she helped her government drastically improve its process for opening up new businesses, a crucial part of economic health and growth, leading to new investments and a World Bank recognition as a top reformer.   \n",
       "4  After fleeing war-torn South Sudan as a child, Mary Maker found security and hope in the school at Kenya's Kakuma Refugee Camp. Now a teacher of young refugees herself, she sees education as an essential tool for rebuilding lives -- and empowering a generation of girls who are too often denied entrance into the classroom. \"For the child of war, an education can turn their tears of loss into a passion for peace,\" Maker says.   \n",
       "\n",
       "                                                                                            URL  \n",
       "0                                 https://www.ted.com/talks/joseph_lekuton_a_parable_for_kenya/  \n",
       "1                           https://www.ted.com/talks/boniface_mwangi_the_day_i_stood_up_alone/  \n",
       "2                          https://www.ted.com/talks/erik_hersman_reporting_crisis_via_texting/  \n",
       "3                       https://www.ted.com/talks/charity_wayua_a_few_ways_to_fix_a_government/  \n",
       "4  https://www.ted.com/talks/mary_maker_why_i_fight_for_the_education_of_refugee_girls_like_me/  "
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "querystr = \"<QUERY>\"\n",
    "\n",
    "df = pd.DataFrame(query_service(querystr)[\"results\"])\n",
    "df[COLUMNS].head()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python3 (cs-quickstart)",
   "language": "python",
   "name": "cs-quickstart"
  },
  "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}