Below is the code necessary to create a Databricks notebook source file that can be imported into Databricks. This file can act as a template for creating ETL logic to build tables in Databricks. Once the notebook is prepared it can be set to run by a Databricks workflow job.
The template is parameterized. This means the developer just needs to provide the destination database, the destination schema, the destination table and the SQL logic.
(Note: this simple example is a full load solution and not a incremental load solution. An incremental load solution can be achieve by writing sufficiently robust SQL that is use case specific.)
The SQL is provided as a variable and the table or table names are stored in a list allowing for a large degree of flexibility for creating a single pipeline that builds multiple database objects.
Another important feature of the code is that it compensates for the fact that Databricks does not have a native acknowledgement of primary keys or restrictions on their violations. A list of primary keys can be provided and if any of those keys are null or not distinct the code will throw an error.
The code will also assign metadata fields to each record created including the job run id as the ETL id, the created date and the updated date.
# Databricks notebook source
# MAGIC %md
# MAGIC https://tidbytez.com/<br />
# MAGIC This is an ETL notebook.<br />
# COMMAND ----------
# Libraries
import os
from pyspark.sql import functions as F
from pyspark.sql.types import StringType
from pyspark.sql.functions import lit, concat_ws, isnan, when, count, col
from datetime import datetime, timedelta
# COMMAND ----------
# Functions
# Generate ETL ID
def get_etl_id():
try:
run_id = (
dbutils.notebook.entry_point.getDbutils()
.notebook()
.getContext()
.currentRunId()
.toString()
)
if run_id.isdigit():
etl_id = bigint(run_id)
return etl_id
else
etl_id = bigint(1)
return etl_id
except:
print("Could not return an etl_id number")
# Build database object
def build_object(dest_db_name, schema_name, table_name, pk_fields, sql_query):
# Destination Database and table
table_location = dest_db_name + "." + table_name
# External table file location
file_location = "/mnt/" + schema_name + "/" + table_name
# Create Dataframe
df = sql_query
# Count nulls in Primary Key
cnt_pk_nulls = df.select(
[count(when(isnan(c) | col(c).isNull(), c)).alias(c) for c in pk_fields]
).collect()[0][0]
# Dataframe record count
cnt_rows = df.count()
# Primary Key distinct count
cnt_dist = df.dropDuplicates(pk_fields).count()
# Error message
message = ""
# Join metadata to dataframe
global meta_df
meta = meta_df
df = df.withColumn("key", lit(1))
# inner join on two dataframes
df = df.join(meta, df.key == meta.key, "inner").drop(df.key).drop(meta.key)
# Write dataframe to table
if cnt_pk_nulls == 0:
if cnt_rows == cnt_dist:
df.write.mode("overwrite").format("delta").option(
"mergeSchema", "false"
).option("path", file_location).saveAsTable(table_location)
else:
message = "Primary Key is not unique"
else:
message = "Primary Key contains nulls"
if message != "":
raise Exception(message)
# COMMAND ----------
# Variables
# Destinations
# File location
schema_name = "YOUR_SCHEMA_NAME_HERE"
# Database location
dest_db_name = "YOUR_DEST_DATABASE_NAME_HERE"
# PK fields
pk_fields = ["EXAMPLE_ID", "EXAMPLE_LOCATION"]
# Metadata
etl_id = get_etl_id()
t = datetime.utcnow()
# Create metadata dataFrame
data = [(1, etl_id, t, t)]
columns = ["key", "ETL_ID", "CREATED_DATE", "UPDATED_DATE"]
meta_df = spark.createDataFrame(data, columns)
meta_df = meta_df.withColumn("ETL_ID", meta_df["ETL_ID"].cast("int"))
# COMMAND ----------
# Table name variable list
table_list = [
{"table_name": "EXAMPLE_TABLE"}
]
# COMMAND ----------
# Iterate through table variables
for i in range(len(table_list)):
table_name = table_list[i].get("table_name")
# SQL query
sql_query = spark.sql(
f"""
SELECT 1 AS EXAMPLE_ID,
'TEXAS' AS "EXAMPLE_LOCATION"
"""
)
build_object(dest_db_name, schema_name, table_name, pk_fields, sql_query)
Tidbytez 