mirror of
https://github.com/dbt-labs/dbt-core
synced 2025-12-17 19:31:34 +00:00
358 lines
12 KiB
Python
358 lines
12 KiB
Python
import importlib.util
|
|
import os
|
|
import time
|
|
import traceback
|
|
from functools import update_wrapper
|
|
from typing import Optional
|
|
|
|
from click import Context
|
|
|
|
import dbt.tracking
|
|
from dbt.adapters.factory import adapter_management, get_adapter, register_adapter
|
|
from dbt.cli.exceptions import ExceptionExit, ResultExit
|
|
from dbt.cli.flags import Flags
|
|
from dbt.config import RuntimeConfig
|
|
from dbt.config.runtime import UnsetProfile, load_profile, load_project
|
|
from dbt.context.providers import generate_runtime_macro_context
|
|
from dbt.context.query_header import generate_query_header_context
|
|
from dbt.events.logging import setup_event_logger
|
|
from dbt.events.types import (
|
|
CommandCompleted,
|
|
MainEncounteredError,
|
|
MainReportArgs,
|
|
MainReportVersion,
|
|
MainStackTrace,
|
|
MainTrackingUserState,
|
|
ResourceReport,
|
|
)
|
|
from dbt.exceptions import DbtProjectError, FailFastError
|
|
from dbt.flags import get_flag_dict, set_flags
|
|
from dbt.mp_context import get_mp_context
|
|
from dbt.parser.manifest import parse_manifest
|
|
from dbt.plugins import set_up_plugin_manager
|
|
from dbt.profiler import profiler
|
|
from dbt.tracking import active_user, initialize_from_flags, track_run
|
|
from dbt.utils import try_get_max_rss_kb
|
|
from dbt.version import installed as installed_version
|
|
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.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 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.
|
|
get_invocation_context()._env = get_env()
|
|
|
|
# Flags
|
|
flags = Flags(ctx)
|
|
ctx.obj["flags"] = flags
|
|
set_flags(flags)
|
|
|
|
# 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()
|
|
|
|
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())
|
|
|
|
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_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_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).
|
|
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
|
|
|
|
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"]
|
|
project = load_project(
|
|
flags.PROJECT_DIR, flags.VERSION_CHECK, ctx.obj["profile"], flags.VARS
|
|
)
|
|
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 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)
|
|
|
|
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"]
|
|
|
|
# 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
|
|
)
|
|
else:
|
|
register_adapter(runtime_config, get_mp_context())
|
|
adapter = get_adapter(runtime_config)
|
|
adapter.set_macro_context_generator(generate_runtime_macro_context)
|
|
adapter.set_macro_resolver(ctx.obj["manifest"])
|
|
query_header_context = generate_query_header_context(
|
|
adapter.config, ctx.obj["manifest"]
|
|
)
|
|
adapter.connections.set_query_header(query_header_context)
|
|
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])
|