run_glue_spark.py¶

Source revision: 75f65139e26eb7079b7a897c877de97859c0120b
This page is the generated source; it shows the complete readable projection of the canonical file.
"""DataCoolie AWS Glue Spark runner reference.

Required Glue job parameters:
  --REGION
  --METADATA_PATH
  --WATERMARK_BASE_PATH
  --LOG_BASE_PATH

Optional Glue job parameters:
  --CONNECTIONS_PATH  Separate connections metadata file
  --SCHEMA_HINTS_PATH  Separate global schema-hints metadata file
  --STAGE  Stage value passed unchanged; omission runs all stages
  --JOB_NUM  Total shard count (default 1)
  --JOB_INDEX  This invocation's zero-based shard index (default 0)

This example uses DataCoolie watermark state and does not enable AWS Glue bookmarks.
Attach ``datacoolie[aws]`` before execution; Glue supplies Spark and the job
role is resolved through the standard boto3 credential chain.
"""

from __future__ import annotations

import logging
import sys

from awsglue.context import GlueContext
from awsglue.utils import getResolvedOptions
from pyspark.context import SparkContext

from datacoolie.core.models.run_config import DataCoolieRunConfig
from datacoolie.engines.spark_engine import SparkEngine
from datacoolie.metadata.file_provider import FileProvider
from datacoolie.orchestration.driver import DataCoolieDriver
from datacoolie.platforms.aws_platform import AWSPlatform


def require_s3_path(value: str, option: str) -> str:
    """Reject launcher-local paths before a Glue Driver is constructed."""
    if not value or not value.startswith(("s3://", "s3a://")):
        raise ValueError(f"{option} must be an s3:// or s3a:// URI; upload the input first")
    return value


def resolve_options(argv: list[str]) -> dict[str, str]:
    required = ["REGION", "METADATA_PATH", "WATERMARK_BASE_PATH", "LOG_BASE_PATH"]
    optional = [
        name
        for name in ("CONNECTIONS_PATH", "SCHEMA_HINTS_PATH", "STAGE", "JOB_NUM", "JOB_INDEX")
        if f"--{name}" in argv
    ]
    resolved = getResolvedOptions(argv, [*required, *optional])
    for name in ("METADATA_PATH", "WATERMARK_BASE_PATH", "LOG_BASE_PATH"):
        require_s3_path(resolved[name], f"--{name}")
    for name in ("CONNECTIONS_PATH", "SCHEMA_HINTS_PATH"):
        if name in resolved:
            require_s3_path(resolved[name], f"--{name}")
    return resolved


args = resolve_options(sys.argv)
stage = args.get("STAGE")

spark = GlueContext(SparkContext.getOrCreate()).spark_session
platform = AWSPlatform(region=args["REGION"])
engine = SparkEngine(spark_session=spark, platform=platform)
metadata = FileProvider(
    config_path=args["METADATA_PATH"],
    connections_path=args.get("CONNECTIONS_PATH"),
    schema_hints_path=args.get("SCHEMA_HINTS_PATH"),
    platform=platform,
    watermark_base_path=args["WATERMARK_BASE_PATH"],
)
config = DataCoolieRunConfig(
    max_workers=4,
    job_num=int(args.get("JOB_NUM", "1")),
    job_index=int(args.get("JOB_INDEX", "0")),
    stop_on_error=True,
    allowed_function_prefixes=[],
)

failed = 0
with DataCoolieDriver(
    engine=engine,
    metadata_provider=metadata,
    log_base_path=args["LOG_BASE_PATH"],
    config=config,
) as driver:
    result = driver.run(stage=stage)
    failed = result.failed

logging.info("DataCoolie run completed with %d failed dataflows", failed)
if failed:
    raise RuntimeError(f"DataCoolie run failed for {failed} dataflows")