{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b7486b98",
   "metadata": {},
   "source": [
    "# Query Opteryx.app with SQLAlchemy\n",
    "\n",
    "Connects to [opteryx.app](https://opteryx.app), the hosted service, with the `opteryx-sqlalchemy` dialect, queries the public **National Vulnerability Database (NVD)** dataset into pandas, then uses the same connection on data only your account can see.\n",
    "\n",
    "This notebook accompanies [Using SQLAlchemy from a Notebook or Script](https://docs.opteryx.app/docs/guides/using-sqlalchemy) on docs.opteryx.app.\n",
    "\n",
    "**You need an opteryx.app account.** The datasets under `public` are readable by any signed-in account, but SQLAlchemy goes through the Jobs API, and that API needs a credential for every query, public data included. Signing up is free at [opteryx.app](https://opteryx.app). Then create a personal access token in Opteryx Studio under **Settings → API Tokens**. It gives you a **client ID** and a **client secret**, and the secret is shown only once."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a23dcf08",
   "metadata": {},
   "source": [
    "## 1. Install"
   ]
  },
  {
   "cell_type": "code",
   "id": "f3bd8ae4",
   "metadata": {},
   "source": [
    "%pip install -q opteryx-sqlalchemy pandas matplotlib"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "5bd35bc7",
   "metadata": {},
   "source": [
    "## 2. Credentials\n",
    "\n",
    "Keep the client secret out of the notebook itself, so it isn't saved with it or shared by accident. This cell looks for the client ID and secret in three places, in order:\n",
    "\n",
    "1. **Colab Secrets.** Click the key icon in Colab's left sidebar, add `OPTERYX_CLIENT_ID` and `OPTERYX_CLIENT_SECRET`, and turn on notebook access for both. Secrets stay with your Colab account, so you set them once and every notebook can use them.\n",
    "2. **Environment variables** with the same names, for running the notebook outside Colab.\n",
    "3. **A prompt.** If neither is set, the cell asks, and the secret is typed into a hidden field."
   ]
  },
  {
   "cell_type": "code",
   "id": "682c6c20",
   "metadata": {},
   "source": [
    "import os\n",
    "from getpass import getpass\n",
    "\n",
    "\n",
    "def credential(name, secret=False):\n",
    "    try:\n",
    "        from google.colab import userdata\n",
    "\n",
    "        return userdata.get(name)\n",
    "    except Exception:\n",
    "        pass   # not in Colab, or the secret isn't set or shared with this notebook\n",
    "    if os.environ.get(name):\n",
    "        return os.environ[name]\n",
    "    return getpass(f\"{name}: \") if secret else input(f\"{name}: \")\n",
    "\n",
    "\n",
    "CLIENT_ID = credential(\"OPTERYX_CLIENT_ID\")\n",
    "CLIENT_SECRET = credential(\"OPTERYX_CLIENT_SECRET\", secret=True)\n",
    "print(\"Client ID:\", CLIENT_ID)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "c50cbaa9",
   "metadata": {},
   "source": [
    "## 3. Connect\n",
    "\n",
    "The client ID and secret go in the username and password slots of the connection URL. `URL.create` builds the URL and escapes any characters in the secret that would otherwise break it. The dialect exchanges the pair for a short-lived access token when the connection opens, so you never call the token endpoint yourself.\n",
    "\n",
    "The first query counts the NVD records, which also checks that the credentials work. If it fails with `Authentication error`, opteryx.app didn't accept the client ID and secret. Check both values, and that the token hasn't been revoked."
   ]
  },
  {
   "cell_type": "code",
   "id": "5683d870",
   "metadata": {},
   "source": [
    "from sqlalchemy import URL, create_engine, text\n",
    "\n",
    "engine = create_engine(\n",
    "    URL.create(\n",
    "        \"opteryx\",\n",
    "        username=CLIENT_ID,\n",
    "        password=CLIENT_SECRET,\n",
    "        host=\"opteryx.app\",\n",
    "        port=443,\n",
    "        database=\"default\",\n",
    "        query={\"ssl\": \"true\"},\n",
    "    )\n",
    ")\n",
    "\n",
    "with engine.connect() as conn:\n",
    "    total = conn.execute(text(\"SELECT COUNT(*) FROM public.security.nvd_vulnerabilities\")).scalar()\n",
    "print(f\"{total:,} CVE records in public.security.nvd_vulnerabilities\")"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "a91a4f84",
   "metadata": {},
   "source": [
    "## 4. Query the public NVD dataset into pandas\n",
    "\n",
    "`public.security.nvd_vulnerabilities` has one row per CVE: its ID, when it was published, its severity and CVSS score, and a description. `pandas.read_sql_query` runs a query and returns a DataFrame.\n",
    "\n",
    "How many CVEs fall into each severity band:"
   ]
  },
  {
   "cell_type": "code",
   "id": "047a4bb8",
   "metadata": {},
   "source": [
    "import pandas as pd\n",
    "\n",
    "with engine.connect() as conn:\n",
    "    by_severity = pd.read_sql_query(\n",
    "        \"\"\"\n",
    "        SELECT severity, COUNT(*) AS cves\n",
    "        FROM public.security.nvd_vulnerabilities\n",
    "        GROUP BY severity\n",
    "        ORDER BY cves DESC\n",
    "        \"\"\",\n",
    "        conn,\n",
    "    )\n",
    "by_severity"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "258b351a",
   "metadata": {},
   "source": [
    "How many CVEs were published each year. The counting happens on opteryx.app, so only one row per year comes back to the notebook:"
   ]
  },
  {
   "cell_type": "code",
   "id": "e56b2b2b",
   "metadata": {},
   "source": [
    "with engine.connect() as conn:\n",
    "    per_year = pd.read_sql_query(\n",
    "        \"\"\"\n",
    "        SELECT EXTRACT(YEAR FROM published_at) AS year, COUNT(*) AS cves\n",
    "        FROM public.security.nvd_vulnerabilities\n",
    "        GROUP BY EXTRACT(YEAR FROM published_at)\n",
    "        ORDER BY year\n",
    "        \"\"\",\n",
    "        conn,\n",
    "    )\n",
    "per_year.plot.bar(x=\"year\", y=\"cves\", legend=False, figsize=(12, 4), title=\"CVEs published per year\")"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "dffe1ca6",
   "metadata": {},
   "source": [
    "The most recently published critical CVEs:"
   ]
  },
  {
   "cell_type": "code",
   "id": "3059d984",
   "metadata": {},
   "source": [
    "with engine.connect() as conn:\n",
    "    recent_critical = pd.read_sql_query(\n",
    "        \"\"\"\n",
    "        SELECT cve_id, published_at, cvss_score, description\n",
    "        FROM public.security.nvd_vulnerabilities\n",
    "        WHERE severity = 'CRITICAL'\n",
    "        ORDER BY published_at DESC\n",
    "        LIMIT 10\n",
    "        \"\"\",\n",
    "        conn,\n",
    "    )\n",
    "recent_critical"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "6fd4a89c",
   "metadata": {},
   "source": [
    "Datasets can be joined like tables in any database. `public.security.cisa_kev` is CISA's catalog of vulnerabilities known to be exploited. Joining it to NVD on the CVE ID adds each one's CVSS score:"
   ]
  },
  {
   "cell_type": "code",
   "id": "01b0de4f",
   "metadata": {},
   "source": [
    "with engine.connect() as conn:\n",
    "    exploited = pd.read_sql_query(\n",
    "        \"\"\"\n",
    "        SELECT kev.cve_id, kev.vendor_project, kev.product, kev.date_added, nvd.severity, nvd.cvss_score\n",
    "        FROM public.security.cisa_kev AS kev\n",
    "        INNER JOIN public.security.nvd_vulnerabilities AS nvd\n",
    "            ON kev.cve_id = nvd.cve_id\n",
    "        ORDER BY kev.date_added DESC\n",
    "        LIMIT 10\n",
    "        \"\"\",\n",
    "        conn,\n",
    "    )\n",
    "exploited"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "042ca601",
   "metadata": {},
   "source": [
    "## 5. The same connection on your own data\n",
    "\n",
    "Nothing about the connection changes for data that isn't public. Each dataset has a dot-separated name, `workspace.collection.dataset`, and the workspace's policies decide what your credential can read. See [Access and Permissions](https://docs.opteryx.app/docs/core-concepts/access-and-permissions) for how.\n",
    "\n",
    "This lists every schema (`workspace.collection`) your credential can see. It reads opteryx.app's metadata rather than running a query:"
   ]
  },
  {
   "cell_type": "code",
   "id": "9cb55283",
   "metadata": {},
   "source": [
    "from sqlalchemy import inspect\n",
    "\n",
    "with engine.connect() as conn:\n",
    "    schemas = inspect(conn).get_schema_names()\n",
    "for schema in schemas:\n",
    "    print(schema)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "3cae786c",
   "metadata": {},
   "source": [
    "Pick one of your own datasets and query it exactly as you queried NVD. Set `MY_DATASET` to its full name, for example `acme.sales.orders`. If you haven't loaded anything yet, [Load and Query Data](https://docs.opteryx.app/docs/getting-started/reading-data) shows how."
   ]
  },
  {
   "cell_type": "code",
   "id": "4610103b",
   "metadata": {},
   "source": [
    "MY_DATASET = \"\"   # e.g. \"acme.sales.orders\"\n",
    "\n",
    "if MY_DATASET:\n",
    "    workspace_collection, dataset = MY_DATASET.rsplit(\".\", 1)\n",
    "    with engine.connect() as conn:\n",
    "        for column in inspect(conn).get_columns(dataset, schema=workspace_collection):\n",
    "            print(f\"{column['name']:<30} {column['type']}\")\n",
    "        mine = pd.read_sql_query(f\"SELECT * FROM {MY_DATASET} LIMIT 10\", conn)\n",
    "    display(mine)\n",
    "else:\n",
    "    print(\"Set MY_DATASET to one of your own datasets, then run this cell again.\")"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "7259a886",
   "metadata": {},
   "source": [
    "## Things to know\n",
    "\n",
    "- **Read-only.** `commit()` and `rollback()` do nothing, and the dialect doesn't support DDL or DML.\n",
    "- **No bound parameters yet.** `text(\"... WHERE cvss_score > :min\")` with a parameters dict fails on the server today. The [guide](https://docs.opteryx.app/docs/guides/using-sqlalchemy#parameterising-queries) shows the `SET` workaround, and when it is safe to use.\n",
    "- **Large results.** `conn.execution_options(stream_results=True, max_row_buffer=500)` fetches a large result in pages, and `result_format=\"parquet\"` fetches those pages as Parquet rather than JSON.\n",
    "- **Rotating the secret.** If a secret leaks, revoke that token under **Settings → API Tokens** and create a new one."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}