dagster-databricks 0.28.7__tar.gz → 0.28.9__tar.gz

{dagster_databricks-0.28.7/dagster_databricks.egg-info → dagster_databricks-0.28.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dagster-databricks
-Version: 0.28.7
+Version: 0.28.9
 Summary: Package for Databricks-specific Dagster framework op and resource components.
 Home-page: https://github.com/dagster-io/dagster/tree/master/python_modules/libraries/dagster-databricks
 Author: Dagster Labs
@@ -11,9 +11,10 @@ Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
 Requires-Python: >=3.10,<3.15
 License-File: LICENSE
-Requires-Dist: dagster==1.12.7
-Requires-Dist: dagster-pipes==1.12.7
-Requires-Dist: dagster-pyspark==0.28.7
+Requires-Dist: dagster==1.12.9
+Requires-Dist: dagster-pipes==1.12.9
+Requires-Dist: dagster-pyspark==0.28.9
+Requires-Dist: aiohttp
 Requires-Dist: databricks-sdk<0.61.0,>=0.41
 Dynamic: author
 Dynamic: author-email

{dagster_databricks-0.28.7 → dagster_databricks-0.28.9}/README.md

@@ -1,7 +1,7 @@
 # dagster-databricks
 
 The docs for `dagster-databricks` can be found
-[here](https://docs.dagster.io/api/python-api/libraries/dagster-databricks).
+[here](https://docs.dagster.io/integrations/libraries/databricks/dagster-databricks).
 
 A guide for integrating Databricks using Dagster Pipes can be found
 [here](https://docs.dagster.io/guides/dagster-pipes/databricks).

{dagster_databricks-0.28.7 → dagster_databricks-0.28.9}/dagster_databricks/__init__.py

@@ -13,6 +13,9 @@ from dagster_shared.libraries import DagsterLibraryRegistry
 from dagster_databricks.components.databricks_asset_bundle.component import (
     DatabricksAssetBundleComponent as DatabricksAssetBundleComponent,
 )
+from dagster_databricks.components.databricks_workspace.component import (
+    DatabricksWorkspaceComponent as DatabricksWorkspaceComponent,
+)
 from dagster_databricks.databricks import (
     DatabricksClient as DatabricksClient,
     DatabricksError as DatabricksError,

{dagster_databricks-0.28.7 → dagster_databricks-0.28.9}/dagster_databricks/components/databricks_asset_bundle/component.py

@@ -1,5 +1,4 @@
 import os
-import re
 from collections import defaultdict
 from dataclasses import dataclass, field
 from functools import cached_property
@@ -35,17 +34,7 @@ from dagster_databricks.components.databricks_asset_bundle.resource import Datab
 from dagster_databricks.components.databricks_asset_bundle.scaffolder import (
     DatabricksAssetBundleScaffolder,
 )
-
-
-def snake_case(name: str) -> str:
-    """Convert a string to snake_case."""
-    # Remove file extension if present
-    name = Path(name).stem
-    # Replace special characters and spaces with underscores
-    name = re.sub(r"[^a-zA-Z0-9]+", "_", name)
-    # Convert CamelCase to snake_case
-    name = re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", name)
-    return name.lower().strip("_")
+from dagster_databricks.utils import snake_case
 
 
 @dataclass
@@ -251,9 +240,9 @@ class DatabricksAssetBundleComponent(Component, Resolvable):
         )
 
     def build_defs(self, context: ComponentLoadContext) -> Definitions:
-        component_defs_path_as_python_str = str(
-            os.path.relpath(context.component_path.file_path, start=context.project_root)
-        ).replace("/", "_")
+        component_defs_path_as_python_str = snake_case(
+            str(os.path.relpath(context.component_path.file_path, start=context.project_root))
+        )
 
         databricks_assets = []
         for task_key, asset_specs in self.asset_specs_by_task_key.items():

{dagster_databricks-0.28.7 → dagster_databricks-0.28.9}/dagster_databricks/components/databricks_asset_bundle/configs.py

@@ -12,6 +12,7 @@ from dagster import (
     get_dagster_logger,
 )
 from dagster._annotations import preview
+from dagster._serdes import whitelist_for_serdes
 from dagster_shared.record import IHaveNew, record, record_custom
 from databricks.sdk.service import jobs
 from typing_extensions import Self, TypeVar
@@ -110,6 +111,7 @@ class DatabricksTaskDependsOnConfig:
     outcome: Optional[str]
 
 
+@whitelist_for_serdes
 @record
 class DatabricksBaseTask(ABC, Generic[T_DatabricksSdkTask]):
     task_key: str
@@ -143,6 +145,7 @@ class DatabricksBaseTask(ABC, Generic[T_DatabricksSdkTask]):
     def to_databricks_sdk_task(self) -> T_DatabricksSdkTask: ...
 
 
+@whitelist_for_serdes
 @record
 class DatabricksNotebookTask(DatabricksBaseTask[jobs.NotebookTask]):
     @property
@@ -186,6 +189,7 @@ class DatabricksNotebookTask(DatabricksBaseTask[jobs.NotebookTask]):
         )
 
 
+@whitelist_for_serdes
 @record
 class DatabricksConditionTask(DatabricksBaseTask[jobs.ConditionTask]):
     @property
@@ -236,6 +240,7 @@ class DatabricksConditionTask(DatabricksBaseTask[jobs.ConditionTask]):
         )
 
 
+@whitelist_for_serdes
 @record
 class DatabricksSparkPythonTask(DatabricksBaseTask[jobs.SparkPythonTask]):
     @property
@@ -282,6 +287,7 @@ class DatabricksSparkPythonTask(DatabricksBaseTask[jobs.SparkPythonTask]):
         )
 
 
+@whitelist_for_serdes
 @record
 class DatabricksPythonWheelTask(DatabricksBaseTask[jobs.PythonWheelTask]):
     @property
@@ -331,6 +337,7 @@ class DatabricksPythonWheelTask(DatabricksBaseTask[jobs.PythonWheelTask]):
         )
 
 
