gluekit 1.0.1.dev1__py3-none-any.whl

gluekit/job_mgmt/resources/magic_map.json

@@ -0,0 +1,83 @@
+{
+  "version": 1,
+  "source": "https://docs.aws.amazon.com/glue/latest/dg/interactive-sessions-magics.html",
+  "stop_session_magic": "%stop_session",
+  "section_order": [
+    "session",
+    "compute",
+    "dependencies",
+    "configure"
+  ],
+  "mappings": [
+    {
+      "field_path": "Role",
+      "emit": "magic",
+      "magic": "%iam_role",
+      "section": "session",
+      "format": "string"
+    },
+    {
+      "field_path": "Connections.Connections",
+      "emit": "magic",
+      "magic": "%connections",
+      "section": "session",
+      "format": "csv"
+    },
+    {
+      "field_path": "GlueVersion",
+      "emit": "magic",
+      "magic": "%glue_version",
+      "section": "compute",
+      "format": "string"
+    },
+    {
+      "field_path": "WorkerType",
+      "emit": "magic",
+      "magic": "%worker_type",
+      "section": "compute",
+      "format": "string"
+    },
+    {
+      "field_path": "NumberOfWorkers",
+      "emit": "magic",
+      "magic": "%number_of_workers",
+      "section": "compute",
+      "format": "string"
+    },
+    {
+      "field_path": "DefaultArguments.--additional-python-modules",
+      "emit": "magic",
+      "magic": "%additional_python_modules",
+      "section": "dependencies",
+      "format": "string"
+    },
+    {
+      "field_path": "DefaultArguments.--extra-py-files",
+      "emit": "magic",
+      "magic": "%extra_py_files",
+      "section": "dependencies",
+      "format": "string"
+    },
+    {
+      "field_path": "DefaultArguments.--extra-jars",
+      "emit": "magic",
+      "magic": "%extra_jars",
+      "section": "dependencies",
+      "format": "string"
+    },
+    {
+      "field_path": "DefaultArguments.--conf",
+      "emit": "magic",
+      "magic": "%spark_conf",
+      "section": "dependencies",
+      "format": "string"
+    },
+    {
+      "field_path": "Command.ScriptLocation",
+      "emit": "configure",
+      "configure_key": "script_location",
+      "section": "configure",
+      "format": "script_location_dir"
+    }
+  ]
+}

