databricks-bundle-decorators 0.4.2__tar.gz → 0.4.3__tar.gz

{databricks_bundle_decorators-0.4.2 → databricks_bundle_decorators-0.4.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: databricks-bundle-decorators
-Version: 0.4.2
+Version: 0.4.3
 Summary: Decorator-based framework for defining Databricks jobs and tasks as Python code.
 Author: boccileonardo
 Author-email: boccileonardo <leonardobocci99@hotmail.com>

{databricks_bundle_decorators-0.4.2 → databricks_bundle_decorators-0.4.3}/pyproject.toml

@@ -4,7 +4,7 @@ requires = [ "uv-build>=0.10.4,<0.11" ]
 
 [project]
 name = "databricks-bundle-decorators"
-version = "0.4.2"
+version = "0.4.3"
 description = "Decorator-based framework for defining Databricks jobs and tasks as Python code."
 readme = "README.md"
 authors = [ { name = "boccileonardo", email = "leonardobocci99@hotmail.com" } ]

{databricks_bundle_decorators-0.4.2 → databricks_bundle_decorators-0.4.3}/src/databricks_bundle_decorators/decorators.py

@@ -31,6 +31,32 @@ from databricks_bundle_decorators.registry import (
 from databricks_bundle_decorators.sdk_types import ClusterConfig, JobConfig, TaskConfig
 
 
+# ---------------------------------------------------------------------------
+# Reserved parameter namespace
+# ---------------------------------------------------------------------------
+
+#: Parameter names that are reserved for internal runtime wiring.
+_RESERVED_PARAM_NAMES: frozenset[str] = frozenset(
+    {"__job_name__", "__task_key__", "__run_id__"}
+)
+
+#: Parameter name prefixes that are reserved for internal runtime wiring.
+_RESERVED_PARAM_PREFIXES: tuple[str, ...] = ("__upstream__",)
+
+
+def _validate_user_params(params: dict[str, str], context: str) -> None:
+    """Raise ``ValueError`` if any *params* key collides with reserved names."""
+    for name in params:
+        if name in _RESERVED_PARAM_NAMES or any(
+            name.startswith(p) for p in _RESERVED_PARAM_PREFIXES
+        ):
+            raise ValueError(
+                f"{context}: parameter name {name!r} is reserved for internal "
+                f"runtime use. Reserved names: {sorted(_RESERVED_PARAM_NAMES)}; "
+                f"reserved prefix: {list(_RESERVED_PARAM_PREFIXES)}."
+            )
+
+
 # ---------------------------------------------------------------------------
 # Job context – tracks which @job body is currently being executed
 # ---------------------------------------------------------------------------
@@ -95,7 +121,8 @@ def task(
 
     When used **outside** a ``@job`` body (e.g. at module level), the
     function is registered under its short name for use in tests or
-    standalone execution.
+    standalone execution.  Duplicate names at module level raise
+    `DuplicateResourceError`.
 
     Parameters
     ----------
@@ -110,7 +137,12 @@ def task(
         directly to the ``databricks.bundles.jobs.Task`` constructor at
         deploy time.  See `TaskConfig`
         for the full list of supported fields.
-    """
+    Notes
+    -----
+    Dependency edges are detected only for `TaskProxy` objects passed as
+    **direct** positional or keyword arguments.  Proxies nested inside
+    lists, dicts, or other container types are **not** inspected and will
+    not register dependency edges."""
 
     def decorator(fn: types.FunctionType) -> Callable[..., Any]:
         task_key = fn.__name__
@@ -128,13 +160,21 @@ def task(
             _current_job_tasks[task_key] = meta
         else:
             # Module-level definition (standalone / test usage).
-            _TASK_REGISTRY[task_key] = meta
+            _register_unique(_TASK_REGISTRY, task_key, meta, "task")
 
         @functools.wraps(fn)
         def wrapper(*args, **kwargs):
             if _current_job_name is not None:
                 # We're being *called* inside a @job body – return a
                 # TaskProxy and record DAG edges from any proxy args.
+                if task_key in _current_job_dag:
+                    raise DuplicateResourceError(
+                        f"Task '{task_key}' is called more than once in job "
+                        f"'{_current_job_name}'. Each @task may only be invoked "
+                        "once per @job body. Use a unique function name for "
+                        "each logical step."
+                    )
+
                 upstream_deps: list[str] = []
                 edge_map: dict[str, str] = {}
 
@@ -314,6 +354,10 @@ def job(
                 f"Duplicate job '{job_name}'. Each job must have a unique name."
             )
 
+        # --- validate param names -----------------------------------------
+        if params:
+            _validate_user_params(params, f"@job('{job_name}')")
+
         # --- validate cluster type -----------------------------------------
         if cluster is not None and not isinstance(cluster, ClusterMeta):
             raise TypeError(

{databricks_bundle_decorators-0.4.2 → databricks_bundle_decorators-0.4.3}/src/databricks_bundle_decorators/io_managers/polars_csv.py

@@ -11,7 +11,7 @@ Requires the ``polars`` optional dependency::
 from __future__ import annotations
 
 from collections.abc import Callable
-from typing import Any
+from typing import Any, cast
 
 from databricks_bundle_decorators.io_manager import (
     InputContext,
@@ -107,7 +107,7 @@ class PolarsCsvIoManager(IoManager):
     def storage_options(self) -> dict[str, str] | None:
         """Resolve *storage_options*, calling it first if it is a callable."""
         if callable(self._storage_options):
-            return self._storage_options()
+            return cast(Callable[[], dict[str, str]], self._storage_options)()
         return self._storage_options
 
     def _uri(self, key: str) -> str:

{databricks_bundle_decorators-0.4.2 → databricks_bundle_decorators-0.4.3}/src/databricks_bundle_decorators/io_managers/polars_delta.py

@@ -12,7 +12,7 @@ Requires the ``polars`` and ``deltalake`` optional dependencies::
 from __future__ import annotations
 
 from collections.abc import Callable
-from typing import Any
+from typing import Any, cast
 
 from databricks_bundle_decorators.io_manager import (
     InputContext,
@@ -81,8 +81,7 @@ class PolarsDeltaIoManager(IoManager):
         they are managed by the IoManager.
     mode : str
         Delta write mode.  One of ``"overwrite"``, ``"append"``,
-        ``"error"``, or ``"ignore"``.  Defaults to ``"overwrite"``
-        for idempotent task outputs.
+        ``"error"``, or ``"ignore"``.  Defaults to ``"error"``.
 
         For **merge** operations, ignore this parameter and return a
         fully-configured `deltalake.table.TableMerger` from your task
@@ -129,7 +128,7 @@ class PolarsDeltaIoManager(IoManager):
     def storage_options(self) -> dict[str, str] | None:
         """Resolve *storage_options*, calling it first if it is a callable."""
         if callable(self._storage_options):
-            return self._storage_options()
+            return cast(Callable[[], dict[str, str]], self._storage_options)()
         return self._storage_options
 
     def _uri(self, key: str) -> str:

{databricks_bundle_decorators-0.4.2 → databricks_bundle_decorators-0.4.3}/src/databricks_bundle_decorators/io_managers/polars_json.py

@@ -15,7 +15,7 @@ Requires the ``polars`` optional dependency::
 from __future__ import annotations
 
 from collections.abc import Callable
-from typing import Any
+from typing import Any, cast
 
 from databricks_bundle_decorators.io_manager import (
     InputContext,
@@ -114,7 +114,7 @@ class PolarsJsonIoManager(IoManager):
     def storage_options(self) -> dict[str, str] | None:
         """Resolve *storage_options*, calling it first if it is a callable."""
         if callable(self._storage_options):
-            return self._storage_options()
+            return cast(Callable[[], dict[str, str]], self._storage_options)()
         return self._storage_options
 
     def _uri(self, key: str) -> str:

{databricks_bundle_decorators-0.4.2 → databricks_bundle_decorators-0.4.3}/src/databricks_bundle_decorators/io_managers/polars_parquet.py

@@ -11,7 +11,7 @@ Requires the ``polars`` optional dependency::
 from __future__ import annotations
 
 from collections.abc import Callable
-from typing import Any
+from typing import Any, cast
 
 from databricks_bundle_decorators.io_manager import (
     InputContext,
@@ -114,7 +114,7 @@ class PolarsParquetIoManager(IoManager):
     def storage_options(self) -> dict[str, str] | None:
         """Resolve *storage_options*, calling it first if it is a callable."""
         if callable(self._storage_options):
-            return self._storage_options()
+            return cast(Callable[[], dict[str, str]], self._storage_options)()
         return self._storage_options
 
     def _uri(self, key: str) -> str:

{databricks_bundle_decorators-0.4.2 → databricks_bundle_decorators-0.4.3}/src/databricks_bundle_decorators/io_managers/spark_delta.py

@@ -14,7 +14,7 @@ Requires PySpark, which is pre-installed on Databricks clusters.
 from __future__ import annotations
 
 from collections.abc import Callable
-from typing import Any
+from typing import Any, cast
 
 from databricks_bundle_decorators.io_manager import (
     InputContext,
@@ -194,7 +194,7 @@ class SparkDeltaIoManager(_SparkDeltaBase):
     def spark_configs(self) -> dict[str, str] | None:
         """Resolve *spark_configs*, calling it first if it is a callable."""
         if callable(self._spark_configs):
-            return self._spark_configs()
+            return cast(Callable[[], dict[str, str]], self._spark_configs)()
         return self._spark_configs
 
     def setup(self) -> None:

{databricks_bundle_decorators-0.4.2 → databricks_bundle_decorators-0.4.3}/src/databricks_bundle_decorators/io_managers/spark_parquet.py

@@ -14,7 +14,7 @@ Requires PySpark, which is pre-installed on Databricks clusters.
 from __future__ import annotations
 
 from collections.abc import Callable
-from typing import Any
+from typing import Any, cast
 
 from databricks_bundle_decorators.io_manager import (
     InputContext,
@@ -153,7 +153,7 @@ class SparkParquetIoManager(_SparkParquetBase):
     def spark_configs(self) -> dict[str, str] | None:
         """Resolve *spark_configs*, calling it first if it is a callable."""
         if callable(self._spark_configs):
-            return self._spark_configs()
+            return cast(Callable[[], dict[str, str]], self._spark_configs)()
         return self._spark_configs
 
     def setup(self) -> None:

{databricks_bundle_decorators-0.4.2 → databricks_bundle_decorators-0.4.3}/src/databricks_bundle_decorators/runtime.py

@@ -11,8 +11,8 @@ Databricks invokes the ``dbxdec-run`` console-script, which calls
 """
 
 import argparse
+import logging
 import os
-import sys
 import typing
 from typing import Any
 
@@ -20,6 +20,8 @@ from databricks_bundle_decorators.context import _populate_params
 from databricks_bundle_decorators.io_manager import InputContext, OutputContext
 from databricks_bundle_decorators.registry import _TASK_REGISTRY
 
+_logger = logging.getLogger(__name__)
+
 
 def _parse_named_args(argv: list[str]) -> dict[str, str]:
     """Parse ``--key=value`` pairs from an argv list into a dict."""
@@ -58,14 +60,14 @@ def run_task(task_key: str, cli_params: dict[str, str]) -> None:
     _populate_params(cli_params)
 
     # ---- look up task metadata -------------------------------------------
-    # Try qualified key (job_name.task_key) first, then short key.
     qualified_key = f"{job_name}.{task_key}"
-    task_meta = _TASK_REGISTRY.get(qualified_key) or _TASK_REGISTRY.get(task_key)
+    task_meta = _TASK_REGISTRY.get(qualified_key)
     if task_meta is None:
-        print(f"Error: Task '{task_key}' not found in registry.", file=sys.stderr)
-        print(f"Tried: '{qualified_key}', '{task_key}'", file=sys.stderr)
-        print(f"Available: {list(_TASK_REGISTRY.keys())}", file=sys.stderr)
-        sys.exit(1)
+        available = list(_TASK_REGISTRY.keys())
+        raise RuntimeError(
+            f"Task '{task_key}' not found by qualified key '{qualified_key}'. "
+            f"Available: {available}"
+        )
 
     # ---- resolve type hints for the current task's parameters ------------
     try:
@@ -77,9 +79,7 @@ def run_task(task_key: str, cli_params: dict[str, str]) -> None:
     kwargs: dict[str, Any] = {}
     for param_name, upstream_task_key in upstream_map.items():
         upstream_qualified = f"{job_name}.{upstream_task_key}"
-        upstream_meta = _TASK_REGISTRY.get(upstream_qualified) or _TASK_REGISTRY.get(
-            upstream_task_key
-        )
+        upstream_meta = _TASK_REGISTRY.get(upstream_qualified)
         if upstream_meta and upstream_meta.io_manager:
             upstream_meta.io_manager._ensure_setup()
             context = InputContext(
@@ -91,10 +91,11 @@ def run_task(task_key: str, cli_params: dict[str, str]) -> None:
             )
             kwargs[param_name] = upstream_meta.io_manager.read(context)
         else:
-            print(
-                f"Warning: upstream task '{upstream_task_key}' has no IoManager – "
-                f"cannot auto-load data for parameter '{param_name}'.",
-                file=sys.stderr,
+            _logger.warning(
+                "Upstream task '%s' has no IoManager – "
+                "cannot auto-load data for parameter '%s'.",
+                upstream_task_key,
+                param_name,
             )
 
     # ---- execute the task function ---------------------------------------
@@ -116,10 +117,10 @@ def run_task(task_key: str, cli_params: dict[str, str]) -> None:
         )
         task_meta.io_manager.write(context, result)
     elif result is not None and not task_meta.io_manager:
-        print(
-            f"Warning: task '{task_key}' returned a value but has no IoManager – "
-            f"the return value will be discarded.",
-            file=sys.stderr,
+        _logger.warning(
+            "Task '%s' returned a value but has no IoManager – "
+            "the return value will be discarded.",
+            task_key,
         )
 
 
@@ -142,8 +143,7 @@ def main() -> None:
 
     task_key = cli_params.get("__task_key__")
     if not task_key:
-        print("Error: --__task_key__=<name> is required.", file=sys.stderr)
-        sys.exit(1)
+        raise RuntimeError("--__task_key__=<name> is required.")
 
     # 3. Run the task.
     run_task(task_key, cli_params)

{databricks_bundle_decorators-0.4.2 → databricks_bundle_decorators-0.4.3}/src/databricks_bundle_decorators/task_values.py

@@ -6,6 +6,7 @@ small primitive payloads (~48 KB) and must be opted-in explicitly by calling
 `set_task_value` inside a ``@task`` function.
 """
 
+import os
 from typing import Any
 
 # Module-level fallback store used during local / test execution.
@@ -16,6 +17,11 @@ _local_task_values: dict[str, dict[str, Any]] = {}
 _current_task_key: str | None = None
 
 
+def _is_databricks_runtime() -> bool:
+    """Return ``True`` when running inside a live Databricks job cluster."""
+    return bool(os.environ.get("DATABRICKS_RUNTIME_VERSION"))
+
+
 def set_task_value(key: str, value: str | int | float | bool) -> None:
     """Write a small value into Databricks task values.
 
@@ -37,14 +43,16 @@ def set_task_value(key: str, value: str | int | float | bool) -> None:
             f"got {type(value).__name__}. Use an IoManager for complex data."
         )
 
-    try:
+    if _is_databricks_runtime():
+        # On Databricks: let any API/permission error propagate so that
+        # real runtime failures are visible rather than silently swallowed.
         from pyspark.dbutils import DBUtils  # type: ignore[import-untyped]
         from pyspark.sql import SparkSession  # type: ignore[import-untyped]
 
         spark = SparkSession.builder.getOrCreate()
         dbutils = DBUtils(spark)
         dbutils.jobs.taskValues.set(key=key, value=value)
-    except Exception:
+    else:
         # Local / testing fallback — use the current task key if set by
         # the runtime runner, otherwise fall back to "__current__".
         store_key = _current_task_key or "__current__"
@@ -61,12 +69,13 @@ def get_task_value(task_key: str, key: str) -> Any:
     key:
         The key passed to `set_task_value`.
     """
-    try:
+    if _is_databricks_runtime():
+        # On Databricks: let any API/permission error propagate.
         from pyspark.dbutils import DBUtils  # type: ignore[import-untyped]
         from pyspark.sql import SparkSession  # type: ignore[import-untyped]
 
         spark = SparkSession.builder.getOrCreate()
         dbutils = DBUtils(spark)
         return dbutils.jobs.taskValues.get(taskKey=task_key, key=key)
-    except Exception:
+    else:
         return _local_task_values.get(task_key, {}).get(key)