[core] Add asset_group_name to IO contexts

Author: RaghunandanKumar

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

@@ -16,6 +16,7 @@
 )
 from dagster._core.errors import DagsterInvariantViolationError
 from dagster._core.instance import DagsterInstance
+from dagster._core.definitions.utils import DEFAULT_GROUP_NAME
 from dagster._utils.warnings import normalize_renamed_param
 
 if TYPE_CHECKING:
@@ -206,6 +207,29 @@ def upstream_output(self) -> Optional["OutputContext"]:
         """Info about the output that produced the object we're loading."""
         return self._upstream_output
 
+    @public
+    @property
+    def asset_group_name(self) -> str:
+        """The group name of the asset being loaded."""
+        if self._upstream_output is not None and self._upstream_output._asset_spec is not None:
+            return self._upstream_output.asset_group_name
+
+        if self._asset_key is not None and self._step_context is not None:
+            return self._step_context.job_def.asset_layer.get(self._asset_key).group_name
+
+        if self._upstream_output is not None and self._upstream_output.asset_key is not None:
+            return DEFAULT_GROUP_NAME
+
+        if self._asset_key is not None:
+            raise DagsterInvariantViolationError(
+                "Attempting to access asset_group_name, but it was not provided when constructing"
+                " the InputContext"
+            )
+
+        raise DagsterInvariantViolationError(
+            "Attempting to access asset_group_name, but input does not correspond to an asset"
+        )
+
     @public
     @property
     def dagster_type(self) -> "DagsterType":

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

@@ -25,6 +25,7 @@
     time_window_for_partition_key_range,
 )
 from dagster._core.errors import DagsterInvalidMetadata, DagsterInvariantViolationError
+from dagster._core.definitions.utils import DEFAULT_GROUP_NAME
 from dagster._core.execution.context.input import KeyRangeNoPartitionsDefPartitionsSubset
 from dagster._core.execution.plan.utils import build_resources_for_manager
 from dagster._core.instance import DagsterInstance
@@ -396,6 +397,26 @@ def asset_spec(self) -> AssetSpec:
             )
         return self._asset_spec
 
+    @public
+    @property
+    def asset_group_name(self) -> str:
+        """The group name of the asset that is being stored as an output."""
+        if self._asset_spec is not None:
+            return self._asset_spec.group_name or DEFAULT_GROUP_NAME
+
+        if self._asset_key is not None and self._step_context is not None:
+            return self._step_context.job_def.asset_layer.get(self._asset_key).group_name
+
+        if self._asset_key is not None:
+            raise DagsterInvariantViolationError(
+                "Attempting to access asset_group_name, but it was not provided when constructing"
+                " the OutputContext"
+            )
+
+        raise DagsterInvariantViolationError(
+            "Attempting to access asset_group_name, but output does not correspond to an asset"
+        )
+
     @property
     def step_context(self) -> "StepExecutionContext":
         if self._warn_on_step_context_use:

python_modules/dagster/dagster_tests/execution_tests/context_tests/test_output.py

@@ -30,6 +30,40 @@ def load_input(self, context: InputContext) -> Any:
     dg.materialize([asset_spec], resources={"io_manager": TestIOManager()})
 
 
+def test_build_output_context_asset_group_name():
+    asset_spec = dg.AssetSpec(key="key", group_name="group")
+    assert (
+        dg.build_output_context(asset_key="key", asset_spec=asset_spec).asset_group_name == "group"
+    )
+
+
+def test_build_output_context_asset_group_name_defaults():
+    asset_spec = dg.AssetSpec(key="key")
+    assert (
+        dg.build_output_context(asset_key="key", asset_spec=asset_spec).asset_group_name
+        == "default"
+    )
+
+
+def test_build_input_context_asset_group_name_from_upstream_output():
+    upstream_output = dg.build_output_context(
+        asset_key="key",
+        asset_spec=dg.AssetSpec(key="key", group_name="group"),
+    )
+    assert (
+        dg.build_input_context(asset_key="key", upstream_output=upstream_output).asset_group_name
+        == "group"
+    )
+
+
+def test_build_input_context_asset_group_name_defaults_from_upstream_output():
+    upstream_output = dg.build_output_context(asset_key="key", asset_spec=dg.AssetSpec(key="key"))
+    assert (
+        dg.build_input_context(asset_key="key", upstream_output=upstream_output).asset_group_name
+        == "default"
+    )
+
+
 def test_build_output_context_with_output_metadata():
     expected_metadata = {"foo": "bar"}
     output_context = dg.build_output_context(output_metadata=expected_metadata)