dagster-azure 0.13.19__tar.gz → 0.28.1__tar.gz

{dagster-azure-0.13.19 → dagster_azure-0.28.1}/LICENSE

@@ -186,7 +186,7 @@
       same "printed page" as the copyright notice for easier
       identification within third-party archives.
 
-   Copyright {yyyy} {name of copyright owner}
+   Copyright 2025 Dagster Labs, Inc.
 
    Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.

{dagster-azure-0.13.19 → dagster_azure-0.28.1}/MANIFEST.in

@@ -3,3 +3,4 @@ recursive-include dagster_azure *.yaml
 recursive-include dagster_azure *.txt
 recursive-include dagster_azure *.template
 include LICENSE
+include dagster_azure/py.typed

dagster_azure-0.28.1/PKG-INFO

@@ -0,0 +1,32 @@
+Metadata-Version: 2.4
+Name: dagster-azure
+Version: 0.28.1
+Summary: Package for Azure-specific Dagster framework op and resource components.
+Home-page: https://github.com/dagster-io/dagster/tree/master/python_modules/libraries/dagster-azure
+Author: Dagster Labs
+Author-email: hello@dagsterlabs.com
+License: Apache-2.0
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.9,<3.14
+License-File: LICENSE
+Requires-Dist: azure-core<2.0.0,>=1.7.0
+Requires-Dist: azure-identity<2.0.0,>=1.7.0
+Requires-Dist: azure-ai-ml<2.0.0,>=1.28.0
+Requires-Dist: azure-storage-blob<13.0.0,>=12.5.0
+Requires-Dist: azure-storage-file-datalake<13.0.0,>=12.5
+Requires-Dist: dagster==1.12.1
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: home-page
+Dynamic: license
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary

dagster_azure-0.28.1/README.md

