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

{dagster-azure-0.21.13 → 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 2023 Dagster Labs, Inc".
+   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.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.21.13 → dagster_azure-0.28.1}/dagster_azure/adls2/__init__.py

@@ -1,19 +1,14 @@
-from .fake_adls2_resource import (
-    FakeADLS2Resource as FakeADLS2Resource,
-    FakeADLS2ServiceClient as FakeADLS2ServiceClient,
-    fake_adls2_resource as fake_adls2_resource,
-)
-from .file_manager import (
+from dagster_azure.adls2.file_manager import (
     ADLS2FileHandle as ADLS2FileHandle,
     ADLS2FileManager as ADLS2FileManager,
 )
-from .io_manager import (
+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 .resources import (
+from dagster_azure.adls2.resources import (
     ADLS2DefaultAzureCredential as ADLS2DefaultAzureCredential,
     ADLS2Key as ADLS2Key,
     ADLS2Resource as ADLS2Resource,
@@ -21,4 +16,4 @@ from .resources import (
     adls2_file_manager as adls2_file_manager,
     adls2_resource as adls2_resource,
 )
-from .utils import create_adls2_client as create_adls2_client
+from dagster_azure.adls2.utils import create_adls2_client as create_adls2_client

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

@@ -1,6 +1,7 @@
 import io
 import uuid
 from contextlib import contextmanager
+from typing import Any, Optional
 
 import dagster._check as check
 from dagster._core.storage.file_manager import (
@@ -10,6 +11,8 @@ from dagster._core.storage.file_manager import (
     check_file_like_obj,
 )
 
+from dagster_azure.adls2.utils import DataLakeServiceClient
+
 
 class ADLS2FileHandle(FileHandle):
     """A reference to a file on ADLS2."""
@@ -20,44 +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 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()
@@ -74,7 +77,7 @@ 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")
@@ -85,31 +88,34 @@ class ADLS2FileManager(FileManager):
         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):
+    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):
+    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.21.13 → dagster_azure-0.28.1}/dagster_azure/adls2/io_manager.py

@@ -1,6 +1,7 @@
 import pickle
+from collections.abc import Iterator
 from contextlib import contextmanager
-from typing import Any, Iterator, Union
+from typing import Any, Union
 
 from dagster import (
     InputContext,
@@ -11,6 +12,7 @@ from dagster import (
 )
 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
@@ -19,20 +21,27 @@ from pydantic import Field
 from upath import UPath
 
 from dagster_azure.adls2.resources import ADLS2Resource
-from dagster_azure.adls2.utils import ResourceNotFoundError
-
-_LEASE_DURATION = 60  # One minute
+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: Any,
-        adls2_client: Any,
-        blob_client: Any,
-        lease_client_constructor: Any,
+        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
@@ -41,7 +50,7 @@ class PickledObjectADLS2IOManager(UPathIOManager):
         self.prefix = check.str_param(prefix, "prefix")
 
         self.lease_client_constructor = lease_client_constructor
-        self.lease_duration = _LEASE_DURATION
+        self.lease_duration = lease_duration
         self.file_system_client.get_file_system_properties()
         super().__init__(base_path=UPath(self.prefix))
 
@@ -74,17 +83,13 @@ class PickledObjectADLS2IOManager(UPathIOManager):
         return True
 
     def _uri_for_path(self, path: UPath, protocol: str = "abfss://") -> str:
-        return "{protocol}{filesystem}@{account}.dfs.core.windows.net/{key}".format(
-            protocol=protocol,
-            filesystem=self.file_system_client.file_system_name,
-            account=self.file_system_client.account_name,
-            key=path.as_posix(),
-        )
+        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:
@@ -134,7 +139,7 @@ class ADLS2PickleIOManager(ConfigurableIOManager):
     .. code-block:: python
 
         from dagster import Definitions, asset
-        from dagster_azure.adls2 import ADLS2PickleIOManager, adls2_resource
+        from dagster_azure.adls2 import ADLS2PickleIOManager, ADLS2Resource, ADLS2SASToken
 
         @asset
         def asset1():
@@ -145,14 +150,17 @@ class ADLS2PickleIOManager(ConfigurableIOManager):
         def asset2(asset1):
             return df[:5]
 
-        defs = Definitions(
+        Definitions(
             assets=[asset1, asset2],
             resources={
                 "io_manager": ADLS2PickleIOManager(
                     adls2_file_system="my-cool-fs",
-                    adls2_prefix="my-cool-prefix"
+                    adls2_prefix="my-cool-prefix",
+                    adls2=ADLS2Resource(
+                        storage_account="my-storage-account",
+                        credential=ADLS2SASToken(token="my-sas-token"),
+                    ),
                 ),
-                "adls2": adls2_resource,
             },
         )
 
@@ -162,15 +170,18 @@ class ADLS2PickleIOManager(ConfigurableIOManager):
     .. code-block:: python
 
         from dagster import job
-        from dagster_azure.adls2 import ADLS2PickleIOManager, adls2_resource
+        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_prefix="my-cool-prefix",
+                    adls2=ADLS2Resource(
+                        storage_account="my-storage-account",
+                        credential=ADLS2SASToken(token="my-sas-token"),
+                    ),
                 ),
-                "adls2": adls2_resource,
             },
         )
         def my_job():
@@ -182,6 +193,10 @@ class ADLS2PickleIOManager(ConfigurableIOManager):
     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:
@@ -196,6 +211,7 @@ class ADLS2PickleIOManager(ConfigurableIOManager):
             self.adls2.blob_client,
             self.adls2.lease_client_constructor,
             self.adls2_prefix,
+            self.lease_duration,
         )
 
     def load_input(self, context: "InputContext") -> Any:
