examples/notebooks/delta2snowflake-incremental-load.ipynb

{
	"cells": [
		{
			"cell_type": "markdown",
			"metadata": {
				"editable": true,
				"trusted": true
			},
			"source": [
				"# Load data incrementally from Delta table to Snowflake\n"
			]
		},
		{
			"cell_type": "markdown",
			"metadata": {
				"editable": true,
				"trusted": true
			},
			"source": [
				"####  Run this cell to set up and start your interactive session.\n"
			]
		},
		{
			"cell_type": "code",
			"execution_count": null,
			"metadata": {
				"editable": true,
				"trusted": true
			},
			"outputs": [],
			"source": [
				"%session_id_prefix delta-snowflake-incremental-\n",
				"%idle_timeout 2880\n",
				"%glue_version 4.0\n",
				"%worker_type G.1X\n",
				"%number_of_workers 5\n",
				"%connections snowflake\n",
				"%%configure\n",
				"{\n",
				"    \"--datalake-formats\": \"delta\",\n",
				"    \"--conf\": \"spark.sql.extensions=io.delta.sql.DeltaSparkSessionExtension --conf spark.sql.catalog.spark_catalog=org.apache.spark.sql.delta.catalog.DeltaCatalog\",\n",
				"    \"--extra-py-files\": \"/opt/aws_glue_connectors/selected/datalake/delta-core_2.12-2.1.0.jar\"\n",
				"}"
			]
		},
		{
			"cell_type": "markdown",
			"metadata": {},
			"source": [
				"#### Configure your resources"
			]
		},
		{
			"cell_type": "code",
			"execution_count": null,
			"metadata": {
				"tags": [],
				"trusted": true
			},
			"outputs": [],
			"source": [
				"AWS_ACCOUNT_ID = \"123456789101\"\n",
				"REGION = \"us-east-1\"\n",
				"\n",
				"DELTA_DATASET_PATH = \"s3://<Your S3 bucket>/<Your S3 prefix>/delta_incremental/ghcn/\"\n",
				"\n",
				"SNOWFLAKE_CONNECTION_NAME = \"snowflake\"\n",
				"SNOWFLAKE_URL = \"YOUR_SNOWFLAKE_URL\"\n",
				"SNOWFLAKE_SECRET_ID = \"snowflake_credentials\"\n",
				"SNOWFLAKE_SCHEMA = \"public\"\n",
				"SNOWFLAKE_WAREHOUSE_NAME = \"YOUR_SNOWFLAKE_WAREHOUSE\"\n",
				"SNOWFLAKE_DATABASE_NAME = \"YOUR_SNOWFLAKE_DATABASE\"\n",
				"SNOWFLAKE_TABLE_NAME = \"ghcn\"\n",
				"SNOWFLAKE_TABLE_PRIMARY_KEYS = [\"ID\", \"DATE\", \"ELEMENT\"]"
			]
		},
		{
			"cell_type": "markdown",
			"metadata": {},
			"source": [
				"#### Initialize SparkSession and GlueContext"
			]
		},
		{
			"cell_type": "code",
			"execution_count": null,
			"metadata": {
				"tags": [],
				"trusted": true
			},
			"outputs": [],
			"source": [
				"import sys\n",
				"import json\n",
				"import boto3\n",
				"from botocore.exceptions import ClientError\n",
				"from awsglue.transforms import *\n",
				"from awsglue.utils import getResolvedOptions\n",
				"from pyspark.context import SparkContext\n",
				"from awsglue.context import GlueContext\n",
				"from awsglue.job import Job\n",
				"from delta.tables import *\n",
				"\n",
				"sc = SparkContext.getOrCreate()\n",
				"glueContext = GlueContext(sc)\n",
				"spark = glueContext.spark_session\n",
				"\n",
				"params = []\n",
				"if '--JOB_NAME' in sys.argv:\n",
				"    params.append('JOB_NAME')\n",
				"if '--TempDir' in sys.argv:\n",
				"    params.append('TempDir')\n",
				"args = getResolvedOptions(sys.argv, params)\n",
				"\n",
				"job_name= None\n",
				"if 'JOB_NAME' in args:\n",
				"    job_name = args['JOB_NAME']\n",
				"if not job_name:\n",
				"    job_name = \"delta-ghcn-incremental-load-notebook\"\n",
				"\n",
				"if 'TempDir' in args:\n",
				"    temp_dir = args['TempDir']\n",
				"if not temp_dir:\n",
				"    temp_dir = f\"s3://aws-glue-assets-{AWS_ACCOUNT_ID}-{REGION}/temporary/\"\n"
			]
		},
		{
			"cell_type": "markdown",
			"metadata": {
				"tags": []
			},
			"source": [
				"#### Determine target time range for CDC"
			]
		},
		{
			"cell_type": "code",
			"execution_count": null,
			"metadata": {
				"tags": [],
				"trusted": true
			},
			"outputs": [],
			"source": [
				"glue = boto3.client('glue')\n",
				"\n",
				"try:\n",
				"    res = glue.get_tags(ResourceArn=f\"arn:aws:glue:{REGION}:{AWS_ACCOUNT_ID}:job/{job_name}\")\n",
				"    if 'Tags' in res and 'lastQueryEndTime' in res['Tags']:\n",
				"        beginTime = res['Tags']['lastQueryEndTime']\n",
				"    else:\n",
				"        beginTime = \"1970-01-01 00:00:00\" ### retrieve all\n",
				"except Exception as e:\n",
				"    raise Exception(\"Failed to retrieve lastQueryEndTime tag via get_tags: \" + e.__str__())\n",
				"\n",
				"beginTime += \".001\" # to exclude the last commit processed at the previous run\n",
				"print(f\"beginTime: {beginTime}\")\n"
			]
		},
		{
			"cell_type": "code",
			"execution_count": null,
			"metadata": {
				"tags": [],
				"trusted": true
			},
			"outputs": [],
			"source": [
				"deltaTable = DeltaTable.forPath(spark, DELTA_DATASET_PATH)\n",
				"lastOperationTimestamp = deltaTable.history(1).select(\"timestamp\").collect()[0][0]\n",
				"\n",
				"endTime=lastOperationTimestamp\n",
				"print(f\"endTime: {endTime}\")"
			]
		},
		{
			"cell_type": "markdown",
			"metadata": {},
			"source": [
				"#### Run query"
			]
		},
		{
			"cell_type": "code",
			"execution_count": null,
			"metadata": {
				"tags": [],
				"trusted": true
			},
			"outputs": [],
			"source": [
				"df = spark.read.format(\"delta\") \\\n",
				"    .option(\"readChangeFeed\", \"true\") \\\n",
				"    .option(\"startingTimestamp\", beginTime) \\\n",
				"    .option(\"endingTimestamp\", endTime) \\\n",
				"    .load(DELTA_DATASET_PATH)"
			]
		},
		{
			"cell_type": "code",
			"execution_count": null,
			"metadata": {
				"tags": [],
				"trusted": true
			},
			"outputs": [],
			"source": [
				"df.show(20)"
			]
		},
		{
			"cell_type": "markdown",
			"metadata": {},
			"source": [
				"#### Merge changes into destination table"
			]
		},
		{
			"cell_type": "code",
			"execution_count": null,
			"metadata": {
				"tags": [],
				"trusted": true
			},
			"outputs": [],
			"source": [
				"column_names = [f.name for f in df.schema.fields]\n",
				"print(column_names)"
			]
		},
		{
			"cell_type": "code",
			"execution_count": null,
			"metadata": {
				"tags": [],
				"trusted": true
			},
			"outputs": [],
			"source": [
				"tmp_table_name = f\"{SNOWFLAKE_TABLE_NAME}_tmp\"\n",
				"\n",
				"post_actions = f\"BEGIN TRANSACTION; CREATE TABLE IF NOT EXISTS {SNOWFLAKE_SCHEMA}.{SNOWFLAKE_TABLE_NAME} AS SELECT * FROM {SNOWFLAKE_SCHEMA}.{tmp_table_name} WHERE 1=0; \"\n",
				"post_actions += f\"MERGE INTO {SNOWFLAKE_SCHEMA}.{SNOWFLAKE_TABLE_NAME} USING {SNOWFLAKE_SCHEMA}.{tmp_table_name} ON \"\n",
				"post_actions += ' AND '.join(f\"{SNOWFLAKE_SCHEMA}.{SNOWFLAKE_TABLE_NAME}.{pk} = {SNOWFLAKE_SCHEMA}.{tmp_table_name}.{pk}\" for pk in SNOWFLAKE_TABLE_PRIMARY_KEYS)\n",
				"\n",
				"post_actions += f\" WHEN MATCHED AND {SNOWFLAKE_SCHEMA}.{tmp_table_name}._change_type = 'update_postimage' THEN UPDATE SET \"\n",
				"post_actions += ', '.join(f\"{col} = {SNOWFLAKE_SCHEMA}.{tmp_table_name}.{col}\" for col in column_names)\n",
				"\n",
				"post_actions += f\" WHEN MATCHED AND {SNOWFLAKE_SCHEMA}.{tmp_table_name}._change_type = 'delete' THEN DELETE\"\n",
				"\n",
				"post_actions += \" WHEN NOT MATCHED THEN INSERT VALUES (\"\n",
				"post_actions += ', '.join(f\"{SNOWFLAKE_SCHEMA}.{tmp_table_name}.{col}\" for col in column_names)\n",
				"\n",
				"post_actions += f\"); DROP TABLE {SNOWFLAKE_SCHEMA}.{tmp_table_name}; COMMIT;\"\n",
				"\n",
				"print(f\"post_actions: {post_actions}\")"
			]
		},
		{
			"cell_type": "code",
			"execution_count": null,
			"metadata": {
				"tags": [],
				"trusted": true
			},
			"outputs": [],
			"source": [
				"secretsmanager = boto3.client('secretsmanager')\n",
				"res = secretsmanager.get_secret_value(SecretId=\"snowflake_credentials\")\n",
				"secret = json.loads(res['SecretString'])\n",
				"sfUser = secret['sfUser']\n",
				"sfPassword = secret['sfPassword']"
			]
		},
		{
			"cell_type": "code",
			"execution_count": null,
			"metadata": {
				"tags": [],
				"trusted": true
			},
			"outputs": [],
			"source": [
				"df.write \\\n",
				"  .format(\"net.snowflake.spark.snowflake\") \\\n",
				"  .option(\"sfUrl\", SNOWFLAKE_URL) \\\n",
				"  .option(\"sfWarehouse\", SNOWFLAKE_WAREHOUSE_NAME) \\\n",
				"  .option(\"sfDatabase\", SNOWFLAKE_DATABASE_NAME) \\\n",
				"  .option(\"sfUser\", sfUser) \\\n",
				"  .option(\"sfPassword\", sfPassword) \\\n",
				"  .option(\"dbtable\", tmp_table_name) \\\n",
				"  .option(\"postactions\", post_actions) \\\n",
				"  .mode(\"error\") \\\n",
				"  .save()"
			]
		},
		{
			"cell_type": "markdown",
			"metadata": {},
			"source": [
				"#### Update the last query end time"
			]
		},
		{
			"cell_type": "code",
			"execution_count": null,
			"metadata": {
				"tags": [],
				"trusted": true
			},
			"outputs": [],
			"source": [
				"tag = {\"lastQueryEndTime\": f\"{endTime}\"}\n",
				"\n",
				"try:\n",
				"    glue.tag_resource(ResourceArn=f\"arn:aws:glue:{REGION}:{AWS_ACCOUNT_ID}:job/{job_name}\",TagsToAdd=tag)\n",
				"except Exception as e:\n",
				"    raise Exception(\"Failed to update lastQueryEndTime tag via tags_resource: \" + e.__str__())"
			]
		},
		{
			"cell_type": "markdown",
			"metadata": {},
			"source": [
				"#### Update the record"
			]
		},
		{
			"cell_type": "code",
			"execution_count": null,
			"metadata": {
				"tags": [],
				"trusted": true
			},
			"outputs": [],
			"source": [
				"deltaTable.update(\n",
				"  condition = \"ID = 'AE000041196' AND DATE = '20221231' AND ELEMENT = 'PRCP'\",\n",
				"  set = { \"DATA_VALUE\": \"12345\" }\n",
				")"
			]
		},
		{
			"cell_type": "markdown",
			"metadata": {
				"tags": []
			},
			"source": [
				"#### Delete the record"
			]
		},
		{
			"cell_type": "code",
			"execution_count": null,
			"metadata": {
				"tags": [],
				"trusted": true
			},
			"outputs": [],
			"source": [
				"deltaTable.delete(\"ID = 'AE000041196' AND DATE = '20221231' AND ELEMENT = 'TMAX'\")"
			]
		},
		{
			"cell_type": "code",
			"execution_count": null,
			"metadata": {},
			"outputs": [],
			"source": []
		}
	],
	"metadata": {
		"kernelspec": {
			"display_name": "Glue PySpark",
			"language": "python",
			"name": "glue_pyspark"
		},
		"language_info": {
			"codemirror_mode": {
				"name": "python",
				"version": 3
			},
			"file_extension": ".py",
			"mimetype": "text/x-python",
			"name": "Python_Glue_Session",
			"pygments_lexer": "python3"
		},
		"toc-autonumbering": false
	},
	"nbformat": 4,
	"nbformat_minor": 4
}