@@ -0,0 +1,4 @@
+# dagster-azure
+
+The docs for `dagster-azure` can be found
+[here](https://docs.dagster.io/api/python-api/libraries/dagster-azure).

dagster_azure-0.28.1/dagster_azure/__init__.py

@@ -0,0 +1,5 @@
+from dagster_shared.libraries import DagsterLibraryRegistry
+
+from dagster_azure.version import __version__
+
+DagsterLibraryRegistry.register("dagster-azure", __version__)

dagster_azure-0.28.1/dagster_azure/adls2/__init__.py

@@ -0,0 +1,19 @@
+from dagster_azure.adls2.file_manager import (
+    ADLS2FileHandle as ADLS2FileHandle,
+    ADLS2FileManager as ADLS2FileManager,
+)
+from dagster_azure.adls2.io_manager import (
+    ADLS2PickleIOManager as ADLS2PickleIOManager,
+    ConfigurablePickledObjectADLS2IOManager as ConfigurablePickledObjectADLS2IOManager,
+    PickledObjectADLS2IOManager as PickledObjectADLS2IOManager,
+    adls2_pickle_io_manager as adls2_pickle_io_manager,
+)
+from dagster_azure.adls2.resources import (
+    ADLS2DefaultAzureCredential as ADLS2DefaultAzureCredential,
+    ADLS2Key as ADLS2Key,
+    ADLS2Resource as ADLS2Resource,
+    ADLS2SASToken as ADLS2SASToken,
+    adls2_file_manager as adls2_file_manager,
+    adls2_resource as adls2_resource,
+)
+from dagster_azure.adls2.utils import create_adls2_client as create_adls2_client

{dagster-azure-0.13.19 → dagster_azure-0.28.1}/dagster_azure/adls2/file_manager.py

@@ -1,17 +1,19 @@
 import io
 import uuid
 from contextlib import contextmanager
+from typing import Any, Optional
 
-from dagster import check, usable_as_dagster_type
-from dagster.core.storage.file_manager import (
+import dagster._check as check
+from dagster._core.storage.file_manager import (
     FileHandle,
     FileManager,
     TempfileManager,
     check_file_like_obj,
 )
 
+from dagster_azure.adls2.utils import DataLakeServiceClient
+
 
-@usable_as_dagster_type
 class ADLS2FileHandle(FileHandle):
     """A reference to a file on ADLS2."""
 
@@ -21,48 +23,44 @@ class ADLS2FileHandle(FileHandle):
         self._key = check.str_param(key, "key")
 
     @property
-    def account(self):
+    def account(self) -> str:
         """str: The name of the ADLS2 account."""
         return self._account
 
     @property
-    def file_system(self):
+    def file_system(self) -> str:
         """str: The name of the ADLS2 file system."""
         return self._file_system
 
     @property
-    def key(self):
+    def key(self) -> str:
         """str: The ADLS2 key."""
         return self._key
 
     @property
-    def path_desc(self):
+    def path_desc(self) -> str:
         """str: The file's ADLS2 URL."""
         return self.adls2_path
 
     @property
-    def adls2_path(self):
+    def adls2_path(self) -> str:
         """str: The file's ADLS2 URL."""
-        return "adfss://{file_system}@{account}.dfs.core.windows.net/{key}".format(
-            file_system=self.file_system,
-            account=self.account,
-            key=self.key,
-        )
+        return f"adfss://{self.file_system}@{self.account}.dfs.core.windows.net/{self.key}"
 
 
 class ADLS2FileManager(FileManager):
-    def __init__(self, adls2_client, file_system, prefix):
+    def __init__(self, adls2_client: DataLakeServiceClient, file_system: str, prefix: str):
         self._client = adls2_client
         self._file_system = check.str_param(file_system, "file_system")
         self._prefix = check.str_param(prefix, "prefix")
-        self._local_handle_cache = {}
+        self._local_handle_cache: dict[str, str] = {}
         self._temp_file_manager = TempfileManager()
 
-    def copy_handle_to_local_temp(self, file_handle):
+    def copy_handle_to_local_temp(self, file_handle: ADLS2FileHandle):  # pyright: ignore[reportIncompatibleMethodOverride]
         self._download_if_not_cached(file_handle)
         return self._get_local_path(file_handle)
 
-    def _download_if_not_cached(self, file_handle):
+    def _download_if_not_cached(self, file_handle: ADLS2FileHandle):
         if not self._file_handle_cached(file_handle):
             # instigate download
             temp_file_obj = self._temp_file_manager.tempfile()
@@ -79,41 +77,45 @@ class ADLS2FileManager(FileManager):
         return file_handle
 
     @contextmanager
-    def read(self, file_handle, mode="rb"):
+    def read(self, file_handle: ADLS2FileHandle, mode: str = "rb"):  # pyright: ignore[reportIncompatibleMethodOverride]
         check.inst_param(file_handle, "file_handle", ADLS2FileHandle)
         check.str_param(mode, "mode")
         check.param_invariant(mode in {"r", "rb"}, "mode")
 
         self._download_if_not_cached(file_handle)
 
-        with open(self._get_local_path(file_handle), mode) as file_obj:
+        encoding = None if "b" in mode else "utf-8"
+        with open(self._get_local_path(file_handle), mode, encoding=encoding) as file_obj:
             yield file_obj
 
-    def _file_handle_cached(self, file_handle):
+    def _file_handle_cached(self, file_handle: ADLS2FileHandle) -> bool:
         return file_handle.adls2_path in self._local_handle_cache
 
-    def _get_local_path(self, file_handle):
+    def _get_local_path(self, file_handle: ADLS2FileHandle) -> str:
         return self._local_handle_cache[file_handle.adls2_path]
 
-    def read_data(self, file_handle):
+    def read_data(self, file_handle: ADLS2FileHandle) -> Any:  # pyright: ignore[reportIncompatibleMethodOverride]
         with self.read(file_handle, mode="rb") as file_obj:
             return file_obj.read()
 
-    def write_data(self, data, ext=None):
+    def write_data(self, data: bytes, ext: Optional[str] = None) -> ADLS2FileHandle:
         check.inst_param(data, "data", bytes)
         return self.write(io.BytesIO(data), mode="wb", ext=ext)
 
-    def write(self, file_obj, mode="wb", ext=None):  # pylint: disable=unused-argument
+    def write(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self, file_obj: io.BytesIO, mode: str = "wb", ext: Optional[str] = None
+    ) -> ADLS2FileHandle:
         check_file_like_obj(file_obj)
         adls2_key = self.get_full_key(str(uuid.uuid4()) + (("." + ext) if ext is not None else ""))
         adls2_file = self._client.get_file_client(
             file_system=self._file_system, file_path=adls2_key
         )
         adls2_file.upload_data(file_obj, overwrite=True)
-        return ADLS2FileHandle(self._client.account_name, self._file_system, adls2_key)
+        account_name = check.not_none(self._client.account_name, "Expected account name to be set")
+        return ADLS2FileHandle(account_name, self._file_system, adls2_key)
 
-    def get_full_key(self, file_key):
-        return "{base_key}/{file_key}".format(base_key=self._prefix, file_key=file_key)
+    def get_full_key(self, file_key: str) -> str:
+        return f"{self._prefix}/{file_key}"
 
-    def delete_local_temp(self):
+    def delete_local_temp(self) -> None:
         self._temp_file_manager.close()

dagster_azure-0.28.1/dagster_azure/adls2/io_manager.py

@@ -0,0 +1,314 @@
+import pickle
+from collections.abc import Iterator
+from contextlib import contextmanager
+from typing import Any, Union
+
+from dagster import (
+    InputContext,
+    OutputContext,
+    ResourceDependency,
+    _check as check,
+    io_manager,
+)
+from dagster._annotations import deprecated
+from dagster._config.pythonic_config import ConfigurableIOManager
+from dagster._core.execution.context.init import InitResourceContext
+from dagster._core.storage.io_manager import dagster_maintained_io_manager
+from dagster._core.storage.upath_io_manager import UPathIOManager
+from dagster._utils import PICKLE_PROTOCOL
+from dagster._utils.cached_method import cached_method
+from pydantic import Field
+from upath import UPath
+
+from dagster_azure.adls2.resources import ADLS2Resource
+from dagster_azure.adls2.utils import (
+    DataLakeLeaseClient,
+    DataLakeServiceClient,
+    ResourceNotFoundError,
+)
+from dagster_azure.blob.utils import BlobLeaseClient, BlobServiceClient
+
+
+class PickledObjectADLS2IOManager(UPathIOManager):
+    def __init__(
+        self,
+        file_system: str,
+        adls2_client: DataLakeServiceClient,
+        blob_client: BlobServiceClient,
+        lease_client_constructor: Union[type[DataLakeLeaseClient], type[BlobLeaseClient]],
+        prefix: str = "dagster",
+        lease_duration: int = 60,
+    ):
+        if lease_duration != -1 and (lease_duration < 15 or lease_duration > 60):
+            raise ValueError("lease_duration must be -1 (unlimited) or between 15 and 60")
+
+        self.adls2_client = adls2_client
+        self.file_system_client = self.adls2_client.get_file_system_client(file_system)
+        # We also need a blob client to handle copying as ADLS doesn't have a copy API yet
+        self.blob_client = blob_client
+        self.blob_container_client = self.blob_client.get_container_client(file_system)
+        self.prefix = check.str_param(prefix, "prefix")
+
+        self.lease_client_constructor = lease_client_constructor
+        self.lease_duration = lease_duration
+        self.file_system_client.get_file_system_properties()
+        super().__init__(base_path=UPath(self.prefix))
+
+    def get_op_output_relative_path(self, context: Union[InputContext, OutputContext]) -> UPath:
+        parts = context.get_identifier()
+        run_id = parts[0]
+        output_parts = parts[1:]
+        return UPath("storage", run_id, "files", *output_parts)
+
+    def get_loading_input_log_message(self, path: UPath) -> str:
+        return f"Loading ADLS2 object from: {self._uri_for_path(path)}"
+
+    def get_writing_output_log_message(self, path: UPath) -> str:
+        return f"Writing ADLS2 object at: {self._uri_for_path(path)}"
+
+    def unlink(self, path: UPath) -> None:
+        file_client = self.file_system_client.get_file_client(path.as_posix())
+        with self._acquire_lease(file_client, is_rm=True) as lease:
+            file_client.delete_file(lease=lease, recursive=True)
+
+    def make_directory(self, path: UPath) -> None:
+        # It is not necessary to create directories in ADLS2
+        return None
+
+    def path_exists(self, path: UPath) -> bool:
+        try:
+            self.file_system_client.get_file_client(path.as_posix()).get_file_properties()
+        except ResourceNotFoundError:
+            return False
+        return True
+
+    def _uri_for_path(self, path: UPath, protocol: str = "abfss://") -> str:
+        return f"{protocol}{self.file_system_client.file_system_name}@{self.file_system_client.account_name}.dfs.core.windows.net/{path.as_posix()}"
+
+    @contextmanager
+    def _acquire_lease(self, client: Any, is_rm: bool = False) -> Iterator[str]:
+        lease_client = self.lease_client_constructor(client=client)
+        try:
+            # Unclear why this needs to be type-ignored
+            lease_client.acquire(lease_duration=self.lease_duration)
+            yield lease_client.id
+        finally:
+            # cannot release a lease on a file that no longer exists, so need to check
+            if not is_rm:
+                lease_client.release()
+
+    def load_from_path(self, context: InputContext, path: UPath) -> Any:
+        if context.dagster_type.typing_type == type(None):
+            return None
+        file = self.file_system_client.get_file_client(path.as_posix())
+        stream = file.download_file()
+        return pickle.loads(stream.readall())
+
+    def dump_to_path(self, context: OutputContext, obj: Any, path: UPath) -> None:
+        if self.path_exists(path):
+            context.log.warning(f"Removing existing ADLS2 key: {path}")
+            self.unlink(path)
+
+        pickled_obj = pickle.dumps(obj, PICKLE_PROTOCOL)
+        file = self.file_system_client.create_file(path.as_posix())
+        with self._acquire_lease(file) as lease:
+            file.upload_data(pickled_obj, lease=lease, overwrite=True)
+
+
+class ADLS2PickleIOManager(ConfigurableIOManager):
+    """Persistent IO manager using Azure Data Lake Storage Gen2 for storage.
+
+    Serializes objects via pickling. Suitable for objects storage for distributed executors, so long
+    as each execution node has network connectivity and credentials for ADLS and the backing
+    container.
+
+    Assigns each op output to a unique filepath containing run ID, step key, and output name.
+    Assigns each asset to a single filesystem path, at "<base_dir>/<asset_key>". If the asset key
+    has multiple components, the final component is used as the name of the file, and the preceding
+    components as parent directories under the base_dir.
+
+    Subsequent materializations of an asset will overwrite previous materializations of that asset.
+    With a base directory of "/my/base/path", an asset with key
+    `AssetKey(["one", "two", "three"])` would be stored in a file called "three" in a directory
+    with path "/my/base/path/one/two/".
+
+    Example usage:
+
+    1. Attach this IO manager to a set of assets.
+
+    .. code-block:: python
+
+        from dagster import Definitions, asset
+        from dagster_azure.adls2 import ADLS2PickleIOManager, ADLS2Resource, ADLS2SASToken
+
+        @asset
+        def asset1():
+            # create df ...
+            return df
+
+        @asset
+        def asset2(asset1):
+            return df[:5]
+
+        Definitions(
+            assets=[asset1, asset2],
+            resources={
+                "io_manager": ADLS2PickleIOManager(
+                    adls2_file_system="my-cool-fs",
+                    adls2_prefix="my-cool-prefix",
+                    adls2=ADLS2Resource(
+                        storage_account="my-storage-account",
+                        credential=ADLS2SASToken(token="my-sas-token"),
+                    ),
+                ),
+            },
+        )
+
+
+    2. Attach this IO manager to your job to make it available to your ops.
+
+    .. code-block:: python
+
+        from dagster import job
+        from dagster_azure.adls2 import ADLS2PickleIOManager, ADLS2Resource, ADLS2SASToken
+
+        @job(
+            resource_defs={
+                "io_manager": ADLS2PickleIOManager(
+                    adls2_file_system="my-cool-fs",
+                    adls2_prefix="my-cool-prefix",
+                    adls2=ADLS2Resource(
+                        storage_account="my-storage-account",
+                        credential=ADLS2SASToken(token="my-sas-token"),
+                    ),
+                ),
+            },
+        )
+        def my_job():
+            ...
+    """
+
+    adls2: ResourceDependency[ADLS2Resource]
+    adls2_file_system: str = Field(description="ADLS Gen2 file system name.")
+    adls2_prefix: str = Field(
+        default="dagster", description="ADLS Gen2 file system prefix to write to."
+    )
+    lease_duration: int = Field(
+        default=60,
+        description="Lease duration in seconds. Must be between 15 and 60 seconds or -1 for infinite.",
+    )
+
+    @classmethod
+    def _is_dagster_maintained(cls) -> bool:
+        return True
+
+    @property
+    @cached_method
+    def _internal_io_manager(self) -> PickledObjectADLS2IOManager:
+        return PickledObjectADLS2IOManager(
+            self.adls2_file_system,
+            self.adls2.adls2_client,
+            self.adls2.blob_client,
+            self.adls2.lease_client_constructor,
+            self.adls2_prefix,
+            self.lease_duration,
+        )
+
+    def load_input(self, context: "InputContext") -> Any:
+        return self._internal_io_manager.load_input(context)
+
+    def handle_output(self, context: "OutputContext", obj: Any) -> None:
+        self._internal_io_manager.handle_output(context, obj)
+
+
+@deprecated(
+    breaking_version="2.0",
+    additional_warn_text="Please use ADLS2PickleIOManager instead.",
+)
+class ConfigurablePickledObjectADLS2IOManager(ADLS2PickleIOManager):
+    """Renamed to ADLS2PickleIOManager. See ADLS2PickleIOManager for documentation."""
+
+    pass
+
+
+@dagster_maintained_io_manager
+@io_manager(
+    config_schema=ADLS2PickleIOManager.to_config_schema(),
+    required_resource_keys={"adls2"},
+)
+def adls2_pickle_io_manager(init_context: InitResourceContext) -> PickledObjectADLS2IOManager:
+    """Persistent IO manager using Azure Data Lake Storage Gen2 for storage.
+
+    Serializes objects via pickling. Suitable for objects storage for distributed executors, so long
+    as each execution node has network connectivity and credentials for ADLS and the backing
+    container.
+
+    Assigns each op output to a unique filepath containing run ID, step key, and output name.
+    Assigns each asset to a single filesystem path, at "<base_dir>/<asset_key>". If the asset key
+    has multiple components, the final component is used as the name of the file, and the preceding
+    components as parent directories under the base_dir.
+
+    Subsequent materializations of an asset will overwrite previous materializations of that asset.
+    With a base directory of "/my/base/path", an asset with key
+    `AssetKey(["one", "two", "three"])` would be stored in a file called "three" in a directory
+    with path "/my/base/path/one/two/".
+
+    Example usage:
+
+    Attach this IO manager to a set of assets.
+
+    .. code-block:: python
+
+        from dagster import Definitions, asset
+        from dagster_azure.adls2 import adls2_pickle_io_manager, adls2_resource
+
+        @asset
+        def asset1():
+            # create df ...
+            return df
+
+        @asset
+        def asset2(asset1):
+            return df[:5]
+
+        Definitions(
+            assets=[asset1, asset2],
+            resources={
+                "io_manager": adls2_pickle_io_manager.configured(
+                    {"adls2_file_system": "my-cool-fs", "adls2_prefix": "my-cool-prefix"}
+                ),
+                "adls2": adls2_resource,
+            },
+        )
+
+
+    Attach this IO manager to your job to make it available to your ops.
+
+    .. code-block:: python
+
+        from dagster import job
+        from dagster_azure.adls2 import adls2_pickle_io_manager, adls2_resource
+
+        @job(
+            resource_defs={
+                "io_manager": adls2_pickle_io_manager.configured(
+                    {"adls2_file_system": "my-cool-fs", "adls2_prefix": "my-cool-prefix"}
+                ),
+                "adls2": adls2_resource,
+            },
+        )
+        def my_job():
+            ...
+    """
+    adls_resource = init_context.resources.adls2
+    adls2_client = adls_resource.adls2_client
+    blob_client = adls_resource.blob_client
+    lease_client = adls_resource.lease_client_constructor
+    return PickledObjectADLS2IOManager(
+        init_context.resource_config["adls2_file_system"],
+        adls2_client,
+        blob_client,
+        lease_client,
+        init_context.resource_config.get("adls2_prefix"),
+        init_context.resource_config.get("lease_duration"),
+    )