@@ -207,7 +223,7 @@ class ADLS2PickleIOManager(ConfigurableIOManager):
 
 @deprecated(
     breaking_version="2.0",
-    additional_warn_text="Please use GCSPickleIOManager instead.",
+    additional_warn_text="Please use ADLS2PickleIOManager instead.",
 )
 class ConfigurablePickledObjectADLS2IOManager(ADLS2PickleIOManager):
     """Renamed to ADLS2PickleIOManager. See ADLS2PickleIOManager for documentation."""
@@ -220,7 +236,7 @@ class ConfigurablePickledObjectADLS2IOManager(ADLS2PickleIOManager):
     config_schema=ADLS2PickleIOManager.to_config_schema(),
     required_resource_keys={"adls2"},
 )
-def adls2_pickle_io_manager(init_context):
+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
@@ -239,7 +255,7 @@ def adls2_pickle_io_manager(init_context):
 
     Example usage:
 
-    1. Attach this IO manager to a set of assets.
+    Attach this IO manager to a set of assets.
 
     .. code-block:: python
 
@@ -255,7 +271,7 @@ def adls2_pickle_io_manager(init_context):
         def asset2(asset1):
             return df[:5]
 
-        defs = Definitions(
+        Definitions(
             assets=[asset1, asset2],
             resources={
                 "io_manager": adls2_pickle_io_manager.configured(
@@ -266,7 +282,7 @@ def adls2_pickle_io_manager(init_context):
         )
 
 
-    2. Attach this IO manager to your job to make it available to your ops.
+    Attach this IO manager to your job to make it available to your ops.
 
     .. code-block:: python
 
@@ -288,11 +304,11 @@ def adls2_pickle_io_manager(init_context):
     adls2_client = adls_resource.adls2_client
     blob_client = adls_resource.blob_client
     lease_client = adls_resource.lease_client_constructor
-    pickled_io_manager = PickledObjectADLS2IOManager(
+    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"),
     )
-    return pickled_io_manager

{dagster-azure-0.21.13 → dagster_azure-0.28.1}/dagster_azure/adls2/resources.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Union
+from typing import Any, Union
 
 from azure.identity import DefaultAzureCredential
 from azure.storage.filedatalake import DataLakeLeaseClient
@@ -17,11 +17,11 @@ from dagster._utils.merger import merge_dicts
 from pydantic import Field
 from typing_extensions import Literal
 
+from dagster_azure.adls2.file_manager import ADLS2FileManager
+from dagster_azure.adls2.io_manager import InitResourceContext
+from dagster_azure.adls2.utils import DataLakeServiceClient, create_adls2_client
 from dagster_azure.blob.utils import BlobServiceClient, create_blob_client
 
-from .file_manager import ADLS2FileManager
-from .utils import DataLakeServiceClient, create_adls2_client
-
 
 class ADLS2SASToken(Config):
     credential_type: Literal["sas"] = "sas"
@@ -35,7 +35,7 @@ class ADLS2Key(Config):
 
 class ADLS2DefaultAzureCredential(Config):
     credential_type: Literal["default_azure_credential"] = "default_azure_credential"
-    kwargs: Dict[str, Any]
+    kwargs: dict[str, Any]
 
 
 class ADLS2BaseResource(ConfigurableResource):
@@ -71,6 +71,64 @@ class ADLS2Resource(ADLS2BaseResource):
 
     Contains a client for both the Data Lake and Blob APIs, to work around the limitations
     of each.
+
+    Example usage:
+
+    Attach this resource to your Definitions to be used by assets and jobs.
+
+    .. code-block:: python
+
+        from dagster import Definitions, asset, job, op
+        from dagster_azure.adls2 import ADLS2Resource, ADLS2SASToken
+
+        @asset
+        def asset1(adls2: ADLS2Resource):
+            adls2.adls2_client.list_file_systems()
+            ...
+
+        @op
+        def my_op(adls2: ADLS2Resource):
+            adls2.adls2_client.list_file_systems()
+            ...
+
+        @job
+        def my_job():
+            my_op()
+
+        Definitions(
+            assets=[asset1],
+            jobs=[my_job],
+            resources={
+                "adls2": ADLS2Resource(
+                    storage_account="my-storage-account",
+                    credential=ADLS2SASToken(token="my-sas-token"),
+                )
+            },
+        )
+
+
+    Attach this resource to your job to make it available to your ops.
+
+    .. code-block:: python
+
+        from dagster import job, op
+        from dagster_azure.adls2 import ADLS2Resource, ADLS2SASToken
+
+        @op
+        def my_op(adls2: ADLS2Resource):
+            adls2.adls2_client.list_file_systems()
+            ...
+
+        @job(
+            resource_defs={
+                "adls2": ADLS2Resource(
+                    storage_account="my-storage-account",
+                    credential=ADLS2SASToken(token="my-sas-token"),
+                )
+            },
+        )
+        def my_job():
+            my_op()
     """
 
     @classmethod
@@ -107,7 +165,7 @@ class ADLS2Resource(ADLS2BaseResource):
 # to construct the new config and then use that to construct the resource.
 @dagster_maintained_resource
 @resource(ADLS2_CLIENT_CONFIG)
-def adls2_resource(context):
+def adls2_resource(context: InitResourceContext) -> ADLS2Resource:
     """Resource that gives ops access to Azure Data Lake Storage Gen2.
 
     The underlying client is a :py:class:`~azure.storage.filedatalake.DataLakeServiceClient`.
@@ -171,7 +229,7 @@ def adls2_resource(context):
         },
     )
 )
