From d1e75de80f2f15b98bffb476b1d478fe55d11a59 Mon Sep 17 00:00:00 2001 From: lorenzop Date: Tue, 20 Jan 2026 13:09:14 +0100 Subject: [PATCH] Added features for Nessie branch management --- .../dremio/macros/adapters/relation.sql | 51 +++++++++++++++++++ .../table/create_table_as.sql | 9 ++++ .../macros/materializations/table/table.sql | 13 ++++- .../materializations/view/create_view_as.sql | 3 +- 4 files changed, 74 insertions(+), 2 deletions(-) diff --git a/dbt/include/dremio/macros/adapters/relation.sql b/dbt/include/dremio/macros/adapters/relation.sql index 1f7c3f64..3cb93b5e 100644 --- a/dbt/include/dremio/macros/adapters/relation.sql +++ b/dbt/include/dremio/macros/adapters/relation.sql @@ -26,6 +26,57 @@ limitations under the License.*/ {%- endcall %} {% endmacro %} +{% macro drop_relation_with_branch(relation, branch=none) -%} + {% call statement('drop_relation', auto_begin=False) -%} + drop {{ relation.type }} if exists {{ relation }}{%- if branch is not none %} at branch {{ branch }}{%- endif %} + {%- endcall %} +{% endmacro %} + +{% macro create_branch_statement(relation, branch) -%} + {%- set nessie_ref = config.get('nessie_ref', validator=validation.any[string]) -%} + {%- if execute -%} + {%- set create_branch_sql -%} + create branch if not exists {{ branch }}{% if nessie_ref is not none %} at ref {{ nessie_ref }}{% endif %} in {{ relation.database }} + {%- endset -%} + {%- do run_query(create_branch_sql) -%} + {%- endif -%} +{%- endmacro %} + +{% macro get_relation_at_branch(database, schema, identifier, branch) -%} + {%- if execute -%} + {# First, check if the branch exists #} + {%- set branch_sql -%} + show branches in {{ database }} + {%- endset -%} + {%- set branch_result = run_query(branch_sql) -%} + {%- set branch_exists = False -%} + {%- for row in branch_result.rows -%} + {%- if row[0] == branch -%} + {%- set branch_exists = True -%} + {%- endif -%} + {%- endfor -%} + + {# If branch exists, check if table exists in that branch #} + {%- set found = False -%} + {%- if branch_exists -%} + {%- set table_sql -%} + show tables in {{ database }}.{{ schema }} at branch {{ branch }} + {%- endset -%} + {%- set table_result = run_query(table_sql) -%} + {%- if table_result.rows | length > 0 -%} + {%- for row in table_result.rows -%} + {%- if row[1] == identifier -%} + {%- set found = True -%} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + {%- endif -%} + {{ found }} + {%- else -%} + {{ False }} + {%- endif -%} +{%- endmacro %} + {% macro dremio__rename_relation(from_relation, to_relation) -%} {% call statement('rename_relation1/2 - create to_relation from from_relation') -%} {{ get_create_table_as_sql(temporary=False, relation=to_relation, sql="select * from " ~ from_relation)}} diff --git a/dbt/include/dremio/macros/materializations/table/create_table_as.sql b/dbt/include/dremio/macros/materializations/table/create_table_as.sql index db8cd1bb..e7f55cb9 100644 --- a/dbt/include/dremio/macros/materializations/table/create_table_as.sql +++ b/dbt/include/dremio/macros/materializations/table/create_table_as.sql @@ -34,6 +34,7 @@ limitations under the License.*/ {{ sql_header if sql_header is not none }} create table {{ relation }} + {{ branch_clause() }} {{ partition_method() }} {{ config_cols("partition by") }} {{ config_cols("distribute by") }} {{ config_cols("localsort by") }} @@ -64,3 +65,11 @@ limitations under the License.*/ with single writer {%- endif -%} {%- endmacro -%} + +{%- macro branch_clause() -%} + {%- set branch = config.get('branch', validator=validation.any[string]) -%} + {%- if branch is not none -%} + at branch {{ branch }} + {%- endif -%} +{%- endmacro -%} + diff --git a/dbt/include/dremio/macros/materializations/table/table.sql b/dbt/include/dremio/macros/materializations/table/table.sql index 60adc72a..93232590 100644 --- a/dbt/include/dremio/macros/materializations/table/table.sql +++ b/dbt/include/dremio/macros/materializations/table/table.sql @@ -15,16 +15,27 @@ limitations under the License.*/ {% materialization table, adapter = 'dremio' %} {%- set identifier = model['alias'] -%} + {%- set branch = config.get('branch', validator=validation.any[string]) -%} {%- set format = config.get('format', validator=validation.any[basestring]) or 'iceberg' -%} {%- set old_relation = adapter.get_relation(database=database, schema=schema, identifier=identifier) -%} {%- set target_relation = this.incorporate(type='table') -%} {% set grant_config = config.get('grants') %} {{ run_hooks(pre_hooks) }} + -- create branch first if needed + {% if branch is not none %} + {{ create_branch_statement(target_relation, branch) }} + {% endif %} + -- setup: if the target relation already exists, drop it -- in case if the existing and future table is delta, we want to do a -- create or replace table instead of dropping, so we don't have the table unavailable - {% if old_relation is not none -%} + {% if branch is not none %} + {%- set branch_relation_exists = get_relation_at_branch(database, schema, identifier, branch) -%} + {% if branch_relation_exists %} + {{ drop_relation_with_branch(target_relation, branch) }} + {% endif %} + {% elif old_relation is not none -%} {{ adapter.drop_relation(old_relation) }} {%- endif %} diff --git a/dbt/include/dremio/macros/materializations/view/create_view_as.sql b/dbt/include/dremio/macros/materializations/view/create_view_as.sql index bf444f8d..0de00477 100644 --- a/dbt/include/dremio/macros/materializations/view/create_view_as.sql +++ b/dbt/include/dremio/macros/materializations/view/create_view_as.sql @@ -15,6 +15,7 @@ limitations under the License.*/ {% macro dremio__create_view_as(relation, sql) -%} {% set contract_config = config.get('contract') %} {%- set sql_header = config.get('sql_header', none) -%} + {%- set branch = config.get('branch', validator=validation.any[string]) -%} {{ sql_header if sql_header is not none }} @@ -23,6 +24,6 @@ limitations under the License.*/ {{ get_assert_columns_equivalent(sql) }} {% set sql = get_select_subquery(sql) %} {% endif %} - as {{ sql }} + as {{ sql }}{%- if branch is not none %} at branch {{ branch }}{%- endif %} {%- endmacro %}