Make input partitions methods based on asset key not input name

python_modules/dagster/dagster/_core/definitions/op_invocation.py

@@ -32,7 +32,8 @@
 from .result import MaterializeResult
 
 if TYPE_CHECKING:
-    from ..execution.context.invocation import DirectOpExecutionContext
+    from ..execution.context.compute import OpExecutionContext
+    from ..execution.context.invocation import BaseDirectExecutionContext
     from .assets import AssetsDefinition
     from .composition import PendingNodeInvocation
     from .decorators.op_decorator import DecoratedOpFunction
@@ -100,6 +101,16 @@ def _separate_args_and_kwargs(
     )
 
 
+def _get_op_context(
+    context,
+) -> "OpExecutionContext":
+    from dagster._core.execution.context.compute import AssetExecutionContext
+
+    if isinstance(context, AssetExecutionContext):
+        return context.op_execution_context
+    return context
+
+
 def direct_invocation_result(
     def_or_invocation: Union[
         "OpDefinition", "PendingNodeInvocation[OpDefinition]", "AssetsDefinition"
@@ -109,7 +120,7 @@ def direct_invocation_result(
 ) -> Any:
     from dagster._config.pythonic_config import Config
     from dagster._core.execution.context.invocation import (
-        DirectOpExecutionContext,
+        BaseDirectExecutionContext,
         build_op_context,
     )
 
@@ -149,12 +160,12 @@ def direct_invocation_result(
                 " no context was provided when invoking."
             )
         if len(args) > 0:
-            if args[0] is not None and not isinstance(args[0], DirectOpExecutionContext):
+            if args[0] is not None and not isinstance(args[0], BaseDirectExecutionContext):
                 raise DagsterInvalidInvocationError(
                     f"Decorated function '{compute_fn.name}' has context argument, "
                     "but no context was provided when invoking."
                 )
-            context = cast(DirectOpExecutionContext, args[0])
+            context = args[0]
             # update args to omit context
             args = args[1:]
         else:  # context argument is provided under kwargs
@@ -165,14 +176,14 @@ def direct_invocation_result(
                     f"'{context_param_name}', but no value for '{context_param_name}' was "
                     f"found when invoking. Provided kwargs: {kwargs}"
                 )
-            context = cast(DirectOpExecutionContext, kwargs[context_param_name])
+            context = kwargs[context_param_name]
             # update kwargs to remove context
             kwargs = {
                 kwarg: val for kwarg, val in kwargs.items() if not kwarg == context_param_name
             }
     # allow passing context, even if the function doesn't have an arg for it
-    elif len(args) > 0 and isinstance(args[0], DirectOpExecutionContext):
-        context = cast(DirectOpExecutionContext, args[0])
+    elif len(args) > 0 and isinstance(args[0], BaseDirectExecutionContext):
+        context = args[0]
         args = args[1:]
 
     resource_arg_mapping = {arg.name: arg.name for arg in compute_fn.get_resource_args()}
@@ -230,7 +241,7 @@ def direct_invocation_result(
 
 
 def _resolve_inputs(
-    op_def: "OpDefinition", args, kwargs, context: "DirectOpExecutionContext"
+    op_def: "OpDefinition", args, kwargs, context: "BaseDirectExecutionContext"
 ) -> Mapping[str, Any]:
     from dagster._core.execution.plan.execute_step import do_type_check
 
@@ -333,7 +344,7 @@ def _resolve_inputs(
     return input_dict
 
 
-def _key_for_result(result: MaterializeResult, context: "DirectOpExecutionContext") -> AssetKey:
+def _key_for_result(result: MaterializeResult, context: "BaseDirectExecutionContext") -> AssetKey:
     if not context.per_invocation_properties.assets_def:
         raise DagsterInvariantViolationError(
             f"Op {context.per_invocation_properties.alias} does not have an assets definition."
@@ -355,7 +366,7 @@ def _key_for_result(result: MaterializeResult, context: "DirectOpExecutionContex
 
 def _output_name_for_result_obj(
     event: MaterializeResult,
-    context: "DirectOpExecutionContext",
+    context: "BaseDirectExecutionContext",
 ):
     if not context.per_invocation_properties.assets_def:
         raise DagsterInvariantViolationError(
@@ -368,7 +379,7 @@ def _output_name_for_result_obj(
 def _handle_gen_event(
     event: T,
     op_def: "OpDefinition",
-    context: "DirectOpExecutionContext",
+    context: "BaseDirectExecutionContext",
     output_defs: Mapping[str, OutputDefinition],
     outputs_seen: Set[str],
 ) -> T:
@@ -402,7 +413,7 @@ def _handle_gen_event(
 
 
 def _type_check_output_wrapper(
-    op_def: "OpDefinition", result: Any, context: "DirectOpExecutionContext"
+    op_def: "OpDefinition", result: Any, context: "BaseDirectExecutionContext"
 ) -> Any:
     """Type checks and returns the result of a op.
 
@@ -496,12 +507,13 @@ def type_check_gen(gen):
 
 
 def _type_check_function_output(
-    op_def: "OpDefinition", result: T, context: "DirectOpExecutionContext"
+    op_def: "OpDefinition", result: T, context: "BaseDirectExecutionContext"
 ) -> T:
     from ..execution.plan.compute_generator import validate_and_coerce_op_result_to_iterator
 
     output_defs_by_name = {output_def.name: output_def for output_def in op_def.output_defs}
-    for event in validate_and_coerce_op_result_to_iterator(result, context, op_def.output_defs):
+    op_context = _get_op_context(context)
+    for event in validate_and_coerce_op_result_to_iterator(result, op_context, op_def.output_defs):
         if isinstance(event, (Output, DynamicOutput)):
             _type_check_output(output_defs_by_name[event.output_name], event, context)
         elif isinstance(event, (MaterializeResult)):
@@ -515,14 +527,14 @@ def _type_check_function_output(
 def _type_check_output(
     output_def: "OutputDefinition",
     output: Union[Output, DynamicOutput],
-    context: "DirectOpExecutionContext",
+    context: "BaseDirectExecutionContext",
 ) -> None:
     """Validates and performs core type check on a provided output.
 
     Args:
         output_def (OutputDefinition): The output definition to validate against.
         output (Any): The output to validate.
-        context (DirectOpExecutionContext): Context containing resources to be used for type
+        context (BaseDirectExecutionContext): Context containing resources to be used for type
             check.
     """
     from ..execution.plan.execute_step import do_type_check

python_modules/dagster/dagster/_core/execution/context/compute.py

@@ -38,6 +38,7 @@
     AssetKey,
     AssetMaterialization,
     AssetObservation,
+    CoercibleToAssetKey,
     ExpectationResult,
     UserEvent,
 )
@@ -948,7 +949,10 @@ def self_dependent_asset(context: AssetExecutionContext, self_dependent_asset):
 
 
         """
-        return self._step_execution_context.asset_partition_key_range_for_input(input_name)
+        upstream_asset_key = self.asset_key_for_input(input_name)
+        return self._step_execution_context.asset_partition_key_range_for_upstream_asset(
+            upstream_asset_key
+        )
 
     @public
     def asset_partition_key_for_input(self, input_name: str) -> str:
@@ -991,7 +995,10 @@ def self_dependent_asset(context: AssetExecutionContext, self_dependent_asset):
                 #   "2023-08-20"
 
         """
-        return self._step_execution_context.asset_partition_key_for_input(input_name)
+        upstream_asset_key = self.asset_key_for_input(input_name)
+        return self._step_execution_context.asset_partition_key_for_upstream_asset(
+            upstream_asset_key
+        )
 
     @public
     def asset_partitions_def_for_output(self, output_name: str = "result") -> PartitionsDefinition:
@@ -1207,9 +1214,10 @@ def self_dependent_asset(context: AssetExecutionContext, self_dependent_asset):
                 # running a backfill of the 2023-08-21 through 2023-08-25 partitions of this asset will log:
                 #   ["2023-08-20", "2023-08-21", "2023-08-22", "2023-08-23", "2023-08-24"]
         """
+        upstream_asset_key = self.asset_key_for_input(input_name)
         return list(
-            self._step_execution_context.asset_partitions_subset_for_input(
-                input_name
+            self._step_execution_context.asset_partitions_subset_for_upstream_asset(
+                upstream_asset_key
             ).get_partition_keys()
         )
 
@@ -1287,7 +1295,10 @@ def self_dependent_asset(context: AssetExecutionContext, self_dependent_asset):
                 #   TimeWindow("2023-08-20", "2023-08-25")
 
         """
-        return self._step_execution_context.asset_partitions_time_window_for_input(input_name)
+        upstream_asset_key = self.asset_key_for_input(input_name)
+        return self._step_execution_context.asset_partitions_time_window_for_upstream_asset(
+            upstream_asset_key
+        )
 
     @public
     @experimental
@@ -1459,6 +1470,29 @@ def job_def(self) -> JobDefinition:
         """
         return self.op_execution_context.job_def
 
+    @public
+    def latest_materialization_for_upstream_asset(
+        self, key: CoercibleToAssetKey
+    ) -> Optional[AssetMaterialization]:
+        """Get the most recent AssetMaterialization event for the key. The key must be an upstream
+        asset for the currently materializing asset. Information like metadata and tags can be found
+        on the AssetMaterialization. If the key is not an upstream asset of the currently
+        materializing asset, an error will be raised. If no AssetMaterialization exists for key, None
+        will be returned.
+
+        Returns: Optional[AssetMaterialization]
+        """
+        materialization_events = (
+            self.op_execution_context._step_execution_context.upstream_asset_materialization_events  # noqa: SLF001
+        )
+        if AssetKey.from_coercible(key) in materialization_events.keys():
+            return materialization_events.get(AssetKey.from_coercible(key))
+
+        raise DagsterInvariantViolationError(
+            f"Cannot fetch AssetMaterialization for asset {key}. {key} must be an upstream dependency"
+            "in order to call latest_materialization_for_upstream_asset."
+        )
+
     ######## Deprecated methods
 
     @deprecated(**_get_deprecation_kwargs("dagster_run"))

python_modules/dagster/dagster/_core/execution/context/input.py

@@ -405,7 +405,7 @@ def asset_partitions_time_window(self) -> TimeWindow:
                 "Tried to access asset_partitions_time_window, but the asset is not partitioned.",
             )
 
-        return self.step_context.asset_partitions_time_window_for_input(self.name)
+        return self.step_context.asset_partitions_time_window_for_upstream_asset(self.asset_key)
 
     @public
     def get_identifier(self) -> Sequence[str]: