Merge pull request #273 from analyst-collective/fix/tracking-import

scope all tracking state to fix imported dbt runs

dbt/main.py

@@ -43,6 +43,8 @@ def handle(args):
  # correct profiles.yml file
  if not config.send_anonymous_usage_stats(parsed.profiles_dir):
  dbt.tracking.do_not_track()
+ else:
+ dbt.tracking.initialize_tracking()
 
  res = run_from_args(parsed)
  dbt.tracking.flush()
@@ -92,10 +94,13 @@ def run_from_args(parsed):
  log_path = proj.get('log-path', 'logs')
 
  initialize_logger(parsed.debug, log_path)
+ logger.debug("Tracking: {}".format(dbt.tracking.active_user.state()))
 
  dbt.tracking.track_invocation_start(project=proj, args=parsed)
+
+ result = None
  try:
- return task.run()
+ result = task.run()
  dbt.tracking.track_invocation_end(
  project=proj, args=parsed, result_type="ok", result=None
  )
@@ -110,6 +115,8 @@ def run_from_args(parsed):
  )
  raise
 
+ return result
+
 
 def invoke_dbt(parsed):
  task = None

dbt/runner.py

@@ -412,7 +412,7 @@ def call_table_exists(schema, table):
 
  self.context = {
  "run_started_at": datetime.now(),
- "invocation_id": dbt.tracking.invocation_id,
+ "invocation_id": dbt.tracking.active_user.invocation_id,
  "get_columns_in_table": call_get_columns_in_table,
  "get_missing_columns": call_get_missing_columns,
  "already_exists": call_table_exists,
@@ -575,8 +575,9 @@ def on_complete(run_model_results):
  run_model_result.execution_time
  )
 
