Data Transformation and Preparation Questions

from airflow.decorators import dag, task
from airflow.utils.dates import days_ago
from datetime import timedelta
import uuid
import time

DEFAULT_ARGS = {
    "retries": 3,
    "retry_delay": timedelta(minutes=5),
    "depends_on_past": False,
}

@dag(start_date=days_ago(7), schedule_interval='@daily', default_args=DEFAULT_ARGS, catchup=True)
def upsert_partitioned_table():
    @task()
    def extract_transform(execution_date=None, **context):
        # produce dataframe/rows for the specific partition (execution_date)
        # write to staging location/table named with run_id for idempotency
        run_id = context['run_id']
        partition = execution_date.strftime('%Y-%m-%d')
        staging_table = f"staging.analytics_{partition}_{run_id.replace(':','_')}"
        # Example: write parquet to s3://bucket/path/{partition}/{run_id}/
        # ensure write is atomic (write to temp then rename)
        return {"staging_table": staging_table, "partition": partition}

    @task()
    def promote(staging_table: str, partition: str):
        lock_key = f"lock:analytics:{partition}"
        acquired = acquire_lock(lock_key, ttl=600, wait=30)
        if not acquired:
            raise Exception("Could not acquire lock; will retry per Airflow policy")

        try:
            # Option A: Warehouse supports transactional operations
            # BEGIN TRANSACTION
            # DELETE FROM analytics WHERE partition=partition AND exists in staging (deterministic)
            # INSERT INTO analytics SELECT * FROM staging
            # COMMIT
            #
            # Option B: No atomic MERGE across partitions -> use swap:
            # 1) Write staging partition as new physical partition (done earlier)
            # 2) Update catalog to point analytics partition to new location (atomic metadata change if supported)
            #
            # Generic safe delete+insert pattern (idempotent if same staging content):
            run_delete_insert(staging_table, partition)
        finally:
            release_lock(lock_key)

        # cleanup: remove staging_table
        cleanup_staging(staging_table)
        return True

    et = extract_transform()
    promote(et["staging_table"], et["partition"])

def acquire_lock(key, ttl=600, wait=30):
    # pseudo: try Redis SETNX with expiry, loop until acquired or timeout
    import redis
    r = redis.Redis(...)
    deadline = time.time() + wait
    token = str(uuid.uuid4())
    while time.time() < deadline:
        if r.set(key, token, ex=ttl, nx=True):
            # store token somewhere to validate release
            r.set(f"{key}:token", token, ex=ttl)
            return token
        time.sleep(1)
    return None

def release_lock(key):
    # safe release: only delete if token matches
    import redis
    r = redis.Redis(...)
    token = r.get(f"{key}:token")
    if token:
        # use Lua script to compare-and-delete atomically
        lua = """
        if redis.call("get", KEYS[1]) == ARGV[1] then
            return redis.call("del", KEYS[1])
        else
            return 0
        end
        """
        r.eval(lua, 1, key, token)

def run_delete_insert(staging_table, partition):
    # pseudo: use DB client
    db = get_db_client()
    # wrap in transaction if DB supports it
    db.begin()
    try:
        # delete rows for this partition that will be replaced
        db.execute("DELETE FROM analytics WHERE partition = %s", (partition,))
        # insert from staging
        db.execute("INSERT INTO analytics SELECT * FROM %s", (staging_table,))
        db.commit()
    except:
        db.rollback()
        raise

upsert_partitioned_table_dag = upsert_partitioned_table()

from delta.tables import DeltaTable

def foreach_batch(df, batch_id):
    # df: columns (pk, payload..., op, lsn, ts)
    df_filtered = df.withColumn("incoming_lsn", col("lsn"))
    deltaTable = DeltaTable.forPath(spark, "/mnt/delta/my_table")
    deltaTable.alias("t").merge(
        df_filtered.alias("s"),
        "t.pk = s.pk AND s.incoming_lsn > t.last_lsn"
    ).whenMatchedUpdateAll().whenNotMatchedInsertAll().execute()