+@whitelist_for_serdes
 @record
 class DatabricksSparkJarTask(DatabricksBaseTask[jobs.SparkJarTask]):
     @property
@@ -376,6 +383,7 @@ class DatabricksSparkJarTask(DatabricksBaseTask[jobs.SparkJarTask]):
         )
 
 
+@whitelist_for_serdes
 @record
 class DatabricksJobTask(DatabricksBaseTask[jobs.RunJobTask]):
     @property
@@ -420,6 +428,7 @@ class DatabricksJobTask(DatabricksBaseTask[jobs.RunJobTask]):
         )
 
 
+@whitelist_for_serdes
 @record
 class DatabricksUnknownTask(DatabricksBaseTask):
     @property
@@ -609,3 +618,13 @@ class ResolvedDatabricksExistingClusterConfig(Resolvable, Model):
 @preview
 class ResolvedDatabricksServerlessConfig(Resolvable, Model):
     is_serverless: bool = True
+
+
+@whitelist_for_serdes
+@record
+class DatabricksJob:
+    """Represents a Databricks Job structure for serialization."""
+
+    job_id: int
+    name: str
+    tasks: Optional[list[DatabricksBaseTask]] = None

{dagster_databricks-0.28.7 → dagster_databricks-0.28.9}/dagster_databricks/components/databricks_asset_bundle/resource.py

@@ -1,5 +1,12 @@
+import asyncio
 from collections.abc import Iterator, Mapping
