from abc import abstractmethod
from typing import Any, Generic, Mapping, Optional, Type, Union, cast
from typing_extensions import TypeVar
from dagster._config.pythonic_config.attach_other_object_to_context import (
IAttachDifferentObjectToOpContext as IAttachDifferentObjectToOpContext,
)
from dagster._config.pythonic_config.config import Config
from dagster._config.pythonic_config.conversion_utils import TResValue
from dagster._config.pythonic_config.resource import (
CoercibleToResource,
ConfigurableResourceFactory,
NestedResourcesResourceDefinition,
PartialResource,
T_Self,
)
from dagster._config.pythonic_config.type_check_utils import safe_is_subclass
from dagster._core.definitions.definition_config_schema import CoercableToConfigSchema
from dagster._core.definitions.resource_definition import ResourceFunction
from dagster._core.execution.context.init import InitResourceContext
from dagster._core.storage.io_manager import IOManager, IOManagerDefinition
from dagster._utils.cached_method import cached_method
TIOManagerValue = TypeVar("TIOManagerValue", bound=IOManager)
class ConfigurableIOManagerFactoryResourceDefinition(
NestedResourcesResourceDefinition,
IOManagerDefinition,
):
def __init__(
self,
configurable_resource_cls: Type,
resource_fn: ResourceFunction,
config_schema: Any,
description: Optional[str],
nested_resources: Mapping[str, Any],
nested_partial_resources: Mapping[str, Any],
input_config_schema: Optional[Union[CoercableToConfigSchema, Type[Config]]] = None,
output_config_schema: Optional[Union[CoercableToConfigSchema, Type[Config]]] = None,
dagster_maintained: bool = False,
):
input_config_schema_resolved: CoercableToConfigSchema = (
cast(Type[Config], input_config_schema).to_config_schema()
if safe_is_subclass(input_config_schema, Config)
else cast(CoercableToConfigSchema, input_config_schema)
)
output_config_schema_resolved: CoercableToConfigSchema = (
cast(Type[Config], output_config_schema).to_config_schema()
if safe_is_subclass(output_config_schema, Config)
else cast(CoercableToConfigSchema, output_config_schema)
)
super().__init__(
resource_fn=resource_fn,
config_schema=config_schema,
description=description,
input_config_schema=input_config_schema_resolved,
output_config_schema=output_config_schema_resolved,
)
self._nested_partial_resources = nested_partial_resources
self._nested_resources = nested_resources
self._configurable_resource_cls = configurable_resource_cls
self._dagster_maintained = dagster_maintained
@property
def configurable_resource_cls(self) -> Type:
return self._configurable_resource_cls
@property
def nested_resources(
self,
) -> Mapping[str, Any]:
return self._nested_resources
@property
def nested_partial_resources(
self,
) -> Mapping[str, "CoercibleToResource"]:
return self._nested_partial_resources
[docs]
class ConfigurableIOManagerFactory(ConfigurableResourceFactory, Generic[TResValue]):
"""Base class for Dagster IO managers that utilize structured config. This base class
is useful for cases in which the returned IO manager is not the same as the class itself
(e.g. when it is a wrapper around the actual IO manager implementation).
This class is a subclass of both :py:class:`IOManagerDefinition` and :py:class:`Config`.
Implementers should provide an implementation of the :py:meth:`resource_function` method,
which should return an instance of :py:class:`IOManager`.
Example definition:
.. code-block:: python
class ExternalIOManager(IOManager):
def __init__(self, connection):
self._connection = connection
def handle_output(self, context, obj):
...
def load_input(self, context):
...
class ConfigurableExternalIOManager(ConfigurableIOManagerFactory):
username: str
password: str
def create_io_manager(self, context) -> IOManager:
with database.connect(username, password) as connection:
return MyExternalIOManager(connection)
defs = Definitions(
...,
resources={
"io_manager": ConfigurableExternalIOManager(
username="dagster",
password=EnvVar("DB_PASSWORD")
)
}
)
"""
def __init__(self, **data: Any):
ConfigurableResourceFactory.__init__(self, **data)
@abstractmethod
def create_io_manager(self, context) -> TResValue:
"""Implement as one would implement a @io_manager decorator function."""
raise NotImplementedError()
def create_resource(self, context: InitResourceContext) -> TResValue:
return self.create_io_manager(context)
@classmethod
def configure_at_launch(cls: "Type[T_Self]", **kwargs) -> "PartialIOManager[T_Self]":
"""Returns a partially initialized copy of the IO manager, with remaining config fields
set at runtime.
"""
return PartialIOManager(resource_cls=cls, data=kwargs)
@cached_method
def get_resource_definition(self) -> ConfigurableIOManagerFactoryResourceDefinition:
return ConfigurableIOManagerFactoryResourceDefinition(
self.__class__,
resource_fn=self._get_initialize_and_run_fn(),
config_schema=self._config_schema,
description=self.__doc__,
nested_partial_resources=self._nested_partial_resources,
nested_resources=self.nested_resources,
input_config_schema=self.__class__.input_config_schema(),
output_config_schema=self.__class__.output_config_schema(),
dagster_maintained=self._is_dagster_maintained(),
)
@classmethod
def input_config_schema(
cls,
) -> Optional[Union[CoercableToConfigSchema, Type[Config]]]:
return None
@classmethod
def output_config_schema(
cls,
) -> Optional[Union[CoercableToConfigSchema, Type[Config]]]:
return None
class PartialIOManager(
PartialResource[TResValue],
Generic[TResValue],
):
@cached_method
def get_resource_definition(self) -> ConfigurableIOManagerFactoryResourceDefinition:
input_config_schema = None
output_config_schema = None
if safe_is_subclass(self.resource_cls, ConfigurableIOManagerFactory):
factory_cls: Type[ConfigurableIOManagerFactory] = cast(
Type[ConfigurableIOManagerFactory], self.resource_cls
)
input_config_schema = factory_cls.input_config_schema()
output_config_schema = factory_cls.output_config_schema()
return ConfigurableIOManagerFactoryResourceDefinition(
self.resource_cls,
resource_fn=self._state__internal__.resource_fn,
config_schema=self._state__internal__.config_schema,
description=self._state__internal__.description,
nested_resources=self._state__internal__.nested_resources,
nested_partial_resources=self._state__internal__.nested_partial_resources,
input_config_schema=input_config_schema,
output_config_schema=output_config_schema,
dagster_maintained=self.resource_cls._is_dagster_maintained(), # noqa: SLF001
)
[docs]
class ConfigurableIOManager(ConfigurableIOManagerFactory, IOManager):
"""Base class for Dagster IO managers that utilize structured config.
This class is a subclass of both :py:class:`IOManagerDefinition`, :py:class:`Config`,
and :py:class:`IOManager`. Implementers must provide an implementation of the
:py:meth:`handle_output` and :py:meth:`load_input` methods.
Example definition:
.. code-block:: python
class MyIOManager(ConfigurableIOManager):
path_prefix: List[str]
def _get_path(self, context) -> str:
return "/".join(context.asset_key.path)
def handle_output(self, context, obj):
write_csv(self._get_path(context), obj)
def load_input(self, context):
return read_csv(self._get_path(context))
defs = Definitions(
...,
resources={
"io_manager": MyIOManager(path_prefix=["my", "prefix"])
}
)
"""
def create_io_manager(self, context) -> IOManager:
return self
class ConfigurableLegacyIOManagerAdapter(ConfigurableIOManagerFactory):
"""Adapter base class for wrapping a decorated, function-style I/O manager
with structured config.
To use this class, subclass it, define config schema fields using Pydantic,
and implement the ``wrapped_io_manager`` method.
Example:
.. code-block:: python
class OldIOManager(IOManager):
def __init__(self, base_path: str):
...
@io_manager(config_schema={"base_path": str})
def old_io_manager(context):
base_path = context.resource_config["base_path"]
return OldIOManager(base_path)
class MyIOManager(ConfigurableLegacyIOManagerAdapter):
base_path: str
@property
def wrapped_io_manager(self) -> IOManagerDefinition:
return old_io_manager
"""
@property
@abstractmethod
def wrapped_io_manager(self) -> IOManagerDefinition:
raise NotImplementedError()
def create_io_manager(self, context) -> IOManager:
raise NotImplementedError(
"Because we override resource_fn in the adapter, this is never called."
)
@cached_method
def get_resource_definition(self) -> ConfigurableIOManagerFactoryResourceDefinition:
return ConfigurableIOManagerFactoryResourceDefinition(
self.__class__,
resource_fn=self.wrapped_io_manager.resource_fn,
config_schema=self._config_schema,
description=self.__doc__,
nested_partial_resources=self._nested_partial_resources,
nested_resources=self.nested_resources,
dagster_maintained=self._is_dagster_maintained(),
)