From 808ed75858a8678415e8f104abbe3083f38d9dff Mon Sep 17 00:00:00 2001 From: Jacob Beck Date: Wed, 12 Dec 2018 15:33:48 -0700 Subject: [PATCH] make tracking use the profile directory, and suppress errors --- dbt/config.py | 16 ------ dbt/main.py | 37 +++++++++----- dbt/tracking.py | 31 ++++++----- test/unit/test_config.py | 54 -------------------- test/unit/test_main.py | 108 +++++++++++++++++++++++++++++++++++++++ 5 files changed, 152 insertions(+), 94 deletions(-) create mode 100644 test/unit/test_main.py diff --git a/dbt/config.py b/dbt/config.py index af6e471167f..a87722f2819 100644 --- a/dbt/config.py +++ b/dbt/config.py @@ -86,22 +86,6 @@ def read_profiles(profiles_dir=None): return profiles -def read_config(profiles_dir): - profile = read_profile(profiles_dir) - if profile is None: - return {} - else: - return profile.get('config', {}) - - -def send_anonymous_usage_stats(config): - return config.get('send_anonymous_usage_stats', True) - - -def colorize_output(config): - return config.get('use_colors', True) - - class ConfigRenderer(object): """A renderer provides configuration rendering for a given set of cli variables and a render type. diff --git a/dbt/main.py b/dbt/main.py index 03435a340c8..1c358ef21d5 100644 --- a/dbt/main.py +++ b/dbt/main.py @@ -28,9 +28,8 @@ import dbt.profiler from dbt.utils import ExitCodes -from dbt.config import Project, RuntimeConfig, DbtProjectError, \ - DbtProfileError, PROFILES_DIR, read_config, \ - send_anonymous_usage_stats, colorize_output, read_profiles +from dbt.config import Project, Profile, RuntimeConfig, PROFILES_DIR, \ + read_profiles from dbt.exceptions import DbtProfileError, DbtProfileError, RuntimeException @@ -111,6 +110,27 @@ def handle(args): return res +def initialize_config_values(parsed): + """Given the parsed args, initialize the dbt tracking code. + + It would be nice to re-use this profile later on instead of parsing it + twice, but dbt's intialization is not structured in a way that makes that + easy. + """ + try: + profile = Profile.from_args(parsed) + except RuntimeException: + profile = None + + if profile is None or profile.send_anonymous_usage_stats: + dbt.tracking.initialize_tracking(parsed.profiles_dir) + else: + dbt.tracking.do_not_track() + + if profile is None or profile.use_colors: + dbt.ui.printer.use_colors() + + def handle_and_check(args): parsed = parse_args(args) profiler_enabled = False @@ -122,16 +142,8 @@ def handle_and_check(args): enable=profiler_enabled, outfile=parsed.record_timing_info ): - # this needs to happen after args are parsed so we can determine the - # correct profiles.yml file - profile_config = read_config(parsed.profiles_dir) - if not send_anonymous_usage_stats(profile_config): - dbt.tracking.do_not_track() - else: - dbt.tracking.initialize_tracking() - if colorize_output(profile_config): - dbt.ui.printer.use_colors() + initialize_config_values(parsed) reset_adapters() @@ -598,6 +610,7 @@ def parse_args(args): sys.exit(1) parsed = p.parse_args(args) + parsed.profiles_dir = os.path.expanduser(parsed.profiles_dir) if not hasattr(parsed, 'which'): # the user did not provide a valid subcommand. trigger the help message diff --git a/dbt/tracking.py b/dbt/tracking.py index 85834ad5e99..ce565aa95f5 100644 --- a/dbt/tracking.py +++ b/dbt/tracking.py @@ -17,8 +17,6 @@ COLLECTOR_URL = "fishtownanalytics.sinter-collect.com" COLLECTOR_PROTOCOL = "https" -COOKIE_PATH = os.path.join(os.path.expanduser('~'), '.dbt/.user.yml') - INVOCATION_SPEC = 'iglu:com.dbt/invocation/jsonschema/1-0-0' PLATFORM_SPEC = 'iglu:com.dbt/platform/jsonschema/1-0-0' RUN_MODEL_SPEC = 'iglu:com.dbt/run_model/jsonschema/1-0-0' @@ -35,8 +33,9 @@ class User(object): - def __init__(self): + def __init__(self, cookie_dir): self.do_not_track = True + self.cookie_dir = cookie_dir self.id = None self.invocation_id = str(uuid.uuid4()) @@ -45,6 +44,10 @@ def __init__(self): def state(self): return "do not track" if self.do_not_track else "tracking" + @property + def cookie_path(self): + return os.path.join(self.cookie_dir, '.user.yml') + def initialize(self): self.do_not_track = False @@ -56,21 +59,20 @@ def initialize(self): tracker.set_subject(subject) def set_cookie(self): - cookie_dir = os.path.dirname(COOKIE_PATH) user = {"id": str(uuid.uuid4())} - dbt.clients.system.make_directory(cookie_dir) + dbt.clients.system.make_directory(self.cookie_dir) - with open(COOKIE_PATH, "w") as fh: + with open(self.cookie_path, "w") as fh: yaml.dump(user, fh) return user def get_cookie(self): - if not os.path.isfile(COOKIE_PATH): + if not os.path.isfile(self.cookie_path): user = self.set_cookie() else: - with open(COOKIE_PATH, "r") as fh: + with open(self.cookie_path, "r") as fh: try: user = yaml.safe_load(fh) if user is None: @@ -266,10 +268,15 @@ def flush(): def do_not_track(): global active_user - active_user = User() + active_user = User(None) -def initialize_tracking(): +def initialize_tracking(cookie_dir): global active_user - active_user = User() - active_user.initialize() + active_user = User(cookie_dir) + try: + active_user.initialize() + except Exception: + logger.debug('Got an exception trying to initialize tracking', + exc_info=True) + active_user = User(None) diff --git a/test/unit/test_config.py b/test/unit/test_config.py index bf36698de84..cd2cfa1e381 100644 --- a/test/unit/test_config.py +++ b/test/unit/test_config.py @@ -62,60 +62,6 @@ def temp_cd(path): )) -class ConfigTest(unittest.TestCase): - def setUp(self): - self.base_dir = tempfile.mkdtemp() - self.profiles_path = os.path.join(self.base_dir, 'profiles.yml') - - def set_up_empty_config(self): - with open(self.profiles_path, 'w') as f: - f.write(yaml.dump({})) - - def set_up_config_options(self, **kwargs): - config = { - 'config': kwargs - } - - with open(self.profiles_path, 'w') as f: - f.write(yaml.dump(config)) - - def tearDown(self): - try: - shutil.rmtree(self.base_dir) - except: - pass - - def test__implicit_opt_in(self): - self.set_up_empty_config() - config = dbt.config.read_config(self.base_dir) - self.assertTrue(dbt.config.send_anonymous_usage_stats(config)) - - def test__explicit_opt_out(self): - self.set_up_config_options(send_anonymous_usage_stats=False) - config = dbt.config.read_config(self.base_dir) - self.assertFalse(dbt.config.send_anonymous_usage_stats(config)) - - def test__explicit_opt_in(self): - self.set_up_config_options(send_anonymous_usage_stats=True) - config = dbt.config.read_config(self.base_dir) - self.assertTrue(dbt.config.send_anonymous_usage_stats(config)) - - def test__implicit_colors(self): - self.set_up_empty_config() - config = dbt.config.read_config(self.base_dir) - self.assertTrue(dbt.config.colorize_output(config)) - - def test__explicit_opt_out(self): - self.set_up_config_options(use_colors=False) - config = dbt.config.read_config(self.base_dir) - self.assertFalse(dbt.config.colorize_output(config)) - - def test__explicit_opt_in(self): - self.set_up_config_options(use_colors=True) - config = dbt.config.read_config(self.base_dir) - self.assertTrue(dbt.config.colorize_output(config)) - - class Args(object): def __init__(self, profiles_dir=None, threads=None, profile=None, cli_vars=None): self.profile = profile diff --git a/test/unit/test_main.py b/test/unit/test_main.py new file mode 100644 index 00000000000..b05443e9562 --- /dev/null +++ b/test/unit/test_main.py @@ -0,0 +1,108 @@ +import os +import tempfile +import unittest + +import mock +import yaml + +from dbt import main +import dbt.tracking +import dbt.ui.printer + + +class FakeArgs(object): + def __init__(self, profiles_dir): + self.profiles_dir = profiles_dir + self.profile = 'test' + + +@mock.patch('dbt.ui.printer.use_colors') +@mock.patch('dbt.tracking.do_not_track') +@mock.patch('dbt.tracking.initialize_tracking') +class TestInitializeConfig(unittest.TestCase): + def setUp(self): + self.base_dir = tempfile.mkdtemp() + self.profiles_path = os.path.join(self.base_dir, 'profiles.yml') + self.args = FakeArgs(self.base_dir) + + def _base_config(self): + return { + 'test': { + 'outputs': { + 'default': { + 'type': 'postgres', + 'host': 'test', + 'port': 5555, + 'user': 'db_user', + 'pass': 'db_pass', + 'dbname': 'dbname', + 'schema': 'schema', + }, + }, + 'target': 'default', + } + } + + def set_up_empty_config(self): + with open(self.profiles_path, 'w') as f: + f.write(yaml.dump(self._base_config())) + + def set_up_config_options(self, **kwargs): + config = self._base_config() + config.update(config=kwargs) + + with open(self.profiles_path, 'w') as f: + f.write(yaml.dump(config)) + + def tearDown(self): + try: + shutil.rmtree(self.base_dir) + except: + pass + + def test__implicit_missing(self, initialize_tracking, do_not_track, use_colors): + main.initialize_config_values(self.args) + + initialize_tracking.assert_called_once_with(self.base_dir) + do_not_track.assert_not_called() + use_colors.assert_called_once_with() + + def test__implicit_opt_in_colors(self, initialize_tracking, do_not_track, use_colors): + self.set_up_empty_config() + main.initialize_config_values(self.args) + + initialize_tracking.assert_called_once_with(self.base_dir) + do_not_track.assert_not_called() + use_colors.assert_called_once_with() + + def test__explicit_opt_out(self, initialize_tracking, do_not_track, use_colors): + self.set_up_config_options(send_anonymous_usage_stats=False) + main.initialize_config_values(self.args) + + initialize_tracking.assert_not_called() + do_not_track.assert_called_once_with() + use_colors.assert_called_once_with() + + def test__explicit_opt_in(self, initialize_tracking, do_not_track, use_colors): + self.set_up_config_options(send_anonymous_usage_stats=True) + main.initialize_config_values(self.args) + + initialize_tracking.assert_called_once_with(self.base_dir) + do_not_track.assert_not_called() + use_colors.assert_called_once_with() + + def test__explicit_no_colors(self, initialize_tracking, do_not_track, use_colors): + self.set_up_config_options(use_colors=False) + main.initialize_config_values(self.args) + + initialize_tracking.assert_called_once_with(self.base_dir) + do_not_track.assert_not_called() + use_colors.assert_not_called() + + def test__explicit_opt_in(self, initialize_tracking, do_not_track, use_colors): + self.set_up_config_options(use_colors=True) + main.initialize_config_values(self.args) + + initialize_tracking.assert_called_once_with(self.base_dir) + do_not_track.assert_not_called() + use_colors.assert_called_once_with()