Implement pre/post hooks and incremental tables on BigQuery. dbt-labs#779 dbt-labs#712

Author: lewish

.gitignore

@@ -68,3 +68,6 @@ target/
 
 # Sublime Text
 *.sublime-*
+
+#PyCharm/IntelliJ
+.idea/

dbt/adapters/bigquery/impl.py

@@ -372,8 +372,11 @@ def get_table_from_response(cls, resp):
 
     @classmethod
     def add_begin_query(cls, profile, name):
-        raise dbt.exceptions.NotImplementedException(
-            '`add_begin_query` is not implemented for this adapter!')
+        pass
+
+    @classmethod
+    def add_commit_query(cls, profile, name):
+        pass
 
     @classmethod
     def create_schema(cls, profile, project_cfg, schema, model_name=None):
@@ -465,21 +468,6 @@ def get_dataset(cls, profile, project_cfg, dataset_name, model_name=None):
         dataset_ref = client.dataset(dataset_name)
         return google.cloud.bigquery.Dataset(dataset_ref)
 
-    @classmethod
-    def warning_on_hooks(cls, hook_type):
-        msg = "{} is not supported in bigquery and will be ignored"
-        dbt.ui.printer.print_timestamped_line(msg.format(hook_type),
-                                              dbt.ui.printer.COLOR_FG_YELLOW)
-
-    @classmethod
-    def add_query(cls, profile, sql, model_name=None, auto_begin=True,
-                  bindings=None, abridge_sql_log=False):
-        if model_name in ['on-run-start', 'on-run-end']:
-            cls.warning_on_hooks(model_name)
-        else:
-            raise dbt.exceptions.NotImplementedException(
-                '`add_query` is not implemented for this adapter!')
-
     @classmethod
     def is_cancelable(cls):
         return False

dbt/include/global_project/macros/adapters/bigquery.sql

@@ -16,6 +16,11 @@
 
   create or replace table {{ relation }}
   {{ partition_by(raw_partition_by) }}
+  {% if temporary %}
+  OPTIONS(
+    expiration_timestamp=TIMESTAMP_ADD(CURRENT_TIMESTAMP(), INTERVAL 1 DAY)
+  )
+  {% endif %}
   as (
     {{ sql }}
   );

dbt/include/global_project/macros/materializations/helpers.sql

@@ -11,6 +11,13 @@
   {% endfor %}
 {% endmacro %}
 
+{% macro run_bigquery_hooks(hooks) %}
+  {% for hook in hooks %}
+    {% call statement(auto_begin=inside_transaction) %}
+      {{ hook.get('sql') }}
+    {% endcall %}
+  {% endfor %}
+{% endmacro %}
 
 {% macro column_list(columns) %}
   {%- for col in columns %}

dbt/include/global_project/macros/materializations/incremental/bigquery_incremental.sql

