Repository: dagster-io/dagster

Issue:
Non-partitioned spec in mixed multi-asset raises "Cannot access partition_key for a non-partitioned run" (regression in 1.12.14)
### What's the issue?

Materializing a **non-partitioned** asset spec inside a `@multi_asset(can_subset=True)` raises `DagsterInvariantViolationError` when:

1. The same multi-asset contains at least one **partitioned** spec
2. The non-partitioned spec depends on an external partitioned asset
3. The external partitioned asset has been materialized at least once

This is the exact pattern produced by `@dbt_assets` with a `DagsterDbtTranslator` that assigns `partitions_def` to some models but not others, where a non-partitioned model references a partitioned source.

Versions prior to 1.12.14 do not exhibit the issue. Worked fine in 1.12.13 and earlier.

### What did you expect to happen?

The non-partitioned asset should materialize without error

### How to reproduce?

See https://github.com/Second-Foundation-Partners/dagster-partition-bug-repro for a simple project that reproduces this error.

### Dagster version

1.12.18

### Deployment type

Dagster Cloud

### Deployment details

This is reproducible locally and in dagster plus.

### Additional information

_No response_

### Message from the maintainers

Impacted by this issue? Give it a 👍! We factor engagement into prioritization.

Discussion:

Fix (patch):
diff --git a/python_modules/dagster/dagster/_core/definitions/assets/definition/assets_definition.py b/python_modules/dagster/dagster/_core/definitions/assets/definition/assets_definition.py
index 5b48e4514ef93..641d5c12f673f 100644
--- a/python_modules/dagster/dagster/_core/definitions/assets/definition/assets_definition.py
+++ b/python_modules/dagster/dagster/_core/definitions/assets/definition/assets_definition.py
@@ -1086,8 +1086,18 @@ def backfill_policy(self) -> BackfillPolicy | None:
     @cached_property
     def partitions_def(self) -> PartitionsDefinition | None:
         """Optional[PartitionsDefinition]: The PartitionsDefinition for this AssetsDefinition (if any)."""
+        selected_entity_keys = self.asset_and_check_keys
         partitions_defs = {
-            spec.partitions_def for spec in self.specs if spec.partitions_def is not None
+            *(
+                spec.partitions_def
+                for spec in self.specs
+                if spec.partitions_def is not None and spec.key in selected_entity_keys
+            ),
+            *(
+                check_spec.partitions_def
+                for check_spec in self.check_specs
+                if check_spec.partitions_def is not None and check_spec.key in selected_entity_keys
+            ),
         }
         if len(partitions_defs) == 1:
             return next(iter(partitions_defs))
diff --git a/python_modules/dagster/dagster/_core/execution/context/system.py b/python_modules/dagster/dagster/_core/execution/context/system.py
index 03f935cb656f9..941d5ce1aa6c2 100644
--- a/python_modules/dagster/dagster/_core/execution/context/system.py
+++ b/python_modules/dagster/dagster/_core/execution/context/system.py
@@ -972,12 +972,7 @@ def entity_partitions_def(self) -> PartitionsDefinition | None:
         expected to all have the same PartitionsDefinition.
         """
         if self.assets_def is not None:
-            for spec in self.assets_def.specs:
-                if spec.partitions_def is not None:
-                    return spec.partitions_def
-            for check_spec in self.assets_def.check_specs:
-                if check_spec.partitions_def is not None:
-                    return check_spec.partitions_def
+            return self.assets_def.partitions_def
         return None
 
     @property


Test changes:
diff --git a/python_modules/dagster/dagster_tests/asset_defs_tests/test_partitioned_assets.py b/python_modules/dagster/dagster_tests/asset_defs_tests/test_partitioned_assets.py
index afc294737e394..a5962b36bb3a4 100644
--- a/python_modules/dagster/dagster_tests/asset_defs_tests/test_partitioned_assets.py
+++ b/python_modules/dagster/dagster_tests/asset_defs_tests/test_partitioned_assets.py
@@ -1116,3 +1116,47 @@ def load_input(self, context):
             resources={"io_manager": IOManagerDefinition.hardcoded_io_manager(CustomIOManager())},
             partition_key=daily_partition_definition.get_partition_key(None),  # pyright: ignore[reportArgumentType]
         )
+
+
+def test_non_partitioned_spec_in_mixed_multi_asset_does_not_raise_partition_key_error():
+    pd = dg.MonthlyPartitionsDefinition(start_date="2024-01-01", end_offset=1)
+
+    @dg.asset(partitions_def=pd)
+    def partitioned_upstream(context: dg.AssetExecutionContext):
+        """External partitioned asset."""
+        context.log.info(f"Materializing partition {context.partition_key}")
+
+    @dg.multi_asset(
+        specs=[
+            # Partitioned spec — its mere existence poisons entity_partitions_def
+            dg.AssetSpec("partitioned_model", partitions_def=pd),
+            # Non-partitioned spec that depends on the partitioned upstream asset
+            dg.AssetSpec("non_partitioned_model", deps=["partitioned_upstream"]),
+        ],
+        can_subset=True,
+    )
+    def my_multi_asset(context: dg.AssetExecutionContext):
+        """Multi-asset with mixed partition specs — reproduces the bug."""
+        selected = context.selected_asset_keys
+        if dg.AssetKey("partitioned_model") in selected:
+            yield dg.MaterializeResult(asset_key="partitioned_model")
+        if dg.AssetKey("non_partitioned_model") in selected:
+            yield dg.MaterializeResult(asset_key="non_partitioned_model")
+
+    defs = dg.Definitions(assets=[partitioned_upstream, my_multi_asset])
+    job = defs.resolve_implicit_global_asset_job_def()
+
+    instance = dg.DagsterInstance.ephemeral()
+
+    result = job.execute_in_process(
+        asset_selection=[dg.AssetKey("partitioned_upstream")],
+        partition_key="2024-01-01",
+        instance=instance,
+    )
+    assert result.success
+
+    result = job.execute_in_process(
+        asset_selection=[dg.AssetKey("non_partitioned_model")],
+        instance=instance,
+    )
+    assert result.success