Tracking works with Click (#5972)

Co-authored-by: Ian Knox <ian.knox@dbtlabs.com>

core/dbt/cli/flags.py

@@ -1,10 +1,12 @@
 # TODO  Move this to /core/dbt/flags.py when we're ready to break things
 import os
+import sys
 from dataclasses import dataclass
+from importlib import import_module
 from multiprocessing import get_context
 from pprint import pformat as pf
 
-from click import get_current_context
+from click import Context, get_current_context
 
 if os.name != "nt":
     # https://bugs.python.org/issue41567
@@ -13,7 +15,7 @@
 
 @dataclass(frozen=True)
 class Flags:
-    def __init__(self, ctx=None) -> None:
+    def __init__(self, ctx: Context = None) -> None:
 
         if ctx is None:
             ctx = get_current_context()
@@ -32,13 +34,26 @@ def assign_params(ctx):
 
         assign_params(ctx)
 
+        # Get the invoked command flags
+        if hasattr(ctx, "invoked_subcommand") and ctx.invoked_subcommand is not None:
+            invoked_subcommand = getattr(import_module("dbt.cli.main"), ctx.invoked_subcommand)
+            invoked_subcommand.allow_extra_args = True
+            invoked_subcommand.ignore_unknown_options = True
+            invoked_subcommand_ctx = invoked_subcommand.make_context(None, sys.argv)
+            assign_params(invoked_subcommand_ctx)
+
         # Hard coded flags
         object.__setattr__(self, "WHICH", ctx.info_name)
         object.__setattr__(self, "MP_CONTEXT", get_context("spawn"))
 
         # Support console DO NOT TRACK initiave
-        if os.getenv("DO_NOT_TRACK", "").lower() in (1, "t", "true", "y", "yes"):
-            object.__setattr__(self, "ANONYMOUS_USAGE_STATS", False)
+        object.__setattr__(
+            self,
+            "ANONYMOUS_USAGE_STATS",
+            False
+            if os.getenv("DO_NOT_TRACK", "").lower() in (1, "t", "true", "y", "yes")
+            else True,
+        )
 
     def __str__(self) -> str:
         return str(pf(self.__dict__))

core/dbt/cli/main.py

@@ -7,6 +7,7 @@
 from dbt.cli import params as p
 from dbt.cli.flags import Flags
 from dbt.profiler import profiler
+from dbt.tracking import initialize_from_flags, track_run
 
 
 def cli_runner():
@@ -52,17 +53,21 @@ def cli(ctx, **kwargs):
     """An ELT tool for managing your SQL transformations and data models.
     For more documentation on these commands, visit: docs.getdbt.com
     """
-    incomplete_flags = Flags()
+    flags = Flags()
+
+    # Tracking
+    initialize_from_flags(flags.ANONYMOUS_USAGE_STATS, flags.PROFILES_DIR)
+    ctx.with_resource(track_run(run_command=ctx.invoked_subcommand))
 
     # Profiling
-    if incomplete_flags.RECORD_TIMING_INFO:
-        ctx.with_resource(profiler(enable=True, outfile=incomplete_flags.RECORD_TIMING_INFO))
+    if flags.RECORD_TIMING_INFO:
+        ctx.with_resource(profiler(enable=True, outfile=flags.RECORD_TIMING_INFO))
 
     # Adapter management
     ctx.with_resource(adapter_management())
 
     # Version info
-    if incomplete_flags.VERSION:
+    if flags.VERSION:
         click.echo(f"`version` called\n ctx.params: {pf(ctx.params)}")
         return
     else:

core/dbt/flags.py

@@ -24,27 +24,28 @@
 STORE_FAILURES = False  # subcommand
 
 # Global CLI commands
-USE_EXPERIMENTAL_PARSER = None
-STATIC_PARSER = None
-WARN_ERROR = None
-WRITE_JSON = None
-PARTIAL_PARSE = None
-USE_COLORS = None
+ANONYMOUS_USAGE_STATS = None
+CACHE_SELECTED_ONLY = None
 DEBUG = None
-LOG_FORMAT = None
-VERSION_CHECK = None
+EVENT_BUFFER_SIZE = 100000
 FAIL_FAST = None
-SEND_ANONYMOUS_USAGE_STATS = None
-PRINTER_WIDTH = 80
-WHICH = None
 INDIRECT_SELECTION = None
 LOG_CACHE_EVENTS = None
-EVENT_BUFFER_SIZE = 100000
-QUIET = None
+LOG_FORMAT = None
+LOG_PATH = None
 NO_PRINT = None
-CACHE_SELECTED_ONLY = None
+PARTIAL_PARSE = None
+PRINTER_WIDTH = 80
+QUIET = None
+SEND_ANONYMOUS_USAGE_STATS = None
+STATIC_PARSER = None
 TARGET_PATH = None
-LOG_PATH = None
+USE_COLORS = None
+USE_EXPERIMENTAL_PARSER = None
+VERSION_CHECK = None
+WARN_ERROR = None
+WHICH = None
+WRITE_JSON = None
 
 _NON_BOOLEAN_FLAGS = [
     "LOG_FORMAT",
@@ -63,27 +64,28 @@
 # CLI args, environment variables, and user_config (profiles.yml).
 # Environment variables use the pattern 'DBT_{flag name}', like DBT_PROFILES_DIR
 flag_defaults = {
-    "USE_EXPERIMENTAL_PARSER": False,
-    "STATIC_PARSER": True,
-    "WARN_ERROR": False,
-    "WRITE_JSON": True,
-    "PARTIAL_PARSE": True,
-    "USE_COLORS": True,
-    "PROFILES_DIR": DEFAULT_PROFILES_DIR,
+    "ANONYMOUS_USAGE_STATS": True,
+    "CACHE_SELECTED_ONLY": False,
     "DEBUG": False,
-    "LOG_FORMAT": None,
-    "VERSION_CHECK": True,
+    "EVENT_BUFFER_SIZE": 100000,
     "FAIL_FAST": False,
-    "SEND_ANONYMOUS_USAGE_STATS": True,
-    "PRINTER_WIDTH": 80,
     "INDIRECT_SELECTION": "eager",
     "LOG_CACHE_EVENTS": False,
-    "EVENT_BUFFER_SIZE": 100000,
-    "QUIET": False,
+    "LOG_FORMAT": None,
+    "LOG_PATH": None,
     "NO_PRINT": False,
-    "CACHE_SELECTED_ONLY": False,
+    "PARTIAL_PARSE": True,
+    "PRINTER_WIDTH": 80,
+    "PROFILES_DIR": DEFAULT_PROFILES_DIR,
+    "QUIET": False,
+    "SEND_ANONYMOUS_USAGE_STATS": True,
+    "STATIC_PARSER": True,
     "TARGET_PATH": None,
-    "LOG_PATH": None,
+    "USE_COLORS": True,
+    "USE_EXPERIMENTAL_PARSER": False,
+    "VERSION_CHECK": True,
+    "WARN_ERROR": False,
+    "WRITE_JSON": True,
 }
 
 
@@ -132,7 +134,7 @@ def set_from_args(args, user_config):
     # black insists in putting them all on one line
     global STRICT_MODE, FULL_REFRESH, WARN_ERROR, USE_EXPERIMENTAL_PARSER, STATIC_PARSER
     global WRITE_JSON, PARTIAL_PARSE, USE_COLORS, STORE_FAILURES, PROFILES_DIR, DEBUG, LOG_FORMAT
-    global INDIRECT_SELECTION, VERSION_CHECK, FAIL_FAST, SEND_ANONYMOUS_USAGE_STATS
+    global INDIRECT_SELECTION, VERSION_CHECK, FAIL_FAST, SEND_ANONYMOUS_USAGE_STATS, ANONYMOUS_USAGE_STATS
     global PRINTER_WIDTH, WHICH, LOG_CACHE_EVENTS, EVENT_BUFFER_SIZE, QUIET, NO_PRINT, CACHE_SELECTED_ONLY
     global TARGET_PATH, LOG_PATH
 
@@ -143,39 +145,42 @@ def set_from_args(args, user_config):
     WHICH = getattr(args, "which", WHICH)
 
     # global cli flags with env var and user_config alternatives
-    USE_EXPERIMENTAL_PARSER = get_flag_value("USE_EXPERIMENTAL_PARSER", args, user_config)
-    STATIC_PARSER = get_flag_value("STATIC_PARSER", args, user_config)
-    WARN_ERROR = get_flag_value("WARN_ERROR", args, user_config)
-    WRITE_JSON = get_flag_value("WRITE_JSON", args, user_config)
-    PARTIAL_PARSE = get_flag_value("PARTIAL_PARSE", args, user_config)
-    USE_COLORS = get_flag_value("USE_COLORS", args, user_config)
-    PROFILES_DIR = get_flag_value("PROFILES_DIR", args, user_config)
+    ANONYMOUS_USAGE_STATS = get_flag_value("ANONYMOUS_USAGE_STATS", args, user_config)
+    CACHE_SELECTED_ONLY = get_flag_value("CACHE_SELECTED_ONLY", args, user_config)
     DEBUG = get_flag_value("DEBUG", args, user_config)
-    LOG_FORMAT = get_flag_value("LOG_FORMAT", args, user_config)
-    VERSION_CHECK = get_flag_value("VERSION_CHECK", args, user_config)
+    EVENT_BUFFER_SIZE = get_flag_value("EVENT_BUFFER_SIZE", args, user_config)
     FAIL_FAST = get_flag_value("FAIL_FAST", args, user_config)
-    SEND_ANONYMOUS_USAGE_STATS = get_flag_value("SEND_ANONYMOUS_USAGE_STATS", args, user_config)
-    PRINTER_WIDTH = get_flag_value("PRINTER_WIDTH", args, user_config)
     INDIRECT_SELECTION = get_flag_value("INDIRECT_SELECTION", args, user_config)
     LOG_CACHE_EVENTS = get_flag_value("LOG_CACHE_EVENTS", args, user_config)
-    EVENT_BUFFER_SIZE = get_flag_value("EVENT_BUFFER_SIZE", args, user_config)
-    QUIET = get_flag_value("QUIET", args, user_config)
+    LOG_FORMAT = get_flag_value("LOG_FORMAT", args, user_config)
+    LOG_PATH = get_flag_value("LOG_PATH", args, user_config)
     NO_PRINT = get_flag_value("NO_PRINT", args, user_config)
-    CACHE_SELECTED_ONLY = get_flag_value("CACHE_SELECTED_ONLY", args, user_config)
+    PARTIAL_PARSE = get_flag_value("PARTIAL_PARSE", args, user_config)
+    PRINTER_WIDTH = get_flag_value("PRINTER_WIDTH", args, user_config)
+    PROFILES_DIR = get_flag_value("PROFILES_DIR", args, user_config)
+    QUIET = get_flag_value("QUIET", args, user_config)
+    SEND_ANONYMOUS_USAGE_STATS = get_flag_value("SEND_ANONYMOUS_USAGE_STATS", args, user_config)
+    STATIC_PARSER = get_flag_value("STATIC_PARSER", args, user_config)
     TARGET_PATH = get_flag_value("TARGET_PATH", args, user_config)
-    LOG_PATH = get_flag_value("LOG_PATH", args, user_config)
+    USE_COLORS = get_flag_value("USE_COLORS", args, user_config)
+    USE_EXPERIMENTAL_PARSER = get_flag_value("USE_EXPERIMENTAL_PARSER", args, user_config)
+    VERSION_CHECK = get_flag_value("VERSION_CHECK", args, user_config)
+    WARN_ERROR = get_flag_value("WARN_ERROR", args, user_config)
+    WRITE_JSON = get_flag_value("WRITE_JSON", args, user_config)
 
     _set_overrides_from_env()
 
 
 def _set_overrides_from_env():
     global SEND_ANONYMOUS_USAGE_STATS
+    global ANONYMOUS_USAGE_STATS
 
     flag_value = _get_flag_value_from_env("DO_NOT_TRACK")
     if flag_value is None:
         return
 
     SEND_ANONYMOUS_USAGE_STATS = not flag_value
+    ANONYMOUS_USAGE_STATS = not flag_value
 
 
 def get_flag_value(flag, args, user_config):
@@ -239,6 +244,7 @@ def get_flag_dict():
         "version_check": VERSION_CHECK,
         "fail_fast": FAIL_FAST,
         "send_anonymous_usage_stats": SEND_ANONYMOUS_USAGE_STATS,
+        "anonymous_usage_stats": ANONYMOUS_USAGE_STATS,
         "printer_width": PRINTER_WIDTH,
         "indirect_selection": INDIRECT_SELECTION,
         "log_cache_events": LOG_CACHE_EVENTS,

core/dbt/main.py

@@ -47,8 +47,6 @@
 from dbt.exceptions import (
     Exception as dbtException,
     InternalException,
-    NotImplementedException,
-    FailedToConnectException,
 )
 
 
@@ -178,7 +176,7 @@ def handle_and_check(args):
         # Set flags from args, user config, and env vars
         user_config = read_user_config(flags.PROFILES_DIR)  # This is read again later
         flags.set_from_args(parsed, user_config)
-        dbt.tracking.initialize_from_flags()
+        dbt.tracking.initialize_from_flags(flags.ANONYMOUS_USAGE_STATS, flags.PROFILES_DIR)
         # Set log_format from flags
         parsed.cls.set_log_format()
 
@@ -201,22 +199,6 @@ def handle_and_check(args):
             return res, success
 
 
-@contextmanager
-def track_run(task):
-    dbt.tracking.track_invocation_start(config=task.config, args=task.args)
-    try:
-        yield
-        dbt.tracking.track_invocation_end(config=task.config, args=task.args, result_type="ok")
-    except (NotImplementedException, FailedToConnectException) as e:
-        fire_event(MainEncounteredError(exc=str(e)))
-        dbt.tracking.track_invocation_end(config=task.config, args=task.args, result_type="error")
-    except Exception:
-        dbt.tracking.track_invocation_end(config=task.config, args=task.args, result_type="error")
-        raise
-    finally:
-        dbt.tracking.flush()
-
-
 def run_from_args(parsed):
     log_cache_events(getattr(parsed, "log_cache_events", False))
 
@@ -240,8 +222,9 @@ def run_from_args(parsed):
         fire_event(MainTrackingUserState(user_state=dbt.tracking.active_user.state()))
 
     results = None
-
-    with track_run(task):
+    # this has been updated with project_id and adapter info removed, these will be added to new cli work
+    # being tracked at #6097 and #6098
+    with dbt.tracking.track_run(parsed.which):
         results = task.run()
     return task, results
 

core/dbt/task/deps.py

@@ -38,7 +38,6 @@ def track_package_install(self, package_name: str, source_type: str, version: st
         elif source_type != "hub":
             package_name = dbt.utils.md5(package_name)
             version = dbt.utils.md5(version)
-
         dbt.tracking.track_package_install(
             self.config,
             self.config.args,