gluekit/job_mgmt/schema.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+import json
+from functools import lru_cache
+from importlib.resources import files
+from pathlib import Path
+from typing import Any
+
+import boto3
+
+DEPRECATED_JOB_FIELDS = (
+    "AllocatedCapacity",
+    "MaxCapacity",
+)
+
+SCHEMA_RESOURCE_NAME = "glue_job_schema.json"
+MAGIC_MAP_RESOURCE_NAME = "magic_map.json"
+RESOURCE_DIR = Path(__file__).resolve().parent / "resources"
+
+
+def _resource_text(name: str) -> str:
+    resource_root = files("gluekit.job_mgmt").joinpath("resources")
+    return resource_root.joinpath(name).read_text(encoding="utf-8")
+
+
+@lru_cache(maxsize=1)
+def load_glue_schema() -> dict[str, Any]:
+    return json.loads(_resource_text(SCHEMA_RESOURCE_NAME))
+
+
+@lru_cache(maxsize=1)
+def load_magic_map() -> dict[str, Any]:
+    return json.loads(_resource_text(MAGIC_MAP_RESOURCE_NAME))
+
+
+def get_updateable_job_fields(prefer_live: bool = False) -> set[str]:
+    if prefer_live:
+        live_fields = inspect_live_update_job_fields()
+        if live_fields:
+            return live_fields
+    return set(load_glue_schema().get("update_job_fields", []))
+
+
+def inspect_live_update_job_fields() -> set[str]:
+    try:
+        session = boto3.Session()
+        glue = session.client("glue", region_name="us-east-1")
+        model = glue._service_model
+        request_shape = model.shape_for("UpdateJobRequest")
+        job_update_shape = request_shape.members.get("JobUpdate")
+        if not job_update_shape:
+            return set()
+        return set(job_update_shape.members.keys())
+    except Exception:
+        return set()
+
+
+def _clean_doc(value: Any) -> str | None:
+    if not value:
+        return None
+    text = " ".join(str(value).split())
+    return text or None
+
+
+def _serialize_shape(shape: Any, seen: set[str] | None = None) -> dict[str, Any]:
+    seen = seen or set()
+    shape_token = getattr(shape, "name", None) or f"anon:{id(shape)}"
+
+    data: dict[str, Any] = {
+        "shape_name": getattr(shape, "name", None),
+        "type": getattr(shape, "type_name", None),
+    }
+
+    if shape_token in seen:
+        data["ref"] = shape_token
+        return data
+
+    next_seen = seen | {shape_token}
+
+    documentation = _clean_doc(getattr(shape, "documentation", None))
+    if documentation:
+        data["documentation"] = documentation
+
+    enum_values = list(getattr(shape, "enum", []) or [])
+    if enum_values:
+        data["enum"] = enum_values
+
+    for attr in ("min", "max"):
+        value = getattr(shape, attr, None)
+        if value is not None:
+            data[attr] = value
+
+    required_members = list(getattr(shape, "required_members", []) or [])
+    if required_members:
+        data["required_members"] = sorted(required_members)
+
+    if data["type"] == "structure":
+        members = getattr(shape, "members", {}) or {}
+        data["members"] = {
+            member_name: _serialize_shape(members[member_name], next_seen)
+            for member_name in sorted(members.keys())
+        }
+    elif data["type"] == "list":
+        member = getattr(shape, "member", None)
+        if member is not None:
+            data["member"] = _serialize_shape(member, next_seen)
+    elif data["type"] == "map":
+        key = getattr(shape, "key", None)
+        value = getattr(shape, "value", None)
+        if key is not None:
+            data["key"] = _serialize_shape(key, next_seen)
+        if value is not None:
+            data["value"] = _serialize_shape(value, next_seen)
+
+    return data
+
+
+def build_glue_schema() -> dict[str, Any]:
+    session = boto3.Session()
+    glue = session.client("glue", region_name="us-east-1")
+    model = glue._service_model
+
+    root_shape_names = ["UpdateJobRequest", "JobUpdate", "Job"]
+    root_shapes = {
+        name: _serialize_shape(model.shape_for(name)) for name in root_shape_names
+    }
+
+    job_shape = root_shapes["Job"]
+    job_update_shape = root_shapes["JobUpdate"]
+
+    return {
+        "schema_version": 1,
+        "service": "glue",
+        "api_version": getattr(model, "api_version", None),
+        "deprecated_fields": sorted(DEPRECATED_JOB_FIELDS),
+        "update_job_fields": sorted(job_update_shape.get("members", {}).keys()),
+        "job_fields": job_shape.get("members", {}),
+        "job_update_fields": job_update_shape.get("members", {}),
+        "root_shapes": root_shapes,
+    }
+
+
+def write_glue_schema(output_path: Path | None = None) -> Path:
+    path = output_path or (RESOURCE_DIR / SCHEMA_RESOURCE_NAME)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    payload = build_glue_schema()
+    path.write_text(
+        json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8"
+    )
+    load_glue_schema.cache_clear()
+    return path
+
+
+__all__ = [
+    "DEPRECATED_JOB_FIELDS",
+    "MAGIC_MAP_RESOURCE_NAME",
+    "RESOURCE_DIR",
+    "SCHEMA_RESOURCE_NAME",
+    "build_glue_schema",
+    "get_updateable_job_fields",
+    "inspect_live_update_job_fields",
+    "load_glue_schema",
+    "load_magic_map",
+    "write_glue_schema",
+]

gluekit/local/__init__.py

@@ -0,0 +1,6 @@
+from . import local_fixtures,local
+
+run_glue_script = local.run_glue_script
+
+
+__all__ = ["local","local_fixtures","run_glue_script"]

gluekit/local/awsglue/__init__.py

@@ -0,0 +1 @@
+"""awsglue top-level module"""

gluekit/local/awsglue/context.py

@@ -0,0 +1,30 @@
+from gluekit.local.pyspark.context import SparkContext
+
+
+class GlueContext:
+    """
+    A local emulation of AWS GlueContext for testing purposes.
+    To allow local runs on aws glue jobs and notebooks during development:
+
+
+    ```python
+
+        from awsglue.context import GlueContext
+        from awsglue.job import Job
+        from awsglue.utils import getResolvedOptions
+        from pyspark.context import SparkContext
+
+        sc = SparkContext()
+        glueContext = GlueContext(sc)
+        spark = glueContext.spark_session
+        args = getResolvedOptions(sys.argv, ['JOB_NAME'])
+        job = Job(glueContext)
+        job.init(args['JOB_NAME'], args)
+
+    ```
+
+
+    """
+
+    def __init__(self, _spark_context: SparkContext):
+        self.spark_session = _spark_context.getOrCreate()

gluekit/local/awsglue/job.py

@@ -0,0 +1,9 @@
+class Job:
+    def __init__(self, _glue_context):
+        self.name = None
+
+    def init(self, name, _args):
+        self.name = name
+
+    def commit(self):
+        return None

gluekit/local/awsglue/utils.py

@@ -0,0 +1,17 @@
+def getResolvedOptions(argv, keys):
+    resolved = {}
+    missing = []
+    for key in keys:
+        flag = f"--{key}"
+        if flag in argv:
+            index = argv.index(flag)
+            if index + 1 < len(argv):
+                resolved[key] = argv[index + 1]
+            else:
+                missing.append(key)
+        else:
+            missing.append(key)
+
+    if missing:
+        raise ValueError(f"Missing required options: {missing}")
+    return resolved