Skip to content

Commit

Permalink
Add tests for get_partitions_metadata (#1)
Browse files Browse the repository at this point in the history
* Add tests using get_partitions_metadata

* Readd asterisk to raw_execute
  • Loading branch information
jtcohen6 authored Oct 21, 2020
1 parent 78bd7c9 commit c31ba10
Show file tree
Hide file tree
Showing 5 changed files with 45 additions and 2 deletions.
2 changes: 1 addition & 1 deletion plugins/bigquery/dbt/adapters/bigquery/connections.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
22 changes: 22 additions & 0 deletions test/integration/022_bigquery_test/macros/partition_metadata.sql
Original file line number Diff line number Diff line change
@@ -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 %}
Original file line number Diff line number Diff line change
Expand Up @@ -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
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
version: 2
models:
- name: my_model
tests:
- number_partitions:
expected: "{{ var('expected', 1) }}"
Original file line number Diff line number Diff line change
Expand Up @@ -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}
Expand All @@ -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):
Expand Down

0 comments on commit c31ba10

Please sign in to comment.