Append query comment

CHANGELOG.md

@@ -3,6 +3,9 @@
 ### Fixes
 - When a jinja value is undefined, give a helpful error instead of failing with cryptic "cannot pickle ParserMacroCapture" errors ([#2110](https://github.com/fishtown-analytics/dbt/issues/2110), [#2184](https://github.com/fishtown-analytics/dbt/pull/2184))
 
+### Features
+- Support for appending query comments to SQL queries. ([#2199](https://github.com/fishtown-analytics/dbt/pull/2199))
+
 ## dbt 0.16.0rc2 (March 4, 2020)
 
 ### Under the hood

core/dbt/adapters/base/query_headers.py

@@ -4,31 +4,10 @@
 from dbt.clients.jinja import QueryStringGenerator
 
 from dbt.context.configured import generate_query_header_context
-from dbt.contracts.connection import AdapterRequiredConfig
+from dbt.contracts.connection import AdapterRequiredConfig, QueryComment
 from dbt.contracts.graph.compiled import CompileResultNode
 from dbt.contracts.graph.manifest import Manifest
 from dbt.exceptions import RuntimeException
-from dbt.helper_types import NoValue
-
-
-DEFAULT_QUERY_COMMENT = '''
-{%- set comment_dict = {} -%}
-{%- do comment_dict.update(
-    app='dbt',
-    dbt_version=dbt_version,
-    profile_name=target.get('profile_name'),
-    target_name=target.get('target_name'),
-) -%}
-{%- if node is not none -%}
-  {%- do comment_dict.update(
-    node_id=node.unique_id,
-  ) -%}
-{% else %}
-  {# in the node context, the connection name is the node_id #}
-  {%- do comment_dict.update(connection_name=connection_name) -%}
-{%- endif -%}
-{{ return(tojson(comment_dict)) }}
-'''
 
 
 class NodeWrapper:
@@ -47,21 +26,32 @@ class _QueryComment(local):
     """
     def __init__(self, initial):
         self.query_comment: Optional[str] = initial
+        self.append = False
 
     def add(self, sql: str) -> str:
         if not self.query_comment:
             return sql
-        else:
-            return '/* {} */\n{}'.format(self.query_comment.strip(), sql)
 
-    def set(self, comment: Optional[str]):
+        if self.append:
+            # replace last ';' with '<comment>;'
+            sql = sql.rstrip()
+            if sql[-1] == ';':
+                sql = sql[:-1]
+                return '{}\n/* {} */;'.format(sql, self.query_comment.strip())
+
+            return '{}\n/* {} */'.format(sql, self.query_comment.strip())
+
+        return '/* {} */\n{}'.format(self.query_comment.strip(), sql)
+
+    def set(self, comment: Optional[str], append: bool):
         if isinstance(comment, str) and '*/' in comment:
             # tell the user "no" so they don't hurt themselves by writing
             # garbage
             raise RuntimeException(
                 f'query comment contains illegal value "*/": {comment}'
             )
         self.query_comment = comment
+        self.append = append
 
 
 QueryStringFunc = Callable[[str, Optional[NodeWrapper]], str]
@@ -87,18 +77,26 @@ def __init__(self, config: AdapterRequiredConfig, manifest: Manifest):
         self.comment = _QueryComment(None)
         self.reset()
 
-    def _get_comment_macro(self):
-        if (
-            self.config.query_comment != NoValue() and
-            self.config.query_comment
-        ):
+    def _get_comment_macro(self) -> Optional[str]:
+        default_comment = QueryComment().comment
+        if not self.config.query_comment:
+            return default_comment
+
+        if isinstance(self.config.query_comment, str):
+            if self.config.query_comment in ('None', ''):
+                # query comment is disabled.
+                return None
             return self.config.query_comment
-        # if the query comment is null/empty string, there is no comment at all
-        elif not self.config.query_comment:
+
+        comment = self.config.query_comment.comment
+
+        if comment in ('None', ''):
             return None
-        else:
-            # else, the default
-            return DEFAULT_QUERY_COMMENT
+
+        if not comment:
+            return default_comment
+
+        return comment
 
     def _get_context(self) -> Dict[str, Any]:
         return generate_query_header_context(self.config, self.manifest)
@@ -114,4 +112,8 @@ def set(self, name: str, node: Optional[CompileResultNode]):
         if node is not None:
             wrapped = NodeWrapper(node)
         comment_str = self.generator(name, wrapped)
-        self.comment.set(comment_str)
+
+        append = False
+        if isinstance(self.config.query_comment, QueryComment):
+            append = self.config.query_comment.append
+        self.comment.set(comment_str, append)

core/dbt/config/project.py

@@ -9,12 +9,12 @@
 from dbt.clients.system import path_exists
 from dbt.clients.system import load_file_contents
 from dbt.clients.yaml_helper import load_yaml_text
+from dbt.contracts.connection import QueryComment
 from dbt.exceptions import DbtProjectError
 from dbt.exceptions import RecursionException
 from dbt.exceptions import SemverException
 from dbt.exceptions import validator_error_message
 from dbt.exceptions import warn_or_error
-from dbt.helper_types import NoValue
 from dbt.semver import VersionSpecifier
 from dbt.semver import versions_compatible
 from dbt.version import get_installed_version
@@ -202,6 +202,17 @@ def _raw_project_from(project_root: str) -> Dict[str, Any]:
     return project_dict
 
 
+def _query_comment_from_cfg(
+        cfg_query_comment: Union[QueryComment, str]
+) -> QueryComment:
+    if isinstance(cfg_query_comment, str):
+        if cfg_query_comment in ('None', ''):
+            return QueryComment(comment='')
+        return QueryComment(comment=cfg_query_comment)
+
+    return cfg_query_comment
+
+
 @dataclass
 class PartialProject:
     profile_name: Optional[str]
@@ -244,7 +255,7 @@ class Project:
     snapshots: Dict[str, Any]
     dbt_version: List[VersionSpecifier]
     packages: Dict[str, Any]
-    query_comment: Optional[Union[str, NoValue]]
+    query_comment: Optional[Union[QueryComment, str]]
 
     @property
     def all_source_paths(self) -> List[str]:
@@ -356,7 +367,10 @@ def from_project_config(
         dbt_raw_version: Union[List[str], str] = '>=0.0.0'
         if cfg.require_dbt_version is not None:
             dbt_raw_version = cfg.require_dbt_version
-        query_comment = cfg.query_comment
+
+        query_comment = None
+        if cfg.query_comment is not None:
+            query_comment = _query_comment_from_cfg(cfg.query_comment)
 
         try:
             dbt_version = _parse_versions(dbt_raw_version)
@@ -442,10 +456,11 @@ def to_project_config(self, with_packages=False):
                 v.to_version_string() for v in self.dbt_version
             ],
         })
+        if self.query_comment is not None:
+            result['query-comment'] = self.query_comment.to_dict()
+
         if with_packages:
             result.update(self.packages.to_dict())
-        if self.query_comment != NoValue():
-            result['query-comment'] = self.query_comment
 
         return result
 

core/dbt/contracts/connection.py

@@ -3,8 +3,7 @@
 from dataclasses import dataclass, field
 from typing import (
     Any, ClassVar, Dict, Tuple, Iterable, Optional, NewType, List, Callable,
-    Union
-)
+    Union)
 from typing_extensions import Protocol
 
 from hologram import JsonSchemaMixin
@@ -14,7 +13,6 @@
 
 from dbt.contracts.util import Replaceable
 from dbt.exceptions import InternalException
-from dbt.helper_types import NoValue
 from dbt.utils import translate_aliases
 
 
@@ -175,8 +173,34 @@ class HasCredentials(Protocol):
     threads: int
 
 
+DEFAULT_QUERY_COMMENT = '''
+{%- set comment_dict = {} -%}
+{%- do comment_dict.update(
+    app='dbt',
+    dbt_version=dbt_version,
+    profile_name=target.get('profile_name'),
+    target_name=target.get('target_name'),
+) -%}
+{%- if node is not none -%}
+  {%- do comment_dict.update(
+    node_id=node.unique_id,
+  ) -%}
+{% else %}
+  {# in the node context, the connection name is the node_id #}
+  {%- do comment_dict.update(connection_name=connection_name) -%}
+{%- endif -%}
+{{ return(tojson(comment_dict)) }}
+'''
+
+
+@dataclass
+class QueryComment(JsonSchemaMixin):
+    comment: str = DEFAULT_QUERY_COMMENT
+    append: bool = False
+
+
 class AdapterRequiredConfig(HasCredentials, Protocol):
     project_name: str
-    query_comment: Optional[Union[str, NoValue]]
+    query_comment: Optional[Union[QueryComment, str]]
     cli_vars: Dict[str, Any]
     target_path: str

core/dbt/contracts/project.py

@@ -1,9 +1,8 @@
 from dbt.contracts.util import Replaceable, Mergeable, list_str
-from dbt.contracts.connection import UserConfigContract
+from dbt.contracts.connection import UserConfigContract, QueryComment
 from dbt.logger import GLOBAL_LOGGER as logger  # noqa
 from dbt import tracking
 from dbt.ui import printer
-from dbt.helper_types import NoValue
 
 from hologram import JsonSchemaMixin, ValidationError
 from hologram.helpers import HyphenatedJsonSchemaMixin, register_pattern, \
@@ -162,7 +161,7 @@ class Project(HyphenatedJsonSchemaMixin, Replaceable):
     seeds: Dict[str, Any] = field(default_factory=dict)
     snapshots: Dict[str, Any] = field(default_factory=dict)
     packages: List[PackageSpec] = field(default_factory=list)
-    query_comment: Optional[Union[str, NoValue]] = NoValue()
+    query_comment: Optional[Union[QueryComment, str]] = None
 
     @classmethod
     def from_dict(cls, data, validate=True):

test/integration/051_query_comments_test/test_query_comments.py

@@ -167,7 +167,7 @@ class TestNullQueryComments(TestDefaultQueryComments):
     @property
     def project_config(self):
         cfg = super().project_config
-        cfg.update({'query-comment': None})
+        cfg.update({'query-comment': 'None'})
         return cfg
 
     def matches_comment(self, msg) -> bool:

test/unit/test_config.py

@@ -14,12 +14,12 @@
 from dbt.adapters.postgres import PostgresCredentials
 from dbt.adapters.redshift import RedshiftCredentials
 from dbt.context.base import generate_base_context
+from dbt.contracts.connection import QueryComment, DEFAULT_QUERY_COMMENT
 from dbt.contracts.project import PackageConfig, LocalPackage, GitPackage
 from dbt.semver import VersionSpecifier
 from dbt.task.run_operation import RunOperationTask
 
-from .utils import normalize
-
+from .utils import normalize, config_from_parts_or_dicts
 
 INITIAL_ROOT = os.getcwd()
 
@@ -844,6 +844,45 @@ def test_cycle(self):
                 self.default_project_data, None
             )
 
+    def test_query_comment_disabled(self):
+        self.default_project_data.update({
+            'query-comment': 'None',
+        })
+        project = dbt.config.Project.from_project_config(self.default_project_data, None)
+        self.assertEqual(project.query_comment.comment, '')
+        self.assertEqual(project.query_comment.append, False)
+
+        self.default_project_data.update({
+            'query-comment': '',
+        })
+        project = dbt.config.Project.from_project_config(self.default_project_data, None)
+        self.assertEqual(project.query_comment.comment, '')
+        self.assertEqual(project.query_comment.append, False)
+
+    def test_default_query_comment(self):
+        project = dbt.config.Project.from_project_config(self.default_project_data, None)
+        self.assertEqual(project.query_comment, None)
+
+    def test_default_query_comment_append(self):
+        self.default_project_data.update({
+            'query-comment': {
+                'append': True
+            },
+        })
+        project = dbt.config.Project.from_project_config(self.default_project_data, None)
+        self.assertEqual(project.query_comment.comment, DEFAULT_QUERY_COMMENT)
+        self.assertEqual(project.query_comment.append, True)
+
+    def test_custom_query_comment_append(self):
+        self.default_project_data.update({
+            'query-comment': {
+                'comment': 'run by user test',
+                'append': True
+            },
+        })
+        project = dbt.config.Project.from_project_config(self.default_project_data, None)
+        self.assertEqual(project.query_comment.comment, 'run by user test')
+        self.assertEqual(project.query_comment.append, True)
 
 class TestProjectWithConfigs(BaseConfigTest):
     def setUp(self):

test/unit/test_query_headers.py

@@ -0,0 +1,58 @@
+import re
+from unittest import TestCase, mock
+
+from dbt.adapters.base.query_headers import MacroQueryStringSetter
+
+from test.unit.utils import config_from_parts_or_dicts
+
+
+class TestQueryHeaders(TestCase):
+
+    def setUp(self):
+        self.profile_cfg = {
+            'outputs': {
+                'test': {
+                    'type': 'postgres',
+                    'dbname': 'postgres',
+                    'user': 'test',
+                    'host': 'test',
+                    'pass': 'test',
+                    'port': 5432,
+                    'schema': 'test'
+                },
+            },
+            'target': 'test'
+        }
+        self.project_cfg = {
+            'name': 'query_headers',
+            'version': '0.1',
+            'profile': 'test',
+        }
+        self.query = "SELECT 1;"
+
+    def test_comment_should_prepend_query_by_default(self):
+        config = config_from_parts_or_dicts(self.project_cfg, self.profile_cfg)
+        query_header = MacroQueryStringSetter(config, mock.MagicMock(macros={}))
+        sql = query_header.add(self.query)
+        self.assertTrue(re.match(f'^\/\*.*\*\/\n{self.query}$', sql))
+
+
+    def test_append_comment(self):
+        self.project_cfg.update({
+            'query-comment': {
+                'comment': 'executed by dbt',
+                'append': True
+            }
+        })
+        config = config_from_parts_or_dicts(self.project_cfg, self.profile_cfg)
+        query_header = MacroQueryStringSetter(config, mock.MagicMock(macros={}))
+        sql = query_header.add(self.query)
+        self.assertEqual(sql, f'{self.query[:-1]}\n/* executed by dbt */;')
+
+    def test_disable_query_comment(self):
+        self.project_cfg.update({
+            'query-comment': ''
+        })
+        config = config_from_parts_or_dicts(self.project_cfg, self.profile_cfg)
+        query_header = MacroQueryStringSetter(config, mock.MagicMock(macros={}))
+        self.assertEqual(query_header.add(self.query), self.query)