dvt-core 1.11.0b4__py3-none-any.whl

dvt/cli/requires.py

@@ -0,0 +1,478 @@
+import importlib.util
+import os
+import time
+import traceback
+from functools import update_wrapper
+from typing import Dict, Optional
+
+import dvt.tracking
+from click import Context
+from dvt.cli.exceptions import ExceptionExit, ResultExit
+from dvt.cli.flags import Flags
+from dvt.config import RuntimeConfig
+from dvt.config.catalogs import get_active_write_integration, load_catalogs
+from dvt.config.profiles_v2 import (
+    ProfileRegistry,
+    UnifiedProfileConfig,
+    load_unified_profiles,
+)
+from dvt.config.runtime import UnsetProfile, load_profile, load_project
+from dvt.context.providers import generate_runtime_macro_context
+from dvt.context.query_header import generate_query_header_context
+from dvt.deprecations import show_deprecations_summary
+from dvt.env_vars import KNOWN_ENGINE_ENV_VARS, validate_engine_env_vars
+from dvt.events.logging import setup_event_logger
+from dvt.events.types import (
+    ArtifactUploadError,
+    CommandCompleted,
+    MainEncounteredError,
+    MainReportArgs,
+    MainReportVersion,
+    MainStackTrace,
+    MainTrackingUserState,
+    ResourceReport,
+)
+from dvt.exceptions import DbtProjectError, FailFastError
+from dvt.flags import get_flag_dict, get_flags, set_flags
+from dvt.mp_context import get_mp_context
+from dvt.parser.manifest import parse_manifest
+from dvt.plugins import set_up_plugin_manager
+from dvt.profiler import profiler
+from dvt.tracking import active_user, initialize_from_flags, track_run
+from dvt.utils import try_get_max_rss_kb
+from dvt.utils.artifact_upload import upload_artifacts
+from dvt.version import installed as installed_version
+
+from dbt.adapters.factory import adapter_management, get_adapter, register_adapter
+from dbt_common.clients.system import get_env
+from dbt_common.context import get_invocation_context, set_invocation_context
+from dbt_common.events.base_types import EventLevel
+from dbt_common.events.event_manager_client import get_event_manager
+from dbt_common.events.functions import LOG_VERSION, fire_event
+from dbt_common.events.helpers import get_json_string_utcnow
+from dbt_common.exceptions import DbtBaseException as DbtException
+from dbt_common.invocation import reset_invocation_id
+from dbt_common.record import (
+    Recorder,
+    RecorderMode,
+    get_record_mode_from_env,
+    get_record_types_from_dict,
+    get_record_types_from_env,
+)
+from dbt_common.utils import cast_dict_to_dict_of_strings
+
+
+def _cross_propagate_engine_env_vars(env_dict: Dict[str, str]) -> None:
+    for env_var in KNOWN_ENGINE_ENV_VARS:
+        if env_var.old_name is not None:
+            # If the old name is in the env dict, and not the new name, set the new name based on the old name
+            if env_var.old_name in env_dict and env_var.name not in env_dict:
+                env_dict[env_var.name] = env_dict[env_var.old_name]
+            # If the new name is in the env dict, override the old name with it
+            elif env_var.name in env_dict:
+                env_dict[env_var.old_name] = env_dict[env_var.name]
+
+
+def preflight(func):
+    def wrapper(*args, **kwargs):
+        ctx = args[0]
+        assert isinstance(ctx, Context)
+        ctx.obj = ctx.obj or {}
+
+        set_invocation_context({})
+
+        # Record/Replay
+        setup_record_replay()
+
+        # Must be set after record/replay is set up so that the env can be
+        # recorded or replayed if needed.
+        env_dict = get_env()
+        _cross_propagate_engine_env_vars(env_dict)
+        get_invocation_context()._env = env_dict
+
+        # Flags
+        flags = Flags(ctx)
+        ctx.obj["flags"] = flags
+        set_flags(flags)
+        get_event_manager().require_warn_or_error_handling = (
+            flags.require_all_warnings_handled_by_warn_error
+        )
+
+        # Reset invocation_id for each 'invocation' of a dbt command (can happen multiple times in a single process)
+        reset_invocation_id()
+
+        # Logging
+        callbacks = ctx.obj.get("callbacks", [])
+        setup_event_logger(flags=flags, callbacks=callbacks)
+
+        # Tracking
+        initialize_from_flags(flags.SEND_ANONYMOUS_USAGE_STATS, flags.PROFILES_DIR)
+        ctx.with_resource(track_run(run_command=flags.WHICH))
+
+        # Now that we have our logger, fire away!
+        fire_event(MainReportVersion(version=str(installed_version), log_version=LOG_VERSION))
+        flags_dict_str = cast_dict_to_dict_of_strings(get_flag_dict())
+        fire_event(MainReportArgs(args=flags_dict_str))
+
+        # Deprecation warnings
+        flags.fire_deprecations(ctx)
+
+        if active_user is not None:  # mypy appeasement, always true
+            fire_event(MainTrackingUserState(user_state=active_user.state()))
+
+        # Profiling
+        if flags.RECORD_TIMING_INFO:
+            ctx.with_resource(profiler(enable=True, outfile=flags.RECORD_TIMING_INFO))
+
+        # Adapter management
+        ctx.with_resource(adapter_management())
+
+        # Validate engine env var restricted name space
+        validate_engine_env_vars()
+
+        return func(*args, **kwargs)
+
+    return update_wrapper(wrapper, func)
+
+
+def setup_record_replay():
+    rec_mode = get_record_mode_from_env()
+    rec_types = get_record_types_from_env()
+
+    recorder: Optional[Recorder] = None
+    if rec_mode == RecorderMode.REPLAY:
+        previous_recording_path = os.environ.get(
+            "DBT_ENGINE_RECORDER_FILE_PATH"
+        ) or os.environ.get("DBT_RECORDER_FILE_PATH")
+        recorder = Recorder(
+            RecorderMode.REPLAY, types=rec_types, previous_recording_path=previous_recording_path
+        )
+    elif rec_mode == RecorderMode.DIFF:
+        previous_recording_path = os.environ.get(
+            "DBT_ENGINE_RECORDER_FILE_PATH"
+        ) or os.environ.get("DBT_RECORDER_FILE_PATH")
+        # ensure types match the previous recording
+        types = get_record_types_from_dict(previous_recording_path)
+        recorder = Recorder(
+            RecorderMode.DIFF, types=types, previous_recording_path=previous_recording_path
+        )
+    elif rec_mode == RecorderMode.RECORD:
+        recorder = Recorder(RecorderMode.RECORD, types=rec_types)
+
+    get_invocation_context().recorder = recorder
+
+
+def tear_down_record_replay():
+    recorder = get_invocation_context().recorder
+    if recorder is not None:
+        if recorder.mode == RecorderMode.RECORD:
+            recorder.write()
+        if recorder.mode == RecorderMode.DIFF:
+            recorder.write()
+            recorder.write_diffs(diff_file_name="recording_diffs.json")
+        elif recorder.mode == RecorderMode.REPLAY:
+            recorder.write_diffs("replay_diffs.json")
+
+
+def postflight(func):
+    """The decorator that handles all exception handling for the click commands.
+    This decorator must be used before any other decorators that may throw an exception."""
+
+    def wrapper(*args, **kwargs):
+        ctx = args[0]
+        start_func = time.perf_counter()
+        success = False
+
+        try:
+            result, success = func(*args, **kwargs)
+        except FailFastError as e:
+            fire_event(MainEncounteredError(exc=str(e)))
+            raise ResultExit(e.result)
+        except DbtException as e:
+            fire_event(MainEncounteredError(exc=str(e)))
+            raise ExceptionExit(e)
+        except BaseException as e:
+            fire_event(MainEncounteredError(exc=str(e)))
+            fire_event(MainStackTrace(stack_trace=traceback.format_exc()))
+            raise ExceptionExit(e)
+        finally:
+            # Fire ResourceReport, but only on systems which support the resource
+            # module. (Skip it on Windows).
+            try:
+                if get_flags().upload_to_artifacts_ingest_api:
+                    upload_artifacts(
+                        get_flags().project_dir, get_flags().target_path, ctx.command.name
+                    )
+
+            except Exception as e:
+                fire_event(ArtifactUploadError(msg=str(e)))
+
+            show_deprecations_summary()
+
+            if importlib.util.find_spec("resource") is not None:
+                import resource
+
+                rusage = resource.getrusage(resource.RUSAGE_SELF)
+                fire_event(
+                    ResourceReport(
+                        command_name=ctx.command.name,
+                        command_success=success,
+                        command_wall_clock_time=time.perf_counter() - start_func,
+                        process_user_time=rusage.ru_utime,
+                        process_kernel_time=rusage.ru_stime,
+                        process_mem_max_rss=try_get_max_rss_kb() or rusage.ru_maxrss,
+                        process_in_blocks=rusage.ru_inblock,
+                        process_out_blocks=rusage.ru_oublock,
+                    ),
+                    (
+                        EventLevel.INFO
+                        if "flags" in ctx.obj and ctx.obj["flags"].SHOW_RESOURCE_REPORT
+                        else None
+                    ),
+                )
+
+            fire_event(
+                CommandCompleted(
+                    command=ctx.command_path,
+                    success=success,
+                    completed_at=get_json_string_utcnow(),
+                    elapsed=time.perf_counter() - start_func,
+                )
+            )
+
+            tear_down_record_replay()
+
+        if not success:
+            raise ResultExit(result)
+
+        return (result, success)
+
+    return update_wrapper(wrapper, func)
+
+
+# TODO: UnsetProfile is necessary for deps and clean to load a project.
+# This decorator and its usage can be removed once https://github.com/dbt-labs/dbt-core/issues/6257 is closed.
+def unset_profile(func):
+    def wrapper(*args, **kwargs):
+        ctx = args[0]
+        assert isinstance(ctx, Context)
+
+        profile = UnsetProfile()
+        ctx.obj["profile"] = profile
+
+        return func(*args, **kwargs)
+
+    return update_wrapper(wrapper, func)
+
+
+def profile(func):
+    def wrapper(*args, **kwargs):
+        ctx = args[0]
+        assert isinstance(ctx, Context)
+
+        flags = ctx.obj["flags"]
+        # TODO: Generalize safe access to flags.THREADS:
+        # https://github.com/dbt-labs/dbt-core/issues/6259
+        threads = getattr(flags, "THREADS", None)
+        profile = load_profile(flags.PROJECT_DIR, flags.VARS, flags.PROFILE, flags.TARGET, threads)
+        ctx.obj["profile"] = profile
+        get_invocation_context().uses_adapter(profile.credentials.type)
+
+        return func(*args, **kwargs)
+
+    return update_wrapper(wrapper, func)
+
+
+def project(func):
+    def wrapper(*args, **kwargs):
+        ctx = args[0]
+        assert isinstance(ctx, Context)
+
+        # TODO: Decouple target from profile, and remove the need for profile here:
+        # https://github.com/dbt-labs/dbt-core/issues/6257
+        if not ctx.obj.get("profile"):
+            raise DbtProjectError("profile required for project")
+
+        flags = ctx.obj["flags"]
+        # TODO deprecations warnings fired from loading the project will lack
+        # the project_id in the snowplow event.
+        project = load_project(
+            flags.PROJECT_DIR, flags.VERSION_CHECK, ctx.obj["profile"], flags.VARS, validate=True
+        )
+        ctx.obj["project"] = project
+
+        # Plugins
+        set_up_plugin_manager(project_name=project.project_name)
+
+        if dbt.tracking.active_user is not None:
+            project_id = None if project is None else project.hashed_name()
+
+            dbt.tracking.track_project_id({"project_id": project_id})
+
+        return func(*args, **kwargs)
+
+    return update_wrapper(wrapper, func)
+
+
+def runtime_config(func):
+    """A decorator used by click command functions for generating a runtime
+    config given a profile and project.
+    """
+
+    def wrapper(*args, **kwargs):
+        ctx = args[0]
+        assert isinstance(ctx, Context)
+
+        req_strs = ["profile", "project"]
+        reqs = [ctx.obj.get(req_str) for req_str in req_strs]
+
+        if None in reqs:
+            raise DbtProjectError("profile and project required for runtime_config")
+
+        config = RuntimeConfig.from_parts(
+            ctx.obj["project"],
+            ctx.obj["profile"],
+            ctx.obj["flags"],
+        )
+
+        ctx.obj["runtime_config"] = config
+
+        if dbt.tracking.active_user is not None:
+            adapter_type = (
+                getattr(config.credentials, "type", None)
+                if hasattr(config, "credentials")
+                else None
+            )
+            adapter_unique_id = (
+                config.credentials.hashed_unique_field()
+                if hasattr(config, "credentials")
+                else None
+            )
+
+            dbt.tracking.track_adapter_info(
+                {
+                    "adapter_type": adapter_type,
+                    "adapter_unique_id": adapter_unique_id,
+                }
+            )
+
+        return func(*args, **kwargs)
+
+    return update_wrapper(wrapper, func)
+
+
+def profiles_registry(func):
+    """
+    A decorator for commands that need to work with multiple profiles.
+
+    This decorator loads all profiles from profiles.yml and creates a ProfileRegistry
+    that can be used to access multiple database connections simultaneously.
+
+    Use this for DVT-specific commands like:
+    - dvt debug --all-profiles
+    - dvt profiles list/test
+    - dvt compute test
+
+    The decorator adds 'profiles_registry' to ctx.obj.
+    """
+
+    def wrapper(*args, **kwargs):
+        ctx = args[0]
+        assert isinstance(ctx, Context)
+
+        flags = ctx.obj.get("flags")
+        if not flags:
+            raise DbtProjectError("flags required for profiles_registry")
+
+        # Load unified profiles from standard locations
+        from pathlib import Path
+
+        project_dir = Path(flags.PROJECT_DIR) if flags.PROJECT_DIR else None
+        unified_profiles = load_unified_profiles(project_dir)
+
+        # Create ProfileRegistry
+        registry = ProfileRegistry(unified_profiles)
+        ctx.obj["profiles_registry"] = registry
+        ctx.obj["unified_profiles"] = unified_profiles
+
+        return func(*args, **kwargs)
+
+    return update_wrapper(wrapper, func)
+
+
+def catalogs(func):
+    """A decorator used by click command functions for loading catalogs"""
+
+    def wrapper(*args, **kwargs):
+        ctx = args[0]
+        assert isinstance(ctx, Context)
+
+        req_strs = ["flags", "profile", "project"]
+        reqs = [ctx.obj.get(req_str) for req_str in req_strs]
+        if None in reqs:
+            raise DbtProjectError("profile and flags required to load catalogs")
+
+        flags = ctx.obj["flags"]
+        ctx_project = ctx.obj["project"]
+
+        _catalogs = load_catalogs(flags.PROJECT_DIR, ctx_project.project_name, flags.VARS)
+        ctx.obj["catalogs"] = _catalogs
+
+        return func(*args, **kwargs)
+
+    return update_wrapper(wrapper, func)
+
+
+def manifest(*args0, write=True, write_perf_info=False):
+    """A decorator used by click command functions for generating a manifest
+    given a profile, project, and runtime config. This also registers the adapter
+    from the runtime config and conditionally writes the manifest to disk.
+    """
+
+    def outer_wrapper(func):
+        def wrapper(*args, **kwargs):
+            ctx = args[0]
+            assert isinstance(ctx, Context)
+
+            setup_manifest(ctx, write=write, write_perf_info=write_perf_info)
+            return func(*args, **kwargs)
+
+        return update_wrapper(wrapper, func)
+
+    # if there are no args, the decorator was used without params @decorator
+    # otherwise, the decorator was called with params @decorator(arg)
+    if len(args0) == 0:
+        return outer_wrapper
+    return outer_wrapper(args0[0])
+
+
+def setup_manifest(ctx: Context, write: bool = True, write_perf_info: bool = False):
+    """Load the manifest and add it to the context."""
+    req_strs = ["profile", "project", "runtime_config"]
+    reqs = [ctx.obj.get(dep) for dep in req_strs]
+
+    if None in reqs:
+        raise DbtProjectError("profile, project, and runtime_config required for manifest")
+
+    runtime_config = ctx.obj["runtime_config"]
+
+    catalogs = ctx.obj["catalogs"] if "catalogs" in ctx.obj else []
+    active_integrations = [get_active_write_integration(catalog) for catalog in catalogs]
+
+    # if a manifest has already been set on the context, don't overwrite it
+    if ctx.obj.get("manifest") is None:
+        ctx.obj["manifest"] = parse_manifest(
+            runtime_config,
+            write_perf_info,
+            write,
+            ctx.obj["flags"].write_json,
+            active_integrations,
+        )
+        adapter = get_adapter(runtime_config)
+    else:
+        register_adapter(runtime_config, get_mp_context())
+        adapter = get_adapter(runtime_config)
+        adapter.set_macro_context_generator(generate_runtime_macro_context)  # type: ignore[arg-type]
+        adapter.set_macro_resolver(ctx.obj["manifest"])
+        query_header_context = generate_query_header_context(adapter.config, ctx.obj["manifest"])  # type: ignore[attr-defined]
+        adapter.connections.set_query_header(query_header_context)

