Add tests for get_partitions_metadata

plugins/bigquery/dbt/adapters/bigquery/connections.py

@@ -269,7 +269,7 @@ def get_table_from_response(cls, resp):
         column_names = [field.name for field in resp.schema]
         return agate_helper.table_from_data_flat(resp, column_names)
 
-    def raw_execute(self, sql, fetch=False, use_legacy_sql=False):
+    def raw_execute(self, sql, fetch=False, *, use_legacy_sql=False):
         conn = self.get_thread_connection()
         client = conn.handle
 

test/integration/022_bigquery_test/macros/partition_metadata.sql

@@ -0,0 +1,22 @@
+
+{% macro test_number_partitions(model, expected) %}
+
+    {%- set result = get_partitions_metadata(model) %}
+
+    {% if result %}
+        {% set partitions = result.columns['partition_id'].values() %}
+    {% else %}
+        {% set partitions = () %}
+    {% endif %}
+
+    {% set actual = partitions | length %}
+
+    {% if model and actual == expected %}
+        select 0 as success
+    {% else %}
+        -- actual: {{ actual }}
+        -- expected: {{ expected }}
+        select 1 as error
+    {% endif %}
+
+{% endmacro %}

test/integration/022_bigquery_test/partition-models/my_model.sql

@@ -9,3 +9,5 @@
 }}
 
 select 1 as id, 'dr. bigquery' as name, current_timestamp() as cur_time, current_date() as cur_date
+union all
+select 2 as id, 'prof. bigquery' as name, current_timestamp() as cur_time, current_date() as cur_date

test/integration/022_bigquery_test/partition-models/schema.yml

@@ -0,0 +1,6 @@
+version: 2
+models:
+- name: my_model
+  tests:
+  - number_partitions:
+      expected: "{{ var('expected', 1) }}"

test/integration/022_bigquery_test/test_bigquery_changing_partitions.py

@@ -19,13 +19,22 @@ def run_changes(self, before, after):
         results = self.run_dbt(['run', '--vars', json.dumps(after)])
         self.assertEqual(len(results), 1)
 
+    def test_partitions(self, expected):
+        test_results = self.run_dbt(['test', '--vars', json.dumps(expected)])
+
+        for result in test_results:
+            self.assertIsNone(result.error)
+            self.assertFalse(result.skipped)
+            # status = # of failing rows
+            self.assertEqual(result.status, 0)
 
     @use_profile('bigquery')
     def test_bigquery_add_partition(self):
         before = {"partition_by": None, "cluster_by": None}
         after = {"partition_by": {'field': 'cur_time', 'data_type': 'timestamp'}, "cluster_by": None}
         self.run_changes(before, after)
-
+        self.test_partitions({"expected": 1})
+
     @use_profile('bigquery')
     def test_bigquery_remove_partition(self):
         before = {"partition_by": {'field': 'cur_time', 'data_type': 'timestamp'}, "cluster_by": None}
@@ -37,14 +46,18 @@ def test_bigquery_change_partitions(self):
         before = {"partition_by": {'field': 'cur_time', 'data_type': 'timestamp'}, "cluster_by": None}
         after = {"partition_by": {'field': "cur_date"}, "cluster_by": None}
         self.run_changes(before, after)
+        self.test_partitions({"expected": 1})
         self.run_changes(after, before)
+        self.test_partitions({"expected": 1})
 
     @use_profile('bigquery')
     def test_bigquery_change_partitions_from_int(self):
         before = {"partition_by": {"field": "id", "data_type": "int64", "range": {"start": 0, "end": 10, "interval": 1}}, "cluster_by": None}
         after = {"partition_by": {"field": "cur_date", "data_type": "date"}, "cluster_by": None}
         self.run_changes(before, after)
+        self.test_partitions({"expected": 1})
         self.run_changes(after, before)
+        self.test_partitions({"expected": 2})
 
     @use_profile('bigquery')
     def test_bigquery_add_clustering(self):