partition methods on AssetExecutionContext

.buildkite/dagster-buildkite/dagster_buildkite/steps/dagster.py

@@ -140,7 +140,10 @@ def build_sql_schema_check_steps() -> List[CommandStep]:
     return [
         CommandStepBuilder(":mysql: mysql-schema")
         .on_test_image(AvailablePythonVersion.get_default())
-        .run("pip install -e python_modules/dagster", "python scripts/check_schemas.py")
+        .run(
+            "pip install -e python_modules/dagster -e python_modules/dagster-ext",
+            "python scripts/check_schemas.py",
+        )
         .with_skip(skip_mysql_if_no_changes_to_dependencies(["dagster"]))
         .build()
     ]
@@ -151,8 +154,8 @@ def build_graphql_python_client_backcompat_steps() -> List[CommandStep]:
         CommandStepBuilder(":graphql: GraphQL Python Client backcompat")
         .on_test_image(AvailablePythonVersion.get_default())
         .run(
-            "pip install -e python_modules/dagster[test] -e python_modules/dagster-graphql -e"
-            " python_modules/automation",
+            "pip install -e python_modules/dagster[test] -e python_modules/dagster-graphql -e "
+            " python_modules/automation -e python_modules/dagster-ext",
             "dagster-graphql-client query check",
         )
         .with_skip(

python_modules/dagster-ext/dagster_ext/__init__.py

@@ -621,6 +621,70 @@ def make_channel(
 # ########################
 
 
+class IContext(ABC):
+    """Base class for asset context implemented by AssetExecutionContext and ExtContext."""
+
+    @property
+    @abstractmethod
+    def asset_key(self):
+        """The AssetKey for the asset being materialized. If no asset is being materialized, errors. If
+        multiple assets are being materialized (as in a @multi_asset), errors.
+        """
+
+    @property
+    @abstractmethod
+    def asset_keys(self):
+        """The AssetKeys for the asset being materialized. If no asset is being materialized, errors."""
+
+    @property
+    @abstractmethod
+    def provenance(self):
+        """The data provenance for the asset being materialized. If no asset is being materialized, errors. If
+        multiple assets are being materialized (as in a @multi_asset), errors.
+        """
+
+    @property
+    @abstractmethod
+    def provenance_by_asset_key(self):
+        """A dictionary of data provenance for the assets being materialized, keyed by asset key.
+        If no asset is being materialized, errors.
+        """
+
+    @property
+    @abstractmethod
+    def code_version(self):
+        """The code version for the asset being materialized. If no asset is being materialized, errors. If
+        multiple assets are being materialized (as in a @multi_asset), errors.
+        """
+
+    @property
+    @abstractmethod
+    def code_version_by_asset_key(self):
+        """A dictionary of code versions for the assets being materialized, keyed by asset key.
+        If no asset is being materialized, errors.
+        """
+
+    @property
+    @abstractmethod
+    def is_partitioned(self) -> bool:
+        """True if the current execution is partitioned."""
+
+    @property
+    @abstractmethod
+    def run_id(self) -> str:
+        """The run id of the current execution."""
+
+    @property
+    @abstractmethod
+    def job_name(self) -> Optional[str]:
+        """The name of the job that is executing."""
+
+    @property
+    @abstractmethod
+    def retry_number(self) -> int:
+        """The number of retries of this execution."""
+
+
 def init_dagster_ext(
     *,
     context_loader: Optional[ExtContextLoader] = None,
@@ -648,7 +712,7 @@ def init_dagster_ext(
     return context
 
 
-class ExtContext:
+class ExtContext(IContext):
     _instance: ClassVar[Optional["ExtContext"]] = None
 
     @classmethod
@@ -730,7 +794,7 @@ def code_version_by_asset_key(self) -> Mapping[str, Optional[str]]:
         return code_version_by_asset_key
 
     @property
-    def is_partition_step(self) -> bool:
+    def is_partitioned(self) -> bool:
         return self._data["partition_key_range"] is not None
 
     @property

python_modules/dagster-ext/dagster_ext_tests/test_context.py

@@ -120,7 +120,7 @@ def test_multi_asset_context():
 def test_no_partition_context():
     context = _make_external_execution_context()
 
-    assert not context.is_partition_step
+    assert not context.is_partitioned
     _assert_undefined(context, "partition_key")
     _assert_undefined(context, "partition_key_range")
     _assert_undefined(context, "partition_time_window")
@@ -135,7 +135,7 @@ def test_single_partition_context():
         partition_time_window=None,
     )
 
-    assert context.is_partition_step
+    assert context.is_partitioned
     assert context.partition_key == "foo"
     assert context.partition_key_range == partition_key_range
     assert context.partition_time_window is None
@@ -151,7 +151,7 @@ def test_multiple_partition_context():
         partition_time_window=time_window,
     )
 
-    assert context.is_partition_step
+    assert context.is_partitioned
     _assert_undefined(context, "partition_key")
     assert context.partition_key_range == partition_key_range
     assert context.partition_time_window == time_window