Merge pull request #2708 from rsenseman/enhancement/color_output_comm…

…and_line_flag_second_attempt Add --use-colors cli option (second attempt)
dbt-labs · Aug 19, 2020 · 7ef7a8f · 7ef7a8f
2 parents 674bd8f + db8eea2
commit 7ef7a8f
Show file tree

Hide file tree

Showing 11 changed files with 123 additions and 32 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -16,6 +16,7 @@
 - Add support for impersonating a service account using `impersonate_service_account` in the BigQuery profile configuration ([#2677](https://github.com/fishtown-analytics/dbt/issues/2677)) ([docs](https://docs.getdbt.com/reference/warehouse-profiles/bigquery-profile#service-account-impersonation))
 - Macros in the current project can override internal dbt macros that are called through `execute_macros`. ([#2301](https://github.com/fishtown-analytics/dbt/issues/2301), [#2686](https://github.com/fishtown-analytics/dbt/pull/2686))
 - Add state:modified and state:new selectors ([#2641](https://github.com/fishtown-analytics/dbt/issues/2641), [#2695](https://github.com/fishtown-analytics/dbt/pull/2695))
+- Add two new flags `--use-colors` and `--no-use-colors` to `dbt run` command to enable or disable log colorization from the command line ([#2708](https://github.com/fishtown-analytics/dbt/pull/2708))
 
 ### Fixes
 - Fix Redshift table size estimation; e.g. 44 GB tables are no longer reported as 44 KB. [#2702](https://github.com/fishtown-analytics/dbt/issues/2702)
@@ -31,6 +32,8 @@ Contributors:
 - [@kconvey](https://github.com/kconvey) ([#2694](https://github.com/fishtown-analytics/dbt/pull/2694))
 - [@vogt4nick](https://github.com/vogt4nick) ([#2702](https://github.com/fishtown-analytics/dbt/issues/2702))
 - [@stephen8chang](https://github.com/stephen8chang) ([docs#106](https://github.com/fishtown-analytics/dbt-docs/pull/106), [docs#108](https://github.com/fishtown-analytics/dbt-docs/pull/108), [docs#113](https://github.com/fishtown-analytics/dbt-docs/pull/113))
+- [@rsenseman](https://github.com/rsenseman) ([#2708](https://github.com/fishtown-analytics/dbt/pull/2708))
+
 
 ## dbt 0.18.0b2 (July 30, 2020)
 

diff --git a/core/dbt/contracts/connection.py b/core/dbt/contracts/connection.py
@@ -160,7 +160,7 @@ def to_dict(self, omit_none=True, validate=False, *, with_aliases=False):
 
 class UserConfigContract(Protocol):
     send_anonymous_usage_stats: bool
-    use_colors: bool
+    use_colors: Optional[bool]
     partial_parse: Optional[bool]
     printer_width: Optional[int]
 

diff --git a/core/dbt/contracts/project.py b/core/dbt/contracts/project.py
@@ -14,7 +14,6 @@
 
 PIN_PACKAGE_URL = 'https://docs.getdbt.com/docs/package-management#section-specifying-package-versions' # noqa
 DEFAULT_SEND_ANONYMOUS_USAGE_STATS = True
-DEFAULT_USE_COLORS = True
 
 
 Name = NewType('Name', str)
@@ -259,7 +258,7 @@ def parse_project_config(
 @dataclass
 class UserConfig(ExtensibleJsonSchemaMixin, Replaceable, UserConfigContract):
     send_anonymous_usage_stats: bool = DEFAULT_SEND_ANONYMOUS_USAGE_STATS
-    use_colors: bool = DEFAULT_USE_COLORS
+    use_colors: Optional[bool] = None
     partial_parse: Optional[bool] = None
     printer_width: Optional[int] = None
 
@@ -269,8 +268,8 @@ def set_values(self, cookie_dir):
         else:
             tracking.do_not_track()
 
-        if self.use_colors:
-            ui.use_colors()
+        if self.use_colors is not None:
+            ui.use_colors(self.use_colors)
 
         if self.printer_width:
             ui.printer_width(self.printer_width)

diff --git a/core/dbt/flags.py b/core/dbt/flags.py
@@ -2,6 +2,7 @@
 import multiprocessing
 from pathlib import Path
 from typing import Optional
+
 # initially all flags are set to None, the on-load call of reset() will set
 # them for their first time.
 STRICT_MODE = None
@@ -11,6 +12,7 @@
 TEST_NEW_PARSER = None
 WRITE_JSON = None
 PARTIAL_PARSE = None
+USE_COLORS = None
 
 
 def env_set_truthy(key: str) -> Optional[str]:
@@ -48,7 +50,7 @@ def _get_context():
 
 def reset():
     global STRICT_MODE, FULL_REFRESH, USE_CACHE, WARN_ERROR, TEST_NEW_PARSER, \
-        WRITE_JSON, PARTIAL_PARSE, MP_CONTEXT
+        WRITE_JSON, PARTIAL_PARSE, MP_CONTEXT, USE_COLORS
 
     STRICT_MODE = False
     FULL_REFRESH = False
@@ -58,11 +60,12 @@ def reset():
     WRITE_JSON = True
     PARTIAL_PARSE = False
     MP_CONTEXT = _get_context()
+    USE_COLORS = True
 
 
 def set_from_args(args):
     global STRICT_MODE, FULL_REFRESH, USE_CACHE, WARN_ERROR, TEST_NEW_PARSER, \
-        WRITE_JSON, PARTIAL_PARSE, MP_CONTEXT
+        WRITE_JSON, PARTIAL_PARSE, MP_CONTEXT, USE_COLORS
 
     USE_CACHE = getattr(args, 'use_cache', USE_CACHE)
 
@@ -78,6 +81,13 @@ def set_from_args(args):
     PARTIAL_PARSE = getattr(args, 'partial_parse', None)
     MP_CONTEXT = _get_context()
 
+    # The use_colors attribute will always have a value because it is assigned
+    # None by default from the add_mutually_exclusive_group function
+    use_colors_override = getattr(args, 'use_colors')
+
+    if use_colors_override is not None:
+        USE_COLORS = use_colors_override
+
 
 # initialize everything to the defaults on module load
 reset()
diff --git a/core/dbt/main.py b/core/dbt/main.py
@@ -902,6 +902,30 @@ def parse_args(args, cls=DBTArgumentParser):
         If set, skip writing the manifest and run_results.json files to disk
         '''
     )
+    colors_flag = p.add_mutually_exclusive_group()
+    colors_flag.add_argument(
+        '--use-colors',
+        action='store_const',
+        const=True,
+        dest='use_colors',
+        help='''
+        Colorize the output DBT prints to the terminal. Output is colorized by
+        default and may also be set in a profile or at the command line.
+        Mutually exclusive with --no-use-colors
+        '''
+    )
+    colors_flag.add_argument(
+        '--no-use-colors',
+        action='store_const',
+        const=False,
+        dest='use_colors',
+        help='''
+        Do not colorize the output DBT prints to the terminal. Output is
+        colorized by default and may also be set in a profile or at the
+        command line.
+        Mutually exclusive with --use-colors
+        '''
+    )
 
     p.add_argument(
         '-S',

diff --git a/core/dbt/ui.py b/core/dbt/ui.py
@@ -1,3 +1,4 @@
+import dbt.flags as flags
 import textwrap
 from typing import Dict
 
@@ -11,8 +12,6 @@
 }
 
 
-USE_COLORS = False
-
 COLOR_FG_RED = COLORS['red']
 COLOR_FG_GREEN = COLORS['green']
 COLOR_FG_YELLOW = COLORS['yellow']
@@ -21,9 +20,8 @@
 PRINTER_WIDTH = 80
 
 
-def use_colors():
-    global USE_COLORS
-    USE_COLORS = True
+def use_colors(use_colors_val=True):
+    flags.USE_COLORS = use_colors_val
 
 
 def printer_width(printer_width):
@@ -32,7 +30,7 @@ def printer_width(printer_width):
 
 
 def color(text: str, color_code: str):
-    if USE_COLORS:
+    if flags.USE_COLORS:
         return "{}{}{}".format(color_code, text, COLOR_RESET_ALL)
     else:
         return text

diff --git a/test/integration/061_use_colors_tests/models/do_nothing_then_fail.sql b/test/integration/061_use_colors_tests/models/do_nothing_then_fail.sql
@@ -0,0 +1 @@
+select 1,
diff --git a/test/integration/061_use_colors_tests/test_no_use_colors.py b/test/integration/061_use_colors_tests/test_no_use_colors.py
@@ -0,0 +1,29 @@
+
+from test.integration.base import DBTIntegrationTest, use_profile
+import logging
+import re
+import sys
+
+class TestNoUseColors(DBTIntegrationTest):
+
+    @property
+    def project_config(self):
+        return {'config-version': 2}
+
+    @property
+    def schema(self):
+        return "use_colors_tests_061"
+
+    @property
+    def models(self):
+        return "models"
+
+    @use_profile('postgres')
+    def test_postgres_no_use_colors(self):
+        # pattern to match formatted log output
+        pattern = re.compile(r'\[31m.*|\[33m.*')
+
+        results, stdout = self.run_dbt_and_capture(args=['--no-use-colors', 'run'], expect_pass=False)
+
+        stdout_contains_formatting_characters = bool(pattern.search(stdout))
+        self.assertFalse(stdout_contains_formatting_characters)
diff --git a/test/integration/061_use_colors_tests/test_use_colors.py b/test/integration/061_use_colors_tests/test_use_colors.py
@@ -0,0 +1,29 @@
+
+from test.integration.base import DBTIntegrationTest, use_profile
+import logging
+import re
+import sys
+
+class TestUseColors(DBTIntegrationTest):
+
+    @property
+    def project_config(self):
+        return {'config-version': 2}
+
+    @property
+    def schema(self):
+        return "use_colors_tests_061"
+
+    @property
+    def models(self):
+        return "models"
+
+    @use_profile('postgres')
+    def test_postgres_use_colors(self):
+        # pattern to match formatted log output
+        pattern = re.compile(r'\[31m.*|\[33m.*')
+
+        results, stdout = self.run_dbt_and_capture(args=['--use-colors', 'run'], expect_pass=False)
+
+        stdout_contains_formatting_characters = bool(pattern.search(stdout))
+        self.assertTrue(stdout_contains_formatting_characters)
diff --git a/test/unit/test_config.py b/test/unit/test_config.py
@@ -241,7 +241,7 @@ def test_from_raw_profiles(self):
         self.assertEqual(profile.target_name, 'postgres')
         self.assertEqual(profile.threads, 7)
         self.assertTrue(profile.config.send_anonymous_usage_stats)
-        self.assertTrue(profile.config.use_colors)
+        self.assertIsNone(profile.config.use_colors)
         self.assertTrue(isinstance(profile.credentials, PostgresCredentials))
         self.assertEqual(profile.credentials.type, 'postgres')
         self.assertEqual(profile.credentials.host, 'postgres-db-hostname')
@@ -254,7 +254,7 @@ def test_from_raw_profiles(self):
     def test_config_override(self):
         self.default_profile_data['config'] = {
             'send_anonymous_usage_stats': False,
-            'use_colors': False
+            'use_colors': False,
         }
         profile = self.from_raw_profiles()
         self.assertEqual(profile.profile_name, 'default')
@@ -271,7 +271,7 @@ def test_partial_config_override(self):
         self.assertEqual(profile.profile_name, 'default')
         self.assertEqual(profile.target_name, 'postgres')
         self.assertFalse(profile.config.send_anonymous_usage_stats)
-        self.assertTrue(profile.config.use_colors)
+        self.assertIsNone(profile.config.use_colors)
         self.assertEqual(profile.config.printer_width, 60)
 
     def test_missing_type(self):
@@ -404,7 +404,7 @@ def test_profile_simple(self):
         self.assertEqual(profile.target_name, 'postgres')
         self.assertEqual(profile.threads, 7)
         self.assertTrue(profile.config.send_anonymous_usage_stats)
-        self.assertTrue(profile.config.use_colors)
+        self.assertIsNone(profile.config.use_colors)
         self.assertTrue(isinstance(profile.credentials, PostgresCredentials))
         self.assertEqual(profile.credentials.type, 'postgres')
         self.assertEqual(profile.credentials.host, 'postgres-db-hostname')
@@ -429,7 +429,7 @@ def test_profile_override(self):
         self.assertEqual(profile.target_name, 'other-postgres')
         self.assertEqual(profile.threads, 3)
         self.assertTrue(profile.config.send_anonymous_usage_stats)
-        self.assertTrue(profile.config.use_colors)
+        self.assertIsNone(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')
@@ -451,7 +451,7 @@ def test_target_override(self):
         self.assertEqual(profile.target_name, 'redshift')
         self.assertEqual(profile.threads, 1)
         self.assertTrue(profile.config.send_anonymous_usage_stats)
-        self.assertTrue(profile.config.use_colors)
+        self.assertIsNone(profile.config.use_colors)
         self.assertTrue(isinstance(profile.credentials, RedshiftCredentials))
         self.assertEqual(profile.credentials.type, 'redshift')
         self.assertEqual(profile.credentials.host, 'redshift-db-hostname')
@@ -460,7 +460,7 @@ def test_target_override(self):
         self.assertEqual(profile.credentials.password, 'db_pass')
         self.assertEqual(profile.credentials.schema, 'redshift-schema')
         self.assertEqual(profile.credentials.database, 'redshift-db-name')
-        self.assertEqual(profile, from_raw)
+        self.assertEqual(profile, from_raw)        
 
     def test_env_vars(self):
         self.args.target = 'with-vars'
@@ -474,7 +474,7 @@ def test_env_vars(self):
         self.assertEqual(profile.target_name, 'with-vars')
         self.assertEqual(profile.threads, 1)
         self.assertTrue(profile.config.send_anonymous_usage_stats)
-        self.assertTrue(profile.config.use_colors)
+        self.assertIsNone(profile.config.use_colors)
         self.assertEqual(profile.credentials.type, 'postgres')
         self.assertEqual(profile.credentials.host, 'env-postgres-host')
         self.assertEqual(profile.credentials.port, 6543)
@@ -496,7 +496,7 @@ def test_env_vars_env_target(self):
         self.assertEqual(profile.target_name, 'with-vars')
         self.assertEqual(profile.threads, 1)
         self.assertTrue(profile.config.send_anonymous_usage_stats)
-        self.assertTrue(profile.config.use_colors)
+        self.assertIsNone(profile.config.use_colors)
         self.assertEqual(profile.credentials.type, 'postgres')
         self.assertEqual(profile.credentials.host, 'env-postgres-host')
         self.assertEqual(profile.credentials.port, 6543)
@@ -528,7 +528,7 @@ def test_cli_and_env_vars(self):
         self.assertEqual(profile.target_name, 'cli-and-env-vars')
         self.assertEqual(profile.threads, 1)
         self.assertTrue(profile.config.send_anonymous_usage_stats)
-        self.assertTrue(profile.config.use_colors)
+        self.assertIsNone(profile.config.use_colors)
         self.assertEqual(profile.credentials.type, 'postgres')
         self.assertEqual(profile.credentials.host, 'cli-postgres-host')
         self.assertEqual(profile.credentials.port, 6543)

diff --git a/test/unit/test_main.py b/test/unit/test_main.py
@@ -7,7 +7,6 @@
 import yaml
 
 from dbt import main
-import dbt.tracking
 
 
 class FakeArgs:
@@ -50,7 +49,6 @@ def set_up_empty_config(self):
     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))
 
@@ -65,44 +63,44 @@ def test__implicit_missing(self, initialize_tracking, do_not_track, use_colors):
 
         initialize_tracking.assert_called_once_with(self.base_dir)
         do_not_track.assert_not_called()
-        use_colors.assert_called_once_with()
+        use_colors.assert_not_called()
 
     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()
+        use_colors.assert_not_called()
 
     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()
+        use_colors.assert_not_called()
 
     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()
+        use_colors.assert_not_called()
 
     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()
+        use_colors.assert_called_once_with(False)
 
-    def test__explicit_opt_in(self, initialize_tracking, do_not_track, use_colors):
+    def test__explicit_yes_colors(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()
+        use_colors.assert_called_once_with(True)