def create_intermediate_storage(
context_creation_data,
intermediate_storage_data,
scoped_resources_builder,
):
check.inst_param(context_creation_data, "context_creation_data",
ContextCreationData)
environment_config, pipeline_def, intermediate_storage_def, pipeline_run = (
context_creation_data.environment_config,
context_creation_data.pipeline_def,
context_creation_data.intermediate_storage_def,
context_creation_data.pipeline_run,
)
intermediate_storage_data = (
intermediate_storage_data
if intermediate_storage_data else construct_intermediate_storage_data(
InitIntermediateStorageContext(
pipeline_def=pipeline_def,
mode_def=context_creation_data.mode_def,
intermediate_storage_def=intermediate_storage_def,
intermediate_storage_config=environment_config.
intermediate_storage.intermediate_storage_config,
pipeline_run=pipeline_run,
instance=context_creation_data.instance,
environment_config=environment_config,
type_storage_plugin_registry=
construct_type_storage_plugin_registry(
pipeline_def, intermediate_storage_def),
resources=scoped_resources_builder.build(
context_creation_data.intermediate_storage_def.
required_resource_keys, ),
)))
return intermediate_storage_data
def create_system_storage_data(context_creation_data, system_storage_data,
scoped_resources_builder):
check.inst_param(context_creation_data, "context_creation_data",
ContextCreationData)
environment_config, pipeline_def, system_storage_def, pipeline_run = (
context_creation_data.environment_config,
context_creation_data.pipeline_def,
context_creation_data.system_storage_def,
context_creation_data.pipeline_run,
)
system_storage_data = (system_storage_data if system_storage_data else
construct_system_storage_data(
InitSystemStorageContext(
pipeline_def=pipeline_def,
mode_def=context_creation_data.mode_def,
system_storage_def=system_storage_def,
system_storage_config=environment_config.
storage.system_storage_config,
pipeline_run=pipeline_run,
instance=context_creation_data.instance,
environment_config=environment_config,
type_storage_plugin_registry=
construct_type_storage_plugin_registry(
pipeline_def, system_storage_def),
resources=scoped_resources_builder.build(
context_creation_data.system_storage_def
.required_resource_keys, ),
)))
return system_storage_data
def create_intermediate_storage(
context_creation_data: ContextCreationData,
intermediate_storage_data: Optional[IntermediateStorage],
scoped_resources_builder: ScopedResourcesBuilder,
) -> IntermediateStorage:
resolved_run_config, pipeline_def, intermediate_storage_def, pipeline_run = (
context_creation_data.resolved_run_config,
context_creation_data.pipeline_def,
context_creation_data.intermediate_storage_def,
context_creation_data.pipeline_run,
)
intermediate_storage_data = (
intermediate_storage_data
if intermediate_storage_data else construct_intermediate_storage_data(
InitIntermediateStorageContext(
pipeline_def=pipeline_def,
mode_def=context_creation_data.mode_def,
intermediate_storage_def=intermediate_storage_def,
intermediate_storage_config=resolved_run_config.
intermediate_storage.intermediate_storage_config,
pipeline_run=pipeline_run,
instance=context_creation_data.instance,
resolved_run_config=resolved_run_config,
type_storage_plugin_registry=
construct_type_storage_plugin_registry(
pipeline_def, intermediate_storage_def),
resources=scoped_resources_builder.build(
context_creation_data.intermediate_storage_def.
required_resource_keys, ),
)))
return intermediate_storage_data
def _object_manager(init_context):
pipeline_run = init_context.pipeline_run
instance = init_context.instance_for_backwards_compat
pipeline_def = init_context.pipeline_def_for_backwards_compat
# depend on InitResourceContext.instance_for_backwards_compat and pipeline_def_for_backwards_compat
environment_config = EnvironmentConfig.build(pipeline_def,
pipeline_run.run_config,
mode=pipeline_run.mode)
mode_def = pipeline_def.get_mode_definition(pipeline_run.mode)
intermediate_storage_context = InitIntermediateStorageContext(
pipeline_def=pipeline_def,
mode_def=mode_def,
intermediate_storage_def=intermediate_storage_def,
pipeline_run=pipeline_run,
instance=instance,
environment_config=environment_config,
type_storage_plugin_registry=construct_type_storage_plugin_registry(
pipeline_def, intermediate_storage_def),
resources=init_context.resources,
intermediate_storage_config=environment_config.
intermediate_storage.intermediate_storage_config,
)
intermediate_storage = intermediate_storage_def.intermediate_storage_creation_fn(
intermediate_storage_context)
return IntermediateStorageAdapter(intermediate_storage)
def _io_manager(init_context):
if not init_context.pipeline_run:
raise DagsterInvariantViolationError(
"Attempted to construct intermediate storage outside of execution context. "
"Intermediate storage can only be constructed within the context of an execution."
)
pipeline_run = init_context.pipeline_run
instance = init_context.instance
pipeline_def = init_context.pipeline_def_for_backwards_compat
# depend on InitResourceContext.instance and pipeline_def_for_backwards_compat
environment_config = EnvironmentConfig.build(pipeline_def,
pipeline_run.run_config,
mode=pipeline_run.mode)
mode_def = pipeline_def.get_mode_definition(pipeline_run.mode)
intermediate_storage_context = InitIntermediateStorageContext(
pipeline_def=pipeline_def,
mode_def=mode_def,
intermediate_storage_def=intermediate_storage_def,
pipeline_run=pipeline_run,
instance=instance,
environment_config=environment_config,
type_storage_plugin_registry=construct_type_storage_plugin_registry(
pipeline_def, intermediate_storage_def),
resources=init_context.resources,
intermediate_storage_config=environment_config.
intermediate_storage.intermediate_storage_config,
)
intermediate_storage = intermediate_storage_def.intermediate_storage_creation_fn(
intermediate_storage_context)
return IntermediateStorageAdapter(intermediate_storage)
def create_system_storage_data(context_creation_data, system_storage_data,
scoped_resources_builder):
check.inst_param(context_creation_data, 'context_creation_data',
ContextCreationData)
environment_config, pipeline_def, system_storage_def, run_config = (
context_creation_data.environment_config,
context_creation_data.pipeline_def,
context_creation_data.system_storage_def,
context_creation_data.run_config,
)
system_storage_data = (
system_storage_data
if system_storage_data else construct_system_storage_data(
InitSystemStorageContext(
pipeline_def=pipeline_def,
mode_def=context_creation_data.mode_def,
system_storage_def=system_storage_def,
system_storage_config=environment_config.storage.
system_storage_config,
run_config=run_config,
environment_config=environment_config,
type_storage_plugin_registry
=construct_type_storage_plugin_registry(
pipeline_def, system_storage_def),
resources=scoped_resources_builder.build(
# currently provide default for resource mapping
lambda resources, resources_deps:
{r: resources.get(r)
for r in resources_deps},
context_creation_data.system_storage_def.
required_resource_keys,
),
)))
system_storage_data.run_storage.create_run(
pipeline_name=pipeline_def.name,
run_id=run_config.run_id,
env_config=environment_config.original_config_dict,
mode=context_creation_data.mode_def.name,
)
return system_storage_data