@@ -0,0 +1,101 @@
+{% macro dbt_bigquery__incremental_delete(target_relation, tmp_relation) -%}
+
+  {%- set unique_key = config.require('unique_key') -%}
+  {%- set identifier = model['name'] -%}
+
+  delete
+  from {{ target_relation }}
+  where ({{ unique_key }}) in (
+    select ({{ unique_key }})
+    from {{ tmp_relation.include(schema=True) }}
+  );
+
+{%- endmacro %}
+
+{% materialization incremental, default -%}
+  {%- set sql_where = config.require('sql_where') -%}
+  {%- set unique_key = config.get('unique_key') -%}
+
+  {%- set identifier = model['name'] -%}
+  {%- set tmp_identifier = model['name'] + '__dbt_incremental_tmp' -%}
+
+  {%- set existing_relations = adapter.list_relations(schema=schema) -%}
+  {%- set old_relation = adapter.get_relation(relations_list=existing_relations,
+                                              schema=schema, identifier=identifier) -%}
+  {%- set target_relation = api.Relation.create(identifier=identifier, schema=schema, type='table') -%}
+  {%- set tmp_relation = api.Relation.create(identifier=tmp_identifier,
+                                                 schema=schema, type='table') -%}
+
+  {%- set non_destructive_mode = (flags.NON_DESTRUCTIVE == True) -%}
+  {%- set full_refresh_mode = (flags.FULL_REFRESH == True) -%}
+
+  {%- set exists_as_table = (old_relation is not none and old_relation.is_table) -%}
+  {%- set exists_not_as_table = (old_relation is not none and not old_relation.is_table) -%}
+
+  {%- set should_truncate = (non_destructive_mode and full_refresh_mode and exists_as_table) -%}
+  {%- set should_drop = (not should_truncate and (full_refresh_mode or exists_not_as_table)) -%}
+  {%- set force_create = (flags.FULL_REFRESH and not flags.NON_DESTRUCTIVE) -%}
+
+  -- setup
+  {% if old_relation is none -%}
+    -- noop
+  {%- elif should_truncate -%}
+    {{ adapter.truncate_relation(old_relation) }}
+  {%- elif should_drop -%}
+    {{ adapter.drop_relation(old_relation) }}
+    {%- set old_relation = none -%}
+  {%- endif %}
+
+  {{ run_bigquery_hooks(pre_hooks) }}
+
+  -- build model
+  {% if force_create or old_relation is none -%}
+    {%- call statement('main') -%}
+      {{ create_table_as(False, target_relation, sql) }}
+    {%- endcall -%}
+  {%- else -%}
+     {%- call statement() -%}
+
+       {% set tmp_table_sql -%}
+         with dbt_incr_sbq as (
+           {{ sql }}
+         )
+         select * from dbt_incr_sbq
+         where ({{ sql_where }})
+           or ({{ sql_where }}) is null
+       {%- endset %}
+
+       {{ dbt.create_table_as(True, tmp_relation, tmp_table_sql) }}
+
+     {%- endcall -%}
+
+     {{ adapter.expand_target_column_types(temp_table=tmp_identifier,
+                                           to_schema=schema,
+                                           to_table=identifier) }}
+
+
+     {% if unique_key is not none -%}
+       {%- call statement() -%}
+         {{ dbt_bigquery__incremental_delete(target_relation, tmp_relation) }}
+     {%- endcall -%}
+
+     {%- endif %}
+     {%- call statement('main') -%}
+       {% set dest_columns = adapter.get_columns_in_table(schema, identifier) %}
+       {% set dest_cols_csv = dest_columns | map(attribute='quoted') | join(', ') %}
+
+       insert into {{ target_relation }} ({{ dest_cols_csv }})
+       (
+         select {{ dest_cols_csv }}
+         from {{ tmp_relation.include(schema=True) }}
+       );
+     {% endcall %}
+
+     -- clean up temp table in bigquery
+      {{ drop_relation_if_exists(tmp_relation) }}
+
+  {%- endif %}
+
+  {{ run_bigquery_hooks(post_hooks) }}
+
+{%- endmaterialization %}

dbt/include/global_project/macros/materializations/table/bigquery_table.sql

@@ -52,6 +52,8 @@
       {% endif %}
   {% endif %}
 
+  {{ run_bigquery_hooks(pre_hooks) }}
+
   {#
       Since dbt uses WRITE_TRUNCATE mode for tables, we only need to drop this thing
       if it is not a table. If it _is_ already a table, then we can overwrite it without downtime
@@ -71,5 +73,6 @@
     {% endcall -%}
   {% endif %}
 
+  {{ run_bigquery_hooks(pre_hooks) }}
 
 {% endmaterialization %}

dbt/include/global_project/macros/materializations/view/bigquery_view.sql

@@ -15,6 +15,8 @@
       identifier=identifier, schema=schema,
       type='view') -%}
 
+  {{ run_bigquery_hooks(pre_hooks) }}
+
   -- drop if exists
   {%- if old_relation is not none -%}
     {%- if old_relation.is_table and not flags.FULL_REFRESH -%}
@@ -36,4 +38,6 @@
     {%- endcall %}
   {%- endif %}
 
+  {{ run_bigquery_hooks(post_hooks) }}
+
 {%- endmaterialization %}