-def adls2_file_manager(context):
+def adls2_file_manager(context: InitResourceContext) -> ADLS2FileManager:
     """FileManager that provides abstract access to ADLS2.
 
     Implements the :py:class:`~dagster._core.storage.file_manager.FileManager` API.

{dagster-azure-0.21.13 → dagster_azure-0.28.1}/dagster_azure/adls2/utils.py

@@ -3,7 +3,7 @@ import warnings
 try:
     # Centralise Azure imports here so we only need to warn in one place
     from azure.core.exceptions import ResourceNotFoundError
-    from azure.storage.filedatalake import DataLakeServiceClient
+    from azure.storage.filedatalake import DataLakeLeaseClient, DataLakeServiceClient
 except ImportError:
     msg = (
         "Could not import required Azure objects. This probably means you have an old version "
@@ -16,7 +16,7 @@ except ImportError:
     raise
 
 
-def _create_url(storage_account, subdomain):
+def _create_url(storage_account: str, subdomain: str) -> str:
     return f"https://{storage_account}.{subdomain}.core.windows.net/"
 
 
@@ -26,4 +26,9 @@ def create_adls2_client(storage_account: str, credential) -> DataLakeServiceClie
     return DataLakeServiceClient(account_url, credential)
 
 
-__all__ = ["create_adls2_client", "DataLakeServiceClient", "ResourceNotFoundError"]
+__all__ = [
+    "DataLakeLeaseClient",
+    "DataLakeServiceClient",
+    "ResourceNotFoundError",
+    "create_adls2_client",
+]

dagster_azure-0.28.1/dagster_azure/blob/__init__.py

@@ -0,0 +1,12 @@
+from dagster_azure.blob.compute_log_manager import (
+    AzureBlobComputeLogManager as AzureBlobComputeLogManager,
+)
+from dagster_azure.blob.fake_blob_client import FakeBlobServiceClient as FakeBlobServiceClient
+from dagster_azure.blob.resources import (
+    AzureBlobStorageAnonymousCredential as AzureBlobStorageAnonymousCredential,
+    AzureBlobStorageDefaultCredential as AzureBlobStorageDefaultCredential,
+    AzureBlobStorageKeyCredential as AzureBlobStorageKeyCredential,
+    AzureBlobStorageResource as AzureBlobStorageResource,
+    AzureBlobStorageSASTokenCredential as AzureBlobStorageSASTokenCredential,
+)
+from dagster_azure.blob.utils import create_blob_client as create_blob_client