from typing import Any, Mapping, Optional, Union
import dagster._check as check
from dagster._annotations import public
from dagster._core.definitions.resource_definition import (
IContainsGenerator,
ResourceDefinition,
Resources,
)
from dagster._core.errors import DagsterInvariantViolationError
from dagster._core.instance import DagsterInstance
from dagster._core.log_manager import DagsterLogManager
from dagster._core.storage.dagster_run import DagsterRun
[docs]class InitResourceContext:
"""The context object available as the argument to the initialization function of a :py:class:`dagster.ResourceDefinition`.
Users should not instantiate this object directly. To construct an `InitResourceContext` for testing purposes, use :py:func:`dagster.build_init_resource_context`.
Example:
.. code-block:: python
from dagster import resource, InitResourceContext
@resource
def the_resource(init_context: InitResourceContext):
init_context.log.info("Hello, world!")
"""
def __init__(
self,
resource_config: Any,
resources: Resources,
resource_def: Optional[ResourceDefinition] = None,
instance: Optional[DagsterInstance] = None,
dagster_run: Optional[DagsterRun] = None,
log_manager: Optional[DagsterLogManager] = None,
):
self._resource_config = resource_config
self._resource_def = resource_def
self._log_manager = log_manager
self._instance = instance
self._resources = resources
self._dagster_run = dagster_run
@public
@property
def resource_config(self) -> Any:
"""The configuration data provided by the run config. The schema
for this data is defined by the ``config_field`` argument to
:py:class:`ResourceDefinition`.
"""
return self._resource_config
@public
@property
def resource_def(self) -> Optional[ResourceDefinition]:
"""The definition of the resource currently being constructed."""
return self._resource_def
@public
@property
def resources(self) -> Resources:
"""The resources that are available to the resource that we are initalizing."""
return self._resources
@public
@property
def instance(self) -> Optional[DagsterInstance]:
"""The Dagster instance configured for the current execution context."""
return self._instance
@property
def dagster_run(self) -> Optional[DagsterRun]:
"""The dagster run to use. When initializing resources outside of execution context, this will be None."""
return self._dagster_run
@public
@property
def log(self) -> Optional[DagsterLogManager]:
"""The Dagster log manager configured for the current execution context."""
return self._log_manager
# backcompat: keep around this property from when InitResourceContext used to be a NamedTuple
@public
@property
def log_manager(self) -> Optional[DagsterLogManager]:
"""The log manager for this run of the job."""
return self._log_manager
@public
@property
def run_id(self) -> Optional[str]:
"""The id for this run of the job or pipeline. When initializing resources outside of
execution context, this will be None.
"""
return self.dagster_run.run_id if self.dagster_run else None
def replace_config(self, config: Any) -> "InitResourceContext":
return InitResourceContext(
resource_config=config,
resources=self.resources,
instance=self.instance,
resource_def=self.resource_def,
dagster_run=self.dagster_run,
log_manager=self.log,
)
class UnboundInitResourceContext(InitResourceContext):
"""Resource initialization context outputted by ``build_init_resource_context``.
Represents a context whose config has not yet been validated against a resource definition,
hence the inability to access the `resource_def` attribute. When an instance of
``UnboundInitResourceContext`` is passed to a resource invocation, config is validated,
and it is subsumed into an `InitResourceContext`, which contains the resource_def validated
against.
"""
def __init__(
self,
resource_config: Any,
resources: Optional[Union[Resources, Mapping[str, Any]]],
instance: Optional[DagsterInstance],
):
from dagster._core.execution.api import ephemeral_instance_if_missing
from dagster._core.execution.build_resources import (
build_resources,
wrap_resources_for_execution,
)
from dagster._core.execution.context_creation_job import initialize_console_manager
self._instance_provided = (
check.opt_inst_param(instance, "instance", DagsterInstance) is not None
)
# Construct ephemeral instance if missing
self._instance_cm = ephemeral_instance_if_missing(instance)
# Pylint can't infer that the ephemeral_instance context manager has an __enter__ method,
# so ignore lint error
instance = self._instance_cm.__enter__()
if isinstance(resources, Resources):
check.failed("Should not have a Resources object directly from this initialization")
self._resource_defs = wrap_resources_for_execution(
check.opt_mapping_param(resources, "resources")
)
self._resources_cm = build_resources(self._resource_defs, instance=instance)
resources = self._resources_cm.__enter__()
self._resources_contain_cm = isinstance(resources, IContainsGenerator)
self._cm_scope_entered = False
super(UnboundInitResourceContext, self).__init__(
resource_config=resource_config,
resources=resources,
resource_def=None,
instance=instance,
dagster_run=None,
log_manager=initialize_console_manager(None),
)
def __enter__(self):
self._cm_scope_entered = True
return self
def __exit__(self, *exc):
self._resources_cm.__exit__(*exc)
if self._instance_provided:
self._instance_cm.__exit__(*exc)
def __del__(self):
if self._resources_cm and self._resources_contain_cm and not self._cm_scope_entered:
self._resources_cm.__exit__(None, None, None)
if self._instance_provided and not self._cm_scope_entered:
self._instance_cm.__exit__(None, None, None)
@property
def resource_config(self) -> Any:
return self._resource_config
@property
def resource_def(self) -> Optional[ResourceDefinition]:
raise DagsterInvariantViolationError(
"UnboundInitLoggerContext has not been validated against a logger definition."
)
@property
def resources(self) -> Resources:
"""The resources that are available to the resource that we are initalizing."""
if self._resources_cm and self._resources_contain_cm and not self._cm_scope_entered:
raise DagsterInvariantViolationError(
"At least one provided resource is a generator, but attempting to access "
"resources outside of context manager scope. You can use the following syntax to "
"open a context manager: `with build_init_resource_context(...) as context:`"
)
return self._resources
@property
def instance(self) -> Optional[DagsterInstance]:
return self._instance
@property
def log(self) -> Optional[DagsterLogManager]:
return self._log_manager
# backcompat: keep around this property from when InitResourceContext used to be a NamedTuple
@property
def log_manager(self) -> Optional[DagsterLogManager]:
return self._log_manager
@property
def run_id(self) -> Optional[str]:
return None
[docs]def build_init_resource_context(
config: Optional[Mapping[str, Any]] = None,
resources: Optional[Mapping[str, Any]] = None,
instance: Optional[DagsterInstance] = None,
) -> InitResourceContext:
"""Builds resource initialization context from provided parameters.
``build_init_resource_context`` can be used as either a function or context manager. If there is a
provided resource to ``build_init_resource_context`` that is a context manager, then it must be
used as a context manager. This function can be used to provide the context argument to the
invocation of a resource.
Args:
resources (Optional[Dict[str, Any]]): The resources to provide to the context. These can be
either values or resource definitions.
config (Optional[Any]): The resource config to provide to the context.
instance (Optional[DagsterInstance]): The dagster instance configured for the context.
Defaults to DagsterInstance.ephemeral().
Examples:
.. code-block:: python
context = build_init_resource_context()
resource_to_init(context)
with build_init_resource_context(
resources={"foo": context_manager_resource}
) as context:
resource_to_init(context)
"""
return UnboundInitResourceContext(
resource_config=check.opt_mapping_param(config, "config", key_type=str),
instance=check.opt_inst_param(instance, "instance", DagsterInstance),
resources=check.opt_mapping_param(resources, "resources", key_type=str),
)