From 7e3da7cc54ea6331e4f3dc68fdbdfb315aca05d0 Mon Sep 17 00:00:00 2001
From: Lewis Davies <lewis2106@gmail.com>
Date: Thu, 12 May 2022 21:24:06 +0100
Subject: [PATCH] Add `where` parameter to `union_relations` (#554)

* Add where parameter

* Add union where test

* Update changelog
---
 CHANGELOG.md                                           |  2 ++
 README.md                                              |  1 +
 integration_tests/models/sql/schema.yml                | 10 ++++++++++
 integration_tests/models/sql/test_union_where.sql      |  5 +++++
 integration_tests/models/sql/test_union_where_base.sql |  4 ++++
 macros/sql/union.sql                                   | 10 +++++++---
 6 files changed, 29 insertions(+), 3 deletions(-)
 create mode 100644 integration_tests/models/sql/test_union_where.sql
 create mode 100644 integration_tests/models/sql/test_union_where_base.sql

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 263e6315..f15cb319 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,8 @@
 ## New features
 - Add an optional `where` clause parameter to `get_column_values()` to filter values returned ([#511](https://github.com/dbt-labs/dbt-utils/issues/511), [#583](https://github.com/dbt-labs/dbt-utils/pull/583)) 
 
+## New features
+- Add `where` parameter to `union_relations` macro ([#554](https://github.com/dbt-labs/dbt-utils/pull/554))
 ## Quality of life
 - Documentation about listagg macro ([#544](https://github.com/dbt-labs/dbt-utils/issues/544), [#560](https://github.com/dbt-labs/dbt-utils/pull/560))
 - Fix links to macro section in table of contents ([#555](https://github.com/dbt-labs/dbt-utils/pull/555))
diff --git a/README.md b/README.md
index ac238b5f..4c149a82 100644
--- a/README.md
+++ b/README.md
@@ -854,6 +854,7 @@ final query. Note the `include` and `exclude` arguments are mutually exclusive.
 e.g. `{"some_field": "varchar(100)"}`.``
 * `source_column_name` (optional, `default="_dbt_source_relation"`): The name of
 the column that records the source of this row.
+* `where` (optional): Filter conditions to include in the `where` clause.
 
 #### generate_series ([source](macros/sql/generate_series.sql))
 This macro implements a cross-database mechanism to generate an arbitrarily long list of numbers. Specify the maximum number you'd like in your list and it will create a 1-indexed SQL result set.
diff --git a/integration_tests/models/sql/schema.yml b/integration_tests/models/sql/schema.yml
index 83733596..e5927fbf 100644
--- a/integration_tests/models/sql/schema.yml
+++ b/integration_tests/models/sql/schema.yml
@@ -153,6 +153,16 @@ models:
       - dbt_utils.equality:
           compare_model: ref('data_union_expected')
 
+  - name: test_union_where
+    columns:
+      - name: id
+        tests:
+          - dbt_utils.expression_is_true:
+              expression: "= 1"
+      - name: favorite_number
+        tests:
+          - dbt_utils.not_constant
+
   - name: test_get_relations_by_pattern
     tests:
       - dbt_utils.equality:
diff --git a/integration_tests/models/sql/test_union_where.sql b/integration_tests/models/sql/test_union_where.sql
new file mode 100644
index 00000000..a5986d41
--- /dev/null
+++ b/integration_tests/models/sql/test_union_where.sql
@@ -0,0 +1,5 @@
+select 
+    id, 
+    favorite_number
+from 
+    {{ ref('test_union_where_base') }}
diff --git a/integration_tests/models/sql/test_union_where_base.sql b/integration_tests/models/sql/test_union_where_base.sql
new file mode 100644
index 00000000..177db51a
--- /dev/null
+++ b/integration_tests/models/sql/test_union_where_base.sql
@@ -0,0 +1,4 @@
+{{ dbt_utils.union_relations(
+    [ref('data_union_table_1'), ref('data_union_table_2')],
+    where="id = 1"
+) }}
diff --git a/macros/sql/union.sql b/macros/sql/union.sql
index a7bf1d95..0d4ebccb 100644
--- a/macros/sql/union.sql
+++ b/macros/sql/union.sql
@@ -1,8 +1,8 @@
-{%- macro union_relations(relations, column_override=none, include=[], exclude=[], source_column_name='_dbt_source_relation') -%}
-    {{ return(adapter.dispatch('union_relations', 'dbt_utils')(relations, column_override, include, exclude, source_column_name)) }}
+{%- macro union_relations(relations, column_override=none, include=[], exclude=[], source_column_name='_dbt_source_relation', where=none) -%}
+    {{ return(adapter.dispatch('union_relations', 'dbt_utils')(relations, column_override, include, exclude, source_column_name, where)) }}
 {% endmacro %}
 
-{%- macro default__union_relations(relations, column_override=none, include=[], exclude=[], source_column_name='_dbt_source_relation') -%}
+{%- macro default__union_relations(relations, column_override=none, include=[], exclude=[], source_column_name='_dbt_source_relation', where=none) -%}
 
     {%- if exclude and include -%}
         {{ exceptions.raise_compiler_error("Both an exclude and include list were provided to the `union` macro. Only one is allowed") }}
@@ -92,6 +92,10 @@
                 {%- endfor %}
 
             from {{ relation }}
+
+            {% if where -%}
+            where {{ where }}
+            {%- endif %}
         )
 
         {% if not loop.last -%}