athena_boto3_reference_guide.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ee7d970-5874-4533-9a7d-e150c1be9fc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82042597-cf92-436f-9bc8-9a6706c8cad9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "600173a1-d0a3-475e-878d-1d23ad87be57",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ.setdefault('AWS_DEFAULT_REGION', 'us-east-1')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e5ec549-0e20-43e5-a9b6-848042576554",
   "metadata": {},
   "outputs": [],
   "source": [
    "athena_client = boto3.client('athena')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bbdf986-b9a0-4a74-b1e8-7025f5a57202",
   "metadata": {},
   "outputs": [],
   "source": [
    "athena_client.list_work_groups?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "312ab940-f6c8-47a2-af2e-2f1276402e86",
   "metadata": {},
   "outputs": [],
   "source": [
    "athena_client.list_work_groups()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bec02f1d-658f-492c-b909-dfddc28cfab2",
   "metadata": {},
   "source": [
    "* Listing Amazon Athena Databases using Python boto3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a52fcbb-da87-44bc-befc-e530b969ecd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "athena_client.list_databases?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5f9c2d3-14d8-4bc6-8c26-5242ba72d7d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "databases = athena_client.list_databases(CatalogName='AwsDataCatalog')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00675c8f-4625-4383-89b0-7e69f60d5ef1",
   "metadata": {},
   "outputs": [],
   "source": [
    "type(databases)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d138f4f5-e74c-4ab5-b412-aac4046a1232",
   "metadata": {},
   "outputs": [],
   "source": [
    "type(databases['DatabaseList'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4673d94b-a9bc-423f-92ef-69140c6912ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "[database['Name'] for database in databases['DatabaseList']]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3093ad8-e00e-4ecf-b71a-8c80c03fddbc",
   "metadata": {},
   "source": [
    "* Listing Amazon Athena Database Tables using Python boto3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbb92cab-4ba3-4352-b9f9-7f5bdf60027e",
   "metadata": {},
   "outputs": [],
   "source": [
    "athena_client.list_table_metadata?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f77bcee-c0b3-4f64-bad3-e5ce999f874e",
   "metadata": {},
   "outputs": [],
   "source": [
    "athena_client.list_table_metadata(\n",
    "    CatalogName='AwsDataCatalog',\n",
    "    DatabaseName='retail_db'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95bc1df2-6897-4d7a-801c-70b27b2fde05",
   "metadata": {},
   "outputs": [],
   "source": [
    "tables = athena_client.list_table_metadata(\n",
    "    CatalogName='AwsDataCatalog',\n",
    "    DatabaseName='retail_db'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1072f65-814b-464a-aca1-7d1fac446825",
   "metadata": {},
   "outputs": [],
   "source": [
    "type(tables)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37df8e4f-f2fb-4b8a-b599-cca7630bb305",
   "metadata": {},
   "outputs": [],
   "source": [
    "tables['TableMetadataList']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f368f05-9be7-487c-afed-3d3072dfba27",
   "metadata": {},
   "outputs": [],
   "source": [
    "[table['Name'] for table in tables['TableMetadataList']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b73f605a-bf82-436f-b9ca-0afda555289f",
   "metadata": {},
   "outputs": [],
   "source": [
    "athena_client.get_table_metadata?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "719dc3d6-e9a7-4834-a6b2-659c2f59491a",
   "metadata": {},
   "outputs": [],
   "source": [
    "athena_client.get_table_metadata(\n",
    "    CatalogName='AwsDataCatalog',\n",
    "    DatabaseName='retail_db',\n",
    "    TableName='orders'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8741bbd8-0054-4b41-8f8b-90d1a178f213",
   "metadata": {},
   "outputs": [],
   "source": [
    "table_metadata = athena_client.get_table_metadata(\n",
    "    CatalogName='AwsDataCatalog',\n",
    "    DatabaseName='retail_db',\n",
    "    TableName='orders'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6c0a9cf-18a9-4ac2-96cc-16b7ae713915",
   "metadata": {},
   "outputs": [],
   "source": [
    "table_metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bf940a1-f204-4f03-a4b9-090a578f2aab",
   "metadata": {},
   "outputs": [],
   "source": [
    "table_metadata['TableMetadata']['TableType']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53b70dda-270b-4dec-9b40-58dae1c92dae",
   "metadata": {},
   "outputs": [],
   "source": [
    "table_metadata['TableMetadata']['Columns']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cc0f69f-c1da-4133-a25f-b58b5ea53727",
   "metadata": {},
   "outputs": [],
   "source": [
    "table_metadata['TableMetadata']['Parameters']['location']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7efbc1af-1fe8-41ca-8462-98bbac8d6107",
   "metadata": {},
   "outputs": [],
   "source": [
    "[table['Parameters']['location'] for table in tables['TableMetadataList']]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9da9ad26-b4a2-4175-845d-5ec5d5104a8b",
   "metadata": {},
   "source": [
    "* Run Queries on Athena Tables using Python boto3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9301e566-edae-4bb2-b3ae-2e22bf692612",
   "metadata": {},
   "outputs": [],
   "source": [
    "athena_client.start_query_execution?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cb86252-fa24-4c72-9a18-974f65a3366e",
   "metadata": {},
   "outputs": [],
   "source": [
    "athena_client.list_work_groups()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a381d8ac-16e6-4e8a-b9e8-5a722e37c734",
   "metadata": {},
   "outputs": [],
   "source": [
    "athena_client.get_work_group(WorkGroup='primary')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62b25d34-6b33-4e65-9083-b7394bae7a1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_execution = athena_client.start_query_execution(\n",
    "    QueryString='SELECT count(*) FROM myretail.orders'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e43adf9b-b1f2-4f67-8a61-d220c165e091",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_execution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09c7452d-e5c2-4eaf-9718-b37cac03eb2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "athena_client.get_query_execution(\n",
    "    QueryExecutionId=query_execution['QueryExecutionId']\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "519bd7e1-5e89-4ac9-9e17-e010f58928d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "!aws s3 ls s3://itvathena/wgprimary/79e570cb-6d51-4408-84f4-bc81556ec9d0.csv"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9dc93c6f-e9bc-4bfe-9f93-e69bfd853195",
   "metadata": {},
   "source": [
    "* Process Athena Query Results using Python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76e2568d-9a2c-468f-8be3-c1001138a6f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "athena_client.get_table_metadata(\n",
    "    CatalogName='AwsDataCatalog',\n",
    "    DatabaseName='myretail',\n",
    "    TableName='orders'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "160a5142-fb74-4c50-8eb5-ab14664036d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_str = \"\"\"\n",
    "    SELECT order_status, count(*) AS order_count\n",
    "    FROM myretail.orders\n",
    "    GROUP BY order_status\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7d47572-fe14-47b7-9093-14e495bad4d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_execution = athena_client.start_query_execution(\n",
    "    QueryString=query_str\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6c98ec3-e11b-4794-b2aa-2bf3cd7415bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "athena_client.get_query_execution(\n",
    "    QueryExecutionId=query_execution['QueryExecutionId']\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4391594d-0681-4372-ad46-309d67804023",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_results = athena_client.get_query_results(\n",
    "    QueryExecutionId=query_execution['QueryExecutionId']\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7081a45-a66d-424f-9828-2e943ae4ce58",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7567f644-d1bd-4120-9d70-a10f5c9ee132",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_results['ResultSet']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b17a862-aa45-40f0-8482-2b96e59797dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_results['ResultSet']['Rows']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d66b6f8e-a386-4d60-8d2e-d766edd14b53",
   "metadata": {},
   "source": [
    "* Custom Location for Athena query results using Python boto3\n",
    "\n",
    "```shell\n",
    "pip install pandas\n",
    "pip install fsspec\n",
    "pip install s3fs\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85b564ee-a3a1-43d8-814b-c89b3c6ffefe",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_str = \"\"\"\n",
    "    SELECT order_status, count(*) AS order_count\n",
    "    FROM myretail.orders\n",
    "    GROUP BY order_status\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d55de29-2380-4c2e-bb16-0a9645900750",
   "metadata": {},
   "outputs": [],
   "source": [
    "athena_client.start_query_execution?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95b1d239-b34e-42a0-b724-50a4f36dc6ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_execution = athena_client.start_query_execution(\n",
    "    QueryString=query_str,\n",
    "    ResultConfiguration={'OutputLocation': 's3://itv-retail/myretail/order_status_count_boto3'}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43521e89-5674-4dbf-9ac0-3f40e367ae80",
   "metadata": {},
   "outputs": [],
   "source": [
    "athena_client.get_query_execution(\n",
    "    QueryExecutionId=query_execution['QueryExecutionId']\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36dafc71-c605-47cf-9b30-4fd34f0bb972",
   "metadata": {},
   "outputs": [],
   "source": [
    "!aws s3 ls s3://itv-retail/myretail/order_status_count_boto3/8d388c59-cd56-4a1f-850e-c4225fcaed74.csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2513a725-bda6-47b5-91ae-430263027891",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_results = athena_client.get_query_results(\n",
    "    QueryExecutionId=query_execution['QueryExecutionId']\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6060cda1-3ff7-4708-abe6-b5d5b78ee0de",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_results['ResultSet']['Rows']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b23cd9bb-3d69-4356-b53a-e1239980822e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5b1bcc2-9947-4a74-9041-e9749328695f",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = pd.read_csv('s3://itv-retail/myretail/order_status_count_boto3/8d388c59-cd56-4a1f-850e-c4225fcaed74.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3b0c9cb-07f1-4c98-a772-e554a47ea408",
   "metadata": {},
   "outputs": [],
   "source": [
    "type(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a436c894-29e5-4d86-8ebc-c8fe0a16389c",
   "metadata": {},
   "outputs": [],
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f8c8629-89e5-4a25-b23a-1580902cc896",
   "metadata": {},
   "source": [
    "* Running CTAS against Athena using Python boto3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "id": "fdb2be67-87c9-464d-9440-716cdfc5f932",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_str = \"\"\"\n",
    "CREATE TABLE myretail.orders_part\n",
    "WITH (\n",
    "    format = 'TEXTFILE',\n",
    "    external_location = 's3://itv-retail/myretail/orders_part/',\n",
    "    field_delimiter = ',',\n",
    "    partitioned_by = ARRAY['order_month']\n",
    ")\n",
    "AS\n",
    "SELECT o.*,\n",
    "    cast(replace(substring(order_date, 1, 7), '-', '') AS INT) AS order_month\n",
    "FROM myretail.orders AS o\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "id": "ae1db471-b645-4738-bba2-2c1cc6b4bb50",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_execution = athena_client.start_query_execution(\n",
    "    QueryString=query_str\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "id": "61221cb8-d541-438a-96ad-5d5bf6b1fab0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'QueryExecution': {'QueryExecutionId': '323c6f93-f175-4ee3-807f-9bd91f490e07',\n",
       "  'Query': \"CREATE TABLE myretail.orders_part\\nWITH (\\n    format = 'TEXTFILE',\\n    external_location = 's3://itv-retail/myretail/orders_part/',\\n    field_delimiter = ',',\\n    partitioned_by = ARRAY['order_month']\\n)\\nAS\\nSELECT o.*,\\n    cast(replace(substring(order_date, 1, 7), '-', '') AS INT) AS order_month\\nFROM myretail.orders AS o\",\n",
       "  'StatementType': 'DDL',\n",
       "  'ResultConfiguration': {'OutputLocation': 's3://itvathena/wgprimary/tables/323c6f93-f175-4ee3-807f-9bd91f490e07'},\n",
       "  'QueryExecutionContext': {},\n",
       "  'Status': {'State': 'SUCCEEDED',\n",
       "   'SubmissionDateTime': datetime.datetime(2021, 9, 26, 9, 43, 30, 629000, tzinfo=tzlocal()),\n",
       "   'CompletionDateTime': datetime.datetime(2021, 9, 26, 9, 43, 35, 287000, tzinfo=tzlocal())},\n",
       "  'Statistics': {'EngineExecutionTimeInMillis': 4333,\n",
       "   'DataScannedInBytes': 327238,\n",
       "   'DataManifestLocation': 's3://itvathena/wgprimary/tables/323c6f93-f175-4ee3-807f-9bd91f490e07-manifest.csv',\n",
       "   'TotalExecutionTimeInMillis': 4658,\n",
       "   'QueryQueueTimeInMillis': 225,\n",
       "   'QueryPlanningTimeInMillis': 261,\n",
       "   'ServiceProcessingTimeInMillis': 100},\n",
       "  'WorkGroup': 'primary',\n",
       "  'EngineVersion': {'SelectedEngineVersion': 'Athena engine version 2',\n",
       "   'EffectiveEngineVersion': 'Athena engine version 2'}},\n",
       " 'ResponseMetadata': {'RequestId': '258af459-996c-4b07-bbdd-a1e68e6f20e6',\n",
       "  'HTTPStatusCode': 200,\n",
       "  'HTTPHeaders': {'content-type': 'application/x-amz-json-1.1',\n",
       "   'date': 'Sun, 26 Sep 2021 14:43:41 GMT',\n",
       "   'x-amzn-requestid': '258af459-996c-4b07-bbdd-a1e68e6f20e6',\n",
       "   'content-length': '2072',\n",
       "   'connection': 'keep-alive'},\n",
       "  'RetryAttempts': 0}}"
      ]
     },
     "execution_count": 96,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "athena_client.get_query_execution(\n",
    "    QueryExecutionId=query_execution['QueryExecutionId']\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fcc2246a-18c2-4e42-8543-8fddeeb1059a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}