Ct 1919/get column schema from query macro core

.changes/unreleased/Features-20221118-141120.yaml

@@ -0,0 +1,8 @@
+kind: Features
+body: Data type constraints are now native to SQL table materializations. Enforce
+  columns are specific data types and not null depending on database functionality.
+time: 2022-11-18T14:11:20.868062-08:00
+custom:
+  Author: sungchun12
+  Issue: "6079"
+  PR: "6271"

core/dbt/adapters/base/impl.py

@@ -268,6 +268,22 @@ def execute(
         """
         return self.connections.execute(sql=sql, auto_begin=auto_begin, fetch=fetch)
 
+    @available.parse(lambda *a, **k: [])
+    def get_column_schema_from_query(
+        self, sql: str, auto_begin: bool = False, fetch: bool = False
+    ) -> Tuple[AdapterResponse, agate.Table]:
+        """Execute the given SQL. This is a thin wrapper around
+        ConnectionManager.execute.
+
+        :param str sql: The sql to execute.
+        :param bool auto_begin: If set, and dbt is not currently inside a
+            transaction, automatically begin one.
+        :param bool fetch: If set, fetch results.
+        :return: A tuple of the query status and results (empty if fetch=False).
+        :rtype: List[(column_name: str, data_type: str]
+        """
+        return self.connections.get_column_schema_from_query(sql=sql)
+
     @available.parse(lambda *a, **k: ("", empty_table()))
     def get_partitions_metadata(self, table: str) -> Tuple[agate.Table]:
         """Obtain partitions metadata for a BigQuery partitioned table.

core/dbt/adapters/sql/connections.py

@@ -128,6 +128,31 @@ def get_result_from_cursor(cls, cursor: Any) -> agate.Table:
 
         return dbt.clients.agate_helper.table_from_data_flat(data, column_names)
 
+    @classmethod
+    def data_type_code_to_name(cls, int) -> str:
+        """Get the string representation of the data type from the type_code."""
+        # https://peps.python.org/pep-0249/#type-objects
+        raise dbt.exceptions.NotImplementedError(
+            "`data_type_code_to_name` is not implemented for this adapter!"
+        )
+
+    @classmethod
+    def get_column_schema_from_cursor(cls, cursor: Any) -> List[Tuple[str, str]]:
+        # (column_name, data_type)
+        columns: List[Tuple[str, str]] = []
+
+        if cursor.description is not None:
+            # https://peps.python.org/pep-0249/#description
+            columns = [
+                # TODO: ignoring size, precision, scale for now
+                # (though it is part of DB-API standard, and our Column class does have these attributes)
+                # IMO user-defined contracts shouldn't have to match an exact size/precision/scale
+                (col[0], cls.data_type_code_to_name(col[1]))
+                for col in cursor.description
+            ]
+
+        return columns
+
     def execute(
         self, sql: str, auto_begin: bool = False, fetch: bool = False
     ) -> Tuple[AdapterResponse, agate.Table]:
@@ -140,6 +165,20 @@ def execute(
             table = dbt.clients.agate_helper.empty_table()
         return response, table
 
+    # TODO: do we need to care about auto_begin here?
+    def get_column_schema_from_query(self, sql: str) -> List[Tuple[str, str]]:
+        sql = self._add_query_comment(sql)
+        _, cursor = self.add_query(sql)
+        return self.get_column_schema_from_cursor(cursor)
+
+    # For dbt-bigquery
+    # def get_column_schema_from_query(cls, sql: str) -> List[Tuple[str, str]]:
+    #    sql = self._add_query_comment(sql)
+    #    # auto_begin is ignored on bigquery, and only included for consistency
+    #    query_job, iterator = self.raw_execute(sql)
+    #    columns = [(field.name, field.field_type) for field in resp.iterator]
+    #    return columns
+
     def add_begin_query(self):
         return self.add_query("BEGIN", auto_begin=False)
 

core/dbt/contracts/graph/model_config.py

@@ -446,6 +446,7 @@ class NodeConfig(NodeAndTestConfig):
         default_factory=Docs,
         metadata=MergeBehavior.Update.meta(),
     )
+    constraints_enabled: Optional[bool] = False
 
     # we validate that node_color has a suitable value to prevent dbt-docs from crashing
     def __post_init__(self):

core/dbt/contracts/graph/nodes.py

@@ -61,6 +61,7 @@
     SnapshotConfig,
 )
 
+
 # =====================================================================
 # This contains the classes for all of the nodes and node-like objects
 # in the manifest. In the "nodes" dictionary of the manifest we find
@@ -146,6 +147,8 @@ class ColumnInfo(AdditionalPropertiesMixin, ExtensibleDbtClassMixin, Replaceable
     description: str = ""
     meta: Dict[str, Any] = field(default_factory=dict)
     data_type: Optional[str] = None
+    constraints: Optional[List[str]] = None
+    constraints_check: Optional[str] = None
     quote: Optional[bool] = None
     tags: List[str] = field(default_factory=list)
     _extra: Dict[str, Any] = field(default_factory=dict)
@@ -400,6 +403,7 @@ class CompiledNode(ParsedNode):
     extra_ctes_injected: bool = False
     extra_ctes: List[InjectedCTE] = field(default_factory=list)
     _pre_injected_sql: Optional[str] = None
+    constraints_enabled: bool = False
 
     @property
     def empty(self):

core/dbt/contracts/graph/unparsed.py

@@ -93,6 +93,8 @@ class HasDocs(AdditionalPropertiesMixin, ExtensibleDbtClassMixin, Replaceable):
     description: str = ""
     meta: Dict[str, Any] = field(default_factory=dict)
     data_type: Optional[str] = None
+    constraints: Optional[List[str]] = None
+    constraints_check: Optional[str] = None
     docs: Docs = field(default_factory=Docs)
     _extra: Dict[str, Any] = field(default_factory=dict)
 

core/dbt/include/global_project/macros/adapters/columns.sql

@@ -17,23 +17,60 @@
 {% endmacro %}
 
 
+{% macro get_empty_subquery_sql(select_sql) -%}
+  {{ return(adapter.dispatch('get_empty_subquery_sql', 'dbt')(select_sql)) }}
+{% endmacro %}
+
+{% macro default__get_empty_subquery_sql(select_sql) %}
+    select * from (
+        {{ select_sql }}
+    ) as __dbt_sbq
+    where false
+    limit 0
+{% endmacro %}
+
+
+{% macro get_empty_schema_sql(columns) -%}
+  {{ return(adapter.dispatch('get_empty_schema_sql', 'dbt')(columns)) }}
+{% endmacro %}
+
+{% macro default__get_empty_schema_sql(columns) %}
+    select
+    {% for i in columns %}
+      {%- set col = columns[i] -%}
+      cast(null as {{ col['data_type'] }}) as {{ col['name'] }}{{ ", " if not loop.last }}
+    {%- endfor -%}
+{% endmacro %}
+
+{% macro get_column_schema_from_query(select_sql) -%}
+  {{ return(adapter.dispatch('get_column_schema_from_query', 'dbt')(select_sql)) }}
+{% endmacro %}
+
+{% macro default__get_column_schema_from_query(select_sql) %}
+    {% set columns = [] %}
+    {% set sql = get_empty_subquery_sql(select_sql) %}
+    {% set column_schema = adapter.get_column_schema_from_query(sql) %}
+    {% for col in column_schema %}
+        -- api.Column.create includes a step for translating data type
+        -- TODO: could include size, precision, scale here
+        {% set column = api.Column.create(col[0], col[1]) %}
+        {% do columns.append(column) %}
+    {% endfor %}
+    {{ return(columns) }}
+{% endmacro %}
+
+-- here for back compat
 {% macro get_columns_in_query(select_sql) -%}
   {{ return(adapter.dispatch('get_columns_in_query', 'dbt')(select_sql)) }}
 {% endmacro %}
 
 {% macro default__get_columns_in_query(select_sql) %}
     {% call statement('get_columns_in_query', fetch_result=True, auto_begin=False) -%}
-        select * from (
-            {{ select_sql }}
-        ) as __dbt_sbq
-        where false
-        limit 0
+        {{ get_empty_subquery_sql(select_sql) }}
     {% endcall %}
-
     {{ return(load_result('get_columns_in_query').table.columns | map(attribute='name') | list) }}
 {% endmacro %}
 
-
 {% macro alter_column_type(relation, column_name, new_column_type) -%}
   {{ return(adapter.dispatch('alter_column_type', 'dbt')(relation, column_name, new_column_type)) }}
 {% endmacro %}

core/dbt/include/global_project/macros/materializations/models/table/columns_spec_ddl.sql

@@ -0,0 +1,54 @@
+{%- macro get_columns_spec_ddl() -%}
+  {{ adapter.dispatch('get_columns_spec_ddl', 'dbt')() }}
+{%- endmacro -%}
+
+{% macro default__get_columns_spec_ddl() -%}
+  {{ return(columns_spec_ddl()) }}
+{%- endmacro %}
+
+{% macro columns_spec_ddl() %}
+  {# loop through user_provided_columns to create DDL with data types and constraints #}
+    {%- set user_provided_columns = model['columns'] -%}
+    (
+    {% for i in user_provided_columns %}
+      {% set col = user_provided_columns[i] %}
+      {% set constraints = col['constraints'] %}
+      {% set constraints_check = col['constraints_check'] %}
+      {{ col['name'] }} {{ col['data_type'] }} {% for x in constraints %} {{ x or "" }} {% endfor %} {% if constraints_check -%} check {{ constraints_check or "" }} {%- endif %} {{ "," if not loop.last }}
+    {% endfor %}
+  )
+{% endmacro %}
+
+{%- macro get_assert_columns_equivalent(sql) -%}
+  {{ adapter.dispatch('get_assert_columns_equivalent', 'dbt')(sql) }}
+{%- endmacro -%}
+
+{% macro default__get_assert_columns_equivalent(sql) -%}
+  {{ return(assert_columns_equivalent(sql)) }}
+{%- endmacro %}
+
+{% macro assert_columns_equivalent(sql) %}
+  {%- set sql_file_provided_columns = get_column_schema_from_query(sql) -%}
+  {%- set schema_file_provided_columns = get_column_schema_from_query(get_empty_schema_sql(model['columns'])) -%}
+
+  {%- set sql_file_provided_columns_formatted = format_columns(sql_file_provided_columns)  -%}
+  {%- set schema_file_provided_columns_formatted = format_columns(schema_file_provided_columns)  -%}
+
+  {%- if sql_file_provided_columns_formatted != schema_file_provided_columns_formatted -%}
+    {%- do exceptions.raise_compiler_error('Please ensure the name, data_type, order, and number of columns in your `yml` file match the columns in your SQL file.\nSchema File Columns: ' ~ sql_file_provided_columns_formatted ~ '\n\nSQL File Columns: ' ~ schema_file_provided_columns_formatted ~ ' ' ) %}
+  {%- endif -%}
+
+{% endmacro %}
+
+{% macro format_columns(columns) %}
+  {{ adapter.dispatch('format_columns', 'dbt')(columns) }}
+{%- endmacro %}
+
+{% macro default__format_columns(columns) -%}
+  {% set formatted_columns = [] %}
+  {% for column in columns %}
+    {%- set formatted_column = column.column.lower() ~ " " ~ column.dtype -%}
+    {%- do formatted_columns.append(formatted_column) -%}
+  {% endfor %}
+  {{ return(formatted_columns|join(', ')) }}
+{%- endmacro -%}

core/dbt/include/global_project/macros/materializations/models/table/create_table_as.sql

@@ -25,6 +25,10 @@
 
   create {% if temporary: -%}temporary{%- endif %} table
     {{ relation.include(database=(not temporary), schema=(not temporary)) }}
+    {% if config.get('constraints_enabled', False) %}
+      {{ get_assert_columns_equivalent(sql) }}
+      {{ get_columns_spec_ddl() }}
+    {% endif %}
   as (
     {{ sql }}
   );

core/dbt/parser/base.py

@@ -18,7 +18,7 @@
 from dbt.contracts.graph.manifest import Manifest
 from dbt.contracts.graph.nodes import ManifestNode, BaseNode
 from dbt.contracts.graph.unparsed import UnparsedNode, Docs
-from dbt.exceptions import DbtInternalError, ConfigUpdateError, DictParseError
+from dbt.exceptions import DbtInternalError, ConfigUpdateError, DictParseError, ParsingError
 from dbt import hooks
 from dbt.node_types import NodeType, ModelLanguage
 from dbt.parser.search import FileBlock
@@ -306,6 +306,19 @@ def update_parsed_node_config(
             else:
                 parsed_node.docs = Docs(show=docs_show)
 
+        # If we have constraints_enabled in the config, copy to node level, for backwards
+        # compatibility with earlier node-only config.
+        if config_dict.get("constraints_enabled", False):
+            parsed_node.constraints_enabled = True
+
+            parser_name = type(self).__name__
+            if parser_name == "ModelParser":
+                original_file_path = parsed_node.original_file_path
+                error_message = "\n    `constraints_enabled=true` can only be configured within `schema.yml` files\n      NOT within a model file(ex: .sql, .py) or `dbt_project.yml`."
+                raise ParsingError(
+                    f"Original File Path: ({original_file_path})\nConstraints must be defined in a `yml` schema configuration file like `schema.yml`.\nOnly the SQL table materialization is supported for constraints. \n`data_type` values must be defined for all columns and NOT be null or blank.{error_message}"
+                )
+
         # unrendered_config is used to compare the original database/schema/alias
         # values and to handle 'same_config' and 'same_contents' calls
         parsed_node.unrendered_config = config.build_config_dict(

core/dbt/parser/schemas.py

@@ -119,6 +119,8 @@ def add(
         column: Union[HasDocs, UnparsedColumn],
         description: str,
         data_type: Optional[str],
+        constraints: Optional[List[str]],
+        constraints_check: Optional[str],
         meta: Dict[str, Any],
     ):
         tags: List[str] = []
@@ -132,6 +134,8 @@ def add(
             name=column.name,
             description=description,
             data_type=data_type,
+            constraints=constraints,
+            constraints_check=constraints_check,
             meta=meta,
             tags=tags,
             quote=quote,
@@ -144,8 +148,10 @@ def from_target(cls, target: Union[HasColumnDocs, HasColumnTests]) -> "ParserRef
         for column in target.columns:
             description = column.description
             data_type = column.data_type
+            constraints = column.constraints
+            constraints_check = column.constraints_check
             meta = column.meta
-            refs.add(column, description, data_type, meta)
+            refs.add(column, description, data_type, constraints, constraints_check, meta)
         return refs
 
 
@@ -914,6 +920,75 @@ def parse_patch(self, block: TargetBlock[NodeTarget], refs: ParserRef) -> None:
                 self.patch_node_config(node, patch)
 
             node.patch(patch)
+            self.validate_constraints(node)
+
+    def validate_constraints(self, patched_node):
+        error_messages = []
+        if (
+            patched_node.resource_type == "model"
+            and patched_node.config.constraints_enabled is True
+        ):
+            validators = [
+                self.constraints_schema_validator(patched_node),
+                self.constraints_materialization_validator(patched_node),
+                self.constraints_language_validator(patched_node),
+                self.constraints_data_type_validator(patched_node),
+            ]
+            error_messages = [validator for validator in validators if validator != "None"]
+
+        if error_messages:
+            original_file_path = patched_node.original_file_path
+            raise ParsingError(
+                f"Original File Path: ({original_file_path})\nConstraints must be defined in a `yml` schema configuration file like `schema.yml`.\nOnly the SQL table materialization is supported for constraints. \n`data_type` values must be defined for all columns and NOT be null or blank.{self.convert_errors_to_string(error_messages)}"
+            )
+
+    def convert_errors_to_string(self, error_messages: List[str]):
+        n = len(error_messages)
+        if not n:
+            return ""
+        if n == 1:
+            return error_messages[0]
+        error_messages_string = "".join(error_messages[:-1]) + f"{error_messages[-1]}"
+        return error_messages_string
+
+    def constraints_schema_validator(self, patched_node):
+        schema_error = False
+        if patched_node.columns == {}:
+            schema_error = True
+        schema_error_msg = "\n    Schema Error: `yml` configuration does NOT exist"
+        schema_error_msg_payload = f"{schema_error_msg if schema_error else None}"
+        return schema_error_msg_payload
+
+    def constraints_materialization_validator(self, patched_node):
+        materialization_error = {}
+        if patched_node.config.materialized != "table":
+            materialization_error = {"materialization": patched_node.config.materialized}
+        materialization_error_msg = f"\n    Materialization Error: {materialization_error}"
+        materialization_error_msg_payload = (
+            f"{materialization_error_msg if materialization_error else None}"
+        )
+        return materialization_error_msg_payload
+
+    def constraints_language_validator(self, patched_node):
+        language_error = {}
+        language = str(patched_node.language)
+        if language != "sql":
+            language_error = {"language": language}
+        language_error_msg = f"\n    Language Error: {language_error}"
+        language_error_msg_payload = f"{language_error_msg if language_error else None}"
+        return language_error_msg_payload
+
+    def constraints_data_type_validator(self, patched_node):
+        data_type_errors = set()
+        for column, column_info in patched_node.columns.items():
+            if column_info.data_type is None:
+                data_type_error = {column}
+                data_type_errors.update(data_type_error)
+        data_type_errors_msg = (
+            f"\n    Columns with `data_type` Blank/Null Errors: {data_type_errors}"
+        )
+        data_type_errors_msg_payload = f"{data_type_errors_msg if data_type_errors else None}"
+        return data_type_errors_msg_payload
 
 
 class TestablePatchParser(NodePatchParser[UnparsedNodeUpdate]):