gluekit 1.0.1.dev1__py3-none-any.whl

gluekit/local/local.py

@@ -0,0 +1,434 @@
+"""Python helpers for running Glue scripts locally with mocked AWS services.
+
+This module exposes the same mock environment used by ``gluekit run`` so that
+individual scripts — and the test suites that exercise them — can drive the
+setup directly from Python instead of going through the CLI.
+
+Typical usage in a test or ad-hoc script::
+
+    from gluekit.local import mock_glue_aws, run_glue_script
+
+    # ── option A: context manager ──────────────────────────────────────────
+    with mock_glue_aws(
+        create_buckets=["my-input-bucket"],
+        ssm_parameters={"/app/env": "dev"},
+    ):
+        import boto3
+        s3 = boto3.client("s3")
+        s3.put_object(Bucket="my-input-bucket", Key="data.csv", Body=b"a,b")
+
+    # ── option B: run a script file ────────────────────────────────────────
+    run_glue_script(
+        "glue/scripts/my_job.py",
+        glue_args={"run_date": "2026-06-18"},
+        create_buckets=["my-input-bucket"],
+        ssm_parameters={"/app/env": "dev"},
+    )
+
+    # ── option C: persist fixture mappings once, then run without repeating them ──
+    # 1. Run once in a terminal:
+    #    gluekit local setup --s3-bucket my-input-bucket=tests/fixtures
+    #    gluekit local setup --ssm-param /app/env=dev
+    #    gluekit local s3 cp tests/fixtures/data.csv s3://my-input-bucket/data.csv
+    #
+    # 2. Create a thin launcher (e.g. run_local.py) and hit "play" in VSCode:
+    from gluekit.local import run_glue_script
+    run_glue_script("glue/scripts/my_job.py", config_file=".gluekit/local.json")
+"""
+
+from __future__ import annotations
+
+import importlib
+import os
+import runpy
+import sys
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Generator
+
+from .local_fixtures import (
+    LOCAL_RUN_CONFIG_FILE,
+    load_runtime_config,
+    parse_s3_uri,
+)
+
+#: Default path for the local run configuration file written by ``gluekit local`` commands.
+
+
+def load_local_run_config(
+    config_file: str | Path | None = None,
+) -> dict[str, object]:
+    """Load a local run configuration from a JSON file.
+
+    Reads the file written by ``gluekit local s3`` and ``gluekit local ssm``.  If *config_file* is
+    ``None`` the default path ``LOCAL_RUN_CONFIG_FILE`` (``.gluekit/local.json``
+    relative to the current working directory) is used.
+
+    Returns an empty dict when the file does not exist or cannot be parsed, so
+    callers can safely spread the result into keyword arguments without
+    additional guards.
+
+    Args:
+        config_file: Path to the JSON config file.  Defaults to
+            ``LOCAL_RUN_CONFIG_FILE`` when ``None``.
+
+    Returns:
+        A dict containing zero or more runtime keyword arguments: ``create_buckets``
+        (``list[str]``), ``s3_objects`` (``dict[str, str]``), and
+        ``ssm_parameters`` (``dict[str, str]``).
+
+    Example:
+        ::
+
+            from gluekit.local import load_local_run_config, mock_glue_aws
+
+            cfg = load_local_run_config()
+            with mock_glue_aws(**cfg):
+                ...
+    """
+    path = Path(config_file) if config_file is not None else LOCAL_RUN_CONFIG_FILE
+    return load_runtime_config(path)
+
+
+def _create_mock_bucket(s3_client, *, bucket_name: str, region: str) -> None:
+    """Create a bucket in the active moto mock, handling the us-east-1 special case."""
+    if region == "us-east-1":
+        s3_client.create_bucket(Bucket=bucket_name)
+        return
+    s3_client.create_bucket(
+        Bucket=bucket_name,
+        CreateBucketConfiguration={"LocationConstraint": region},
+    )
+
+
+def _contains_glue_arg(argv: list[str], name: str) -> bool:
+    """Return True if *name* already appears as a ``--name`` flag in *argv*."""
+    flag = f"--{name}"
+    for index, token in enumerate(argv):
+        if token == flag:
+            return True
+        if token.startswith(f"{flag}="):
+            return True
+        if token == "--":
+            return flag in argv[index + 1 :]
+    return False
+
+
+class LocalSparkSession:
+    """Fallback SparkSession used when SQLFrame is not installed."""
+
+    class _Builder:
+        def appName(self, _name: str) -> "LocalSparkSession._Builder":
+            return self
+
+        def config(self, _key: str, _value: object) -> "LocalSparkSession._Builder":
+            return self
+
+        def getOrCreate(self) -> "LocalSparkSession":
+            return LocalSparkSession()
+
+    builder = _Builder()
+
+
+class LocalSparkContext:
+    """Minimal SparkContext compatibility shim for local Glue bootstrap code."""
+
+    def getOrCreate(self):
+        from gluekit.local.pyspark.sql import SparkSession
+
+        return SparkSession.builder.getOrCreate()
+
+
+class LocalGlueContext:
+    """Minimal GlueContext compatibility shim backed by the active local Spark API."""
+
+    def __init__(self, spark_context: LocalSparkContext):
+        self.spark_session = spark_context.getOrCreate()
+
+
+class LocalGlueJob:
+    """Minimal Glue Job shim for scripts that call ``init`` and ``commit``."""
+
+    def __init__(self, _glue_context: LocalGlueContext):
+        self.name: str | None = None
+
+    def init(self, name: str, _args: dict[str, str]) -> None:
+        self.name = name
+
+    def commit(self) -> None:
+        return None
+
+
+def get_resolved_options(argv: list[str], keys: list[str]) -> dict[str, str]:
+    """Local implementation of ``awsglue.utils.getResolvedOptions``."""
+    resolved = {}
+    missing = []
+    for key in keys:
+        flag = f"--{key}"
+        value = None
+        for index, token in enumerate(argv):
+            if token == flag and index + 1 < len(argv):
+                value = argv[index + 1]
+                break
+            if token.startswith(f"{flag}="):
+                value = token[len(flag) + 1 :]
+                break
+        if value is None:
+            missing.append(key)
+        else:
+            resolved[key] = value
+
+    if missing:
+        raise ValueError(f"Missing required options: {missing}")
+    return resolved
+
+
+@contextmanager
+def mock_glue_runtime(
+    *, sqlframe_engine: str = "duckdb"
+) -> Generator[None, None, None]:
+    """Temporarily expose Glue runtime import paths for local script execution.
+
+    AWS Glue scripts commonly import ``awsglue`` and ``pyspark`` as top-level
+    packages that only exist inside the Glue runtime.  Gluekit keeps those
+    emulations scoped to this context so installing the CLI does not shadow real
+    ``awsglue`` or ``pyspark`` packages in unrelated Python processes.
+    """
+    module_names = [
+        "awsglue",
+        "pyspark",
+    ]
+    original_modules = {name: sys.modules.get(name) for name in module_names}
+
+    from sqlframe import activate
+    activate(engine=sqlframe_engine)
+    sys.modules.update(
+        {
+            "awsglue": importlib.import_module("gluekit.local.awsglue"),
+            "pyspark": importlib.import_module("gluekit.local.pyspark")        
+        }
+    )
+
+    try:
+        yield
+    finally:
+        for name, module in original_modules.items():
+            if module is None:
+                sys.modules.pop(name, None)
+            else:
+                sys.modules[name] = module
+
+
+@contextmanager
+def mock_glue_aws(
+    *,
+    create_buckets: list[str] | None = None,
+    s3_objects: dict[str, str] | None = None,
+    ssm_parameters: dict[str, str] | None = None,
+    ssm_parameter_types: dict[str, str] | None = None,
+    aws_region: str | None = None,
+    config_file: str | Path | None = None,
+) -> Generator[None, None, None]:
+    """Context manager that activates a mocked AWS environment for local Glue development.
+
+    Inside the ``with`` block every ``boto3`` call is intercepted by
+    `moto <https://docs.getmoto.org/>`_ so no real AWS credentials are needed
+    and no real AWS resources are touched.
+
+    AWS credentials are temporarily set to ``"testing"`` values required by
+    moto, and are restored to their original values on exit.
+
+    Args:
+        create_buckets: S3 bucket names to pre-create in the mock environment.
+        s3_objects: ``{"s3://bucket/key": "local/file"}`` mappings to upload into
+            mocked S3 before entering the block.
+        ssm_parameters: ``{"/parameter/name": "value"}`` pairs to seed into
+            the mock SSM Parameter Store before entering the block.
+        ssm_parameter_types: Optional ``{"/parameter/name": "String"}`` type map
+            for seeded SSM parameters. Defaults to ``"String"``.
+        aws_region: AWS region for mock service clients. Defaults to
+            ``"us-east-1"``.
+        config_file: Path to a JSON config file written by ``gluekit local
+            fixture commands. Values in the file fill in any parameter that was not
+            provided explicitly (``create_buckets``, ``s3_objects``,
+            and ``ssm_parameters``).
+            When ``None`` no config file is loaded.
+
+    Raises:
+        ImportError: If ``moto[s3,ssm]`` is not installed.
+
+    Example:
+        ::
+
+            from gluekit.local import mock_glue_aws
+            import boto3
+
+            with mock_glue_aws(
+                create_buckets=["my-bucket"],
+                ssm_parameters={"/app/env": "dev"},
+            ):
+                s3 = boto3.client("s3")
+                s3.put_object(Bucket="my-bucket", Key="out.txt", Body=b"ok")
+                body = s3.get_object(Bucket="my-bucket", Key="out.txt")["Body"].read()
+                print(body)  # b"ok"
+
+            # Or load persisted fixtures written by ``gluekit local s3`` and ``gluekit local ssm``:
+            with mock_glue_aws(config_file=".gluekit/local.json"):
+                ...
+    """
+    if config_file is not None:
+        file_cfg = load_local_run_config(config_file)
+        if create_buckets is None:
+            create_buckets = file_cfg.get("create_buckets")  # type: ignore[assignment]
+        if s3_objects is None:
+            s3_objects = file_cfg.get("s3_objects")  # type: ignore[assignment]
+        if ssm_parameters is None:
+            ssm_parameters = file_cfg.get("ssm_parameters")  # type: ignore[assignment]
+
+    # Resolve the final region after config file fallback.
+    effective_region: str = aws_region if isinstance(aws_region, str) else "us-east-1"
+    try:
+        from moto import mock_aws
+    except ImportError as exc:
+        raise ImportError(
+            "mock_glue_aws requires moto with S3/SSM support. "
+            'Install with: uv add "moto[s3,ssm]" or pip install "moto[s3,ssm]"'
+        ) from exc
+
+    env_overrides = {
+        "AWS_ACCESS_KEY_ID": "testing",
+        "AWS_SECRET_ACCESS_KEY": "testing",
+        "AWS_SESSION_TOKEN": "testing",
+        "AWS_DEFAULT_REGION": effective_region,
+        "AWS_REGION": effective_region,
+    }
+    original_env = {key: os.environ.get(key) for key in env_overrides}
+    os.environ.update(env_overrides)
+
+    try:
+        with mock_aws(), mock_glue_runtime():
+            import boto3
+
+            s3_client = boto3.client("s3", region_name=effective_region)
+            ssm_client = boto3.client("ssm", region_name=effective_region)
+
+            bucket_names = set(create_buckets or [])
+            for object_uri in s3_objects or {}:
+                bucket_name, _ = parse_s3_uri(object_uri)
+                bucket_names.add(bucket_name)
+
+            for bucket_name in sorted(bucket_names):
+                _create_mock_bucket(
+                    s3_client, bucket_name=bucket_name, region=effective_region
+                )
+
+            for object_uri, local_path in (s3_objects or {}).items():
+                bucket_name, object_key = parse_s3_uri(object_uri)
+                s3_client.upload_file(str(Path(local_path)), bucket_name, object_key)
+
+            for name, value in (ssm_parameters or {}).items():
+                ssm_client.put_parameter(
+                    Name=name,
+                    Value=value,
+                    Type=(ssm_parameter_types or {}).get(name, "String"),
+                    Overwrite=True,
+                )
+
+            yield
+    finally:
+        for key, value in original_env.items():
+            if value is None:
+                os.environ.pop(key, None)
+            else:
+                os.environ[key] = value
+
+
+def run_glue_script(
+    script_path: str | Path,
+    *,
+    script_args: list[str] | None = None,
+    glue_args: dict[str, str] | None = None,
+    job_name: str = "local-glue-job",
+    create_buckets: list[str] | None = None,
+    s3_objects: dict[str, str] | None = None,
+    ssm_parameters: dict[str, str] | None = None,
+    ssm_parameter_types: dict[str, str] | None = None,
+    aws_region: str | None = None,
+    config_file: str | Path | None = None,
+) -> None:
+    """Run a Glue script locally with emulated ``awsglue``/``pyspark`` and mocked AWS.
+
+    This mirrors the behaviour of ``gluekit run`` but is callable from Python
+    code — useful for integration tests or ad-hoc local runs where you want
+    programmatic control over the mock fixtures.
+
+    Args:
+        script_path: Path to the Glue Python script to execute.
+        script_args: Raw ``sys.argv``-style tokens (e.g.
+            ``["--JOB_NAME", "my-job"]``) appended **before** any
+            ``glue_args``.  Use this to pass arguments verbatim when you
+            already have them in flag form.
+        glue_args: Keyword-style Glue arguments as a dict
+            (e.g. ``{"run_date": "2026-06-18"}``).  Each entry is translated
+            to ``--key value`` and appended to ``sys.argv``.
+        job_name: Value injected as ``--JOB_NAME`` when the argument is not
+            already present in *script_args* or *glue_args*.
+        create_buckets: S3 bucket names to pre-create in the mock environment.
+        s3_objects: ``{"s3://bucket/key": "local/file"}`` mappings to upload into
+            mocked S3 before running the script.
+        ssm_parameters: ``{"/parameter/name": "value"}`` pairs to seed into
+            the mock SSM Parameter Store.
+        ssm_parameter_types: Optional ``{"/parameter/name": "String"}`` type map
+            for seeded SSM parameters. Defaults to ``"String"``.
+        aws_region: AWS region for mock service clients. Defaults to
+            ``"us-east-1"``.
+        config_file: Path to a JSON config file written by ``gluekit local
+            fixture commands. Values in the file fill in any parameter that was not
+            provided explicitly (``create_buckets``, ``s3_objects``,
+            and ``ssm_parameters``).
+            When ``None`` no config file is loaded.
+
+    Raises:
+        ImportError: If ``moto[s3,ssm]`` is not installed.
+
+    Example:
+        ::
+
+            from gluekit.local import run_glue_script
+
+            run_glue_script(
+                "glue/scripts/my_job.py",
+                glue_args={"run_date": "2026-06-18"},
+                create_buckets=["my-input-bucket"],
+                ssm_parameters={"/app/env": "dev"},
+            )
+
+            # Or drive everything from a persisted config file:
+            run_glue_script(
+                "glue/scripts/my_job.py",
+                config_file=".gluekit/local.json",
+            )
+    """
+    script_path = Path(script_path).resolve()
+
+    argv: list[str] = list(script_args or [])
+    for key, value in (glue_args or {}).items():
+        argv.extend([f"--{key}", value])
+
+    if not _contains_glue_arg(argv, "JOB_NAME"):
+        argv.extend(["--JOB_NAME", job_name])
+
+    original_argv = sys.argv[:]
+    try:
+        with mock_glue_aws(
+            create_buckets=create_buckets,
+            s3_objects=s3_objects,
+            ssm_parameters=ssm_parameters,
+            ssm_parameter_types=ssm_parameter_types,
+            aws_region=aws_region,
+            config_file=config_file,
+        ):
+            sys.argv = [str(script_path), *argv]
+            runpy.run_path(str(script_path), run_name="__main__")
+    finally:
+        sys.argv = original_argv