-from typing import TYPE_CHECKING, Any, Union
+from typing import TYPE_CHECKING, Any, Optional, Union
+
+import aiohttp
+
+DATABRICKS_JOBS_API_PATH = "/api/2.1/jobs"
+MAX_CONCURRENT_REQUESTS = 10
+RATE_LIMIT_STATUS_CODE = 429
 
 from dagster import (
     AssetExecutionContext,
@@ -16,6 +23,7 @@ from databricks.sdk.service import compute, jobs
 from pydantic import Field
 
 from dagster_databricks.components.databricks_asset_bundle.configs import (
+    DatabricksJob,
     ResolvedDatabricksExistingClusterConfig,
     ResolvedDatabricksNewClusterConfig,
     ResolvedDatabricksServerlessConfig,
@@ -44,6 +52,61 @@ class DatabricksWorkspace(ConfigurableResource):
             token=self.token,
         )
 
+    async def fetch_jobs(self, databricks_filter: Any) -> list[DatabricksJob]:
+        """Fetches jobs efficiently using async I/O directly from the resource."""
+        headers = {"Authorization": f"Bearer {self.token}"}
+        base_url = self.host.rstrip("/")
+
+        async with aiohttp.ClientSession(headers=headers) as session:
+            list_url = f"{base_url}{DATABRICKS_JOBS_API_PATH}list"
+            async with session.get(list_url) as resp:
+                resp.raise_for_status()
+                data = await resp.json()
+                all_jobs_lite = data.get("jobs", [])
+
+        job_ids_to_fetch = []
+        for j in all_jobs_lite:
+            if databricks_filter and not databricks_filter.include_job(j):
+                continue
+            job_ids_to_fetch.append(j["job_id"])
+
+        if not job_ids_to_fetch:
+            return []
+
+        semaphore = asyncio.Semaphore(MAX_CONCURRENT_REQUESTS)
+
+        async def _fetch_single_job(job_id: int) -> Optional[dict]:
+            async with semaphore:
+                async with aiohttp.ClientSession(headers=headers) as session:
+                    url = f"{base_url}{DATABRICKS_JOBS_API_PATH}/get?job_id={job_id}"
+                    async with session.get(url) as resp:
+                        if resp.status == MAX_CONCURRENT_REQUESTS:
+                            await asyncio.sleep(1)
+                            return await _fetch_single_job(job_id)
+
+                        if resp.status != 200:
+                            resp.raise_for_status()
+
+                        return await resp.json()
+
+        tasks_coroutines = [_fetch_single_job(jid) for jid in job_ids_to_fetch]
+        raw_jobs = await asyncio.gather(*tasks_coroutines)
+
+        final_jobs = []
+        for rj in raw_jobs:
+            if not rj:
+                continue
+
+            settings = rj.get("settings", {})
+            job = DatabricksJob(
+                job_id=rj["job_id"],
+                name=settings.get("name", "Unnamed Job"),
+                tasks=settings.get("tasks", []),
+            )
+            final_jobs.append(job)
+
+        return final_jobs
+
     def submit_and_poll(
         self, component: "DatabricksAssetBundleComponent", context: AssetExecutionContext
     ) -> Iterator[Union[AssetMaterialization, MaterializeResult]]:

dagster_databricks-0.28.9/dagster_databricks/components/databricks_workspace/component.py

@@ -0,0 +1,207 @@
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Annotated, Any, Optional
+
+from dagster import (
+    AssetExecutionContext,
+    AssetKey,
+    AssetsDefinition,
+    AssetSpec,
+    Definitions,
+    MaterializeResult,
+    MetadataValue,
+    Resolvable,
+    ResolvedAssetSpec,
+    multi_asset,
+)
+from dagster._serdes import whitelist_for_serdes
+from dagster._symbol_annotations.lifecycle import preview
+from dagster.components import Resolver
+from dagster.components.component.state_backed_component import StateBackedComponent
+from dagster.components.utils.defs_state import (
+    DefsStateConfig,
+    DefsStateConfigArgs,
+    ResolvedDefsStateConfig,
+)
+from dagster_shared.record import record
+from dagster_shared.serdes.serdes import deserialize_value, serialize_value
+from databricks.sdk.service.jobs import RunResultState
+
+from dagster_databricks.components.databricks_asset_bundle.component import (
+    DatabricksWorkspaceArgs,
+    resolve_databricks_workspace,
+)
+from dagster_databricks.components.databricks_asset_bundle.configs import (
+    DatabricksBaseTask,
+    DatabricksJob,
+)
+from dagster_databricks.components.databricks_asset_bundle.resource import DatabricksWorkspace
+from dagster_databricks.components.databricks_workspace.schema import ResolvedDatabricksFilter
+from dagster_databricks.utils import snake_case
+
+
+@whitelist_for_serdes
+@record
+class DatabricksWorkspaceData:
+    """Container for serialized Databricks workspace state."""
+
+    jobs: list[DatabricksJob]
+
+
+@preview
+@dataclass
+class DatabricksWorkspaceComponent(StateBackedComponent, Resolvable):
+    """Component that fetches Databricks workspace jobs and exposes them as assets."""
+
+    workspace: Annotated[
+        DatabricksWorkspace,
+        Resolver(
+            resolve_databricks_workspace,
+            model_field_type=DatabricksWorkspaceArgs.model(),
+            description="The mapping defining a DatabricksWorkspace.",
+        ),
+    ]
+
+    databricks_filter: Annotated[
+        Optional[ResolvedDatabricksFilter],
+        Resolver.default(description="Filter which Databricks jobs to include"),
+    ] = None
+
+    assets_by_task_key: Annotated[
+        Optional[dict[str, list[ResolvedAssetSpec]]],
+        Resolver.default(
+            description="Optional mapping of Databricks task keys to lists of Dagster AssetSpecs.",
+        ),
+    ] = None
+
+    defs_state: ResolvedDefsStateConfig = field(
+        default_factory=DefsStateConfigArgs.legacy_code_server_snapshots
+    )
+
+    @property
+    def defs_state_config(self) -> DefsStateConfig:
+        default_key = f"{self.__class__.__name__}[{self.workspace.host}]"
+        return DefsStateConfig.from_args(self.defs_state, default_key=default_key)
+
+    async def write_state_to_path(self, state_path: Path) -> None:
+        """Async implementation of state fetching."""
+        jobs = await self.workspace.fetch_jobs(self.databricks_filter)
+
+        data = DatabricksWorkspaceData(jobs=jobs)
+        state_path.write_text(serialize_value(data))
+
+    def build_defs_from_state(self, context: Any, state_path: Optional[Path]) -> Definitions:
+        """Build Dagster Definitions from the cached state."""
+        if not state_path or not state_path.exists():
+            return Definitions()
+
+        workspace_data = deserialize_value(state_path.read_text(), DatabricksWorkspaceData)
+        jobs_state = workspace_data.jobs
+
+        databricks_assets = []
+
+        for job in jobs_state:
+            job_specs = []
+            task_key_map = {}
+            tasks = job.tasks or []
+
+            for task in tasks:
+                specs = self.get_asset_specs(task=task, job_name=job.name)
+
+                for spec in specs:
+                    job_specs.append(spec)
+                    task_key_map[spec.key] = task.task_key
+
+            if job_specs:
+                asset_def = self._create_job_asset_def(job, job_specs, task_key_map)
+                databricks_assets.append(asset_def)
+
+        return Definitions(assets=databricks_assets)
+
+    def _create_job_asset_def(
+        self, job: DatabricksJob, specs: list[Any], task_key_map: dict
+    ) -> AssetsDefinition:
+        asset_name = f"databricks_job_{job.job_id}"
+
+        @multi_asset(name=asset_name, specs=specs, can_subset=True)
+        def _execution_fn(context: AssetExecutionContext):
+            client = self.workspace.get_client()
+            selected_keys = context.selected_asset_keys
+
+            tasks_to_run = [
+                task_key
+                for task_key, specs in (self.assets_by_task_key or {}).items()
+                if any(spec.key in selected_keys for spec in specs)
+            ]
+            context.log.info(f"Triggering Databricks job {job.job_id} for tasks: {tasks_to_run}")
+
+            run = client.jobs.run_now(
+                job_id=job.job_id, only=tasks_to_run if tasks_to_run else None
+            )
+            if run.run_page_url:
+                context.log.info(f"Run URL: {run.run_page_url}")
+
+            client.jobs.wait_get_run_job_terminated_or_skipped(run.run_id)
+
+            final_run = client.jobs.get_run(run.run_id)
+            state_obj = final_run.state
+            result_state = state_obj.result_state if state_obj else None
+
+            if result_state != RunResultState.SUCCESS:
+                status_str = result_state.value if result_state else "UNKNOWN"
+                error_msg = f"Job {job.job_id} failed: {status_str}. URL: {run.run_page_url}"
+                context.log.error(error_msg)
+                raise Exception(error_msg)
+
+            for spec in specs:
+                if spec.key in selected_keys:
+                    current_task_key = next(
+                        (
+                            t_key
+                            for t_key, t_specs in (self.assets_by_task_key or {}).items()
+                            if any(s.key == spec.key for s in t_specs)
+                        ),
+                        "unknown",
+                    )
+                    yield MaterializeResult(
+                        asset_key=spec.key,
+                        metadata={
+                            "dagster-databricks/job_id": MetadataValue.int(job.job_id),
+                            "dagster-databricks/run_id": MetadataValue.int(run.run_id),
+                            "dagster-databricks/run_url": MetadataValue.url(run.run_page_url or ""),
+                            "dagster-databricks/task_key": current_task_key,
+                        },
+                    )
+
+        return _execution_fn
+
+    def get_asset_specs(self, task: DatabricksBaseTask, job_name: str) -> list[AssetSpec]:
+        """Return a list of AssetSpec objects for the given task."""
+        task_key = task.task_key
+
+        if self.assets_by_task_key and task_key in self.assets_by_task_key:
+            return [
+                spec.merge_attributes(
+                    kinds={"databricks"},
+                    metadata={
+                        "dagster-databricks/task_key": task_key,
+                        "dagster-databricks/job_name": job_name,
+                    },
+                )
+                for spec in self.assets_by_task_key[task_key]
+            ]
+
+        clean_job = snake_case(job_name)
+        clean_task = snake_case(task_key)
+
+        return [
+            AssetSpec(
+                key=AssetKey([clean_job, clean_task]),
+                description=f"Databricks task {task_key} in job {job_name}",
+                kinds={"databricks"},
+                metadata={
+                    "dagster-databricks/task_key": task_key,
+                    "dagster-databricks/job_name": job_name,
+                },
+            )
+        ]

dagster_databricks-0.28.9/dagster_databricks/components/databricks_workspace/schema.py

@@ -0,0 +1,49 @@
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Annotated, Any, Optional, Union
+
+from dagster.components import Resolver
+from pydantic import BaseModel
+
+from dagster_databricks.components.databricks_asset_bundle.configs import DatabricksJob
+
+DatabricksJobInfo = Union[dict[str, Any], DatabricksJob]
+
+
+@dataclass
+class DatabricksFilter:
+    include_job: Callable[[DatabricksJobInfo], bool]
+
+
+class IncludeJobsConfig(BaseModel):
+    job_ids: list[int]
+
+
+class DatabricksFilterConfig(BaseModel):
+    include_jobs: Optional[IncludeJobsConfig] = None
+
+
+def resolve_databricks_filter(context, config: DatabricksFilterConfig) -> DatabricksFilter:
+    """Convert a DatabricksFilterConfig into a DatabricksFilter."""
+    if config and config.include_jobs and getattr(config.include_jobs, "job_ids", None):
+        allowed_ids = set(config.include_jobs.job_ids)
+
+        def include_job(job: DatabricksJobInfo) -> bool:
+            job_id = job.get("job_id") if isinstance(job, dict) else job.job_id
+            return job_id in allowed_ids
+    else:
+
+        def include_job(job: DatabricksJobInfo) -> bool:
+            return True
+
+    return DatabricksFilter(include_job=include_job)
+
+
+ResolvedDatabricksFilter = Annotated[
+    DatabricksFilter,
+    Resolver(
+        resolve_databricks_filter,
+        model_field_type=DatabricksFilterConfig,
+        description="Filter which Databricks jobs to include",
+    ),
+]

dagster_databricks-0.28.9/dagster_databricks/utils.py

@@ -0,0 +1,8 @@
+import re
+
+
+def snake_case(name: str) -> str:
+    """Standard snake_case utility for Databricks components."""
+    name = re.sub(r"[^a-zA-Z0-9]+", "_", str(name))
+    name = re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", name)
+    return name.lower().strip("_")

dagster_databricks-0.28.9/dagster_databricks/version.py

@@ -0,0 +1 @@
+__version__ = "0.28.9"

{dagster_databricks-0.28.7 → dagster_databricks-0.28.9/dagster_databricks.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dagster-databricks
-Version: 0.28.7
+Version: 0.28.9
 Summary: Package for Databricks-specific Dagster framework op and resource components.
 Home-page: https://github.com/dagster-io/dagster/tree/master/python_modules/libraries/dagster-databricks
 Author: Dagster Labs
@@ -11,9 +11,10 @@ Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
 Requires-Python: >=3.10,<3.15
 License-File: LICENSE
-Requires-Dist: dagster==1.12.7
-Requires-Dist: dagster-pipes==1.12.7
-Requires-Dist: dagster-pyspark==0.28.7
+Requires-Dist: dagster==1.12.9
+Requires-Dist: dagster-pipes==1.12.9
+Requires-Dist: dagster-pyspark==0.28.9
+Requires-Dist: aiohttp
 Requires-Dist: databricks-sdk<0.61.0,>=0.41
 Dynamic: author
 Dynamic: author-email

{dagster_databricks-0.28.7 → dagster_databricks-0.28.9}/dagster_databricks.egg-info/SOURCES.txt

@@ -14,6 +14,7 @@ dagster_databricks/pipes.py
 dagster_databricks/py.typed
 dagster_databricks/resources.py
 dagster_databricks/types.py
+dagster_databricks/utils.py
 dagster_databricks/version.py
 dagster_databricks.egg-info/PKG-INFO
 dagster_databricks.egg-info/SOURCES.txt
@@ -26,4 +27,7 @@ dagster_databricks/components/databricks_asset_bundle/__init__.py
 dagster_databricks/components/databricks_asset_bundle/component.py
 dagster_databricks/components/databricks_asset_bundle/configs.py
 dagster_databricks/components/databricks_asset_bundle/resource.py
-dagster_databricks/components/databricks_asset_bundle/scaffolder.py
+dagster_databricks/components/databricks_asset_bundle/scaffolder.py
+dagster_databricks/components/databricks_workspace/__init__.py
+dagster_databricks/components/databricks_workspace/component.py
+dagster_databricks/components/databricks_workspace/schema.py

dagster_databricks-0.28.9/dagster_databricks.egg-info/requires.txt

@@ -0,0 +1,5 @@
+dagster==1.12.9
+dagster-pipes==1.12.9
+dagster-pyspark==0.28.9
+aiohttp
+databricks-sdk<0.61.0,>=0.41

{dagster_databricks-0.28.7 → dagster_databricks-0.28.9}/setup.py

@@ -31,9 +31,10 @@ setup(
     include_package_data=True,
     python_requires=">=3.10,<3.15",
     install_requires=[
-        "dagster==1.12.7",
-        "dagster-pipes==1.12.7",
-        "dagster-pyspark==0.28.7",
+        "dagster==1.12.9",
+        "dagster-pipes==1.12.9",
+        "dagster-pyspark==0.28.9",
+        "aiohttp",
         "databricks-sdk>=0.41,<0.61.0",  # dbt-databricks is pinned to this version
     ],
     zip_safe=False,

dagster_databricks-0.28.7/dagster_databricks/version.py

@@ -1 +0,0 @@
-__version__ = "0.28.7"

dagster_databricks-0.28.7/dagster_databricks.egg-info/requires.txt

@@ -1,4 +0,0 @@
-dagster==1.12.7
-dagster-pipes==1.12.7
-dagster-pyspark==0.28.7
-databricks-sdk<0.61.0,>=0.41