docs/examples/testing_new_awswrangler.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fb00cddf-d747-46c4-a2cb-c3c495e46d85",
   "metadata": {},
   "source": [
    "# Testing pydbtools with awswrangler 2.12+\n",
    "\n",
    "Later versions of awswrangler have a fix that allows large dates, e.g. the year 2500, to be processed correctly. This needs the application of additional parameters which `pydbtools` should apply automatically.\n",
    "\n",
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8e26667d-a825-4191-8cbc-0848181a22b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "import awswrangler as wr\n",
    "import pydbtools as pydb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2f75d48f-693a-4217-aa40-6df30a61fdee",
   "metadata": {},
   "outputs": [],
   "source": [
    "# setup your own testing area (set foldername = GH username)\n",
    "foldername = \"mratford\"  # GH username\n",
    "foldername = foldername.lower().replace(\"-\", \"_\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "832bdd69-231b-4d59-8147-ab81a9cb0cee",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "deleting objs\n",
      "aws_example_mratford found deleting\n"
     ]
    }
   ],
   "source": [
    "bucketname = \"alpha-everyone\"\n",
    "s3_base_path = f\"s3://{bucketname}/{foldername}/\"\n",
    "\n",
    "db_name = f\"aws_example_{foldername}\"\n",
    "source_db_base_path = f\"s3://{bucketname}/{foldername}/source_db/\"\n",
    "\n",
    "# Delete all the s3 files in a given path\n",
    "if wr.s3.list_objects(s3_base_path):\n",
    "    print(\"deleting objs\")\n",
    "    wr.s3.delete_objects(s3_base_path)\n",
    "\n",
    "# Delete the database if it exists\n",
    "df_dbs = wr.catalog.databases(None)\n",
    "if db_name in df_dbs[\"Database\"].to_list():\n",
    "    print(f\"{db_name} found deleting\")\n",
    "    wr.catalog.delete_database(name=db_name)\n",
    "\n",
    "# Setup source database\n",
    "# Create the database\n",
    "wr.catalog.create_database(db_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c855d4f4-f404-48b4-ab9b-94285093e529",
   "metadata": {},
   "source": [
    "## Create a dataframe with difficult dates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7f518116-edc1-4132-9721-d1284c27d976",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame(\n",
    "    {\n",
    "        \"event\": [\n",
    "            \"Foundation of Mega-City One\",\n",
    "            \"Newspeak eclipses oldspeak\",\n",
    "            \"The year 3000\",\n",
    "        ],\n",
    "        \"future_date\": [\"2031-03-17\", \"2050-10-05\", \"3000-01-01\"],\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5bea3cd4-8784-484e-bee2-0b4e1038d5c9",
   "metadata": {},
   "source": [
    "Write to a table. We need to use strings as types to get the data into the table."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e2871323-52d7-44f0-be1d-48e37f820065",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'paths': ['s3://alpha-everyone/mratford/future/28e04de56d6240c0b03185d11dce343f.snappy.parquet'],\n",
       " 'partitions_values': {}}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "table_name = \"future\"\n",
    "wr.s3.to_parquet(\n",
    "    df=df,\n",
    "    path=os.path.join(s3_base_path, table_name),\n",
    "    index=False,\n",
    "    dataset=True,\n",
    "    database=db_name,\n",
    "    table=table_name,\n",
    "    mode=\"overwrite\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d79e21a5-f56e-4928-adcb-285ebdb69a21",
   "metadata": {},
   "source": [
    "Now use the existing table to create a new table with the date strings converted to dates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "727e14c6-b4c0-413f-9267-d38c52964eab",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "CREATE TABLE aws_example_mratford.new_future\n",
      "WITH (\n",
      "      external_location = 's3://alpha-everyone/mratford/source_db/new_future'\n",
      ")\n",
      "AS SELECT event, date(future_date) as futuredate\n",
      "FROM aws_example_mratford.future;\n",
      "\n"
     ]
    }
   ],
   "source": [
    "sql = f\"\"\"\n",
    "CREATE TABLE {db_name}.new_future\n",
    "WITH (\n",
    "      external_location = '{source_db_base_path}new_future'\n",
    ")\n",
    "AS SELECT event, date(future_date) as futuredate\n",
    "FROM {db_name}.future;\n",
    "\"\"\"\n",
    "print(sql)\n",
    "_ = pydb.start_query_execution_and_wait(sql)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a812cf2d-a11f-4537-aa0e-484d5fca850e",
   "metadata": {},
   "source": [
    "If we query the table using a standard awswrangler query it will fail."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ff1586e8-b03b-4344-a7de-f78f0d804f08",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Failed as expected.\n"
     ]
    }
   ],
   "source": [
    "sql = f\"\"\"\n",
    "select event, futuredate\n",
    "from {db_name}.new_future\n",
    "\"\"\"\n",
    "\n",
    "try:\n",
    "    wr.athena.read_sql_query(sql, database=db_name, ctas_approach=False)\n",
    "except AttributeError:\n",
    "    print(\"Failed as expected.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4edc318b-b981-4268-b922-aff25198edb6",
   "metadata": {},
   "source": [
    "A pydbtools query passes the additional parameters to awswrangler to handle the dates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "be0d4ffc-96b7-4863-901f-41fef7c2bce4",
   "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>event</th>\n",
       "      <th>futuredate</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Foundation of Mega-City One</td>\n",
       "      <td>2031-03-17</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Newspeak eclipses oldspeak</td>\n",
       "      <td>2050-10-05</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>The year 3000</td>\n",
       "      <td>3000-01-01</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                         event  futuredate\n",
       "0  Foundation of Mega-City One  2031-03-17\n",
       "1   Newspeak eclipses oldspeak  2050-10-05\n",
       "2                The year 3000  3000-01-01"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pydb.read_sql_query(sql)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1755821-d7eb-48c6-ade6-d489b39ae968",
   "metadata": {},
   "source": [
    "Clean up the database and table."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "9b02e820-d623-4200-88bf-4cb9ede656be",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pydb.delete_database_and_data(db_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5bbf90d0-e36b-4697-b4ac-8bc773f1a602",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pydbtools",
   "language": "python",
   "name": "pydbtools"
  },
  "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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}