colorized dbt output

dbt/config.py

@@ -25,14 +25,8 @@ def read_profile(profiles_dir):
 
 def read_config(profiles_dir):
     profile = read_profile(profiles_dir)
-    return profile.get('config')
+    return profile.get('config', {})
 
 
-def send_anonymous_usage_stats(profiles_dir):
-    config = read_config(profiles_dir)
-
-    if config is not None \
-       and not config.get("send_anonymous_usage_stats", True):
-        return False
-
-    return True
+def send_anonymous_usage_stats(config):
+    return config.get('send_anonymous_usage_stats', True)

dbt/logger.py

@@ -3,15 +3,35 @@
 import os
 import sys
 
+import colorama
+
 # disable logs from other modules, excepting CRITICAL logs
 logging.getLogger('botocore').setLevel(logging.CRITICAL)
 logging.getLogger('contracts').setLevel(logging.CRITICAL)
 logging.getLogger('requests').setLevel(logging.CRITICAL)
 logging.getLogger('urllib3').setLevel(logging.CRITICAL)
 logging.getLogger('snowflake.connector').setLevel(logging.CRITICAL)
 
+
+# Colorama needs some help on windows because we're using logger.info
+# intead of print(). If the Windows env doesn't have a TERM var set,
+# then we should override the logging stream to use the colorama
+# converter. If the TERM var is set (as with Git Bash), then it's safe
+# to send escape characters and no log handler injection is needed.
+colorama_stdout = sys.stdout
+colorama_wrap = True
+
+if sys.platform == 'win32' and not os.environ.get('TERM'):
+    colorama_wrap = False
+    colorama_stdout = colorama.AnsiToWin32(sys.stdout).stream
+
+elif sys.platform == 'win32':
+    colorama_wrap = False
+
+colorama.init(wrap=colorama_wrap)
+
 # create a global console logger for dbt
-stdout_handler = logging.StreamHandler(sys.stdout)
+stdout_handler = logging.StreamHandler(colorama_stdout)
 stdout_handler.setFormatter(logging.Formatter('%(message)s'))
 stdout_handler.setLevel(logging.INFO)
 

dbt/main.py

@@ -20,6 +20,7 @@
 import dbt.config as config
 import dbt.adapters.cache as adapter_cache
 
+import dbt.printer
 
 def main(args=None):
     if args is None:
@@ -39,11 +40,15 @@ def handle(args):
 
     # this needs to happen after args are parsed so we can determine the
     # correct profiles.yml file
-    if not config.send_anonymous_usage_stats(parsed.profiles_dir):
+    profile_config = config.read_config(parsed.profiles_dir)
+    if not config.send_anonymous_usage_stats(profile_config):
         dbt.tracking.do_not_track()
     else:
         dbt.tracking.initialize_tracking()
 
+    if profile_config.get('use_colors', True):
+        dbt.printer.use_colors()
+
     res = run_from_args(parsed)
     dbt.tracking.flush()
 

dbt/printer.py

