Support latest dagster-polars dependencies

Author: cmpadden

libraries/dagster-polars/dagster_polars/io_managers/base.py

@@ -85,6 +85,21 @@ def setup_for_execution(self, context: InitResourceContext) -> None:
             else UPath(check.not_none(context.instance).storage_directory())
         )
 
+    @property
+    def storage_options(self) -> dict[str, Any]:
+        """Return UPath storage options across universal-pathlib versions.
+
+        Dagster's ``UPathIOManager`` currently reads the private ``_kwargs``
+        attribute, which was removed in universal-pathlib 0.3. Prefer the
+        public ``storage_options`` attribute when present and fall back to the
+        legacy private attribute for universal-pathlib 0.2.
+        """
+        if hasattr(self._base_path, "storage_options"):
+            return dict(self._base_path.storage_options or {})
+        if hasattr(self._base_path, "_kwargs"):
+            return self._base_path._kwargs.copy()  # noqa: SLF001
+        return {}
+
     @abstractmethod
     def write_df_to_path(
         self,

libraries/dagster-polars/dagster_polars/io_managers/delta.py

@@ -5,7 +5,14 @@
 from typing import TYPE_CHECKING, Any, ClassVar, Optional
 
 import polars as pl
-from dagster import InputContext, MetadataValue, MultiPartitionKey, OutputContext
+from dagster import (
+    InputContext,
+    MetadataValue,
+    MultiPartitionKey,
+    MultiPartitionsDefinition,
+    OutputContext,
+    TimeWindowPartitionsDefinition,
+)
 from dagster._core.errors import DagsterInvariantViolationError
 from dagster._core.storage.upath_io_manager import is_dict_type
 from packaging.version import parse as parse_version
@@ -484,7 +491,18 @@ def get_predicate(
                 f"Invalid context type: {type(context)}"
             )
 
-        def key_to_predicate(key):
+        def key_to_predicate(key: str, dim: str | None = None) -> str:
+            partitions_def = context.asset_partitions_def
+            if dim is not None and isinstance(
+                partitions_def, MultiPartitionsDefinition
+            ):
+                dim_partitions_def = partitions_def.get_partitions_def_for_dimension(
+                    dim
+                )
+                if isinstance(dim_partitions_def, TimeWindowPartitionsDefinition):
+                    return f"DATE '{key}'"
+            elif isinstance(partitions_def, TimeWindowPartitionsDefinition):
+                return f"DATE '{key}'"
             return f"'{key}'"
 
         if partition_by is None or not context.has_asset_partitions:
@@ -502,7 +520,7 @@ def key_to_predicate(key):
 
             predicate = " AND ".join(
                 [
-                    f"{partition_by[dim]} in ({', '.join(map(key_to_predicate, keys))})"
+                    f"{partition_by[dim]} in ({', '.join(key_to_predicate(key, dim) for key in keys)})"
                     for dim, keys in all_keys_by_dim.items()
                 ]
             )
@@ -514,7 +532,7 @@ def key_to_predicate(key):
             )
 
             if len(context.asset_partition_keys) == 1:
-                predicate = f"{partition_by} = '{context.asset_partition_keys[0]}'"
+                predicate = f"{partition_by} = {key_to_predicate(context.asset_partition_keys[0])}"
             else:
                 predicate = f"{partition_by} in ({', '.join(map(key_to_predicate, context.asset_partition_keys))})"
 

libraries/dagster-polars/pyproject.toml

@@ -18,12 +18,12 @@ dependencies = [
     "polars>=1.0.0",
     "pyarrow>=8.0.0",
     "typing-extensions>=4.7.0",
-    "universal-pathlib>=0.1.4,<0.3",
+    "universal-pathlib>=0.1.4",
 ]
 dynamic = ["version"]
 
 [project.optional-dependencies]
-deltalake = ["deltalake>=0.25.0,<1.2"]
+deltalake = ["deltalake>=0.25.0"]
 gcp = ["dagster-gcp>=0.19.5"]
 patito = [
     "patito>=0.8.3",