Generic destination / sink decorator

dlt/__init__.py

@@ -29,6 +29,8 @@
 
 from dlt import sources
 from dlt.extract.decorators import source, resource, transformer, defer
+from dlt.destinations.decorators import destination
+
 from dlt.pipeline import (
     pipeline as _pipeline,
     run,
@@ -62,6 +64,7 @@
     "resource",
     "transformer",
     "defer",
+    "destination",
     "pipeline",
     "run",
     "attach",

dlt/common/configuration/__init__.py

@@ -1,7 +1,7 @@
 from .specs.base_configuration import configspec, is_valid_hint, is_secret_hint, resolve_type
 from .specs import known_sections
 from .resolve import resolve_configuration, inject_section
-from .inject import with_config, last_config, get_fun_spec
+from .inject import with_config, last_config, get_fun_spec, create_resolved_partial
 
 from .exceptions import (
     ConfigFieldMissingException,

dlt/common/configuration/container.py

@@ -1,7 +1,7 @@
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext, AbstractContextManager
 import re
 import threading
-from typing import ClassVar, Dict, Iterator, Tuple, Type, TypeVar
+from typing import ClassVar, Dict, Iterator, Tuple, Type, TypeVar, Any
 
 from dlt.common.configuration.specs.base_configuration import ContainerInjectableContext
 from dlt.common.configuration.exceptions import (
@@ -34,13 +34,17 @@ class Container:
 
     thread_contexts: Dict[int, Dict[Type[ContainerInjectableContext], ContainerInjectableContext]]
     """A thread aware mapping of injection context """
+    _context_container_locks: Dict[str, threading.Lock]
+    """Locks for container types on threads."""
+
     main_context: Dict[Type[ContainerInjectableContext], ContainerInjectableContext]
     """Injection context for the main thread"""
 
     def __new__(cls: Type["Container"]) -> "Container":
         if not cls._INSTANCE:
             cls._INSTANCE = super().__new__(cls)
             cls._INSTANCE.thread_contexts = {}
+            cls._INSTANCE._context_container_locks = {}
             cls._INSTANCE.main_context = cls._INSTANCE.thread_contexts[
                 Container._MAIN_THREAD_ID
             ] = {}
@@ -84,22 +88,22 @@ def _thread_context(
         self, spec: Type[TConfiguration]
     ) -> Dict[Type[ContainerInjectableContext], ContainerInjectableContext]:
         if spec.global_affinity:
-            context = self.main_context
+            return self.main_context
         else:
             # thread pool names used in dlt contain originating thread id. use this id over pool id
             if m := re.match(r"dlt-pool-(\d+)-", threading.currentThread().getName()):
                 thread_id = int(m.group(1))
             else:
                 thread_id = threading.get_ident()
+
             # return main context for main thread
             if thread_id == Container._MAIN_THREAD_ID:
                 return self.main_context
             # we may add a new empty thread context so lock here
             with Container._LOCK:
-                context = self.thread_contexts.get(thread_id)
-                if context is None:
+                if (context := self.thread_contexts.get(thread_id)) is None:
                     context = self.thread_contexts[thread_id] = {}
-        return context
+                return context
 
     def _thread_getitem(
         self, spec: Type[TConfiguration]
@@ -127,29 +131,44 @@ def _thread_delitem(
         del context[spec]
 
     @contextmanager
-    def injectable_context(self, config: TConfiguration) -> Iterator[TConfiguration]:
+    def injectable_context(
+        self, config: TConfiguration, lock_context: bool = False
+    ) -> Iterator[TConfiguration]:
         """A context manager that will insert `config` into the container and restore the previous value when it gets out of scope."""
+
         config.resolve()
         spec = type(config)
         previous_config: ContainerInjectableContext = None
-        context, previous_config = self._thread_getitem(spec)
-
-        # set new config and yield context
-        self._thread_setitem(context, spec, config)
-        try:
-            yield config
-        finally:
-            # before setting the previous config for given spec, check if there was no overlapping modification
-            context, current_config = self._thread_getitem(spec)
-            if current_config is config:
-                # config is injected for spec so restore previous
-                if previous_config is None:
-                    self._thread_delitem(context, spec)
+        context = self._thread_context(spec)
+        lock: AbstractContextManager[Any]
+
+        # if there is a lock_id, we need a lock for the lock_id in the scope of the current context
+        if lock_context:
+            lock_key = f"{id(context)}"
+            if (lock := self._context_container_locks.get(lock_key)) is None:
+                with Container._LOCK:
+                    self._context_container_locks[lock_key] = lock = threading.Lock()
+        else:
+            lock = nullcontext()
+
+        with lock:
+            # remember context and set item
+            previous_config = context.get(spec)
+            self._thread_setitem(context, spec, config)
+            try:
+                yield config
+            finally:
+                # before setting the previous config for given spec, check if there was no overlapping modification
+                context, current_config = self._thread_getitem(spec)
+                if current_config is config:
+                    # config is injected for spec so restore previous
+                    if previous_config is None:
+                        self._thread_delitem(context, spec)
+                    else:
+                        self._thread_setitem(context, spec, previous_config)
                 else:
-                    self._thread_setitem(context, spec, previous_config)
-            else:
-                # value was modified in the meantime and not restored
-                raise ContainerInjectableContextMangled(spec, context[spec], config)
+                    # value was modified in the meantime and not restored
+                    raise ContainerInjectableContextMangled(spec, context[spec], config)
 
     @staticmethod
     def thread_pool_prefix() -> str:

dlt/common/configuration/inject.py

@@ -1,12 +1,15 @@
 import inspect
+
 from functools import wraps
-from typing import Callable, Dict, Type, Any, Optional, Tuple, TypeVar, overload
+from typing import Callable, Dict, Type, Any, Optional, Tuple, TypeVar, overload, cast
 from inspect import Signature, Parameter
+from contextlib import nullcontext
 
 from dlt.common.typing import DictStrAny, StrAny, TFun, AnyFun
 from dlt.common.configuration.resolve import resolve_configuration, inject_section
 from dlt.common.configuration.specs.base_configuration import BaseConfiguration
 from dlt.common.configuration.specs.config_section_context import ConfigSectionContext
+
 from dlt.common.reflection.spec import spec_from_signature
 
 
@@ -32,6 +35,8 @@ def with_config(
     auto_pipeline_section: bool = False,
     include_defaults: bool = True,
     accept_partial: bool = False,
+    initial_config: BaseConfiguration = None,
+    base: Type[BaseConfiguration] = BaseConfiguration,
 ) -> TFun: ...
 
 
@@ -45,6 +50,8 @@ def with_config(
     auto_pipeline_section: bool = False,
     include_defaults: bool = True,
     accept_partial: bool = False,
+    initial_config: Optional[BaseConfiguration] = None,
+    base: Type[BaseConfiguration] = BaseConfiguration,
 ) -> Callable[[TFun], TFun]: ...
 
 
@@ -58,6 +65,7 @@ def with_config(
     include_defaults: bool = True,
     accept_partial: bool = False,
     initial_config: Optional[BaseConfiguration] = None,
+    base: Type[BaseConfiguration] = BaseConfiguration,
 ) -> Callable[[TFun], TFun]:
     """Injects values into decorated function arguments following the specification in `spec` or by deriving one from function's signature.
 
@@ -71,10 +79,11 @@ def with_config(
         prefer_existing_sections: (bool, optional): When joining existing section context, the existing context will be preferred to the one in `sections`. Default: False
         auto_pipeline_section (bool, optional): If True, a top level pipeline section will be added if `pipeline_name` argument is present . Defaults to False.
         include_defaults (bool, optional): If True then arguments with default values will be included in synthesized spec. If False only the required arguments marked with `dlt.secrets.value` and `dlt.config.value` are included
-
+        base (Type[BaseConfiguration], optional): A base class for synthesized spec. Defaults to BaseConfiguration.
     Returns:
         Callable[[TFun], TFun]: A decorated function
     """
+
     section_f: Callable[[StrAny], str] = None
     # section may be a function from function arguments to section
     if callable(sections):
@@ -88,9 +97,8 @@ def decorator(f: TFun) -> TFun:
         )
         spec_arg: Parameter = None
         pipeline_name_arg: Parameter = None
-
         if spec is None:
-            SPEC = spec_from_signature(f, sig, include_defaults)
+            SPEC = spec_from_signature(f, sig, include_defaults, base=base)
         else:
             SPEC = spec
 
@@ -109,49 +117,52 @@ def decorator(f: TFun) -> TFun:
                 pipeline_name_arg = p
                 pipeline_name_arg_default = None if p.default == Parameter.empty else p.default
 
-        @wraps(f)
-        def _wrap(*args: Any, **kwargs: Any) -> Any:
+        def resolve_config(bound_args: inspect.BoundArguments) -> BaseConfiguration:
+            """Resolve arguments using the provided spec"""
             # bind parameters to signature
-            bound_args = sig.bind(*args, **kwargs)
             # for calls containing resolved spec in the kwargs, we do not need to resolve again
             config: BaseConfiguration = None
-            if _LAST_DLT_CONFIG in kwargs:
-                config = last_config(**kwargs)
+
+            # if section derivation function was provided then call it
+            if section_f:
+                curr_sections: Tuple[str, ...] = (section_f(bound_args.arguments),)
+                # sections may be a string
+            elif isinstance(sections, str):
+                curr_sections = (sections,)
             else:
-                # if section derivation function was provided then call it
-                if section_f:
-                    curr_sections: Tuple[str, ...] = (section_f(bound_args.arguments),)
-                    # sections may be a string
-                elif isinstance(sections, str):
-                    curr_sections = (sections,)
-                else:
-                    curr_sections = sections
-
-                # if one of arguments is spec the use it as initial value
-                if initial_config:
-                    config = initial_config
-                elif spec_arg:
-                    config = bound_args.arguments.get(spec_arg.name, None)
-                # resolve SPEC, also provide section_context with pipeline_name
-                if pipeline_name_arg:
-                    curr_pipeline_name = bound_args.arguments.get(
-                        pipeline_name_arg.name, pipeline_name_arg_default
-                    )
-                else:
-                    curr_pipeline_name = None
-                section_context = ConfigSectionContext(
-                    pipeline_name=curr_pipeline_name,
-                    sections=curr_sections,
-                    merge_style=sections_merge_style,
+                curr_sections = sections
+
+            # if one of arguments is spec the use it as initial value
+            if initial_config:
+                config = initial_config
+            elif spec_arg:
+                config = bound_args.arguments.get(spec_arg.name, None)
+            # resolve SPEC, also provide section_context with pipeline_name
+            if pipeline_name_arg:
+                curr_pipeline_name = bound_args.arguments.get(
+                    pipeline_name_arg.name, pipeline_name_arg_default
                 )
-                # this may be called from many threads so section_context is thread affine
-                with inject_section(section_context):
-                    # print(f"RESOLVE CONF in inject: {f.__name__}: {section_context.sections} vs {sections}")
-                    config = resolve_configuration(
-                        config or SPEC(),
-                        explicit_value=bound_args.arguments,
-                        accept_partial=accept_partial,
-                    )
+            else:
+                curr_pipeline_name = None
+            section_context = ConfigSectionContext(
+                pipeline_name=curr_pipeline_name,
+                sections=curr_sections,
+                merge_style=sections_merge_style,
+            )
+
+            # this may be called from many threads so section_context is thread affine
+            with inject_section(section_context, lock_context=True):
+                # print(f"RESOLVE CONF in inject: {f.__name__}: {section_context.sections} vs {sections}")
+                return resolve_configuration(
+                    config or SPEC(),
+                    explicit_value=bound_args.arguments,
+                    accept_partial=accept_partial,
+                )
+
+        def update_bound_args(
+            bound_args: inspect.BoundArguments, config: BaseConfiguration, args: Any, kwargs: Any
+        ) -> None:
+            # overwrite or add resolved params
             resolved_params = dict(config)
             # overwrite or add resolved params
             for p in sig.parameters.values():
@@ -167,12 +178,56 @@ def _wrap(*args: Any, **kwargs: Any) -> Any:
                 bound_args.arguments[kwargs_arg.name].update(resolved_params)
                 bound_args.arguments[kwargs_arg.name][_LAST_DLT_CONFIG] = config
                 bound_args.arguments[kwargs_arg.name][_ORIGINAL_ARGS] = (args, kwargs)
+
+        def with_partially_resolved_config(config: Optional[BaseConfiguration] = None) -> Any:
+            # creates a pre-resolved partial of the decorated function
+            empty_bound_args = sig.bind_partial()
+            if not config:
+                config = resolve_config(empty_bound_args)
+
+            def wrapped(*args: Any, **kwargs: Any) -> Any:
+                nonlocal config
+
+                # Do we need an exception here?
+                if spec_arg and spec_arg.name in kwargs:
+                    from dlt.common import logger
+
+                    logger.warning(
+                        "Spec argument is provided in kwargs, ignoring it for resolved partial"
+                        " function."
+                    )
+
+                # we can still overwrite the config
+                if _LAST_DLT_CONFIG in kwargs:
+                    config = last_config(**kwargs)
+
+                # call the function with the pre-resolved config
+                bound_args = sig.bind(*args, **kwargs)
+                update_bound_args(bound_args, config, args, kwargs)
+                return f(*bound_args.args, **bound_args.kwargs)
+
+            return wrapped
+
+        @wraps(f)
+        def _wrap(*args: Any, **kwargs: Any) -> Any:
+            # Resolve config
+            config: BaseConfiguration = None
+            bound_args = sig.bind(*args, **kwargs)
+            if _LAST_DLT_CONFIG in kwargs:
+                config = last_config(**kwargs)
+            else:
+                config = resolve_config(bound_args)
+
             # call the function with resolved config
+            update_bound_args(bound_args, config, args, kwargs)
             return f(*bound_args.args, **bound_args.kwargs)
 
         # register the spec for a wrapped function
         _FUNC_SPECS[id(_wrap)] = SPEC
 
+        # add a method to create a pre-resolved partial
+        setattr(_wrap, "__RESOLVED_PARTIAL_FUNC__", with_partially_resolved_config)  # noqa: B010
+
         return _wrap  # type: ignore
 
     # See if we're being called as @with_config or @with_config().
@@ -197,3 +252,10 @@ def last_config(**kwargs: Any) -> Any:
 
 def get_orig_args(**kwargs: Any) -> Tuple[Tuple[Any], DictStrAny]:
     return kwargs[_ORIGINAL_ARGS]  # type: ignore
+
+
+def create_resolved_partial(f: AnyFun, config: Optional[BaseConfiguration] = None) -> AnyFun:
+    """Create a pre-resolved partial of the with_config decorated function"""
+    if partial_func := getattr(f, "__RESOLVED_PARTIAL_FUNC__", None):
+        return cast(AnyFun, partial_func(config))
+    return f

dlt/common/configuration/resolve.py

@@ -92,13 +92,14 @@ def initialize_credentials(hint: Any, initial_value: Any) -> CredentialsConfigur
 
 
 def inject_section(
-    section_context: ConfigSectionContext, merge_existing: bool = True
+    section_context: ConfigSectionContext, merge_existing: bool = True, lock_context: bool = False
 ) -> ContextManager[ConfigSectionContext]:
     """Context manager that sets section specified in `section_context` to be used during configuration resolution. Optionally merges the context already in the container with the one provided
 
     Args:
         section_context (ConfigSectionContext): Instance providing a pipeline name and section context
         merge_existing (bool, optional): Merges existing section context with `section_context` in the arguments by executing `merge_style` function on `section_context`. Defaults to True.
+        lock_context (bool, optional): Instruct to threadlock the current thread to prevent race conditions in context injection.
 
     Default Merge Style:
         Gets `pipeline_name` and `sections` from existing context if they are not provided in `section_context` argument.
@@ -112,7 +113,7 @@ def inject_section(
     if merge_existing:
         section_context.merge(existing_context)
 
-    return container.injectable_context(section_context)
+    return container.injectable_context(section_context, lock_context=lock_context)
 
 
 def _maybe_parse_native_value(