def _get_code_pointer_dict_from_kwargs(kwargs):
python_file = kwargs.get("python_file")
module_name = kwargs.get("module_name")
working_directory = get_working_directory_from_kwargs(kwargs)
attribute = kwargs.get("attribute")
if python_file:
_check_cli_arguments_none(kwargs, "module_name")
return {
repository_def_from_target_def(
loadable_target.target_definition).name:
CodePointer.from_python_file(python_file,
loadable_target.attribute,
working_directory)
for loadable_target in get_loadable_targets(
python_file, module_name, working_directory, attribute)
}
elif module_name:
_check_cli_arguments_none(kwargs, "python_file", "working_directory")
return {
repository_def_from_target_def(
loadable_target.target_definition).name:
CodePointer.from_module(module_name, loadable_target.attribute)
for loadable_target in get_loadable_targets(
python_file, module_name, working_directory, attribute)
}
else:
check.failed("Must specify a Python file or module name")
def ListRepositories(self, request, _context):
repository_code_pointer_dict = {}
for loadable_repository_symbol in self._loadable_repository_symbols:
if self._loadable_target_origin.python_file:
repository_code_pointer_dict[
loadable_repository_symbol.
repository_name] = CodePointer.from_python_file(
self._loadable_target_origin.python_file,
loadable_repository_symbol.attribute,
self._loadable_target_origin.working_directory,
)
if self._loadable_target_origin.module_name:
repository_code_pointer_dict[
loadable_repository_symbol.
repository_name] = CodePointer.from_module(
self._loadable_target_origin.module_name,
loadable_repository_symbol.attribute,
)
return api_pb2.ListRepositoriesReply(
serialized_list_repositories_response=serialize_dagster_namedtuple(
ListRepositoriesResponse(
self._loadable_repository_symbols,
executable_path=self._loadable_target_origin.
executable_path if self._loadable_target_origin else None,
repository_code_pointer_dict=repository_code_pointer_dict,
)))
def get_repository_origin_from_kwargs(kwargs):
load_target = created_workspace_load_target(kwargs)
if isinstance(load_target, PythonFileTarget):
return RepositoryPythonOrigin(
executable_path=sys.executable,
code_pointer=CodePointer.from_python_file(
load_target.python_file,
load_target.attribute,
working_directory=load_target.working_directory,
),
)
elif isinstance(load_target, ModuleTarget):
return RepositoryPythonOrigin(
executable_path=sys.executable,
code_pointer=CodePointer.from_module(load_target.module_name,
load_target.attribute),
)
elif isinstance(load_target, GrpcServerTarget):
return RepositoryGrpcServerOrigin(
host=load_target.host,
port=load_target.port,
socket=load_target.socket,
repository_name=kwargs['repository'],
)
else:
check.failed('invalid')
def build_code_pointers_by_repo_name(loadable_target_origin, loadable_repository_symbols):
repository_code_pointer_dict = {}
for loadable_repository_symbol in loadable_repository_symbols:
if loadable_target_origin.python_file:
repository_code_pointer_dict[
loadable_repository_symbol.repository_name
] = CodePointer.from_python_file(
loadable_target_origin.python_file,
loadable_repository_symbol.attribute,
loadable_target_origin.working_directory,
)
elif loadable_target_origin.package_name:
repository_code_pointer_dict[
loadable_repository_symbol.repository_name
] = CodePointer.from_python_package(
loadable_target_origin.package_name, loadable_repository_symbol.attribute,
)
else:
repository_code_pointer_dict[
loadable_repository_symbol.repository_name
] = CodePointer.from_module(
loadable_target_origin.module_name, loadable_repository_symbol.attribute,
)
return repository_code_pointer_dict
def location_handle_from_python_file(python_file,
attribute,
location_name=None,
working_directory=None):
check.str_param(python_file, 'python_file')
check.opt_str_param(attribute, 'attribute')
check.opt_str_param(location_name, 'location_name')
check.opt_str_param(working_directory, 'working_directory')
loadable_targets = ([
LoadableTarget(
attribute,
load_def_in_python_file(python_file, attribute, working_directory))
] if attribute else loadable_targets_from_python_file(
python_file, working_directory))
repository_code_pointer_dict = {}
for loadable_target in loadable_targets:
repository_code_pointer_dict[loadable_target.target_definition.
name] = CodePointer.from_python_file(
python_file,
loadable_target.attribute,
working_directory)
return RepositoryLocationHandle.create_out_of_process_location(
repository_code_pointer_dict=repository_code_pointer_dict,
# default to the name of the repository symbol for now
location_name=assign_location_name(location_name,
repository_code_pointer_dict),
)
def _get_code_pointer_dict_from_kwargs(kwargs):
python_file = kwargs.get("python_file")
module_name = kwargs.get("module_name")
working_directory = get_working_directory_from_kwargs(kwargs)
attribute = kwargs.get("attribute")
loadable_targets = get_loadable_targets(python_file, module_name,
working_directory, attribute)
if python_file:
return {
repository_def_from_target_def(
loadable_target.target_definition).name:
CodePointer.from_python_file(python_file,
loadable_target.attribute,
working_directory)
for loadable_target in loadable_targets
}
elif module_name:
return {
repository_def_from_target_def(
loadable_target.target_definition).name:
CodePointer.from_module(module_name, loadable_target.attribute)
for loadable_target in loadable_targets
}
else:
check.failed("invalid")
def get_repository_python_origin_from_kwargs(kwargs):
provided_repo_name = kwargs.get("repository")
if not (kwargs.get("python_file") or kwargs.get("module_name")
or kwargs.get("package_name")):
raise click.UsageError("Must specify a python file or module name")
# Short-circuit the case where an attribute and no repository name is passed in,
# giving us enough information to return an origin without loading any target
# definitions - we may need to return an origin for a non-existent repository
# (e.g. to log an origin ID for an error message)
if kwargs.get("attribute") and not provided_repo_name:
if kwargs.get("python_file"):
_check_cli_arguments_none(kwargs, "module_name", "package_name")
code_pointer = CodePointer.from_python_file(
kwargs.get("python_file"),
kwargs.get("attribute"),
get_working_directory_from_kwargs(kwargs),
)
elif kwargs.get("module_name"):
_check_cli_arguments_none(kwargs, "python_file",
"working_directory", "package_name")
code_pointer = CodePointer.from_module(
kwargs.get("module_name"),
kwargs.get("attribute"),
)
elif kwargs.get("package_name"):
_check_cli_arguments_none(kwargs, "python_file",
"working_directory", "module_name")
code_pointer = CodePointer.from_python_package(
kwargs.get("package_name"),
kwargs.get("attribute"),
)
else:
check.failed("Must specify a Python file or module name")
return RepositoryPythonOrigin(executable_path=sys.executable,
code_pointer=code_pointer)
code_pointer_dict = _get_code_pointer_dict_from_kwargs(kwargs)
if provided_repo_name is None and len(code_pointer_dict) == 1:
code_pointer = next(iter(code_pointer_dict.values()))
elif provided_repo_name is None:
raise click.UsageError(
("Must provide --repository as there is more than one repository. "
"Options are: {repos}.").format(
repos=_sorted_quoted(code_pointer_dict.keys())))
elif not provided_repo_name in code_pointer_dict:
raise click.UsageError(
'Repository "{provided_repo_name}" not found. Found {found_names} instead.'
.format(
provided_repo_name=provided_repo_name,
found_names=_sorted_quoted(code_pointer_dict.keys()),
))
else:
code_pointer = code_pointer_dict[provided_repo_name]
return RepositoryPythonOrigin(executable_path=sys.executable,
code_pointer=code_pointer)
def define_context_for_file(python_file, fn_name, instance):
check.inst_param(instance, "instance", DagsterInstance)
return DagsterGraphQLContext(
workspace=Workspace([
RepositoryLocationHandle.create_in_process_location(
CodePointer.from_python_file(python_file, fn_name, None))
]),
instance=instance,
)
def test_single_repository():
single_repo_path = file_relative_path(__file__, 'single_repository.py')
loadable_targets = loadable_targets_from_python_file(single_repo_path)
assert len(loadable_targets) == 1
symbol = loadable_targets[0].attribute
assert symbol == 'single_repository'
repo_def = CodePointer.from_python_file(single_repo_path, symbol).load_target()
isinstance(repo_def, RepositoryDefinition)
assert repo_def.name == 'single_repository'
def __init__(self, shutdown_server_event, loadable_target_origin=None):
super(DagsterApiServer, self).__init__()
self._shutdown_server_event = check.inst_param(
shutdown_server_event, 'shutdown_server_event',
seven.ThreadingEventType)
self._loadable_target_origin = check.opt_inst_param(
loadable_target_origin, 'loadable_target_origin',
LoadableTargetOrigin)
if loadable_target_origin:
loadable_targets = get_loadable_targets(
loadable_target_origin.python_file,
loadable_target_origin.module_name,
loadable_target_origin.working_directory,
loadable_target_origin.attribute,
)
self._loadable_repository_symbols = [
LoadableRepositorySymbol(
attribute=loadable_target.attribute,
repository_name=repository_def_from_target_def(
loadable_target.target_definition).name,
) for loadable_target in loadable_targets
]
else:
self._loadable_repository_symbols = []
self._shutdown_server_event = check.inst_param(
shutdown_server_event, 'shutdown_server_event',
seven.ThreadingEventType)
# Dict[str, multiprocessing.Process] of run_id to execute_run process
self._executions = {}
# Dict[str, multiprocessing.Event]
self._termination_events = {}
self._execution_lock = threading.Lock()
self._repository_code_pointer_dict = {}
for loadable_repository_symbol in self._loadable_repository_symbols:
if self._loadable_target_origin.python_file:
self._repository_code_pointer_dict[
loadable_repository_symbol.
repository_name] = CodePointer.from_python_file(
self._loadable_target_origin.python_file,
loadable_repository_symbol.attribute,
self._loadable_target_origin.working_directory,
)
if self._loadable_target_origin.module_name:
self._repository_code_pointer_dict[
loadable_repository_symbol.
repository_name] = CodePointer.from_module(
self._loadable_target_origin.module_name,
loadable_repository_symbol.attribute,
)
def create_python_env_location(
loadable_target_origin,
location_name=None,
user_process_api=UserProcessApi.GRPC,
use_python_package=False,
):
check.inst_param(loadable_target_origin, "loadable_target_origin", LoadableTargetOrigin)
check.opt_str_param(location_name, "location_name")
check.bool_param(use_python_package, "use_python_package")
if user_process_api == UserProcessApi.GRPC:
return RepositoryLocationHandle.create_process_bound_grpc_server_location(
loadable_target_origin=loadable_target_origin, location_name=location_name
)
response = sync_list_repositories(
executable_path=loadable_target_origin.executable_path,
python_file=loadable_target_origin.python_file,
module_name=loadable_target_origin.module_name,
working_directory=loadable_target_origin.working_directory,
attribute=loadable_target_origin.attribute,
)
if loadable_target_origin.python_file:
repository_code_pointer_dict = {
lrs.repository_name: CodePointer.from_python_file(
loadable_target_origin.python_file,
lrs.attribute,
loadable_target_origin.working_directory,
)
for lrs in response.repository_symbols
}
elif use_python_package:
repository_code_pointer_dict = {
lrs.repository_name: CodePointer.from_python_package(
loadable_target_origin.module_name, lrs.attribute
)
for lrs in response.repository_symbols
}
else:
repository_code_pointer_dict = {
lrs.repository_name: CodePointer.from_module(
loadable_target_origin.module_name, lrs.attribute
)
for lrs in response.repository_symbols
}
return PythonEnvRepositoryLocationHandle(
location_name=location_name
if location_name
else _assign_python_env_location_name(repository_code_pointer_dict),
loadable_target_origin=loadable_target_origin,
repository_code_pointer_dict=repository_code_pointer_dict,
)
def test_single_graph():
single_graph_path = file_relative_path(__file__, "single_graph.py")
loadable_targets = loadable_targets_from_python_file(single_graph_path)
assert len(loadable_targets) == 1
symbol = loadable_targets[0].attribute
assert symbol == "graph_one"
repo_def = repository_def_from_pointer(
CodePointer.from_python_file(single_graph_path, symbol, None))
isinstance(repo_def, RepositoryDefinition)
assert repo_def.get_pipeline("graph_one")
def test_single_pipeline():
single_pipeline_path = file_relative_path(__file__, 'single_pipeline.py')
loadable_targets = loadable_targets_from_python_file(single_pipeline_path)
assert len(loadable_targets) == 1
symbol = loadable_targets[0].attribute
assert symbol == 'a_pipeline'
repo_def = repository_def_from_pointer(
CodePointer.from_python_file(single_pipeline_path, symbol)
)
isinstance(repo_def, RepositoryDefinition)
assert repo_def.get_pipeline('a_pipeline')
def test_single_asset_group():
path = file_relative_path(__file__, "single_asset_group.py")
loadable_targets = loadable_targets_from_python_file(path)
assert len(loadable_targets) == 1
symbol = loadable_targets[0].attribute
assert symbol == "my_asset_group"
repo_def = repository_def_from_pointer(
CodePointer.from_python_file(path, symbol, None))
isinstance(repo_def, RepositoryDefinition)
the_job = repo_def.get_job("__ASSET_GROUP")
assert len(the_job.graph.node_defs) == 2
def location_handle_from_python_file(python_file,
attribute,
user_process_api,
location_name=None,
working_directory=None):
check.str_param(python_file, 'python_file')
check.opt_str_param(attribute, 'attribute')
check.inst_param(user_process_api, 'user_process_api', UserProcessApi)
check.opt_str_param(location_name, 'location_name')
check.opt_str_param(working_directory, 'working_directory')
if user_process_api == UserProcessApi.GRPC:
return RepositoryLocationHandle.create_process_bound_grpc_server_location(
loadable_target_origin=LoadableTargetOrigin(
executable_path=sys.executable,
python_file=python_file,
module_name=None,
working_directory=working_directory,
attribute=attribute,
),
location_name=location_name,
)
loadable_targets = ([
LoadableTarget(
attribute,
load_def_in_python_file(python_file, attribute, working_directory))
] if attribute else loadable_targets_from_python_file(
python_file, working_directory))
repository_code_pointer_dict = {}
for loadable_target in loadable_targets:
repository_code_pointer_dict[loadable_target.target_definition.
name] = CodePointer.from_python_file(
python_file,
loadable_target.attribute,
working_directory)
return RepositoryLocationHandle.create_out_of_process_location(
repository_code_pointer_dict=repository_code_pointer_dict,
# default to the name of the repository symbol for now
location_name=assign_location_name(location_name,
repository_code_pointer_dict),
)
def location_handle_from_python_file(
python_file,
attribute,
user_process_api,
location_name=None,
working_directory=None,
executable_path=sys.executable,
):
check.str_param(python_file, 'python_file')
check.opt_str_param(attribute, 'attribute')
check.inst_param(user_process_api, 'user_process_api', UserProcessApi)
check.opt_str_param(location_name, 'location_name')
check.opt_str_param(working_directory, 'working_directory')
if user_process_api == UserProcessApi.GRPC:
return RepositoryLocationHandle.create_process_bound_grpc_server_location(
loadable_target_origin=LoadableTargetOrigin(
executable_path=executable_path,
python_file=python_file,
module_name=None,
working_directory=working_directory,
attribute=attribute,
),
location_name=location_name,
)
else:
response = sync_list_repositories(
executable_path=executable_path,
python_file=python_file,
module_name=None,
working_directory=working_directory,
attribute=attribute,
)
return RepositoryLocationHandle.create_python_env_location(
executable_path=executable_path,
location_name=location_name,
repository_code_pointer_dict={
lrs.repository_name:
CodePointer.from_python_file(python_file, lrs.attribute,
working_directory)
for lrs in response.repository_symbols
},
)
def _get_code_pointer(loadable_target_origin, loadable_repository_symbol):
if loadable_target_origin.python_file:
return CodePointer.from_python_file(
loadable_target_origin.python_file,
loadable_repository_symbol.attribute,
loadable_target_origin.working_directory,
)
elif loadable_target_origin.package_name:
return CodePointer.from_python_package(
loadable_target_origin.package_name,
loadable_repository_symbol.attribute,
loadable_target_origin.working_directory,
)
else:
return CodePointer.from_module(
loadable_target_origin.module_name,
loadable_repository_symbol.attribute,
loadable_target_origin.working_directory,
)
def _location_handle_from_python_environment_config(python_environment_config, yaml_path):
check.dict_param(python_environment_config, 'python_environment_config')
check.str_param(yaml_path, 'yaml_path')
executable_path, target_config = (
# do shell expansion on path
os.path.expanduser(python_environment_config['executable_path']),
python_environment_config['target'],
)
check.invariant(is_target_config(target_config))
python_file_config, python_module_config = (
target_config.get('python_file'),
target_config.get('python_module'),
)
if python_file_config:
absolute_path, attribute, location_name = _get_python_file_config_data(
python_file_config, yaml_path
)
if not attribute:
response = sync_list_repositories(
executable_path=executable_path, python_file=absolute_path, module_name=None,
)
return RepositoryLocationHandle.create_python_env_location(
executable_path=executable_path,
location_name=location_name,
repository_code_pointer_dict={
lrs.attribute: CodePointer.from_python_file(absolute_path, lrs.attribute)
for lrs in response.repository_symbols
},
)
else:
return RepositoryLocationHandle.create_python_env_location(
executable_path=executable_path,
location_name=location_name,
repository_code_pointer_dict={
attribute: CodePointer.from_python_file(absolute_path, attribute)
},
)
else:
check.invariant(python_module_config)
module_name, attribute, location_name = _get_module_config_data(python_module_config)
if not attribute:
response = sync_list_repositories(
executable_path=executable_path, python_file=None, module_name=module_name,
)
return RepositoryLocationHandle.create_python_env_location(
executable_path=executable_path,
location_name=location_name,
repository_code_pointer_dict={
lrs.attribute: CodePointer.from_module(module_name, lrs.attribute)
for lrs in response.repository_symbols
},
)
else:
return RepositoryLocationHandle.create_python_env_location(
executable_path=executable_path,
location_name=location_name,
repository_code_pointer_dict={
attribute: CodePointer.from_module(module_name, attribute)
},
)
def _location_handle_from_python_environment_config(python_environment_config,
yaml_path,
user_process_api):
check.dict_param(python_environment_config, 'python_environment_config')
check.str_param(yaml_path, 'yaml_path')
check.inst_param(user_process_api, 'user_process_api', UserProcessApi)
executable_path, target_config = (
# do shell expansion on path
os.path.expanduser(python_environment_config['executable_path']),
python_environment_config['target'],
)
check.invariant(is_target_config(target_config))
python_file_config, python_module_config, python_package_config = (
target_config.get('python_file'),
target_config.get('python_module'),
target_config.get('python_package'),
)
if python_file_config:
absolute_path, attribute, location_name, working_directory = _get_python_file_config_data(
python_file_config, yaml_path)
if user_process_api == UserProcessApi.GRPC:
return RepositoryLocationHandle.create_process_bound_grpc_server_location(
loadable_target_origin=LoadableTargetOrigin(
executable_path=executable_path,
python_file=absolute_path,
module_name=None,
working_directory=None,
attribute=attribute,
),
location_name=location_name,
)
elif not attribute:
response = sync_list_repositories(
executable_path=executable_path,
python_file=absolute_path,
module_name=None,
working_directory=None,
)
return RepositoryLocationHandle.create_python_env_location(
executable_path=executable_path,
location_name=location_name,
repository_code_pointer_dict={
lrs.repository_name:
CodePointer.from_python_file(absolute_path, lrs.attribute,
working_directory)
for lrs in response.repository_symbols
},
)
else:
return RepositoryLocationHandle.create_python_env_location(
executable_path=executable_path,
location_name=location_name,
repository_code_pointer_dict={
attribute:
CodePointer.from_python_file(absolute_path, attribute,
working_directory)
},
)
elif python_module_config:
check.invariant(python_module_config)
module_name, attribute, location_name = _get_module_config_data(
python_module_config)
if user_process_api == UserProcessApi.GRPC:
return RepositoryLocationHandle.create_process_bound_grpc_server_location(
loadable_target_origin=LoadableTargetOrigin(
executable_path=executable_path,
python_file=None,
module_name=module_name,
working_directory=None,
attribute=attribute,
),
location_name=location_name,
)
elif not attribute:
response = sync_list_repositories(
executable_path=executable_path,
python_file=None,
module_name=module_name,
working_directory=None,
)
return RepositoryLocationHandle.create_python_env_location(
executable_path=executable_path,
location_name=location_name,
repository_code_pointer_dict={
lrs.repository_name:
CodePointer.from_module(module_name, lrs.attribute)
for lrs in response.repository_symbols
},
)
else:
return RepositoryLocationHandle.create_python_env_location(
executable_path=executable_path,
location_name=location_name,
repository_code_pointer_dict={
attribute: CodePointer.from_module(module_name, attribute)
},
)
else:
check.invariant(python_package_config)
package_name, attribute, location_name = _get_package_config_data(
python_package_config)
if user_process_api == UserProcessApi.GRPC:
return RepositoryLocationHandle.create_process_bound_grpc_server_location(
loadable_target_origin=LoadableTargetOrigin(
executable_path=executable_path,
python_file=None,
module_name=package_name,
working_directory=None,
attribute=attribute,
),
location_name=location_name,
)
elif not attribute:
response = sync_list_repositories(
executable_path=executable_path,
python_file=None,
module_name=package_name,
working_directory=None,
)
return RepositoryLocationHandle.create_python_env_location(
executable_path=executable_path,
location_name=location_name,
repository_code_pointer_dict={
lrs.repository_name:
CodePointer.from_python_package(package_name,
lrs.attribute)
for lrs in response.repository_symbols
},
)
else:
return RepositoryLocationHandle.create_python_env_location(
executable_path=executable_path,
location_name=location_name,
repository_code_pointer_dict={
attribute:
CodePointer.from_python_package(package_name, attribute)
},
)
def load_def_in_python_file(python_file, attribute):
return def_from_pointer(CodePointer.from_python_file(python_file, attribute))
def load_def_in_python_file(python_file, attribute, working_directory):
return def_from_pointer(CodePointer.from_python_file(python_file, attribute, working_directory))
def test_out_of_process_reload_location(self, graphql_context):
result = execute_dagster_graphql(graphql_context,
RELOAD_REPOSITORY_LOCATION_QUERY,
{"repositoryLocationName": "test"})
assert result
assert result.data
assert result.data["reloadRepositoryLocation"]
assert result.data["reloadRepositoryLocation"][
"__typename"] == "RepositoryLocation"
assert result.data["reloadRepositoryLocation"]["name"] == "test"
assert result.data["reloadRepositoryLocation"]["repositories"] == [{
"name":
"test_repo"
}]
assert result.data["reloadRepositoryLocation"][
"isReloadSupported"] is True
with mock.patch(
# note it where the function is *used* that needs to mocked, not
# where it is defined.
# see https://docs.python.org/3/library/unittest.mock.html#where-to-patch
"dagster.core.host_representation.handle.sync_list_repositories_grpc"
) as cli_command_mock:
with mock.patch(
# note it where the function is *used* that needs to mocked, not
# where it is defined.
# see https://docs.python.org/3/library/unittest.mock.html#where-to-patch
"dagster.core.host_representation.repository_location.sync_get_streaming_external_repositories_grpc"
) as external_repository_mock:
@repository
def new_repo():
return []
new_repo_data = external_repository_data_from_def(new_repo)
external_repository_mock.return_value = [
ExternalRepository(
new_repo_data,
RepositoryHandle(
"new_repo", graphql_context.
repository_locations[0].location_handle),
)
]
cli_command_mock.return_value = ListRepositoriesResponse(
repository_symbols=[],
executable_path=sys.executable,
repository_code_pointer_dict={
"new_repo":
CodePointer.from_python_file(__file__, "new_repo",
None)
},
)
result = execute_dagster_graphql(
graphql_context,
RELOAD_REPOSITORY_LOCATION_QUERY,
{"repositoryLocationName": "test"},
)
assert cli_command_mock.call_count == 1
assert external_repository_mock.call_count == 1
assert result.data["reloadRepositoryLocation"][
"repositories"] == [{
"name": "new_repo"
}]
def test_all_step_events(): # pylint: disable=too-many-locals
instance = DagsterInstance.ephemeral()
pipeline_def = define_test_events_pipeline()
workspace = create_in_process_ephemeral_workspace(
pointer=CodePointer.from_python_file(
__file__,
define_test_events_pipeline.__name__,
working_directory=None))
mode = pipeline_def.get_default_mode_name()
execution_plan = create_execution_plan(pipeline_def, mode=mode)
pipeline_run = instance.create_run_for_pipeline(
pipeline_def=pipeline_def, execution_plan=execution_plan, mode=mode)
step_levels = execution_plan.topological_step_levels()
unhandled_events = STEP_EVENTS.copy()
# Exclude types that are not step events
ignored_events = {
"LogMessageEvent",
"PipelineStartEvent",
"PipelineSuccessEvent",
"PipelineInitFailureEvent",
"PipelineFailureEvent",
}
event_counts = defaultdict(int)
for step_level in step_levels:
for step in step_level:
variables = {
"executionParams": {
"selector": {
"repositoryLocationName":
IN_PROCESS_NAME,
"repositoryName":
get_ephemeral_repository_name(pipeline_def.name),
"pipelineName":
pipeline_def.name,
},
"runConfigData": {
"storage": {
"filesystem": {}
}
},
"mode": mode,
"executionMetadata": {
"runId": pipeline_run.run_id
},
"stepKeys": [step.key],
},
}
res = execute_query(
workspace,
EXECUTE_PLAN_MUTATION,
variables,
instance=instance,
)
# go through the same dict, decrement all the event records we've seen from the GraphQL
# response
if not res.get("errors"):
assert "data" in res, res
assert "executePlan" in res["data"], res
assert "stepEvents" in res["data"]["executePlan"], res
step_events = res["data"]["executePlan"]["stepEvents"]
events = [
dagster_event_from_dict(e, pipeline_def.name)
for e in step_events
if e["__typename"] not in ignored_events
]
for event in events:
if event.step_key:
key = event.step_key + "." + event.event_type_value
else:
key = event.event_type_value
event_counts[key] -= 1
unhandled_events -= {
DagsterEventType(e.event_type_value)
for e in events
}
else:
raise Exception(res["errors"])
# build up a dict, incrementing all the event records we've produced in the run storage
logs = instance.all_logs(pipeline_run.run_id)
for log in logs:
if not log.dagster_event or (DagsterEventType(
log.dagster_event.event_type_value) not in STEP_EVENTS.union(
set([DagsterEventType.ENGINE_EVENT]))):
continue
if log.dagster_event.step_key:
key = log.dagster_event.step_key + "." + log.dagster_event.event_type_value
else:
key = log.dagster_event.event_type_value
event_counts[key] += 1
# Ensure we've processed all the events that were generated in the run storage
assert sum(event_counts.values()) == 0
# Ensure we've handled the universe of event types
# Why are these retry events not handled? Because right now there is no way to configure retries
# on executePlan -- this needs to change, and we should separate the ExecutionParams that get
# sent to executePlan fromm those that get sent to startPipelineExecution and friends
assert unhandled_events == {
DagsterEventType.STEP_UP_FOR_RETRY, DagsterEventType.STEP_RESTARTED
}
