FROM python:3.13-slim

# Spark 4.x runs on the JVM (Course 2 Module 1 Lesson 3) and needs Java 17 or
# 21 -- a headless JRE is enough (no compiler, no javadoc, no javafx), which
# keeps this image far slimmer than a full JDK would. python:3.13-slim's
# current Debian base (trixie) only ships JDK 21 in its repos (no standalone
# "openjdk-17-jre-headless" package there), so this image uses 21 -- still
# squarely inside Spark 4's supported 17/21 range.
RUN apt-get update \
    && apt-get install --no-install-recommends -y openjdk-21-jre-headless \
    && rm -rf /var/lib/apt/lists/*
ENV JAVA_HOME=/usr/lib/jvm/java-21-openjdk-arm64

WORKDIR /app

COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY cityflow_pyspark_job.py .
COPY cityflow_trips_2024-01-01.parquet /opt/cityflow/data/cityflow_trips_2024-01-01.parquet
COPY taxi_zone_lookup.csv /opt/cityflow/data/taxi_zone_lookup.csv

CMD ["python", "cityflow_pyspark_job.py"]

# gate: skip
"""CityFlow's real PySpark-on-Kubernetes job (Module 7 Lessons 3 and 5).

Runs inside the cityflow-pyspark-job:1.0 image, on a real Kind Kubernetes Job
pod, submitted by an Airflow DAG. Reads the real CityFlow data baked into the
image at build time (a single-day slice, Jan 1 2024 -- the same real
81,013-row slice Module 6 Lesson 1's regression test already reproduced from
this course's own pandas logic), joins it to the real taxi zone lookup table,
and aggregates real trip counts and revenue by borough with a real Spark
groupBy -- small, real, CityFlow-relevant work, not a placeholder computation.

The result is printed to stdout as one line prefixed CITYFLOW_PYSPARK_RESULT_JSON:
so the Airflow task (and anyone reading `kubectl logs`) can find and parse it
without scraping Spark's own log noise.
"""
import json

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

TRIPS_PATH = "/opt/cityflow/data/cityflow_trips_2024-01-01.parquet"
ZONES_PATH = "/opt/cityflow/data/taxi_zone_lookup.csv"


def main():
    spark = (
        SparkSession.builder
        .appName("cityflow-k8s-pyspark-job")
        .master("local[*]")
        .config("spark.ui.showConsoleProgress", "false")
        .getOrCreate()
    )
    spark.sparkContext.setLogLevel("ERROR")

    trips = spark.read.parquet(TRIPS_PATH)
    total_rows = trips.count()

    zones = spark.read.csv(ZONES_PATH, header=True, inferSchema=True)

    joined = trips.join(zones, trips.PULocationID == zones.LocationID, "left")
    # A left join leaves a real NULL Borough for any PULocationID that matched
    # no row in the zone lookup at all -- a *different*, more honest signal
    # than the lookup table's own "Unknown" Borough category. pandas'
    # groupby() drops NULL-key rows silently by default; Spark's groupBy()
    # keeps them as their own real group, which is why this coalesce is here
    # instead of being silently lost the way a naive pandas merge would lose it.
    joined = joined.withColumn("Borough", F.coalesce(F.col("Borough"), F.lit("Unmatched Zone")))
    by_borough = (
        joined.groupBy("Borough")
        .agg(
            F.count("*").alias("trip_count"),
            F.round(F.sum("total_amount"), 2).alias("total_revenue"),
        )
        .orderBy(F.col("trip_count").desc())
    )
    rows = [r.asDict() for r in by_borough.collect()]
    top = rows[0]

    result = {
        "total_rows": total_rows,
        "by_borough": rows,
        "top_borough": top["Borough"],
        "top_borough_trip_count": top["trip_count"],
    }

    print("CityFlow PySpark job on Kubernetes -- real result:")
    print(f"  total rows (real, baked-in CityFlow slice): {total_rows}")
    for r in rows:
        print(f"  {r['Borough']:>14}: {r['trip_count']:>6} trips, ${r['total_revenue']:>12,.2f} revenue")
    print(f"  busiest borough: {top['Borough']} ({top['trip_count']} trips)")
    print("CITYFLOW_PYSPARK_RESULT_JSON:" + json.dumps(result))

    spark.stop()


if __name__ == "__main__":
    main()

# gate: skip
"""CityFlow's real Airflow-submits-a-real-PySpark-job-to-Kubernetes DAG
(Module 7 Lesson 3).

Same submit -> poll -> read-logs shape Lesson 2's cityflow_k8s_submit_demo_dag
proved with a trivial pod, with the trivial python:3.13-slim command replaced
by the real cityflow-pyspark-job:1.0 image: python:3.13-slim + a headless
JDK 21 JRE + PySpark, running a real Spark job against a real 81,013-row
CityFlow slice baked into the image. Still no
apache-airflow-providers-cncf-kubernetes -- the same real kubectl subprocess
pattern, against the same real kubectl binary + kubeconfig this course
host-mounted into the Airflow containers."""
import re
from datetime import datetime

from airflow.decorators import dag, task

KUBECTL_BIN = "/opt/airflow/dags/bin/kubectl"
KUBECONFIG_PATH = "/opt/airflow/dags/data/kubeconfig-kind.yaml"
NAMESPACE = "default"
JOB_NAME = "cityflow-pyspark-on-k8s"
IMAGE = "cityflow-pyspark-job:1.0"
JOB_MANIFEST = f"""\
apiVersion: batch/v1
kind: Job
metadata:
  name: {JOB_NAME}
  namespace: {NAMESPACE}
  labels:
    app: {JOB_NAME}
spec:
  backoffLimit: 0
  ttlSecondsAfterFinished: 3600
  template:
    metadata:
      labels:
        app: {JOB_NAME}
    spec:
      restartPolicy: Never
      containers:
        - name: cityflow-pyspark
          image: {IMAGE}
          imagePullPolicy: IfNotPresent
"""


@dag(
    dag_id="cityflow_pyspark_on_k8s",
    schedule=None,
    start_date=datetime(2026, 1, 1),
    catchup=False,
    tags=["cityflow", "module-7", "kubernetes", "pyspark"],
)
def cityflow_pyspark_on_k8s_dag():
    @task
    def submit_job() -> str:
        import subprocess

        subprocess.run(
            [KUBECTL_BIN, "--kubeconfig", KUBECONFIG_PATH, "delete", "job", JOB_NAME,
             "-n", NAMESPACE, "--ignore-not-found"],
            capture_output=True, text=True, check=True,
        )
        result = subprocess.run(
            [KUBECTL_BIN, "--kubeconfig", KUBECONFIG_PATH, "apply", "-f", "-"],
            input=JOB_MANIFEST, capture_output=True, text=True,
        )
        print(f"submit_job: kubectl apply -> {result.stdout.strip()}")
        if result.returncode != 0:
            raise RuntimeError(f"kubectl apply failed: {result.stderr}")
        return JOB_NAME

    @task
    def poll_job_completion(job_name: str) -> dict:
        import subprocess
        import time

        deadline = time.time() + 180
        polls = 0
        while time.time() < deadline:
            polls += 1
            succeeded = subprocess.run(
                [KUBECTL_BIN, "--kubeconfig", KUBECONFIG_PATH, "get", "job", job_name,
                 "-n", NAMESPACE, "-o", "jsonpath={.status.succeeded}"],
                capture_output=True, text=True, check=True,
            ).stdout.strip()
            failed = subprocess.run(
                [KUBECTL_BIN, "--kubeconfig", KUBECONFIG_PATH, "get", "job", job_name,
                 "-n", NAMESPACE, "-o", "jsonpath={.status.failed}"],
                capture_output=True, text=True, check=True,
            ).stdout.strip()
            print(f"poll_job_completion: poll #{polls} succeeded={succeeded!r} failed={failed!r}")
            if succeeded == "1":
                return {"job_name": job_name, "polls": polls, "outcome": "succeeded"}
            if failed and int(failed) > 0:
                raise RuntimeError(f"real Job {job_name} reported {failed} failed pod(s)")
            time.sleep(2)
        raise TimeoutError(f"real Job {job_name} did not succeed within the poll deadline")

    @task
    def report_spark_result(poll_result: dict) -> dict:
        """Reads the real pod's real logs and parses the real
        CITYFLOW_PYSPARK_RESULT_JSON: line the Spark job itself printed --
        proof the actual PySpark computation ran on the actual Kind cluster,
        not just that the Job's status turned green."""
        import json
        import subprocess

        job_name = poll_result["job_name"]
        pod_name = subprocess.run(
            [KUBECTL_BIN, "--kubeconfig", KUBECONFIG_PATH, "get", "pods", "-n", NAMESPACE,
             "-l", f"app={job_name}", "-o", "jsonpath={.items[0].metadata.name}"],
            capture_output=True, text=True, check=True,
        ).stdout.strip()
        logs = subprocess.run(
            [KUBECTL_BIN, "--kubeconfig", KUBECONFIG_PATH, "logs", pod_name, "-n", NAMESPACE],
            capture_output=True, text=True, check=True,
        ).stdout

        match = re.search(r"CITYFLOW_PYSPARK_RESULT_JSON:(\{.*\})", logs)
        if not match:
            raise RuntimeError(f"no CITYFLOW_PYSPARK_RESULT_JSON line found in real pod logs:\n{logs}")
        result = json.loads(match.group(1))
        print(f"report_spark_result: real pod {pod_name!r} computed: {result}")
        assert result["total_rows"] == 81013, result["total_rows"]
        assert result["top_borough"] == "Manhattan", result["top_borough"]
        return result

    report_spark_result(poll_job_completion(submit_job()))


cityflow_pyspark_on_k8s_dag()

import json
import subprocess
import time

dag_id = "cityflow_pyspark_on_k8s"

subprocess.run(["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "unpause", dag_id],
                capture_output=True, text=True, check=True)

run_id = f"lesson3_m7_verify_{int(time.time())}"
trigger = subprocess.run(
    ["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "trigger", dag_id, "--run-id", run_id],
    capture_output=True, text=True,
)
if trigger.returncode != 0:
    subprocess.run(["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "reserialize"],
                    capture_output=True, text=True, check=True)
    subprocess.run(
        ["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "trigger", dag_id, "--run-id", run_id],
        capture_output=True, text=True, check=True,
    )
print(f"triggered {dag_id!r} as {run_id!r}")

deadline = time.time() + 180
state = "queued"
while state not in ("success", "failed") and time.time() < deadline:
    time.sleep(3)
    r = subprocess.run(
        ["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "list-runs", "-d", dag_id, "--output", "json"],
        capture_output=True, text=True, check=True,
    )
    runs = json.loads(r.stdout)
    match = [x for x in runs if x["run_id"] == run_id]
    state = match[0]["state"] if match else "queued"
print(f"run {run_id!r} reached state: {state}")
assert state == "success", f"expected success, got {state}"

r = subprocess.run(
    ["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "tasks", "states-for-dag-run", dag_id, run_id],
    capture_output=True, text=True, check=True,
)
print(r.stdout)