Module 26: Python Project on Databricks

orders_pipeline/
├── pyproject.toml
├── src/orders_pipeline/
│   ├── __init__.py
│   ├── config.py        # parameters -> Config
│   ├── transforms.py    # pure DataFrame functions
│   └── main.py          # entry points (one per job task)
├── tests/
│   ├── conftest.py
│   └── test_transforms.py
├── databricks.yml       # bundle + job definition (Module 27)
└── .github/workflows/   # CI (Module 27)

# src/orders_pipeline/transforms.py
from pyspark.sql import DataFrame, functions as F


def clean_orders(df: DataFrame) -> DataFrame:
    """Typed, lower-case status, no invalid rows, one row per order."""
    return (
        df.withColumn("status", F.lower(F.trim("status")))
          .withColumn("order_date", F.to_date("order_ts"))
          .filter("customer_id IS NOT NULL AND quantity > 0")
          .dropDuplicates(["order_id"])
    )


def daily_revenue(orders: DataFrame, products: DataFrame) -> DataFrame:
    return (
        orders.filter("status IN ('completed', 'shipped')")
              .join(products.select("product_id", "unit_price"), "product_id")
              .groupBy("order_date")
              .agg(F.count("*").alias("orders"),
                   F.round(F.sum(F.col("quantity") * F.col("unit_price")), 2).alias("revenue"))
    )

# src/orders_pipeline/config.py
import argparse
from dataclasses import dataclass
from datetime import date, timedelta


@dataclass(frozen=True)
class Config:
    catalog: str
    schema: str
    run_date: str
    days: int

    def table(self, name: str) -> str:
        return f"{self.catalog}.{self.schema}.{name}"

    @property
    def window_start(self) -> str:
        return (date.fromisoformat(self.run_date) - timedelta(days=self.days - 1)).isoformat()


def parse_args(argv=None) -> Config:
    p = argparse.ArgumentParser()
    p.add_argument("--catalog", required=True)
    p.add_argument("--schema", default="dbx_project")
    p.add_argument("--run-date", "--run_date", dest="run_date", default="")
    p.add_argument("--days", type=int, default=7)
    a, _unknown = p.parse_known_args(argv)
    run_date = a.run_date or date.today().isoformat()
    date.fromisoformat(run_date)                      # fail early on a bad date
    return Config(a.catalog, a.schema, run_date, a.days)

# src/orders_pipeline/main.py
import logging
import sys
from pyspark.sql import SparkSession, functions as F
from .config import parse_args
from .transforms import clean_orders, daily_revenue

log = logging.getLogger("orders_pipeline")


def _spark() -> SparkSession:
    return SparkSession.builder.getOrCreate()


def load_silver() -> None:
    cfg = parse_args()
    spark = _spark()
    raw = spark.table(f"{cfg.catalog}.dbx_course.orders_raw")
    clean_orders(raw).write.mode("overwrite").saveAsTable(cfg.table("silver_orders"))
    log.info("silver_orders written for %s", cfg.run_date)


def check_quality() -> None:
    cfg = parse_args()
    silver = _spark().table(cfg.table("silver_orders"))
    bad = silver.filter("order_id IS NULL OR quantity <= 0").count()
    if bad:
        raise ValueError(f"{bad} invalid rows in silver_orders")


def build_gold() -> None:
    cfg = parse_args()
    spark = _spark()
    silver = (spark.table(cfg.table("silver_orders"))
                   .filter(F.col("order_date").between(cfg.window_start, cfg.run_date)))
    products = spark.table(f"{cfg.catalog}.dbx_course.products")
    out = daily_revenue(silver, products)
    (out.write.mode("overwrite")
        .option("replaceWhere", f"order_date >= '{cfg.window_start}' AND order_date <= '{cfg.run_date}'")
        .saveAsTable(cfg.table("gold_daily_revenue")))

# pyproject.toml (fragment)
[project.scripts]
load_silver   = "orders_pipeline.main:load_silver"
check_quality = "orders_pipeline.main:check_quality"
build_gold    = "orders_pipeline.main:build_gold"

# tests/conftest.py
import pytest


@pytest.fixture(scope="session")
def spark():
    try:
        from databricks.connect import DatabricksSession
        return DatabricksSession.builder.getOrCreate()      # remote: serverless or cluster
    except ImportError:
        from pyspark.sql import SparkSession
        return SparkSession.builder.master("local[1]").getOrCreate()

# tests/test_transforms.py
from pyspark.testing import assertDataFrameEqual
from orders_pipeline.transforms import clean_orders

def test_clean_orders_drops_invalid_and_dedupes(spark):
    df = spark.createDataFrame(
        [(1, 10, 2, " Shipped ", "2025-03-01 10:00:00"),
         (1, 10, 2, "shipped",   "2025-03-01 10:00:00"),     # duplicate id
         (2, None, 1, "completed", "2025-03-01 11:00:00"),   # no customer
         (3, 11, 0, "completed", "2025-03-01 12:00:00")],    # bad quantity
        "order_id INT, customer_id INT, quantity INT, status STRING, order_ts STRING")
    out = clean_orders(df)
    assert out.count() == 1
    assert out.first()["status"] == "shipped"

pip uninstall -y pyspark                          # databricks-connect and pyspark conflict
pip install "databricks-connect==17.3.*"          # match the compute version, see the docs
databricks auth login --host https://<workspace-url>
pytest -q

python -m pip install build
python -m build --wheel                 # creates dist/orders_pipeline-0.1.0-py3-none-any.whl

%pip install /Volumes/main/dbx_project/artifacts/orders_pipeline-0.1.0-py3-none-any.whl

dbutils.library.restartPython()
from orders_pipeline.main import load_silver

# src/orders_pipeline/logging_setup.py
import logging
import sys


def setup_logging(level: int = logging.INFO) -> None:
    logging.basicConfig(
        level=level,
        stream=sys.stdout,
        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
        force=True,
    )

# in main.py
import logging, time
from .logging_setup import setup_logging

def run_step(name, fn):
    setup_logging()
    log = logging.getLogger("orders_pipeline")
    t0 = time.time()
    try:
        fn()
        log.info("step=%s status=ok seconds=%.1f", name, time.time() - t0)
    except Exception:
        log.exception("step=%s status=failed", name)
        raise                                   # the job task must fail