diff --git a/integration_tests/tests/adapter_query_runner.py b/integration_tests/tests/adapter_query_runner.py index c45cbb83a..fe8e75b56 100644 --- a/integration_tests/tests/adapter_query_runner.py +++ b/integration_tests/tests/adapter_query_runner.py @@ -15,6 +15,7 @@ from typing import Any, Dict, List, Optional from dbt.adapters.base import BaseAdapter +from elementary.clients.dbt.factory import RunnerMethod, get_dbt_runner_method from logger import get_logger logger = get_logger(__name__) @@ -76,11 +77,23 @@ class AdapterQueryRunner: Path to the dbt project directory. target : str Name of the dbt target / profile output to use. + runner_method : RunnerMethod, optional + How the tests run dbt. Only out-of-process runners need the DuckDB file + lock released between queries (see ``_release_duckdb_file_lock``). """ - def __init__(self, project_dir: str, target: str) -> None: + def __init__( + self, + project_dir: str, + target: str, + runner_method: Optional[RunnerMethod] = None, + ) -> None: self._project_dir = project_dir self._target = target + # Every runner except API shells out, so it needs the file lock released. + self._dbt_runs_out_of_process = ( + runner_method or get_dbt_runner_method() + ) != RunnerMethod.API self._adapter: BaseAdapter = self._create_adapter(project_dir, target) self._ref_map: Optional[Dict[str, str]] = None self._source_map: Optional[Dict[tuple, str]] = None @@ -246,8 +259,45 @@ def has_non_ref_jinja(query: str) -> bool: def execute_sql(self, sql: str) -> None: """Execute a SQL statement that does not return results (DDL/DML).""" - with self._adapter.connection_named("execute_sql"): - self._adapter.execute(sql, fetch=False) + try: + with self._adapter.connection_named("execute_sql"): + self._adapter.execute(sql, fetch=False) + finally: + self._release_duckdb_file_lock() + + def _release_duckdb_file_lock(self) -> None: + """Release the DuckDB connection so an out-of-process dbt run can take the file lock. + + A file-backed DuckDB database allows a single read-write process at a time, so + the environment cached on ``DuckDBConnectionManager`` has to be dropped before + dbt runs in another process (any ``--runner-method`` other than ``api``). Dropping + the last reference closes the DuckDB connection via ``LocalEnvironment.__del__``; + the next query builds a new environment. + + No-op unless dbt actually runs out of process, and for non-DuckDB adapters, + ``path: ":memory:"`` (closing it would discard the database) and MotherDuck + (hosted, not a local file). + """ + if not self._dbt_runs_out_of_process: + return + + if self._adapter.type() != "duckdb": + return + + credentials = self._adapter.config.credentials + if credentials.path == ":memory:" or credentials.is_motherduck: + return + + from dbt.adapters.duckdb.connections import DuckDBConnectionManager + + # Called from a `finally`, so neither step may mask the original error. + try: + self._adapter.connections.cleanup_all() + DuckDBConnectionManager.close_all_connections() + except Exception as exc: + logger.warning( + f"Ignoring error while releasing the DuckDB connection: {exc}" + ) @property def schema_name(self) -> str: @@ -267,8 +317,11 @@ def run_query(self, prerendered_query: str) -> List[Dict[str, Any]]: if self.has_non_ref_jinja(prerendered_query): raise UnsupportedJinjaError(prerendered_query) sql = self.resolve_refs(prerendered_query) - with self._adapter.connection_named("run_query"): - _response, table = self._adapter.execute(sql, fetch=True) + try: + with self._adapter.connection_named("run_query"): + _response, table = self._adapter.execute(sql, fetch=True) + finally: + self._release_duckdb_file_lock() # Convert agate Table → list[dict] matching agate_to_dicts behaviour columns = [c.lower() for c in table.column_names] diff --git a/integration_tests/tests/conftest.py b/integration_tests/tests/conftest.py index 579c5f70e..d36d83a85 100644 --- a/integration_tests/tests/conftest.py +++ b/integration_tests/tests/conftest.py @@ -1,3 +1,4 @@ +import os import shutil from pathlib import Path from tempfile import mkdtemp @@ -6,7 +7,12 @@ import pytest import yaml from dbt.version import __version__ as dbt_version -from dbt_project import PYTEST_XDIST_WORKER, SCHEMA_NAME_SUFFIX, DbtProject +from dbt_project import ( + DBT2_RUNNERS, + PYTEST_XDIST_WORKER, + SCHEMA_NAME_SUFFIX, + DbtProject, +) from elementary.clients.dbt.factory import RunnerMethod from env import Environment from logger import get_logger @@ -18,6 +24,7 @@ "bigquery", "redshift", "databricks_catalog", + "duckdb", ] logger = get_logger(__name__) @@ -70,7 +77,7 @@ def _edit_packages_yml_to_include_absolute_elementary_package_path( def _remove_python_models_for_dbt_fusion( project_dir_copy: str, runner_method: Optional[RunnerMethod] ): - if runner_method != RunnerMethod.FUSION: + if runner_method not in DBT2_RUNNERS: return logger.info(f"Removing python tests for project {project_dir_copy}") @@ -129,7 +136,7 @@ def only_on_targets(request, target: str): @pytest.fixture(autouse=True) def skip_for_dbt_fusion(request, runner_method: Optional[RunnerMethod]): if request.node.get_closest_marker("skip_for_dbt_fusion"): - if runner_method == RunnerMethod.FUSION: + if runner_method in DBT2_RUNNERS: pytest.skip("Test unsupported for dbt fusion") @@ -169,16 +176,60 @@ def clear_on_end(request) -> bool: return request.config.getoption("--clear-on-end") +def _profile_path(target: str) -> Optional[str]: + """Return the ``path`` dbt resolves for *target*, if it can be resolved. + + Goes through dbt's own profile handling instead of reading ``profiles.yml`` + directly, so ``env_var`` and any other Jinja in ``path`` render the way dbt + will render it, and an omitted ``path`` picks up dbt-duckdb's ``:memory:`` + default. ``env_var`` needs the invocation context, hence the + ``set_invocation_context`` call. + + Returns ``None`` when the profile cannot be resolved, so callers only act on + a value they positively read. + """ + from argparse import Namespace + + from dbt.config.runtime import RuntimeConfig + from dbt.flags import set_from_args + from dbt_common.context import set_invocation_context + + profiles_dir = os.environ.get("DBT_PROFILES_DIR", os.path.expanduser("~/.dbt")) + args = Namespace( + project_dir=str(DBT_PROJECT_PATH), + profiles_dir=profiles_dir, + target=target, + threads=1, + vars={}, + profile=None, + PROFILES_DIR=profiles_dir, + PROJECT_DIR=str(DBT_PROJECT_PATH), + ) + try: + set_invocation_context(os.environ) + set_from_args(args, None) + return getattr(RuntimeConfig.from_args(args).credentials, "path", None) + except Exception: + return None + + @pytest.fixture(scope="session") def runner_method(request, target: str) -> Optional[RunnerMethod]: runner_method_str = request.config.getoption("--runner-method") if runner_method_str: runner_method = RunnerMethod(runner_method_str) + if runner_method in DBT2_RUNNERS and target not in DBT_FUSION_SUPPORTED_TARGETS: + raise ValueError(f"Fusion runner is not supported for target: {target}") if ( - runner_method == RunnerMethod.FUSION - and target not in DBT_FUSION_SUPPORTED_TARGETS + runner_method in DBT2_RUNNERS + and target == "duckdb" + and _profile_path(target) == ":memory:" ): - raise ValueError(f"Fusion runner is not supported for target: {target}") + raise ValueError( + f"The {runner_method.value} runner runs dbt in a subprocess, which " + "cannot share an in-memory DuckDB database with the tests. Set " + "`path` to a file in the duckdb target of your profiles.yml." + ) return runner_method return None diff --git a/integration_tests/tests/dbt_project.py b/integration_tests/tests/dbt_project.py index 020e27db4..507445ac8 100644 --- a/integration_tests/tests/dbt_project.py +++ b/integration_tests/tests/dbt_project.py @@ -23,6 +23,9 @@ PYTEST_XDIST_WORKER = os.environ.get("PYTEST_XDIST_WORKER", None) SCHEMA_NAME_SUFFIX = f"_{PYTEST_XDIST_WORKER}" if PYTEST_XDIST_WORKER else "" +# Both spellings of the dbt 2.0 runner: FUSION is a legacy alias for DBT2. +DBT2_RUNNERS = (RunnerMethod.DBT2, RunnerMethod.FUSION) + _DEFAULT_VARS = { "disable_dbt_invocation_autoupload": True, "disable_dbt_artifacts_autoupload": True, @@ -89,7 +92,7 @@ def _get_query_runner(self) -> AdapterQueryRunner: """Lazily initialize the direct adapter query runner.""" if self._query_runner is None: self._query_runner = AdapterQueryRunner( - str(self.project_dir_path), self.target + str(self.project_dir_path), self.target, self.runner_method ) return self._query_runner @@ -398,7 +401,7 @@ def seed(self, data: List[dict], table_name: str): def _fix_seed_if_needed(self, table_name: str) -> None: # Hack for BigQuery - seems like we get empty strings instead of nulls in seeds, so we # fix them here. - if self.runner_method == RunnerMethod.FUSION and self.target == "bigquery": + if self.runner_method in DBT2_RUNNERS and self.target == "bigquery": self.dbt_runner.run_operation( "elementary_tests.replace_empty_strings_with_nulls", macro_args={"table_name": table_name}, diff --git a/macros/edr/materializations/test/test.sql b/macros/edr/materializations/test/test.sql index fdef2c8fc..5c336aa68 100644 --- a/macros/edr/materializations/test/test.sql +++ b/macros/edr/materializations/test/test.sql @@ -193,6 +193,15 @@ {% set database, schema = elementary.get_package_database_and_schema() %} {% set test_id = model["alias"] %} {% set relation = elementary.create_temp_table(database, schema, test_id, sql) %} + {# On DuckDB under dbt-fusion this is a regular table rather than a temp one, so + register it for the on_run_end cleanup - otherwise it survives until the stale + cleanup runs. #} + {% if target.type == "duckdb" and elementary.is_dbt_fusion() %} + {% set test_entry = elementary.get_cache( + "temp_test_table_relations_map" + ).setdefault(test_id, {}) %} + {% do test_entry.update({"test_result": relation}) %} + {% endif %} {% set new_sql %} select * from {{ relation }} {% endset %} diff --git a/macros/edr/tests/on_run_end/handle_tests_results.sql b/macros/edr/tests/on_run_end/handle_tests_results.sql index b41114080..d7e33a8e6 100644 --- a/macros/edr/tests/on_run_end/handle_tests_results.sql +++ b/macros/edr/tests/on_run_end/handle_tests_results.sql @@ -227,6 +227,8 @@ {% if not elementary.has_temp_table_support() %} {% do elementary.fully_drop_relation(temp_relation) %} + {# DuckDB: commit the drop too, so it survives the same ROLLBACK #} + {% if target.type == "duckdb" %} {% do adapter.commit() %} {% endif %} {% endif %} {% endmacro %} @@ -287,6 +289,8 @@ {% if not elementary.has_temp_table_support() %} {% do elementary.fully_drop_relation(temp_relation) %} + {# DuckDB: commit the drop too, so it survives the same ROLLBACK #} + {% if target.type == "duckdb" %} {% do adapter.commit() %} {% endif %} {% endif %} {% endmacro %} diff --git a/macros/utils/table_operations/create_table_as.sql b/macros/utils/table_operations/create_table_as.sql index 61a16a6a5..f0399009c 100644 --- a/macros/utils/table_operations/create_table_as.sql +++ b/macros/utils/table_operations/create_table_as.sql @@ -120,6 +120,11 @@ {% macro duckdb__edr_get_create_table_as_sql( temporary, relation, sql_query, expiration_hours=none ) %} + {# dbt-fusion uses connection pooling - DuckDB temp tables are session-scoped and + aren't visible in other sessions. DuckDB also rejects TEMPORARY on a qualified + name. Create regular tables instead - these are cleaned up by Elementary's + normal cleanup logic. #} + {% set temporary = temporary and not elementary.is_dbt_fusion() %} create or replace {% if temporary %} temporary {% endif %} table {{ relation }} as {{ sql_query }} {% endmacro %} diff --git a/macros/utils/table_operations/has_temp_table_support.sql b/macros/utils/table_operations/has_temp_table_support.sql index 988267675..3af7cf0d7 100644 --- a/macros/utils/table_operations/has_temp_table_support.sql +++ b/macros/utils/table_operations/has_temp_table_support.sql @@ -25,11 +25,13 @@ {% endmacro %} +{# Session-scoped temp tables are invisible to the next statement under dbt-fusion, + which runs each statement on a pooled connection ("relation does not exist"). + These adapters fall back to regular tables with cleanup instead. #} {% macro redshift__has_temp_table_support() %} - {# dbt-fusion uses connection pooling, so temp tables created in one session - are not visible in other sessions, causing "relation does not exist" errors. - Use regular tables with cleanup instead. #} - {% if elementary.is_dbt_fusion() %} {% do return(false) %} - {% else %} {% do return(true) %} - {% endif %} + {% do return(not elementary.is_dbt_fusion()) %} +{% endmacro %} + +{% macro duckdb__has_temp_table_support() %} + {% do return(not elementary.is_dbt_fusion()) %} {% endmacro %} diff --git a/macros/utils/table_operations/make_temp_relation.sql b/macros/utils/table_operations/make_temp_relation.sql index bcf031937..14050eff0 100644 --- a/macros/utils/table_operations/make_temp_relation.sql +++ b/macros/utils/table_operations/make_temp_relation.sql @@ -72,20 +72,41 @@ {{ return(elementary.spark__edr_make_temp_relation(base_relation, suffix)) }} {% endmacro %} +{% macro _fusion_regular_temp_relation(base_relation, suffix) %} + {# dbt-fusion runs each statement on a pooled connection, so a session-scoped temp + relation is not visible to the statement that reads it. Build a regular relation + with an explicit database/schema instead - Elementary's normal cleanup drops it. #} + {% set tmp_identifier = elementary.table_name_with_suffix( + base_relation.identifier, suffix + ) %} + {% do return( + api.Relation.create( + identifier=tmp_identifier, + schema=base_relation.schema, + database=base_relation.database, + type="table", + ) + ) %} +{% endmacro %} + {% macro redshift__edr_make_temp_relation(base_relation, suffix) %} {% if elementary.is_dbt_fusion() %} {# Workaround for dbt-fusion temp table metadata bug - create regular relations with explicit schema/database instead of temp relations #} - {% set tmp_identifier = elementary.table_name_with_suffix( - base_relation.identifier, suffix + {% do return( + elementary._fusion_regular_temp_relation(base_relation, suffix) ) %} - {% set tmp_relation = api.Relation.create( - identifier=tmp_identifier, - schema=base_relation.schema, - database=base_relation.database, - type="table", + {% else %} {% do return(dbt.make_temp_relation(base_relation, suffix)) %} + {% endif %} +{% endmacro %} + +{% macro duckdb__edr_make_temp_relation(base_relation, suffix) %} + {% if elementary.is_dbt_fusion() %} + {# DuckDB additionally rejects TEMPORARY on a qualified name, so the regular + relation is required rather than merely convenient. #} + {% do return( + elementary._fusion_regular_temp_relation(base_relation, suffix) ) %} - {% do return(tmp_relation) %} {% else %} {% do return(dbt.make_temp_relation(base_relation, suffix)) %} {% endif %} {% endmacro %}