dvt/cli/resolvers.py

@@ -0,0 +1,32 @@
+from pathlib import Path
+
+from dvt.config.project import PartialProject
+from dvt.exceptions import DbtProjectError
+
+
+def default_project_dir() -> Path:
+    paths = list(Path.cwd().parents)
+    paths.insert(0, Path.cwd())
+    return next((x for x in paths if (x / "dbt_project.yml").exists()), Path.cwd())
+
+
+def default_profiles_dir() -> Path:
+    return Path.cwd() if (Path.cwd() / "profiles.yml").exists() else Path.home() / ".dbt"
+
+
+def default_log_path(project_dir: Path, verify_version: bool = False) -> Path:
+    """If available, derive a default log path from dbt_project.yml. Otherwise, default to "logs".
+    Known limitations:
+    1. Using PartialProject here, so no jinja rendering of log-path.
+    2. Programmatic invocations of the cli via dbtRunner may pass a Project object directly,
+       which is not being taken into consideration here to extract a log-path.
+    """
+    default_log_path = Path("logs")
+    try:
+        partial = PartialProject.from_project_root(str(project_dir), verify_version=verify_version)
+        partial_log_path = partial.project_dict.get("log-path") or default_log_path
+        default_log_path = Path(project_dir) / partial_log_path
+    except DbtProjectError:
+        pass
+
+    return default_log_path

