pass optional sql_header to empty subquery sql rendering (#7734)

.changes/unreleased/Fixes-20230531-131919.yaml

@@ -0,0 +1,6 @@
+kind: Fixes
+body: send sql header on contract enforcement
+time: 2023-05-31T13:19:19.801391-04:00
+custom:
+  Author: michelleark
+  Issue: "7714"

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

@@ -17,15 +17,18 @@
 {% endmacro %}
 
 
-{% macro get_empty_subquery_sql(select_sql) -%}
-  {{ return(adapter.dispatch('get_empty_subquery_sql', 'dbt')(select_sql)) }}
+{% macro get_empty_subquery_sql(select_sql, select_sql_header=none) -%}
+  {{ return(adapter.dispatch('get_empty_subquery_sql', 'dbt')(select_sql, select_sql_header)) }}
 {% endmacro %}
 
 {#
   Builds a query that results in the same schema as the given select_sql statement, without necessitating a data scan.
   Useful for running a query in a 'pre-flight' context, such as model contract enforcement (assert_columns_equivalent macro).
 #}
-{% macro default__get_empty_subquery_sql(select_sql) %}
+{% macro default__get_empty_subquery_sql(select_sql, select_sql_header=none) %}
+    {%- if select_sql_header is not none -%}
+    {{ select_sql_header }}
+    {%- endif -%}
     select * from (
         {{ select_sql }}
     ) as __dbt_sbq
@@ -53,10 +56,10 @@
     {%- endif -%}
 {% endmacro %}
 
-{% macro get_column_schema_from_query(select_sql) -%}
+{% macro get_column_schema_from_query(select_sql, select_sql_header=none) -%}
     {% set columns = [] %}
     {# -- Using an 'empty subquery' here to get the same schema as the given select_sql statement, without necessitating a data scan.#}
-    {% set sql = get_empty_subquery_sql(select_sql) %}
+    {% set sql = get_empty_subquery_sql(select_sql, select_sql_header) %}
     {% set column_schema = adapter.get_column_schema_from_query(sql) %}
     {{ return(column_schema) }}
 {% endmacro %}

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

@@ -34,7 +34,7 @@
 #}
 {% macro assert_columns_equivalent(sql) %}
   {#-- Obtain the column schema provided by sql file. #}
-  {%- set sql_file_provided_columns = get_column_schema_from_query(sql) -%}
+  {%- set sql_file_provided_columns = get_column_schema_from_query(sql, config.get('sql_header', none)) -%}
   {#--Obtain the column schema provided by the schema file by generating an 'empty schema' query from the model's columns. #}
   {%- set schema_file_provided_columns = get_column_schema_from_query(get_empty_schema_sql(model['columns'])) -%}
 

tests/adapter/dbt/tests/adapter/constraints/fixtures.py

@@ -133,6 +133,33 @@
   {sql_value} as wrong_data_type_column_name
 """
 
+my_model_contract_sql_header_sql = """
+{{
+  config(
+    materialized = "table"
+  )
+}}
+
+{% call set_sql_header(config) %}
+set session time zone 'Asia/Kolkata';
+{%- endcall %}
+select current_setting('timezone') as column_name
+"""
+
+my_model_incremental_contract_sql_header_sql = """
+{{
+  config(
+    materialized = "incremental",
+    on_schema_change="append_new_columns"
+  )
+}}
+
+{% call set_sql_header(config) %}
+set session time zone 'Asia/Kolkata';
+{%- endcall %}
+select current_setting('timezone') as column_name
+"""
+
 # model breaking constraints
 my_model_with_nulls_sql = """
 {{
@@ -303,3 +330,15 @@
       - name: wrong_data_type_column_name
         data_type: {data_type}
 """
+
+model_contract_header_schema_yml = """
+version: 2
+models:
+  - name: my_model_contract_sql_header
+    config:
+      contract:
+        enforced: true
+    columns:
+      - name: column_name
+        data_type: text
+"""

tests/adapter/dbt/tests/adapter/constraints/test_constraints.py

@@ -25,6 +25,9 @@
     my_model_incremental_with_nulls_sql,
     model_schema_yml,
     constrained_model_schema_yml,
+    my_model_contract_sql_header_sql,
+    my_model_incremental_contract_sql_header_sql,
+    model_contract_header_schema_yml,
 )
 
 
@@ -358,6 +361,45 @@ class TestIncrementalConstraintsRollback(BaseIncrementalConstraintsRollback):
     pass
 
 
+class BaseContractSqlHeader:
+    """Tests a contracted model with a sql header dependency."""
+
+    def test__contract_sql_header(self, project):
+        run_dbt(["run", "-s", "my_model_contract_sql_header"])
+
+        manifest = get_manifest(project.project_root)
+        model_id = "model.test.my_model_contract_sql_header"
+        model_config = manifest.nodes[model_id].config
+
+        assert model_config.contract.enforced
+
+
+class BaseTableContractSqlHeader(BaseContractSqlHeader):
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "my_model_contract_sql_header.sql": my_model_contract_sql_header_sql,
+            "constraints_schema.yml": model_contract_header_schema_yml,
+        }
+
+
+class BaseIncrementalContractSqlHeader(BaseContractSqlHeader):
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "my_model_contract_sql_header.sql": my_model_incremental_contract_sql_header_sql,
+            "constraints_schema.yml": model_contract_header_schema_yml,
+        }
+
+
+class TestTableContractSqlHeader(BaseTableContractSqlHeader):
+    pass
+
+
+class TestIncrementalContractSqlHeader(BaseIncrementalContractSqlHeader):
+    pass
+
+
 class BaseModelConstraintsRuntimeEnforcement:
     """
     These model-level constraints pass muster for dbt's preflight checks. Make sure they're