+ invocation_id = dbt.tracking.active_user.invocation_id
  dbt.tracking.track_model_run({
- "invocation_id": dbt.tracking.invocation_id,
+ "invocation_id": invocation_id,
  "index": index,
  "total": num_models,
  "execution_time": run_model_result.execution_time,

dbt/tracking.py

@@ -31,35 +31,56 @@
 emitter = Emitter(COLLECTOR_URL, protocol=COLLECTOR_PROTOCOL, buffer_size=1)
 tracker = Tracker(emitter, namespace="cf", app_id="dbt")
 
+active_user = None
 
-def __write_user():
- user = {
- "id": str(uuid.uuid4())
- }
 
- cookie_dir = os.path.dirname(COOKIE_PATH)
- if not os.path.exists(cookie_dir):
- os.makedirs(cookie_dir)
+class User(object):
 
- with open(COOKIE_PATH, "w") as fh:
- yaml.dump(user, fh)
+ def __init__(self):
+ self.do_not_track = True
 
- return user
+ self.id = None
+ self.invocation_id = None
 
+ def state(self):
+ return "do not track" if self.do_not_track else "tracking"
 
-def get_user():
- if os.path.isfile(COOKIE_PATH):
- with open(COOKIE_PATH, "r") as fh:
- try:
- user = yaml.safe_load(fh)
- if user is None:
- user = __write_user()
- except yaml.reader.ReaderError as e:
- user = __write_user()
- else:
- user = __write_user()
+ def initialize(self):
+ self.do_not_track = False
+
+ self.invocation_id = str(uuid.uuid4())
+
+ cookie = self.get_cookie()
+ self.id = cookie.get('id')
+
+ subject = Subject()
+ subject.set_user_id(self.id)
+ tracker.set_subject(subject)
+
+ def set_cookie(self):
+ cookie_dir = os.path.dirname(COOKIE_PATH)
+ user = {"id": str(uuid.uuid4())}
+
+ if not os.path.exists(cookie_dir):
+ os.makedirs(cookie_dir)
 
- return user
+ with open(COOKIE_PATH, "w") as fh:
+ yaml.dump(user, fh)
+
+ return user
+
+ def get_cookie(self):
+ if not os.path.isfile(COOKIE_PATH):
+ user = self.set_cookie()
+ else:
+ with open(COOKIE_PATH, "r") as fh:
+ try:
+ user = yaml.safe_load(fh)
+ if user is None:
+ user = self.set_cookie()
+ except yaml.reader.ReaderError as e:
+ user = self.set_cookie()
+ return user
 
 
 def get_options(args):
@@ -75,11 +96,11 @@ def get_run_type(args):
  return 'regular'
 
 
-def get_invocation_context(invocation_id, user, project, args):
+def get_invocation_context(user, project, args):
  return {
  "project_id": None if project is None else project.hashed_name(),
- "user_id": user.get("id", None),
- "invocation_id": invocation_id,
+ "user_id": user.id,
+ "invocation_id": user.invocation_id,
 
  "command": args.which,
  "options": get_options(args),
@@ -89,8 +110,8 @@ def get_invocation_context(invocation_id, user, project, args):
  }
 
 
-def get_invocation_start_context(invocation_id, user, project, args):
- data = get_invocation_context(invocation_id, user, project, args)
+def get_invocation_start_context(user, project, args):
+ data = get_invocation_context(user, project, args)
 
  start_data = {
  "progress": "start",
@@ -102,10 +123,8 @@ def get_invocation_start_context(invocation_id, user, project, args):
  return SelfDescribingJson(INVOCATION_SPEC, data)
 
 
-def get_invocation_end_context(
- invocation_id, user, project, args, result_type, result
-):
- data = get_invocation_context(invocation_id, user, project, args)
+def get_invocation_end_context(user, project, args, result_type, result):
+ data = get_invocation_context(user, project, args)
 
  start_data = {
  "progress": "end",
@@ -117,10 +136,8 @@ def get_invocation_end_context(
  return SelfDescribingJson(INVOCATION_SPEC, data)
 
 
-def get_invocation_invalid_context(
- invocation_id, user, project, args, result_type, result
-):
- data = get_invocation_context(invocation_id, user, project, args)
+def get_invocation_invalid_context(user, project, args, result_type, result):
+ data = get_invocation_context(user, project, args)
 
  start_data = {
  "progress": "invalid",
@@ -155,20 +172,9 @@ def get_dbt_env_context():
 
  return SelfDescribingJson(INVOCATION_ENV_SPEC, data)
 
-invocation_id = str(uuid.uuid4())
-platform_context = get_platform_context()
-env_context = get_dbt_env_context()
-
-user = get_user()
-subject = Subject()
-subject.set_user_id(user.get("id", None))
-tracker.set_subject(subject)
-
-__is_do_not_track = False
-
 
-def track(*args, **kwargs):
- if __is_do_not_track:
+def track(user, *args, **kwargs):
+ if user.do_not_track:
  return
  else:
  logger.debug("Sending event: {}".format(kwargs))
@@ -181,43 +187,76 @@ def track(*args, **kwargs):
 
 
 def track_invocation_start(project=None, args=None):
- invocation_context = get_invocation_start_context(
- invocation_id, user, project, args
+ context = [
+ get_invocation_start_context(active_user, project, args),
+ get_platform_context(),
+ get_dbt_env_context()
+ ]
+
+ track(
+ active_user,
+ category="dbt",
+ action='invocation',
+ label='start',
+ context=context
  )
- context = [invocation_context, platform_context, env_context]
- track(category="dbt", action='invocation', label='start', context=context)
 
 
 def track_model_run(options):
  context = [SelfDescribingJson(RUN_MODEL_SPEC, options)]
  model_id = options['model_id']
  track(
+ active_user,
  category="dbt",
  action='run_model',
- label=invocation_id,
+ label=active_user.invocation_id,
  context=context
  )
 
 
 def track_invocation_end(
  project=None, args=None, result_type=None, result=None
 ):
- invocation_context = get_invocation_end_context(
- invocation_id, user, project, args, result_type, result
+ user = active_user
+ context = [
+ get_invocation_end_context(user, project, args, result_type, result),
+ get_platform_context(),
+ get_dbt_env_context()
+ ]
+ track(
+ active_user,
+ category="dbt",
+ action='invocation',
+ label='end',
+ context=context
  )
- context = [invocation_context, platform_context, env_context]
- track(category="dbt", action='invocation', label='end', context=context)
 
 
 def track_invalid_invocation(
  project=None, args=None, result_type=None, result=None
 ):
+
+ user = active_user
  invocation_context = get_invocation_invalid_context(
- invocation_id, user, project, args, result_type, result
+ user,
+ project,
+ args,
+ result_type,
+ result
  )
- context = [invocation_context, platform_context, env_context]
+
+ context = [
+ invocation_context,
+ get_platform_context(),
+ get_dbt_env_context()
+ ]
+
  track(
- category="dbt", action='invocation', label='invalid', context=context
+ active_user,
+ category="dbt",
+ action='invocation',
+ label='invalid',
+ context=context
  )
 
 
@@ -227,6 +266,11 @@ def flush():
 
 
 def do_not_track():
- global __is_do_not_track
- logger.debug("Not sending anonymous usage events")
- __is_do_not_track = True
+ global active_user
+ active_user = User()
+
+
+def initialize_tracking():
+ global active_user
+ active_user = User()
+ active_user.initialize()