dvt/cli/types.py

@@ -0,0 +1,40 @@
+from enum import Enum
+from typing import List
+
+from dbt_common.exceptions import DbtInternalError
+
+
+class Command(Enum):
+    BUILD = "build"
+    CLEAN = "clean"
+    COMPILE = "compile"
+    CLONE = "clone"
+    DOCS_GENERATE = "generate"
+    DOCS_SERVE = "serve"
+    DEBUG = "debug"
+    DEPS = "deps"
+    INIT = "init"
+    LIST = "list"
+    PARSE = "parse"
+    RUN = "run"
+    RUN_OPERATION = "run-operation"
+    SEED = "seed"
+    SHOW = "show"
+    SNAPSHOT = "snapshot"
+    SOURCE_FRESHNESS = "freshness"
+    TEST = "test"
+    RETRY = "retry"
+
+    @classmethod
+    def from_str(cls, s: str) -> "Command":
+        try:
+            return cls(s)
+        except ValueError:
+            raise DbtInternalError(f"No value '{s}' exists in Command enum")
+
+    def to_list(self) -> List[str]:
+        return {
+            Command.DOCS_GENERATE: ["docs", "generate"],
+            Command.DOCS_SERVE: ["docs", "serve"],
+            Command.SOURCE_FRESHNESS: ["source", "freshness"],
+        }.get(self, [self.value])

