From bfc3d53f224b65aa11a98dd5aa7b2759df6bb8da Mon Sep 17 00:00:00 2001 From: rfeith2 <159465873+rfeith2@users.noreply.github.com> Date: Mon, 21 Jul 2025 20:46:04 -0700 Subject: [PATCH 1/5] Update changelog for UDF materialization --- CHANGELOG.md | 6 ++++ .../dremio/macros/materializations/udf.sql | 35 +++++++++++++++++++ 2 files changed, 41 insertions(+) create mode 100644 dbt/include/dremio/macros/materializations/udf.sql diff --git a/CHANGELOG.md b/CHANGELOG.md index 6a00476b..9872701c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,9 @@ +# dbt-dremio v1.9.1 + +## Features + +- Added built-in `udf` materialization to create user-defined functions + # dbt-dremio v1.9.0 ## Changes diff --git a/dbt/include/dremio/macros/materializations/udf.sql b/dbt/include/dremio/macros/materializations/udf.sql new file mode 100644 index 00000000..1998ba2a --- /dev/null +++ b/dbt/include/dremio/macros/materializations/udf.sql @@ -0,0 +1,35 @@ +/*Copyright (C) 2022 Dremio Corporation + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License.*/ + +{% materialization udf, adapter='dremio' %} + {%- set folder_name = config.get('folder_name', '') -%} + {%- set target = folder_name ~ ('.' if folder_name else '') ~ this.identifier -%} + + {%- set parameter_list = config.get('parameter_list') -%} + {%- set ret = config.get('returns') -%} + + {%- set create_sql -%} +CREATE OR REPLACE FUNCTION {{ target }}({{ parameter_list }}) +RETURNS {{ ret }} + +{{ sql }} +; + {%- endset -%} + + {% call statement('main') -%} + {{ create_sql }} + {%- endcall %} + + {{ return({'relations': []}) }} +{% endmaterialization %} From cb20be04c6ce6ce958a95efa794f97f240e84605 Mon Sep 17 00:00:00 2001 From: rfeith2 <159465873+rfeith2@users.noreply.github.com> Date: Mon, 21 Jul 2025 20:46:54 -0700 Subject: [PATCH 2/5] Add tests for UDF materialization --- tests/unit/test_udf_materialization.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) create mode 100644 tests/unit/test_udf_materialization.py diff --git a/tests/unit/test_udf_materialization.py b/tests/unit/test_udf_materialization.py new file mode 100644 index 00000000..77aa237d --- /dev/null +++ b/tests/unit/test_udf_materialization.py @@ -0,0 +1,18 @@ +import os + +MACRO_PATH = os.path.join( + "dbt", "include", "dremio", "macros", "materializations", "udf.sql" +) + + +def test_udf_materialization_file_exists(): + assert os.path.isfile(MACRO_PATH) + + +def test_udf_materialization_contains_expected_sql(): + with open(MACRO_PATH) as f: + contents = f.read() + + assert "{% materialization udf" in contents + assert "CREATE OR REPLACE FUNCTION" in contents + assert "return({'relations': []})" in contents From 0e035cf2f8d6cd638e6f1b268ce68da30978c55b Mon Sep 17 00:00:00 2001 From: rfeith2 <159465873+rfeith2@users.noreply.github.com> Date: Mon, 21 Jul 2025 20:46:59 -0700 Subject: [PATCH 3/5] Fix UDF materialization and bump version --- dbt/adapters/dremio/__version__.py | 2 +- dbt/include/dremio/macros/materializations/udf.sql | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/dbt/adapters/dremio/__version__.py b/dbt/adapters/dremio/__version__.py index b9b62836..06f1d95f 100644 --- a/dbt/adapters/dremio/__version__.py +++ b/dbt/adapters/dremio/__version__.py @@ -9,4 +9,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -version = "1.9.0" +version = "1.9.1" diff --git a/dbt/include/dremio/macros/materializations/udf.sql b/dbt/include/dremio/macros/materializations/udf.sql index 1998ba2a..7861c77c 100644 --- a/dbt/include/dremio/macros/materializations/udf.sql +++ b/dbt/include/dremio/macros/materializations/udf.sql @@ -13,11 +13,11 @@ See the License for the specific language governing permissions and limitations under the License.*/ {% materialization udf, adapter='dremio' %} - {%- set folder_name = config.get('folder_name', '') -%} + {%- set folder_name = config.get('folder_name', validator=validation.any[basestring]) or '' -%} {%- set target = folder_name ~ ('.' if folder_name else '') ~ this.identifier -%} - {%- set parameter_list = config.get('parameter_list') -%} - {%- set ret = config.get('returns') -%} + {%- set parameter_list = config.get('parameter_list', validator=validation.any[basestring]) -%} + {%- set ret = config.get('returns', validator=validation.any[basestring]) -%} {%- set create_sql -%} CREATE OR REPLACE FUNCTION {{ target }}({{ parameter_list }}) From 0f4b5140bb67aecfa04647532fc3917d8a6518bb Mon Sep 17 00:00:00 2001 From: rfeith2 Date: Thu, 24 Jul 2025 21:15:26 -0700 Subject: [PATCH 4/5] perfecting testing --- .../get_custom_name/is_datalake_node.sql | 2 +- .../dremio/macros/materializations/udf.sql | 19 +- .../adapter/materialization/test_udf.py | 286 ++++++++++++++++++ tests/unit/test_udf_materialization.py | 79 ++++- 4 files changed, 381 insertions(+), 5 deletions(-) create mode 100644 tests/functional/adapter/materialization/test_udf.py diff --git a/dbt/include/dremio/macros/get_custom_name/is_datalake_node.sql b/dbt/include/dremio/macros/get_custom_name/is_datalake_node.sql index b58740c2..adcfb278 100644 --- a/dbt/include/dremio/macros/get_custom_name/is_datalake_node.sql +++ b/dbt/include/dremio/macros/get_custom_name/is_datalake_node.sql @@ -14,5 +14,5 @@ limitations under the License.*/ {% macro is_datalake_node(node) -%} {{ return(node.resource_type in ['test', 'seed', 'snapshot'] - or (node.resource_type == 'model' and node.config.materialized not in ['view', 'reflection'])) }} + or (node.resource_type == 'model' and node.config.materialized not in ['view', 'reflection', 'udf'])) }} {%- endmacro %} diff --git a/dbt/include/dremio/macros/materializations/udf.sql b/dbt/include/dremio/macros/materializations/udf.sql index 7861c77c..ae8cae17 100644 --- a/dbt/include/dremio/macros/materializations/udf.sql +++ b/dbt/include/dremio/macros/materializations/udf.sql @@ -13,12 +13,25 @@ See the License for the specific language governing permissions and limitations under the License.*/ {% materialization udf, adapter='dremio' %} - {%- set folder_name = config.get('folder_name', validator=validation.any[basestring]) or '' -%} - {%- set target = folder_name ~ ('.' if folder_name else '') ~ this.identifier -%} + {%- set identifier = model['alias'] -%} + {%- set target_relation = this.incorporate(type='view') -%} + + {%- set udf_database = config.get('database', validator=validation.any[basestring]) or database -%} + {%- set udf_schema = config.get('schema', validator=validation.any[basestring]) or schema -%} + {%- set target = udf_database ~ '.' ~ udf_schema ~ '.' ~ identifier -%} {%- set parameter_list = config.get('parameter_list', validator=validation.any[basestring]) -%} {%- set ret = config.get('returns', validator=validation.any[basestring]) -%} + {# Validate required configurations #} + {%- if parameter_list is none or parameter_list == '' -%} + {{ exceptions.raise_compiler_error("UDF materialization requires 'parameter_list' configuration (e.g., parameter_list='x INT, y INT').") }} + {%- endif -%} + + {%- if ret is none -%} + {{ exceptions.raise_compiler_error("UDF materialization requires 'returns' configuration specifying the return type (e.g., returns='INT').") }} + {%- endif -%} + {%- set create_sql -%} CREATE OR REPLACE FUNCTION {{ target }}({{ parameter_list }}) RETURNS {{ ret }} @@ -31,5 +44,5 @@ RETURNS {{ ret }} {{ create_sql }} {%- endcall %} - {{ return({'relations': []}) }} + {{ return({'relations': [target_relation]}) }} {% endmaterialization %} diff --git a/tests/functional/adapter/materialization/test_udf.py b/tests/functional/adapter/materialization/test_udf.py new file mode 100644 index 00000000..efbf8a6f --- /dev/null +++ b/tests/functional/adapter/materialization/test_udf.py @@ -0,0 +1,286 @@ +# Copyright (C) 2022 Dremio Corporation + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +from dbt.tests.util import run_dbt, check_result_nodes_by_name + + +# Basic UDF with simple arithmetic +basic_udf_sql = """ +{{ config( + materialized='udf', + parameter_list='x INT, y INT', + returns='INT' +) }} + +RETURN x + y +""" + +# UDF with no parameters - should fail validation +no_params_udf_sql = """ +{{ config( + materialized='udf', + parameter_list='', + returns='VARCHAR' +) }} + +RETURN 'Hello World' +""" + +# UDF with no parameter_list config - should fail +no_param_list_config_sql = """ +{{ config( + materialized='udf', + returns='VARCHAR' +) }} + +RETURN 'Hello World' +""" + +# UDF with complex types +complex_udf_sql = """ +{{ config( + materialized='udf', + parameter_list='name VARCHAR, age INT, salary DECIMAL(10,2)', + returns='VARCHAR' +) }} + +RETURN CONCAT(name, ' is ', CAST(age AS VARCHAR), ' years old and earns ', CAST(salary AS VARCHAR)) +""" + +# UDF with custom database configuration (uses existing test schema) +custom_location_udf_sql = """ +{{ config( + materialized='udf', + database='dbt_test', + schema='{{ target.schema }}', + parameter_list='x DOUBLE', + returns='DOUBLE' +) }} + +RETURN x * x +""" + +# Model that uses a UDF +downstream_model_sql = """ +{{ config(materialized='view') }} + +SELECT + {{ ref('basic_udf') }}(10, 20) as udf_output +FROM (SELECT 1 as dummy) t +""" + +# Invalid UDF - missing returns config +invalid_udf_no_returns_sql = """ +{{ config( + materialized='udf', + parameter_list='x INT' +) }} + +RETURN x * 2 +""" + +# Invalid UDF - missing parameter_list config +invalid_udf_no_params_sql = """ +{{ config( + materialized='udf', + returns='INT' +) }} + +RETURN 42 +""" + + +class TestBasicUDFMaterialization: + @pytest.fixture(scope="class") + def models(self): + return { + "basic_udf.sql": basic_udf_sql, + "downstream_model.sql": downstream_model_sql, + } + + def test_basic_udf_creation(self, project): + # First compile to see what's generated + compile_results = run_dbt(["compile"]) + + # Read the compiled SQL to debug + import os + compiled_path = os.path.join(project.project_root, "target", "compiled", "test", "models", "downstream_model.sql") + if os.path.exists(compiled_path): + with open(compiled_path, 'r') as f: + compiled_sql = f.read() + print(f"\n\nCOMPILED downstream_model.sql:\n{compiled_sql}\n\n") + + # Run both models + results = run_dbt(["run"]) + assert len(results) == 2 + check_result_nodes_by_name(results, ["basic_udf", "downstream_model"]) + + # Test docs generation + results = run_dbt(["docs", "generate"]) + assert results is not None + + +class TestUDFParameterVariations: + @pytest.fixture(scope="class") + def models(self): + return { + "complex_udf.sql": complex_udf_sql, + } + + def test_udf_parameter_variations(self, project): + # Run complex UDF model + results = run_dbt(["run"]) + assert len(results) == 1 + check_result_nodes_by_name(results, ["complex_udf"]) + + # Verify they can be re-run (CREATE OR REPLACE) + results = run_dbt(["run"]) + assert len(results) == 1 + + +class TestUDFWithCustomLocation: + @pytest.fixture(scope="class") + def models(self, unique_schema): + # Use the unique_schema fixture to dynamically set the schema + custom_udf_with_schema = """ +{{ config( + materialized='udf', + database='dbt_test', + schema='""" + unique_schema + """', + parameter_list='x DOUBLE', + returns='DOUBLE' +) }} + +RETURN x * x +""" + return { + "custom_location_udf.sql": custom_udf_with_schema, + } + + def test_udf_with_custom_location(self, project): + # Run the UDF with custom database/schema configuration + results = run_dbt(["run", "--select", "custom_location_udf"]) + assert len(results) == 1 + check_result_nodes_by_name(results, ["custom_location_udf"]) + + +class TestUDFErrorHandling: + @pytest.fixture(scope="class") + def models(self): + return { + "invalid_udf_no_returns.sql": invalid_udf_no_returns_sql, + "invalid_udf_no_params.sql": invalid_udf_no_params_sql, + "no_params_udf.sql": no_params_udf_sql, + "no_param_list_config.sql": no_param_list_config_sql, + } + + def test_missing_returns_config(self, project): + # This should fail due to missing returns config + results = run_dbt(["run", "--select", "invalid_udf_no_returns"], expect_pass=False) + assert len(results) == 1 + assert results[0].status == "error" + assert "UDF materialization requires 'returns' configuration" in results[0].message + + def test_missing_parameter_list_config(self, project): + # This should fail due to missing parameter_list config + results = run_dbt(["run", "--select", "invalid_udf_no_params"], expect_pass=False) + assert len(results) == 1 + assert results[0].status == "error" + assert "UDF materialization requires 'parameter_list' configuration" in results[0].message + + def test_empty_parameter_list(self, project): + # This should fail due to empty parameter_list + results = run_dbt(["run", "--select", "no_params_udf"], expect_pass=False) + assert len(results) == 1 + assert results[0].status == "error" + assert "UDF materialization requires 'parameter_list' configuration" in results[0].message + + def test_no_parameter_list_config_key(self, project): + # This should fail due to no parameter_list config key + results = run_dbt(["run", "--select", "no_param_list_config"], expect_pass=False) + assert len(results) == 1 + assert results[0].status == "error" + assert "UDF materialization requires 'parameter_list' configuration" in results[0].message + + +class TestUDFDownstreamUsage: + @pytest.fixture(scope="class") + def models(self): + return { + "add_numbers.sql": """ + {{ config( + materialized='udf', + parameter_list='a INT, b INT', + returns='INT' + ) }} + + RETURN a + b + """, + "simple_view.sql": """ + {{ config(materialized='view') }} + + SELECT + {{ ref('add_numbers') }}(5, 3) as sum_output + """, + "multiply_numbers.sql": """ + {{ config( + materialized='udf', + parameter_list='x INT, y INT', + returns='INT' + ) }} + + RETURN x * y + """, + "calculations_view.sql": """ + {{ config(materialized='view') }} + + SELECT + {{ ref('add_numbers') }}(5, 3) as sum_result, + {{ ref('multiply_numbers') }}(4, 6) as product_result, + {{ ref('add_numbers') }}( + {{ ref('multiply_numbers') }}(2, 3), + 10 + ) as nested_result + """, + "final_report.sql": """ + {{ config( + materialized='view' + ) }} + + SELECT + sum_result, + product_result, + nested_result, + {{ ref('add_numbers') }}(sum_result, product_result) as total + FROM {{ ref('calculations_view') }} + """, + } + + def test_udf_downstream_usage(self, project): + # Run all models + results = run_dbt(["run"]) + assert len(results) == 5 + check_result_nodes_by_name( + results, ["add_numbers", "multiply_numbers", "simple_view", "calculations_view", "final_report"] + ) + + # Test selective execution + results = run_dbt(["run", "--select", "+final_report"]) + assert len(results) == 4 + + # Test that UDFs are properly tracked in the manifest + results = run_dbt(["ls", "--select", "add_numbers+"]) + assert "calculations_view" in str(results) + assert "final_report" in str(results) \ No newline at end of file diff --git a/tests/unit/test_udf_materialization.py b/tests/unit/test_udf_materialization.py index 77aa237d..161332cb 100644 --- a/tests/unit/test_udf_materialization.py +++ b/tests/unit/test_udf_materialization.py @@ -1,4 +1,5 @@ import os +import re MACRO_PATH = os.path.join( "dbt", "include", "dremio", "macros", "materializations", "udf.sql" @@ -6,13 +7,89 @@ def test_udf_materialization_file_exists(): + """Test that the UDF materialization macro file exists.""" assert os.path.isfile(MACRO_PATH) def test_udf_materialization_contains_expected_sql(): + """Test that the UDF materialization contains expected SQL components.""" with open(MACRO_PATH) as f: contents = f.read() + # Basic structure assert "{% materialization udf" in contents + assert "adapter='dremio'" in contents assert "CREATE OR REPLACE FUNCTION" in contents - assert "return({'relations': []})" in contents + assert "return({'relations': [target_relation]})" in contents + + # Configuration handling + assert "config.get('database'" in contents + assert "config.get('schema'" in contents + assert "config.get('parameter_list'" in contents + assert "config.get('returns'" in contents + + # SQL construction + assert "RETURNS {{ ret }}" in contents + assert "{{ target }}({{ parameter_list }})" in contents + assert "{{ sql }}" in contents + + +def test_udf_materialization_macro_structure(): + """Test the macro structure and Jinja2 syntax.""" + with open(MACRO_PATH) as f: + contents = f.read() + + # Check proper Jinja2 block structure + assert contents.count("{% materialization") == 1 + assert contents.count("{% endmaterialization %}") == 1 + assert contents.count("{%- set") >= 5 # identifier, target_relation, udf_database, udf_schema, parameter_list, ret, target, create_sql + assert contents.count("{%- endset -%}") >= 1 # create_sql + + # Check statement execution + assert "{% call statement('main')" in contents + assert "{%- endcall %}" in contents + + +def test_udf_materialization_target_construction(): + """Test that the target name is properly constructed with database.schema.identifier.""" + with open(MACRO_PATH) as f: + contents = f.read() + + # Check the target construction logic + assert "udf_database ~ '.' ~ udf_schema ~ '.' ~ identifier" in contents + assert "config.get('database'" in contents + assert "config.get('schema'" in contents + + # Verify the pattern creates: database.schema.function_name + target_pattern = r"set target = udf_database ~ '\.' ~ udf_schema ~ '\.' ~ identifier" + assert re.search(target_pattern, contents, re.IGNORECASE) + + +def test_udf_materialization_sql_generation(): + """Test that the CREATE FUNCTION SQL is properly generated.""" + with open(MACRO_PATH) as f: + contents = f.read() + + # Check the SQL template structure + create_sql_block = re.search( + r"{%- set create_sql -%}(.*?){%- endset -%}", + contents, + re.DOTALL + ) + assert create_sql_block is not None + + sql_content = create_sql_block.group(1) + assert "CREATE OR REPLACE FUNCTION" in sql_content + assert "{{ target }}" in sql_content + assert "{{ parameter_list }}" in sql_content + assert "RETURNS {{ ret }}" in sql_content + assert "{{ sql }}" in sql_content + + +def test_udf_materialization_validation(): + """Test that the macro includes proper validation.""" + with open(MACRO_PATH) as f: + contents = f.read() + + # Check that validators are used for config values + assert "validator=validation.any[basestring]" in contents From 340f29e39c511ab509cf83539e3e270fedc370e2 Mon Sep 17 00:00:00 2001 From: rfeith2 Date: Fri, 25 Jul 2025 05:33:43 -0700 Subject: [PATCH 5/5] final changes --- .../dremio/macros/materializations/udf.sql | 34 +++-- .../adapter/materialization/test_udf.py | 116 +++++++++++------- .../adapter/materialization/test_udf_hooks.py | 83 +++++++++++++ tests/unit/test_udf_materialization.py | 59 +++++---- 4 files changed, 205 insertions(+), 87 deletions(-) create mode 100644 tests/functional/adapter/materialization/test_udf_hooks.py diff --git a/dbt/include/dremio/macros/materializations/udf.sql b/dbt/include/dremio/macros/materializations/udf.sql index ae8cae17..130421fa 100644 --- a/dbt/include/dremio/macros/materializations/udf.sql +++ b/dbt/include/dremio/macros/materializations/udf.sql @@ -14,35 +14,49 @@ limitations under the License.*/ {% materialization udf, adapter='dremio' %} {%- set identifier = model['alias'] -%} - {%- set target_relation = this.incorporate(type='view') -%} {%- set udf_database = config.get('database', validator=validation.any[basestring]) or database -%} {%- set udf_schema = config.get('schema', validator=validation.any[basestring]) or schema -%} - {%- set target = udf_database ~ '.' ~ udf_schema ~ '.' ~ identifier -%} + + {# Create proper relation object for UDF #} + {%- set target_relation = api.Relation.create( + identifier=identifier, + schema=udf_schema, + database=udf_database, + type='view') -%} {%- set parameter_list = config.get('parameter_list', validator=validation.any[basestring]) -%} {%- set ret = config.get('returns', validator=validation.any[basestring]) -%} {# Validate required configurations #} {%- if parameter_list is none or parameter_list == '' -%} - {{ exceptions.raise_compiler_error("UDF materialization requires 'parameter_list' configuration (e.g., parameter_list='x INT, y INT').") }} + {{ exceptions.raise_compiler_error("UDF materialization requires 'parameter_list' configuration specifying the function parameters (e.g., parameter_list='x INT, y INT').") }} {%- endif -%} {%- if ret is none -%} {{ exceptions.raise_compiler_error("UDF materialization requires 'returns' configuration specifying the return type (e.g., returns='INT').") }} {%- endif -%} - {%- set create_sql -%} -CREATE OR REPLACE FUNCTION {{ target }}({{ parameter_list }}) -RETURNS {{ ret }} + {% set grant_config = config.get('grants') %} + + {{ run_hooks(pre_hooks) }} -{{ sql }} -; - {%- endset -%} + {# Check for existing relation and handle conflicts #} + {%- set old_relation = adapter.get_relation(database=udf_database, schema=udf_schema, identifier=identifier) -%} + {%- if old_relation is not none -%} + {{ adapter.drop_relation(old_relation) }} + {%- endif -%} + {# Build model #} {% call statement('main') -%} - {{ create_sql }} + create or replace function {{ target_relation }}({{ parameter_list }}) + returns {{ ret }} + {{ sql }} {%- endcall %} + {% do apply_grants(target_relation, grant_config, should_revoke=should_revoke) %} + + {{ run_hooks(post_hooks) }} + {{ return({'relations': [target_relation]}) }} {% endmaterialization %} diff --git a/tests/functional/adapter/materialization/test_udf.py b/tests/functional/adapter/materialization/test_udf.py index efbf8a6f..64b5c0f1 100644 --- a/tests/functional/adapter/materialization/test_udf.py +++ b/tests/functional/adapter/materialization/test_udf.py @@ -102,28 +102,54 @@ """ -class TestBasicUDFMaterialization: +class TestUDFMaterialization: @pytest.fixture(scope="class") - def models(self): + def models(self, unique_schema): + # Use the unique_schema fixture to dynamically set the schema for custom location test + custom_udf_with_schema = ( + """ +{{ config( + materialized='udf', + database='dbt_test', + schema='""" + + unique_schema + + """', + parameter_list='x DOUBLE', + returns='DOUBLE' +) }} + +RETURN x * x +""" + ) return { "basic_udf.sql": basic_udf_sql, "downstream_model.sql": downstream_model_sql, + "complex_udf.sql": complex_udf_sql, + "custom_location_udf.sql": custom_udf_with_schema, } def test_basic_udf_creation(self, project): # First compile to see what's generated compile_results = run_dbt(["compile"]) - + # Read the compiled SQL to debug import os - compiled_path = os.path.join(project.project_root, "target", "compiled", "test", "models", "downstream_model.sql") + + compiled_path = os.path.join( + project.project_root, + "target", + "compiled", + "test", + "models", + "downstream_model.sql", + ) if os.path.exists(compiled_path): - with open(compiled_path, 'r') as f: + with open(compiled_path, "r") as f: compiled_sql = f.read() print(f"\n\nCOMPILED downstream_model.sql:\n{compiled_sql}\n\n") - - # Run both models - results = run_dbt(["run"]) + + # Run both models + results = run_dbt(["run", "--select", "basic_udf downstream_model"]) assert len(results) == 2 check_result_nodes_by_name(results, ["basic_udf", "downstream_model"]) @@ -131,44 +157,16 @@ def test_basic_udf_creation(self, project): results = run_dbt(["docs", "generate"]) assert results is not None - -class TestUDFParameterVariations: - @pytest.fixture(scope="class") - def models(self): - return { - "complex_udf.sql": complex_udf_sql, - } - def test_udf_parameter_variations(self, project): # Run complex UDF model - results = run_dbt(["run"]) + results = run_dbt(["run", "--select", "complex_udf"]) assert len(results) == 1 check_result_nodes_by_name(results, ["complex_udf"]) # Verify they can be re-run (CREATE OR REPLACE) - results = run_dbt(["run"]) + results = run_dbt(["run", "--select", "complex_udf"]) assert len(results) == 1 - -class TestUDFWithCustomLocation: - @pytest.fixture(scope="class") - def models(self, unique_schema): - # Use the unique_schema fixture to dynamically set the schema - custom_udf_with_schema = """ -{{ config( - materialized='udf', - database='dbt_test', - schema='""" + unique_schema + """', - parameter_list='x DOUBLE', - returns='DOUBLE' -) }} - -RETURN x * x -""" - return { - "custom_location_udf.sql": custom_udf_with_schema, - } - def test_udf_with_custom_location(self, project): # Run the UDF with custom database/schema configuration results = run_dbt(["run", "--select", "custom_location_udf"]) @@ -188,31 +186,48 @@ def models(self): def test_missing_returns_config(self, project): # This should fail due to missing returns config - results = run_dbt(["run", "--select", "invalid_udf_no_returns"], expect_pass=False) + results = run_dbt( + ["run", "--select", "invalid_udf_no_returns"], expect_pass=False + ) assert len(results) == 1 assert results[0].status == "error" - assert "UDF materialization requires 'returns' configuration" in results[0].message + assert ( + "UDF materialization requires 'returns' configuration" in results[0].message + ) def test_missing_parameter_list_config(self, project): # This should fail due to missing parameter_list config - results = run_dbt(["run", "--select", "invalid_udf_no_params"], expect_pass=False) + results = run_dbt( + ["run", "--select", "invalid_udf_no_params"], expect_pass=False + ) assert len(results) == 1 assert results[0].status == "error" - assert "UDF materialization requires 'parameter_list' configuration" in results[0].message + assert ( + "UDF materialization requires 'parameter_list' configuration" + in results[0].message + ) def test_empty_parameter_list(self, project): # This should fail due to empty parameter_list results = run_dbt(["run", "--select", "no_params_udf"], expect_pass=False) assert len(results) == 1 assert results[0].status == "error" - assert "UDF materialization requires 'parameter_list' configuration" in results[0].message + assert ( + "UDF materialization requires 'parameter_list' configuration" + in results[0].message + ) def test_no_parameter_list_config_key(self, project): # This should fail due to no parameter_list config key - results = run_dbt(["run", "--select", "no_param_list_config"], expect_pass=False) + results = run_dbt( + ["run", "--select", "no_param_list_config"], expect_pass=False + ) assert len(results) == 1 assert results[0].status == "error" - assert "UDF materialization requires 'parameter_list' configuration" in results[0].message + assert ( + "UDF materialization requires 'parameter_list' configuration" + in results[0].message + ) class TestUDFDownstreamUsage: @@ -273,7 +288,14 @@ def test_udf_downstream_usage(self, project): results = run_dbt(["run"]) assert len(results) == 5 check_result_nodes_by_name( - results, ["add_numbers", "multiply_numbers", "simple_view", "calculations_view", "final_report"] + results, + [ + "add_numbers", + "multiply_numbers", + "simple_view", + "calculations_view", + "final_report", + ], ) # Test selective execution @@ -283,4 +305,4 @@ def test_udf_downstream_usage(self, project): # Test that UDFs are properly tracked in the manifest results = run_dbt(["ls", "--select", "add_numbers+"]) assert "calculations_view" in str(results) - assert "final_report" in str(results) \ No newline at end of file + assert "final_report" in str(results) diff --git a/tests/functional/adapter/materialization/test_udf_hooks.py b/tests/functional/adapter/materialization/test_udf_hooks.py new file mode 100644 index 00000000..4849392b --- /dev/null +++ b/tests/functional/adapter/materialization/test_udf_hooks.py @@ -0,0 +1,83 @@ +# Copyright (C) 2022 Dremio Corporation + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +from dbt.tests.util import run_dbt +from tests.utils.util import SOURCE + + +# UDF with pre-hook that creates a view and post-hook that drops it +udf_with_hooks_sql = """ +{{ config( + materialized='udf', + parameter_list='x INT', + returns='INT', + pre_hook="CREATE VIEW {{ target.database }}.{{ this.schema }}.hook_test_view AS SELECT 1 as id", + post_hook="DROP VIEW {{ target.database }}.{{ this.schema }}.hook_test_view" +) }} + +RETURN x * 2 +""" + +# Multiple hooks test +udf_with_multiple_hooks_sql = """ +{{ config( + materialized='udf', + parameter_list='x INT, y INT', + returns='INT', + pre_hook=[ + "CREATE VIEW {{ target.database }}.{{ this.schema }}.hook_test_view_1 AS SELECT 1 as id", + "CREATE VIEW {{ target.database }}.{{ this.schema }}.hook_test_view_2 AS SELECT 2 as id" + ], + post_hook=[ + "DROP VIEW {{ target.database }}.{{ this.schema }}.hook_test_view_1", + "DROP VIEW {{ target.database }}.{{ this.schema }}.hook_test_view_2" + ] +) }} + +RETURN x + y +""" + + +class TestUDFHooks: + @pytest.fixture(scope="class") + def models(self): + return { + "udf_with_hooks.sql": udf_with_hooks_sql, + "udf_with_multiple_hooks.sql": udf_with_multiple_hooks_sql, + } + + def test_udf_single_hook_execution(self, project): + # Run the UDF with hooks + results = run_dbt(["run", "--select", "udf_with_hooks"]) + assert len(results) == 1 + + # If this succeeded, it means: + # 1. Pre-hook created the table + # 2. UDF was created successfully + # 3. Post-hook dropped the table + + # Run it again to ensure table was properly cleaned up + # (would fail if table still exists) + results = run_dbt(["run", "--select", "udf_with_hooks"]) + assert len(results) == 1 + + def test_udf_multiple_hooks_execution(self, project): + # Run the UDF with multiple hooks + results = run_dbt(["run", "--select", "udf_with_multiple_hooks"]) + assert len(results) == 1 + + # Run it again to ensure both tables were cleaned up + results = run_dbt(["run", "--select", "udf_with_multiple_hooks"]) + assert len(results) == 1 diff --git a/tests/unit/test_udf_materialization.py b/tests/unit/test_udf_materialization.py index 161332cb..45e08f3d 100644 --- a/tests/unit/test_udf_materialization.py +++ b/tests/unit/test_udf_materialization.py @@ -19,18 +19,18 @@ def test_udf_materialization_contains_expected_sql(): # Basic structure assert "{% materialization udf" in contents assert "adapter='dremio'" in contents - assert "CREATE OR REPLACE FUNCTION" in contents + assert "create or replace function" in contents assert "return({'relations': [target_relation]})" in contents - + # Configuration handling assert "config.get('database'" in contents assert "config.get('schema'" in contents assert "config.get('parameter_list'" in contents assert "config.get('returns'" in contents - + # SQL construction - assert "RETURNS {{ ret }}" in contents - assert "{{ target }}({{ parameter_list }})" in contents + assert "returns {{ ret }}" in contents + assert "{{ target_relation }}({{ parameter_list }})" in contents assert "{{ sql }}" in contents @@ -38,51 +38,50 @@ def test_udf_materialization_macro_structure(): """Test the macro structure and Jinja2 syntax.""" with open(MACRO_PATH) as f: contents = f.read() - + # Check proper Jinja2 block structure assert contents.count("{% materialization") == 1 assert contents.count("{% endmaterialization %}") == 1 - assert contents.count("{%- set") >= 5 # identifier, target_relation, udf_database, udf_schema, parameter_list, ret, target, create_sql - assert contents.count("{%- endset -%}") >= 1 # create_sql - + assert ( + contents.count("{%- set") >= 5 + ) # identifier, udf_database, udf_schema, target_relation, parameter_list, ret, old_relation + # No longer using create_sql block - SQL is directly in statement call + # Check statement execution assert "{% call statement('main')" in contents assert "{%- endcall %}" in contents def test_udf_materialization_target_construction(): - """Test that the target name is properly constructed with database.schema.identifier.""" + """Test that the target relation is properly constructed using api.Relation.create.""" with open(MACRO_PATH) as f: contents = f.read() - - # Check the target construction logic - assert "udf_database ~ '.' ~ udf_schema ~ '.' ~ identifier" in contents + + # Check the new relation construction logic using api.Relation.create + assert "api.Relation.create(" in contents + assert "identifier=identifier" in contents + assert "schema=udf_schema" in contents + assert "database=udf_database" in contents assert "config.get('database'" in contents assert "config.get('schema'" in contents - - # Verify the pattern creates: database.schema.function_name - target_pattern = r"set target = udf_database ~ '\.' ~ udf_schema ~ '\.' ~ identifier" - assert re.search(target_pattern, contents, re.IGNORECASE) def test_udf_materialization_sql_generation(): """Test that the CREATE FUNCTION SQL is properly generated.""" with open(MACRO_PATH) as f: contents = f.read() - - # Check the SQL template structure - create_sql_block = re.search( - r"{%- set create_sql -%}(.*?){%- endset -%}", - contents, - re.DOTALL + + # Check that SQL is properly structured within the statement call + statement_block = re.search( + r"{% call statement\('main'\) -%}(.*?){%- endcall %}", contents, re.DOTALL ) - assert create_sql_block is not None - - sql_content = create_sql_block.group(1) - assert "CREATE OR REPLACE FUNCTION" in sql_content - assert "{{ target }}" in sql_content + assert statement_block is not None + + sql_content = statement_block.group(1) + assert "create or replace function" in sql_content + assert "{{ target_relation }}" in sql_content assert "{{ parameter_list }}" in sql_content - assert "RETURNS {{ ret }}" in sql_content + assert "returns {{ ret }}" in sql_content assert "{{ sql }}" in sql_content @@ -90,6 +89,6 @@ def test_udf_materialization_validation(): """Test that the macro includes proper validation.""" with open(MACRO_PATH) as f: contents = f.read() - + # Check that validators are used for config values assert "validator=validation.any[basestring]" in contents