notebooks/verify_setup.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Verify Polaris Setup\n",
    "\n",
    "This notebook allows us to check if the Apache Polaris setup has been successful and we are able to work with the catalog e.g create namespaces, tables etc.,"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import traceback\n",
    "from pathlib import Path\n",
    "\n",
    "import pyarrow as pa\n",
    "from pyiceberg.catalog.rest import RestCatalog\n",
    "from pyiceberg.exceptions import NamespaceAlreadyExistsError, TableAlreadyExistsError\n",
    "from pyiceberg.types import StringType"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Retrieve Principal Credentials\n",
    "As part of the catalog setup script, the Principal(`super_user`) credentials are stored in  `$PROJECT_HOME/work/principal.txt`, let us retrieve it for further operations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "principal_creds = Path(os.getcwd()).parent.joinpath(\"work\", \"principal.txt\")\n",
    "with open(principal_creds, \"r\") as file:\n",
    "    realm, client_id, client_secret = file.readline().split(\",\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define Variables\n",
    "Let us define some variables for us across the notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "namespace = \"demo_db\"\n",
    "table_name = \"fruits\"\n",
    "# IMPORTANT!!! /api/catalog or get the prefix from your OpenCatalog instance\n",
    "CATALOG_URI = \"http://localhost:18181/api/catalog\"\n",
    "catalog_name = \"polardb\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Working with Catalog\n",
    "Let us retrieve the catalog `polardb` that we created earlier using the `catalog_setup.yml` script."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "catalog = RestCatalog(\n",
    "    name=catalog_name,\n",
    "    **{\n",
    "        \"uri\": CATALOG_URI,\n",
    "        \"credential\": f\"{client_id}:{client_secret}\",\n",
    "        \"header.content-type\": \"application/vnd.api+json\",\n",
    "        \"header.X-Iceberg-Access-Delegation\": \"vended-credentials\",\n",
    "        \"header.Polaris-Realm\": realm,\n",
    "        \"warehouse\": catalog_name,\n",
    "        \"scope\": \"PRINCIPAL_ROLE:ALL\",\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Namespace\n",
    "Create a new namespace named `demo_db`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    catalog.create_namespace(namespace)\n",
    "except NamespaceAlreadyExistsError:\n",
    "    print(f\"Namespace '{namespace}' already exists\")\n",
    "except Exception as e:\n",
    "    print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Table\n",
    "Create a table named `fruits` with two columns."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_schema = pa.schema(\n",
    "    [\n",
    "        pa.field(\"id\", pa.int64(), nullable=False),\n",
    "        pa.field(\"name\", pa.string(), nullable=True),\n",
    "    ]\n",
    ")\n",
    "try:\n",
    "    new_tbl = catalog.create_table(\n",
    "        identifier=f\"{namespace}.{table_name}\",\n",
    "        schema=_schema,\n",
    "    )\n",
    "    print(new_tbl)\n",
    "except TableAlreadyExistsError:\n",
    "    print(f\"Table '{table_name}' already exists\")\n",
    "except Exception as e:\n",
    "    print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load Table\n",
    "Let us load the created table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    table = catalog.load_table(f\"{namespace}.{table_name}\")\n",
    "    df = table.scan().to_pandas()\n",
    "    print(df.head())\n",
    "except Exception as e:\n",
    "    print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Insert Data\n",
    "Insert some fruits data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    data = pa.Table.from_pylist(\n",
    "        [\n",
    "            {\"id\": 1, \"name\": \"mango\"},\n",
    "            {\"id\": 2, \"name\": \"banana\"},\n",
    "            {\"id\": 3, \"name\": \"orange\"},\n",
    "        ],\n",
    "        schema=_schema,\n",
    "    )\n",
    "    table.append(data)\n",
    "except Exception:\n",
    "    print(traceback.format_exc())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Query Data\n",
    "Query the inserted data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = table.scan().to_pandas()\n",
    "df.head(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Schema Evolution\n",
    "Let us now add a new column named `season` to the `fruits` table."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "with table.update_schema() as update:\n",
    "    update.add_column(\"season\",StringType(),doc=\"Fruit Season\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Print the table to view its structure and other details. If you go back and query the data again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(table)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you query the table again your query still works and in fact gets the new column with empty/null value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = table.scan().to_pandas()\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us insert season data using the modified new schema."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_schema =  _schema.append( pa.field(\"season\", pa.string(), nullable=True),)\n",
    "# New data with season column\n",
    "new_table = pa.Table.from_pylist([\n",
    "    {\"id\": 1, \"name\": \"mango\",\"season\": \"summer\"},\n",
    "    {\"id\": 2, \"name\": \"banana\",\"season\": \"all\"},\n",
    "    {\"id\": 3, \"name\": \"orange\",\"season\": \"winter\"},\n",
    "   \n",
    "],schema=new_schema)\n",
    "table.overwrite(new_table)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now querying again,will show the updated data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = table.scan().to_pandas()\n",
    "df.head()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}