dvt/clients/checked_load.py

@@ -0,0 +1,82 @@
+import collections
+import dataclasses
+from typing import Any, Dict, List, Optional, Tuple
+
+import yaml
+from dvt import deprecations
+from dvt.clients.yaml_helper import load_yaml_text
+
+# the C version is faster, but it doesn't always exist
+try:
+    from yaml import CSafeLoader as SafeLoader
+except ImportError:
+    from yaml import SafeLoader  # type: ignore  # noqa: F401
+
+
+@dataclasses.dataclass
+class YamlCheckFailure:
+    failure_type: str
+    message: str
+
+
+def checked_load(contents) -> Tuple[Optional[Dict[str, Any]], List[YamlCheckFailure]]:
+    # A hacky (but sadly justified) method for modifying a bit of PyYAML. We create
+    # a new local subclass of SafeLoader, since we need to associate state with
+    # the static class, but static classes do not have non-static state. This allows
+    # us to be sure we have exclusive access to the class.
+    class CheckedLoader(SafeLoader):
+        check_failures: List[YamlCheckFailure] = []
+
+        def construct_mapping(self, node, deep=False):
+            if not isinstance(node, yaml.MappingNode):
+                raise yaml.constructor.ConstructorError(
+                    None, None, "expected a mapping node, but found %s" % node.id, node.start_mark
+                )
+            is_override = (
+                len(node.value) > 0
+                and len(node.value[0]) > 0
+                and getattr(node.value[0][0], "value") == "<<"
+            )
+            self.flatten_mapping(node)
+            mapping = {}
+            for key_node, value_node in node.value:
+                key = self.construct_object(key_node, deep=deep)
+                if not isinstance(key, collections.abc.Hashable):
+                    raise yaml.constructor.ConstructorError(
+                        "while constructing a mapping",
+                        node.start_mark,
+                        "found unhashable key",
+                        key_node.start_mark,
+                    )
+                value = self.construct_object(value_node, deep=deep)
+
+                if not is_override and key in mapping:
+                    start_mark = str(key_node.start_mark)
+                    if start_mark.startswith("  in"):  # this means it was at the top level
+                        message = f"Duplicate key '{key}' {start_mark.lstrip()}"
+                    else:
+                        message = f"Duplicate key '{key}' at {key_node.start_mark}"
+
+                    self.check_failures.append(YamlCheckFailure("duplicate_key", message))
+
+                mapping[key] = value
+            return mapping
+
+    CheckedLoader.add_constructor(
+        yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, CheckedLoader.construct_mapping
+    )
+
+    dct = load_yaml_text(contents, loader=CheckedLoader)
+    check_failures = CheckedLoader.check_failures
+
+    return (dct, check_failures)
+
+
+def issue_deprecation_warnings_for_failures(failures: List[YamlCheckFailure], file: str):
+    for failure in failures:
+        if failure.failure_type == "duplicate_key":
+            deprecations.warn(
+                "duplicate-yaml-keys-deprecation",
+                duplicate_description=failure.message,
+                file=file,
+            )