153 lines
5.0 KiB
Python
153 lines
5.0 KiB
Python
# Databricks notebook source
|
|
# MAGIC %md
|
|
# MAGIC # Bronze ingestion: Redshift foreign catalog -> Delta bronze (Trakt example)
|
|
# MAGIC The Redshift login lives in a Unity Catalog **connection**; the **foreign catalog**
|
|
# MAGIC built on it is what shows up in the Catalog Explorer UI. No secrets, JDBC URL or S3
|
|
# MAGIC tempdir in this notebook. Foreign tables are read-only, so we land into a native Delta catalog.
|
|
|
|
# COMMAND ----------
|
|
|
|
import json
|
|
import uuid
|
|
from datetime import datetime, timezone
|
|
from pyspark.sql import functions as F
|
|
|
|
dbutils.widgets.text("source_catalog", "redshift_bronze") # foreign catalog (Redshift connection)
|
|
dbutils.widgets.text("source_schema", "trakt_raw")
|
|
dbutils.widgets.text("target_catalog", "dev") # native Delta catalog
|
|
dbutils.widgets.dropdown("mode", "incremental", ["incremental", "full"])
|
|
|
|
SRC_CAT = dbutils.widgets.get("source_catalog")
|
|
SRC_SCHEMA = dbutils.widgets.get("source_schema")
|
|
TGT_CAT = dbutils.widgets.get("target_catalog")
|
|
MODE = dbutils.widgets.get("mode")
|
|
|
|
BATCH_ID = str(uuid.uuid4())
|
|
RUN_TS = datetime.now(timezone.utc)
|
|
|
|
BRONZE = f"{TGT_CAT}.bronze"
|
|
WM = f"{TGT_CAT}.ops.ingestion_watermark"
|
|
|
|
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 ## Sanity check the federation (what you see in the UI, via SQL)
|
|
|
|
# COMMAND ----------
|
|
|
|
spark.sql(f"SHOW TABLES IN {SRC_CAT}.{SRC_SCHEMA}").show(truncate=False)
|
|
spark.sql(f"DESCRIBE EXTENDED {SRC_CAT}.{SRC_SCHEMA}.watch_history").show(50, truncate=False)
|
|
|
|
# COMMAND ----------
|
|
# MAGIC %md ## Target objects
|
|
|
|
# COMMAND ----------
|
|
|
|
spark.sql(f"CREATE SCHEMA IF NOT EXISTS {BRONZE}")
|
|
spark.sql(f"CREATE SCHEMA IF NOT EXISTS {TGT_CAT}.ops")
|
|
spark.sql(f"""
|
|
CREATE TABLE IF NOT EXISTS {WM} (
|
|
source_table STRING, last_cursor TIMESTAMP,
|
|
last_batch_id STRING, updated_at TIMESTAMP
|
|
) USING DELTA
|
|
""")
|
|
spark.conf.set("spark.databricks.delta.schema.autoMerge.enabled", "true")
|
|
|
|
# COMMAND ----------
|
|
# MAGIC %md ## Helpers
|
|
|
|
# COMMAND ----------
|
|
|
|
|
|
def get_watermark(name):
|
|
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, value):
|
|
if value is None:
|
|
return
|
|
ts = 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(df, name):
|
|
return (df.withColumn("_source_table", F.lit(f"{SRC_CAT}.{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, keys):
|
|
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)
|
|
spark.sql(f"""
|
|
MERGE INTO {target} t USING incoming s ON {on_clause}
|
|
WHEN NOT MATCHED THEN INSERT *
|
|
""")
|
|
|
|
|
|
def dq_gate(df, name, keys):
|
|
total = df.count()
|
|
null_keys = df.filter(" OR ".join(f"{k} IS NULL" for k in keys)).count()
|
|
if null_keys:
|
|
raise ValueError(f"[{name}] {null_keys} NULL business keys; failing batch {BATCH_ID}")
|
|
return {"rows_read": total, "null_keys": null_keys}
|
|
|
|
# COMMAND ----------
|
|
# MAGIC %md ## Run
|
|
|
|
# COMMAND ----------
|
|
|
|
results = {}
|
|
|
|
for name, cfg in TABLES.items():
|
|
keys, cursor = cfg["key"], cfg["cursor"]
|
|
source = f"{SRC_CAT}.{SRC_SCHEMA}.{name}"
|
|
target = f"{BRONZE}.trakt_{name}"
|
|
|
|
wm = get_watermark(name)
|
|
src_df = spark.table(source) # federated read, no connector options
|
|
if wm is not None:
|
|
src_df = src_df.filter(F.col(cursor) >= F.lit(wm)) # should push down to Redshift
|
|
|
|
df = add_audit(src_df, name).cache()
|
|
stats = dq_gate(df, name, keys)
|
|
new_wm = df.agg(F.max(cursor)).first()[0] if stats["rows_read"] else None
|
|
|
|
merge_insert_only(df, target, keys)
|
|
set_watermark(name, new_wm)
|
|
|
|
stats.update({"target": target, "watermark_after": str(new_wm)})
|
|
results[name] = stats
|
|
df.unpersist()
|
|
|
|
# COMMAND ----------
|
|
# MAGIC %md ## Verify pushdown (look for the filter inside the federated scan, not above it)
|
|
|
|
# COMMAND ----------
|
|
|
|
spark.table(f"{SRC_CAT}.{SRC_SCHEMA}.watch_history") \
|
|
.filter("watched_at >= '2026-01-01'") \
|
|
.select("id", "watched_at") \
|
|
.explain("formatted")
|
|
|
|
# COMMAND ----------
|
|
|
|
dbutils.notebook.exit(json.dumps({"batch_id": BATCH_ID, "mode": MODE, "results": results})) |