diff --git a/python_modules/dagster-graphql/dagster_graphql_tests/graphql/repo.py b/python_modules/dagster-graphql/dagster_graphql_tests/graphql/repo.py index 34355ec88f161..60c96fa563777 100644 --- a/python_modules/dagster-graphql/dagster_graphql_tests/graphql/repo.py +++ b/python_modules/dagster-graphql/dagster_graphql_tests/graphql/repo.py @@ -41,6 +41,7 @@ Map, Noneable, Nothing, + OpExecutionContext, Out, Output, PythonObjectDagsterType, @@ -733,7 +734,7 @@ def spawn() -> int: @op( required_resource_keys={"retry_count"}, ) - def fail(context: AssetExecutionContext, depth: int) -> int: + def fail(context: OpExecutionContext, depth: int) -> int: if context.resources.retry_count <= depth: raise Exception("fail") diff --git a/python_modules/dagster/dagster_tests/execution_tests/test_execute_job.py b/python_modules/dagster/dagster_tests/execution_tests/test_execute_job.py index 878296ba5ff71..a755d62473a25 100644 --- a/python_modules/dagster/dagster_tests/execution_tests/test_execute_job.py +++ b/python_modules/dagster/dagster_tests/execution_tests/test_execute_job.py @@ -1,12 +1,12 @@ import dagster._check as check import pytest from dagster import ( - AssetExecutionContext, AssetKey, DagsterExecutionStepNotFoundError, DagsterInvalidConfigError, DagsterInvariantViolationError, Field, + OpExecutionContext, Out, Output, ReexecutionOptions, @@ -377,7 +377,7 @@ def echo(x): @op -def fail_once(context: AssetExecutionContext, x): +def fail_once(context: OpExecutionContext, x): key = context.op_handle.name if context.instance.run_storage.get_cursor_values({key}).get(key): return x diff --git a/python_modules/libraries/dagster-dbt/dagster_dbt_tests/core/test_resources_v2.py b/python_modules/libraries/dagster-dbt/dagster_dbt_tests/core/test_resources_v2.py index dc836a135ed45..53db3f264230e 100644 --- a/python_modules/libraries/dagster-dbt/dagster_dbt_tests/core/test_resources_v2.py +++ b/python_modules/libraries/dagster-dbt/dagster_dbt_tests/core/test_resources_v2.py @@ -16,7 +16,7 @@ materialize, op, ) -from dagster._core.execution.context.compute import OpExecutionContext +from dagster._core.execution.context.compute import AssetExecutionContext from dagster_dbt import dbt_assets from dagster_dbt.asset_utils import build_dbt_asset_selection from dagster_dbt.core.resources_v2 import ( @@ -209,7 +209,7 @@ def test_dbt_with_partial_parse() -> None: def test_dbt_cli_debug_execution() -> None: @dbt_assets(manifest=manifest) - def my_dbt_assets(context: OpExecutionContext, dbt: DbtCliResource): + def my_dbt_assets(context: AssetExecutionContext, dbt: DbtCliResource): yield from dbt.cli(["--debug", "run"], context=context).stream() result = materialize( @@ -230,7 +230,7 @@ def test_dbt_cli_subsetted_execution() -> None: ) @dbt_assets(manifest=manifest, select=dbt_select) - def my_dbt_assets(context: OpExecutionContext, dbt: DbtCliResource): + def my_dbt_assets(context: AssetExecutionContext, dbt: DbtCliResource): dbt_cli_invocation = dbt.cli(["run"], context=context) assert dbt_cli_invocation.process.args == ["dbt", "run", "--select", dbt_select] @@ -253,7 +253,7 @@ def test_dbt_cli_asset_selection() -> None: ] @dbt_assets(manifest=manifest) - def my_dbt_assets(context: OpExecutionContext, dbt: DbtCliResource): + def my_dbt_assets(context: AssetExecutionContext, dbt: DbtCliResource): dbt_cli_invocation = dbt.cli(["run"], context=context) dbt_cli_args: List[str] = list(dbt_cli_invocation.process.args) # type: ignore @@ -283,7 +283,7 @@ def my_dbt_assets(context: OpExecutionContext, dbt: DbtCliResource): @pytest.mark.parametrize("exclude", [None, "fqn:dagster_dbt_test_project.subdir.least_caloric"]) def test_dbt_cli_default_selection(exclude: Optional[str]) -> None: @dbt_assets(manifest=manifest, exclude=exclude) - def my_dbt_assets(context: OpExecutionContext, dbt: DbtCliResource): + def my_dbt_assets(context: AssetExecutionContext, dbt: DbtCliResource): dbt_cli_invocation = dbt.cli(["run"], context=context) expected_args = ["dbt", "run", "--select", "fqn:*"]