diff --git a/dbt/config.py b/dbt/config.py index af6e471167f..71a913c0849 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. @@ -476,13 +460,48 @@ def warn_for_unused_resource_config_paths(self, resource_fqns, disabled): logger.info(dbt.ui.printer.yellow(msg)) +class UserConfig(object): + def __init__(self, send_anonymous_usage_stats, use_colors): + self.send_anonymous_usage_stats = send_anonymous_usage_stats + self.use_colors = use_colors + + @classmethod + def from_dict(cls, cfg=None): + if cfg is None: + cfg = {} + send_anonymous_usage_stats = cfg.get( + 'send_anonymous_usage_stats', + DEFAULT_SEND_ANONYMOUS_USAGE_STATS + ) + use_colors = cfg.get( + 'use_colors', + DEFAULT_USE_COLORS + ) + return cls(send_anonymous_usage_stats, use_colors) + + def to_dict(self): + return { + 'send_anonymous_usage_stats': self.send_anonymous_usage_stats, + 'use_colors': self.use_colors, + } + + @classmethod + def from_directory(cls, directory): + user_cfg = None + profile = read_profile(directory) + if profile: + user_cfg = profile.get('config', {}) + return cls.from_dict(user_cfg) + + class Profile(object): - def __init__(self, profile_name, target_name, send_anonymous_usage_stats, - use_colors, threads, credentials): + def __init__(self, profile_name, target_name, config, threads, + credentials): self.profile_name = profile_name self.target_name = target_name - self.send_anonymous_usage_stats = send_anonymous_usage_stats - self.use_colors = use_colors + if isinstance(config, dict): + config = UserConfig.from_dict(config) + self.config = config self.threads = threads self.credentials = credentials @@ -498,8 +517,7 @@ def to_profile_info(self, serialize_credentials=False): result = { 'profile_name': self.profile_name, 'target_name': self.target_name, - 'send_anonymous_usage_stats': self.send_anonymous_usage_stats, - 'use_colors': self.use_colors, + 'config': self.config.to_dict(), 'threads': self.threads, 'credentials': self.credentials.incorporate(), } @@ -584,21 +602,11 @@ def from_credentials(cls, credentials, threads, profile_name, target_name, :raises DbtProfileError: If the profile is invalid. :returns Profile: The new Profile object. """ - if user_cfg is None: - user_cfg = {} - send_anonymous_usage_stats = user_cfg.get( - 'send_anonymous_usage_stats', - DEFAULT_SEND_ANONYMOUS_USAGE_STATS - ) - use_colors = user_cfg.get( - 'use_colors', - DEFAULT_USE_COLORS - ) + config = UserConfig.from_dict(user_cfg) profile = cls( profile_name=profile_name, target_name=target_name, - send_anonymous_usage_stats=send_anonymous_usage_stats, - use_colors=use_colors, + config=config, threads=threads, credentials=credentials ) @@ -745,10 +753,8 @@ def from_args(cls, args, project_profile_name=None, cli_vars=None): cli_vars = dbt.utils.parse_cli_vars(getattr(args, 'vars', '{}')) threads_override = getattr(args, 'threads', None) - # TODO(jeb): is it even possible for this to not be set? - profiles_dir = getattr(args, 'profiles_dir', PROFILES_DIR) target_override = getattr(args, 'target', None) - raw_profiles = read_profile(profiles_dir) + raw_profiles = read_profile(args.profiles_dir) profile_name = cls.pick_profile_name(args.profile, project_profile_name) @@ -797,9 +803,8 @@ def __init__(self, project_name, version, project_root, source_paths, macro_paths, data_paths, test_paths, analysis_paths, docs_paths, target_path, clean_targets, log_path, modules_path, quoting, models, on_run_start, on_run_end, - archive, seeds, profile_name, target_name, - send_anonymous_usage_stats, use_colors, threads, credentials, - packages, args): + archive, seeds, profile_name, target_name, config, + threads, credentials, packages, args): # 'vars' self.args = args self.cli_vars = dbt.utils.parse_cli_vars(getattr(args, 'vars', '{}')) @@ -833,8 +838,7 @@ def __init__(self, project_name, version, project_root, source_paths, self, profile_name=profile_name, target_name=target_name, - send_anonymous_usage_stats=send_anonymous_usage_stats, - use_colors=use_colors, + config=config, threads=threads, credentials=credentials ) @@ -877,8 +881,7 @@ def from_parts(cls, project, profile, args): packages=project.packages, profile_name=profile.profile_name, target_name=profile.target_name, - send_anonymous_usage_stats=profile.send_anonymous_usage_stats, - use_colors=profile.use_colors, + config=profile.config, threads=profile.threads, credentials=profile.credentials, args=args diff --git a/dbt/contracts/project.py b/dbt/contracts/project.py index f14e74784e3..ad34dadc8c3 100644 --- a/dbt/contracts/project.py +++ b/dbt/contracts/project.py @@ -284,6 +284,20 @@ class PackageConfig(APIObject): SCHEMA = PACKAGE_FILE_CONTRACT +USER_CONFIG_CONTRACT = { + 'type': 'object', + 'additionalProperties': True, + 'properties': { + 'send_anonymous_usage_stats': { + 'type': 'boolean', + }, + 'use_colors': { + 'type': 'boolean', + }, + }, +} + + PROFILE_INFO_CONTRACT = { 'type': 'object', 'additionalProperties': False, @@ -294,12 +308,7 @@ class PackageConfig(APIObject): 'target_name': { 'type': 'string', }, - 'send_anonymous_usage_stats': { - 'type': 'boolean', - }, - 'use_colors': { - 'type': 'boolean', - }, + 'config': USER_CONFIG_CONTRACT, 'threads': { 'type': 'number', }, @@ -313,8 +322,7 @@ class PackageConfig(APIObject): }, }, 'required': [ - 'profile_name', 'target_name', 'send_anonymous_usage_stats', - 'use_colors', 'threads', 'credentials' + 'profile_name', 'target_name', 'config', 'threads', 'credentials' ], } diff --git a/dbt/main.py b/dbt/main.py index c8d13c18eac..248745c9476 100644 --- a/dbt/main.py +++ b/dbt/main.py @@ -28,10 +28,9 @@ 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.exceptions import DbtProfileError, DbtProfileError, RuntimeException +from dbt.config import Project, UserConfig, RuntimeConfig, PROFILES_DIR, \ + read_profiles +from dbt.exceptions import DbtProjectError, DbtProfileError, RuntimeException PROFILES_HELP_MESSAGE = """ @@ -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: + cfg = UserConfig.from_directory(parsed.profiles_dir) + except RuntimeException: + cfg = UserConfig.from_dict(None) + + if cfg.send_anonymous_usage_stats: + dbt.tracking.initialize_tracking(parsed.profiles_dir) + else: + dbt.tracking.do_not_track() + + if cfg.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() @@ -599,6 +611,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/integration/base.py b/test/integration/base.py index 3f6492298c8..5b12d3b8c99 100644 --- a/test/integration/base.py +++ b/test/integration/base.py @@ -42,6 +42,7 @@ class TestArgs(object): def __init__(self, kwargs): self.which = 'run' self.single_threaded = False + self.profiles_dir = DBT_CONFIG_DIR self.__dict__.update(kwargs) diff --git a/test/unit/test_config.py b/test/unit/test_config.py index bf36698de84..156620df968 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 @@ -263,8 +209,8 @@ def test_from_raw_profiles(self): self.assertEqual(profile.profile_name, 'default') self.assertEqual(profile.target_name, 'postgres') self.assertEqual(profile.threads, 7) - self.assertTrue(profile.send_anonymous_usage_stats) - self.assertTrue(profile.use_colors) + self.assertTrue(profile.config.send_anonymous_usage_stats) + self.assertTrue(profile.config.use_colors) self.assertTrue(isinstance(profile.credentials, PostgresCredentials)) self.assertEqual(profile.credentials.type, 'postgres') self.assertEqual(profile.credentials.host, 'postgres-db-hostname') @@ -282,8 +228,8 @@ def test_config_override(self): profile = self.from_raw_profiles() self.assertEqual(profile.profile_name, 'default') self.assertEqual(profile.target_name, 'postgres') - self.assertFalse(profile.send_anonymous_usage_stats) - self.assertFalse(profile.use_colors) + self.assertFalse(profile.config.send_anonymous_usage_stats) + self.assertFalse(profile.config.use_colors) def test_partial_config_override(self): self.default_profile_data['config'] = { @@ -292,8 +238,8 @@ def test_partial_config_override(self): profile = self.from_raw_profiles() self.assertEqual(profile.profile_name, 'default') self.assertEqual(profile.target_name, 'postgres') - self.assertFalse(profile.send_anonymous_usage_stats) - self.assertTrue(profile.use_colors) + self.assertFalse(profile.config.send_anonymous_usage_stats) + self.assertTrue(profile.config.use_colors) def test_missing_type(self): del self.default_profile_data['default']['outputs']['postgres']['type'] @@ -418,8 +364,8 @@ def test_profile_simple(self): self.assertEqual(profile.profile_name, 'default') self.assertEqual(profile.target_name, 'postgres') self.assertEqual(profile.threads, 7) - self.assertTrue(profile.send_anonymous_usage_stats) - self.assertTrue(profile.use_colors) + self.assertTrue(profile.config.send_anonymous_usage_stats) + self.assertTrue(profile.config.use_colors) self.assertTrue(isinstance(profile.credentials, PostgresCredentials)) self.assertEqual(profile.credentials.type, 'postgres') self.assertEqual(profile.credentials.host, 'postgres-db-hostname') @@ -443,8 +389,8 @@ def test_profile_override(self): self.assertEqual(profile.profile_name, 'other') self.assertEqual(profile.target_name, 'other-postgres') self.assertEqual(profile.threads, 3) - self.assertTrue(profile.send_anonymous_usage_stats) - self.assertTrue(profile.use_colors) + self.assertTrue(profile.config.send_anonymous_usage_stats) + self.assertTrue(profile.config.use_colors) self.assertTrue(isinstance(profile.credentials, PostgresCredentials)) self.assertEqual(profile.credentials.type, 'postgres') self.assertEqual(profile.credentials.host, 'other-postgres-db-hostname') @@ -465,8 +411,8 @@ def test_target_override(self): self.assertEqual(profile.profile_name, 'default') self.assertEqual(profile.target_name, 'redshift') self.assertEqual(profile.threads, 1) - self.assertTrue(profile.send_anonymous_usage_stats) - self.assertTrue(profile.use_colors) + self.assertTrue(profile.config.send_anonymous_usage_stats) + self.assertTrue(profile.config.use_colors) self.assertTrue(isinstance(profile.credentials, RedshiftCredentials)) self.assertEqual(profile.credentials.type, 'redshift') self.assertEqual(profile.credentials.host, 'redshift-db-hostname') @@ -488,8 +434,8 @@ def test_env_vars(self): self.assertEqual(profile.profile_name, 'default') self.assertEqual(profile.target_name, 'with-vars') self.assertEqual(profile.threads, 1) - self.assertTrue(profile.send_anonymous_usage_stats) - self.assertTrue(profile.use_colors) + self.assertTrue(profile.config.send_anonymous_usage_stats) + self.assertTrue(profile.config.use_colors) self.assertEqual(profile.credentials.type, 'postgres') self.assertEqual(profile.credentials.host, 'env-postgres-host') self.assertEqual(profile.credentials.port, 6543) @@ -510,8 +456,8 @@ def test_env_vars_env_target(self): self.assertEqual(profile.profile_name, 'default') self.assertEqual(profile.target_name, 'with-vars') self.assertEqual(profile.threads, 1) - self.assertTrue(profile.send_anonymous_usage_stats) - self.assertTrue(profile.use_colors) + self.assertTrue(profile.config.send_anonymous_usage_stats) + self.assertTrue(profile.config.use_colors) self.assertEqual(profile.credentials.type, 'postgres') self.assertEqual(profile.credentials.host, 'env-postgres-host') self.assertEqual(profile.credentials.port, 6543) @@ -541,8 +487,8 @@ def test_cli_and_env_vars(self): self.assertEqual(profile.profile_name, 'default') self.assertEqual(profile.target_name, 'cli-and-env-vars') self.assertEqual(profile.threads, 1) - self.assertTrue(profile.send_anonymous_usage_stats) - self.assertTrue(profile.use_colors) + self.assertTrue(profile.config.send_anonymous_usage_stats) + self.assertTrue(profile.config.use_colors) self.assertEqual(profile.credentials.type, 'postgres') self.assertEqual(profile.credentials.host, 'cli-postgres-host') self.assertEqual(profile.credentials.port, 6543) @@ -1032,7 +978,7 @@ def test_validate_fails(self): project = self.get_project() profile = self.get_profile() # invalid - must be boolean - profile.use_colors = None + profile.config.use_colors = None with self.assertRaises(dbt.exceptions.DbtProjectError): dbt.config.RuntimeConfig.from_parts(project, 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()