@@ -0,0 +1,214 @@
+
+from dbt.logger import GLOBAL_LOGGER as logger
+from dbt.utils import get_materialization, NodeType
+
+import colorama
+import time
+
+USE_COLORS = False
+
+def use_colors():
+    global USE_COLORS
+    USE_COLORS = True
+
+def get_timestamp():
+    return time.strftime("%H:%M:%S")
+
+
+def color(text, color_code):
+    if USE_COLORS:
+        return "{}{}{}".format(color_code, text, colorama.Style.RESET_ALL)
+    else:
+        return text
+
+
+def green(text):
+    return color(text, colorama.Fore.GREEN)
+
+
+def yellow(text):
+    return color(text, colorama.Fore.YELLOW)
+
+
+def red(text):
+    return color(text, colorama.Fore.RED)
+
+
+def print_timestamped_line(msg):
+    logger.info("{} | {}".format(get_timestamp(), msg))
+
+
+def print_fancy_output_line(msg, status, index, total, execution_time=None):
+    prefix = "{timestamp} | {index} of {total} {message}".format(
+        timestamp=get_timestamp(),
+        index=index,
+        total=total,
+        message=msg)
+
+    justified = prefix.ljust(80, ".")
+
+    if execution_time is None:
+        status_time = ""
+    else:
+        status_time = " in {execution_time:0.2f}s".format(
+            execution_time=execution_time)
+
+    status_txt = status
+
+    output = "{justified} [{status}{status_time}]".format(
+        justified=justified, status=status_txt, status_time=status_time)
+
+    logger.info(output)
+
+
+def print_skip_line(model, schema, relation, index, num_models):
+    msg = 'SKIP relation {}.{}'.format(schema, relation)
+    print_fancy_output_line(msg, yellow('SKIP'), index, num_models)
+
+
+def get_counts(flat_nodes):
+    counts = {}
+
+    for node in flat_nodes:
+        t = node.get('resource_type')
+
+        if node.get('resource_type') == NodeType.Model:
+            t = '{} {}'.format(get_materialization(node), t)
+
+        counts[t] = counts.get(t, 0) + 1
+
+    stat_line = ", ".join(
+        ["{} {}s".format(v, k) for k, v in counts.items()])
+
+    return stat_line
+
+
+def print_test_start_line(model, schema_name, index, total):
+    msg = "START test {name}".format(
+        name=model.get('name'))
+
+    run = 'RUN'
+    print_fancy_output_line(msg, run, index, total)
+
+
+def print_model_start_line(model, schema_name, index, total):
+    msg = "START {model_type} model {schema}.{relation}".format(
+        model_type=get_materialization(model),
+        schema=schema_name,
+        relation=model.get('name'))
+
+    run = 'RUN'
+    print_fancy_output_line(msg, run, index, total)
+
+
+def print_archive_start_line(model, index, total):
+    cfg = model.get('config', {})
+    msg = "START archive {source_schema}.{source_table} --> "\
+          "{target_schema}.{target_table}".format(**cfg)
+
+    run = 'RUN'
+    print_fancy_output_line(msg, run, index, total)
+
+
+def print_test_result_line(result, schema_name, index, total):
+    model = result.node
+    info = 'PASS'
+
+    if result.errored:
+        info = "ERROR"
+        color = red
+
+    elif result.status > 0:
+        info = 'FAIL {}'.format(result.status)
+        color = red
+
+        result.fail = True
+    elif result.status == 0:
+        info = 'PASS'
+        color = green
+
+    else:
+        raise RuntimeError("unexpected status: {}".format(result.status))
+
+    print_fancy_output_line(
+        "{info} {name}".format(info=info, name=model.get('name')),
+        color(info),
+        index,
+        total,
+        result.execution_time)
+
+
+def get_printable_result(result, success, error):
+    if result.errored:
+        info = 'ERROR {}'.format(error)
+        status = red(result.status)
+    else:
+        info = 'OK {}'.format(success)
+        status = green(result.status)
+
+    return info, status
+
+
+def print_archive_result_line(result, index, total):
+    model = result.node
+
+    info, status = get_printable_result(result, 'archived', 'archiving')
+    cfg = model.get('config', {})
+
+    print_fancy_output_line(
+        "{info} {source_schema}.{source_table} --> "
+        "{target_schema}.{target_table}".format(info=info, **cfg),
+        status,
+        index,
+        total,
+        result.execution_time)
+
+
+def print_model_result_line(result, schema_name, index, total):
+    model = result.node
+
+    info, status = get_printable_result(result, 'created', 'creating')
+
+    print_fancy_output_line(
+        "{info} {model_type} model {schema}.{relation}".format(
+            info=info,
+            model_type=get_materialization(model),
+            schema=schema_name,
+            relation=model.get('name')),
+        status,
+        index,
+        total,
+        result.execution_time)
+
+
+def interpret_run_result(result):
+    if result.errored or result.failed:
+        return 'error'
+    elif result.skipped:
+        return 'skip'
+    else:
+        return 'pass'
+
+
+def get_run_status_line(results):
+    stats = {
+        'error': 0,
+        'skip': 0,
+        'pass': 0,
+        'total': 0,
+    }
+
+    for r in results:
+        result_type = interpret_run_result(r)
+        stats[result_type] += 1
+        stats['total'] += 1
+
+    if stats['error'] == 0:
+        message = green('Completed successfully')
+    else:
+        message = red('Completed with errors')
+
+    stats_line = "Done. PASS={pass} ERROR={error} SKIP={skip} TOTAL={total}"
+    stats_line = stats_line.format(**stats)
+
+    return "{}\n{}".format(message, stats_line)