195 lines
6.3 KiB
Python
195 lines
6.3 KiB
Python
|
# Databricks notebook source
|
||
|
|
# MAGIC %md
|
||
|
|
# MAGIC # Bronze ingestion: Redshift -> Unity Catalog (Trakt.tv example)
|
||
|
|
# MAGIC Assumes Trakt personal data (watch history, ratings, watchlist) was already
|
||
|
|
# MAGIC exported from the Trakt API into Redshift schema `trakt_raw`.
|
||
|
|
# MAGIC Pattern: incremental read -> audit columns -> idempotent MERGE -> watermark -> DQ gate.
|
||
|
|
|
||
|
|
# COMMAND ----------
|
||
|
|
|
||
|
|
import json
|
||
|
|
import uuid
|
||
|
|
from datetime import datetime, timezone
|
||
|
|
from pyspark.sql import functions as F
|
||
|
|
|
||
|
|
# ---- Parameters (set per environment by the job / bundle) ----
|
||
|
|
dbutils.widgets.text("catalog", "dev")
|
||
|
|
dbutils.widgets.text("source_schema", "trakt_raw")
|
||
|
|
dbutils.widgets.dropdown("mode", "incremental", ["incremental", "full"])
|
||
|
|
|
||
|
|
CATALOG = dbutils.widgets.get("catalog")
|
||
|
|
SRC_SCHEMA = dbutils.widgets.get("source_schema")
|
||
|
|
MODE = dbutils.widgets.get("mode")
|
||
|
|
|
||
|
|
BATCH_ID = str(uuid.uuid4())
|
||
|
|
RUN_TS = datetime.now(timezone.utc)
|
||
|
|
|
||
|
|
BRONZE = f"{CATALOG}.bronze"
|
||
|
|
OPS = f"{CATALOG}.ops"
|
||
|
|
WM = f"{OPS}.ingestion_watermark"
|
||
|
|
|
||
|
|
# ---- Source table config: business key + incremental cursor column ----
|
||
|
|
TABLES = {
|
||
|
|
"watch_history": {"key": ["id"], "cursor": "watched_at"},
|
||
|
|
"ratings": {"key": ["trakt_id", "rated_at"], "cursor": "rated_at"},
|
||
|
|
"watchlist": {"key": ["trakt_id", "listed_at"], "cursor": "listed_at"},
|
||
|
|
}
|
||
|
|
|
||
|
|
# COMMAND ----------
|
||
|
|
# MAGIC %md ## Secrets and connection (no credentials in code)
|
||
|
|
|
||
|
|
# COMMAND ----------
|
||
|
|
|
||
|
|
JDBC_URL = dbutils.secrets.get("redshift", "jdbc_url") # jdbc:redshift://host:5439/db
|
||
|
|
RS_USER = dbutils.secrets.get("redshift", "user")
|
||
|
|
RS_PASSWORD = dbutils.secrets.get("redshift", "password")
|
||
|
|
IAM_ROLE = dbutils.secrets.get("redshift", "unload_iam_role")
|
||
|
|
TEMPDIR = dbutils.secrets.get("redshift", "tempdir") # s3://bucket/redshift-tmp/ (UC external location)
|
||
|
|
|
||
|
|
|
||
|
|
def read_redshift(query: str):
|
||
|
|
"""Redshift connector unloads via S3 (faster than plain JDBC for large reads)."""
|
||
|
|
return (
|
||
|
|
spark.read.format("redshift")
|
||
|
|
.option("url", JDBC_URL)
|
||
|
|
.option("user", RS_USER)
|
||
|
|
.option("password", RS_PASSWORD)
|
||
|
|
.option("query", query)
|
||
|
|
.option("tempdir", TEMPDIR)
|
||
|
|
.option("aws_iam_role", IAM_ROLE)
|
||
|
|
.load()
|
||
|
|
)
|
||
|
|
|
||
|
|
# COMMAND ----------
|
||
|
|
# MAGIC %md ## Unity Catalog objects (Spark SQL DDL)
|
||
|
|
|
||
|
|
# COMMAND ----------
|
||
|
|
|
||
|
|
spark.sql(f"CREATE SCHEMA IF NOT EXISTS {BRONZE} COMMENT 'Raw landed data, minimal transformation'")
|
||
|
|
spark.sql(f"CREATE SCHEMA IF NOT EXISTS {OPS} COMMENT 'Pipeline control tables'")
|
||
|
|
|
||
|
|
spark.sql(f"""
|
||
|
|
CREATE TABLE IF NOT EXISTS {WM} (
|
||
|
|
source_table STRING,
|
||
|
|
last_cursor TIMESTAMP,
|
||
|
|
last_batch_id STRING,
|
||
|
|
updated_at TIMESTAMP
|
||
|
|
) USING DELTA
|
||
|
|
""")
|
||
|
|
|
||
|
|
# Allow new source columns to flow into bronze instead of failing the job
|
||
|
|
spark.conf.set("spark.databricks.delta.schema.autoMerge.enabled", "true")
|
||
|
|
|
||
|
|
# COMMAND ----------
|
||
|
|
# MAGIC %md ## Helpers: watermark, audit columns, idempotent write, DQ
|
||
|
|
|
||
|
|
# COMMAND ----------
|
||
|
|
|
||
|
|
|
||
|
|
def get_watermark(name: str):
|
||
|
|
if MODE == "full":
|
||
|
|
return None
|
||
|
|
row = spark.sql(
|
||
|
|
f"SELECT last_cursor FROM {WM} WHERE source_table = '{name}'"
|
||
|
|
).first()
|
||
|
|
return row["last_cursor"] if row else None
|
||
|
|
|
||
|
|
|
||
|
|
def set_watermark(name: str, cursor_value):
|
||
|
|
if cursor_value is None:
|
||
|
|
return
|
||
|
|
ts = cursor_value.strftime("%Y-%m-%d %H:%M:%S.%f")
|
||
|
|
spark.sql(f"""
|
||
|
|
MERGE INTO {WM} t
|
||
|
|
USING (SELECT '{name}' AS source_table,
|
||
|
|
TIMESTAMP'{ts}' AS last_cursor,
|
||
|
|
'{BATCH_ID}' AS last_batch_id,
|
||
|
|
current_timestamp() AS updated_at) s
|
||
|
|
ON t.source_table = s.source_table
|
||
|
|
WHEN MATCHED THEN UPDATE SET *
|
||
|
|
WHEN NOT MATCHED THEN INSERT *
|
||
|
|
""")
|
||
|
|
|
||
|
|
|
||
|
|
def add_audit_columns(df, name: str):
|
||
|
|
return (
|
||
|
|
df.withColumn("_source_table", F.lit(f"{SRC_SCHEMA}.{name}"))
|
||
|
|
.withColumn("_batch_id", F.lit(BATCH_ID))
|
||
|
|
.withColumn("_ingested_at", F.lit(RUN_TS))
|
||
|
|
.withColumn("_ingest_date", F.to_date(F.lit(RUN_TS)))
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def merge_insert_only(df, target: str, keys: list):
|
||
|
|
"""Idempotent: re-running a batch (or overlapping windows) never duplicates rows."""
|
||
|
|
if not spark.catalog.tableExists(target):
|
||
|
|
(df.write.format("delta")
|
||
|
|
.option("delta.enableChangeDataFeed", "true")
|
||
|
|
.saveAsTable(target))
|
||
|
|
return
|
||
|
|
df.createOrReplaceTempView("incoming")
|
||
|
|
on_clause = " AND ".join(f"t.{k} <=> s.{k}" for k in keys) # null-safe equality
|
||
|
|
spark.sql(f"""
|
||
|
|
MERGE INTO {target} t
|
||
|
|
USING incoming s
|
||
|
|
ON {on_clause}
|
||
|
|
WHEN NOT MATCHED THEN INSERT *
|
||
|
|
""")
|
||
|
|
|
||
|
|
|
||
|
|
def dq_gate(df, name: str, keys: list):
|
||
|
|
total = df.count()
|
||
|
|
null_keys = df.filter(" OR ".join(f"{k} IS NULL" for k in keys)).count()
|
||
|
|
if null_keys > 0:
|
||
|
|
raise ValueError(f"[{name}] {null_keys} rows with NULL business key; failing batch {BATCH_ID}")
|
||
|
|
return {"rows_read": total, "null_keys": null_keys}
|
||
|
|
|
||
|
|
# COMMAND ----------
|
||
|
|
# MAGIC %md ## Run: one pass per source table
|
||
|
|
|
||
|
|
# COMMAND ----------
|
||
|
|
|
||
|
|
results = {}
|
||
|
|
|
||
|
|
for name, cfg in TABLES.items():
|
||
|
|
keys, cursor = cfg["key"], cfg["cursor"]
|
||
|
|
target = f"{BRONZE}.trakt_{name}"
|
||
|
|
|
||
|
|
wm = get_watermark(name)
|
||
|
|
query = f"SELECT * FROM {SRC_SCHEMA}.{name}"
|
||
|
|
if wm is not None:
|
||
|
|
# >= (not >) so late-arriving rows at the boundary aren't missed;
|
||
|
|
# the MERGE de-duplicates the overlap
|
||
|
|
query += f" WHERE {cursor} >= '{wm.strftime('%Y-%m-%d %H:%M:%S')}'"
|
||
|
|
|
||
|
|
df = read_redshift(query)
|
||
|
|
df = add_audit_columns(df, name).cache()
|
||
|
|
|
||
|
|
stats = dq_gate(df, name, keys)
|
||
|
|
new_wm = df.agg(F.max(cursor)).first()[0] if stats["rows_read"] > 0 else None
|
||
|
|
|
||
|
|
merge_insert_only(df, target, keys)
|
||
|
|
set_watermark(name, new_wm) # only after a successful write
|
||
|
|
|
||
|
|
stats.update({"target": target, "watermark_before": str(wm), "watermark_after": str(new_wm)})
|
||
|
|
results[name] = stats
|
||
|
|
df.unpersist()
|
||
|
|
|
||
|
|
# COMMAND ----------
|
||
|
|
# MAGIC %md ## Governance (Unity Catalog grants) and post-load checks
|
||
|
|
|
||
|
|
# COMMAND ----------
|
||
|
|
|
||
|
|
spark.sql(f"GRANT USE SCHEMA, SELECT ON SCHEMA {BRONZE} TO `data-engineers`")
|
||
|
|
|
||
|
|
spark.sql(f"""
|
||
|
|
SELECT _source_table, _batch_id, COUNT(*) AS rows_loaded
|
||
|
|
FROM {BRONZE}.trakt_watch_history
|
||
|
|
GROUP BY _source_table, _batch_id
|
||
|
|
ORDER BY MAX(_ingested_at) DESC
|
||
|
|
LIMIT 5
|
||
|
|
""").show(truncate=False)
|
||
|
|
|
||
|
|
# COMMAND ----------
|
||
|
|
|
||
|
|
dbutils.notebook.exit(json.dumps({"batch_id": BATCH_ID, "mode": MODE, "results": results}))
|