ml

claude/bronze_redshift.py

@@ -0,0 +1,195 @@
+# 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}))