From 3a80f77a0e6d84f7953829597cd74668eb4b07c1 Mon Sep 17 00:00:00 2001
From: JamieDeMaria <jamie@dagsterlabs.com>
Date: Fri, 26 Jan 2024 16:02:52 -0500
Subject: [PATCH] deprecate op related methods from AssetExecutionContext

---
 .../_core/execution/context/compute.py        | 48 ++++++++++++++-----
 .../test_asset_execution_context.py           |  8 ----
 2 files changed, 35 insertions(+), 21 deletions(-)

diff --git a/python_modules/dagster/dagster/_core/execution/context/compute.py b/python_modules/dagster/dagster/_core/execution/context/compute.py
index bc3ae18cb1f9d..80c0bd86726ed 100644
--- a/python_modules/dagster/dagster/_core/execution/context/compute.py
+++ b/python_modules/dagster/dagster/_core/execution/context/compute.py
@@ -1375,6 +1375,16 @@ def _copy_docs_from_op_execution_context(obj):
     "get_tag": "context.run.tags.get(key)",
 }
 
+USE_OP_CONTEXT = [
+    "op_config",
+    "node_handle",
+    "op_handle",
+    "op",
+    "op_def",
+    "get_mapping_key",
+    "selected_output_names",
+]
+
 
 def _get_deprecation_kwargs(attr: str):
     deprecation_kwargs = {"breaking_version": "1.8.0"}
@@ -1390,6 +1400,11 @@ def _get_deprecation_kwargs(attr: str):
             f"You have called the deprecated method {attr} on AssetExecutionContext. Use"
             f" {ALTERNATE_EXPRESSIONS[attr]} instead."
         )
+    elif attr in USE_OP_CONTEXT:
+        deprecation_kwargs["additional_warn_text"] = (
+            f"You have called the deprecated method {attr} on AssetExecutionContext. Use"
+            f" context.op_execution_context.{attr} instead."
+        )
 
     return deprecation_kwargs
 
@@ -1533,57 +1548,64 @@ def asset_partitions_def_for_output(self, output_name: str = "result") -> Partit
     def asset_partition_keys_for_output(self, output_name: str = "result") -> Sequence[str]:
         return self.op_execution_context.asset_partition_keys_for_output(output_name=output_name)
 
-    ########## pass-through to op context
-
-    #### op related
-
-    @property
-    @_copy_docs_from_op_execution_context
-    def retry_number(self):
-        return self.op_execution_context.retry_number
-
+    @deprecated(**_get_deprecation_kwargs("op_config"))
     @public
     @property
     @_copy_docs_from_op_execution_context
     def op_config(self) -> Any:
         return self.op_execution_context.op_config
 
+    @deprecated(**_get_deprecation_kwargs("node_handle"))
     @property
     @_copy_docs_from_op_execution_context
     def node_handle(self) -> NodeHandle:
         return self.op_execution_context.node_handle
 
+    @deprecated(**_get_deprecation_kwargs("op_handle"))
     @property
     @_copy_docs_from_op_execution_context
     def op_handle(self) -> NodeHandle:
         return self.op_execution_context.op_handle
 
+    @deprecated(**_get_deprecation_kwargs("op"))
     @property
     @_copy_docs_from_op_execution_context
     def op(self) -> Node:
         return self.op_execution_context.op
 
+    @deprecated(**_get_deprecation_kwargs("op_def"))
     @public
     @property
     @_copy_docs_from_op_execution_context
     def op_def(self) -> OpDefinition:
         return self.op_execution_context.op_def
 
-    @_copy_docs_from_op_execution_context
-    def describe_op(self) -> str:
-        return self.op_execution_context.describe_op()
-
+    @deprecated(**_get_deprecation_kwargs("get_mapping_key"))
     @public
     @_copy_docs_from_op_execution_context
     def get_mapping_key(self) -> Optional[str]:
         return self.op_execution_context.get_mapping_key()
 
+    @deprecated(**_get_deprecation_kwargs("selected_output_names"))
     @public
     @property
     @_copy_docs_from_op_execution_context
     def selected_output_names(self) -> AbstractSet[str]:
         return self.op_execution_context.selected_output_names
 
+    ########## pass-through to op context
+
+    #### op related
+
+    @property
+    @_copy_docs_from_op_execution_context
+    def retry_number(self):
+        return self.op_execution_context.retry_number
+
+    @_copy_docs_from_op_execution_context
+    def describe_op(self) -> str:
+        return self.op_execution_context.describe_op()
+
     #### job related
 
     @public
diff --git a/python_modules/dagster/dagster_tests/core_tests/execution_tests/test_asset_execution_context.py b/python_modules/dagster/dagster_tests/core_tests/execution_tests/test_asset_execution_context.py
index 0c0a11b3b71f9..55d127c3ba1a6 100644
--- a/python_modules/dagster/dagster_tests/core_tests/execution_tests/test_asset_execution_context.py
+++ b/python_modules/dagster/dagster_tests/core_tests/execution_tests/test_asset_execution_context.py
@@ -67,7 +67,6 @@ def test_deprecation_warnings():
         "has_partition_key",
         "instance",
         "job_name",
-        "op_handle",
         "output_for_asset_key",
         "partition_key",
         "partition_key_range",
@@ -76,19 +75,12 @@ def test_deprecation_warnings():
         "resources",
         "selected_asset_check_keys",
         "selected_asset_keys",
-        "selected_output_names",
         "set_data_version",
         "set_requires_typed_event_stream",
         "typed_event_stream_error_message",
         "describe_op",
         "has_assets_def",
-        "get_mapping_key",
         "get_step_execution_context",
-        "node_handle",
-        "op",
-        "op_config",
-        "op_def",
-        "op_handle",
         "step_launcher",
         "has_events",
         "consume_events",