From 8ff9ef2c3ea27efe2b09736d6b70a65652f78e12 Mon Sep 17 00:00:00 2001 From: Bill Date: Sat, 6 Jan 2024 10:18:43 +0000 Subject: [PATCH 1/2] =?UTF-8?q?=E2=9C=A8=20feat:=20Enable=20module=20impor?= =?UTF-8?q?ts=20and=20better=20environment=20variable=20handling?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- dbt_py/main.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/dbt_py/main.py b/dbt_py/main.py index c55b4b7..1b4f56a 100644 --- a/dbt_py/main.py +++ b/dbt_py/main.py @@ -15,9 +15,7 @@ PROJECT_ROOT = pathlib.Path(__file__).parent.parent # Python-style ref, e.g. `package.module.submodule` -PACKAGE_ROOT: str = os.environ.get("DBT_PY_PACKAGE_ROOT", "custom") -# The name to associate with the package -PACKAGE_NAME: str = os.environ.get("DBT_PY_PACKAGE_NAME", PACKAGE_ROOT) +PACKAGE_ROOT: str = os.environ.get("DBT_PY_PACKAGE_ROOT") def import_submodules( @@ -30,6 +28,10 @@ def import_submodules( - https://stackoverflow.com/a/25562415/10730311 """ package = importlib.import_module(package_name) + if not hasattr(package, "__path__"): + # `package` is a module, don't recurse any further + return {} + results = {} for loader, name, is_pkg in pkgutil.walk_packages(package.__path__): full_name = f"{package.__name__}.{name}" @@ -44,9 +46,12 @@ def new_get_context_modules() -> dict[str, dict[str, Any]]: """ Append the custom modules into the whitelisted dbt modules. """ - import_submodules(PACKAGE_ROOT) + package_root: str = PACKAGE_ROOT or "custom" + package_name: str = os.environ.get("DBT_PY_PACKAGE_NAME") or package_root + + import_submodules(package_root) modules = _get_context_modules() - modules[PACKAGE_NAME] = importlib.import_module(PACKAGE_ROOT) # type: ignore + modules[package_name] = importlib.import_module(package_root) # type: ignore return modules From 3c11a917906994a6aa555b2189dca44c7a454009 Mon Sep 17 00:00:00 2001 From: Bill Date: Sat, 6 Jan 2024 10:44:05 +0000 Subject: [PATCH 2/2] =?UTF-8?q?=E2=9C=85=20refactor:=20Give=20the=20Jinja?= =?UTF-8?q?=20context=20a=20different=20name=20to=20the=20package?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tests/integration/jaffle-shop/dbt-commands.sh | 2 +- tests/integration/jaffle-shop/models/example.sql | 8 ++++---- tests/integration/test__integration.py | 2 +- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/integration/jaffle-shop/dbt-commands.sh b/tests/integration/jaffle-shop/dbt-commands.sh index 7aad172..8dabdf8 100644 --- a/tests/integration/jaffle-shop/dbt-commands.sh +++ b/tests/integration/jaffle-shop/dbt-commands.sh @@ -1,6 +1,6 @@ # Just for reference # $env:DBT_PY_PACKAGE_ROOT = "tests.integration.jaffle-shop.dbt_py_test" -# $env:DBT_PY_PACKAGE_NAME = "dbt_py_test" +# $env:DBT_PY_PACKAGE_NAME = "custom_py" dbt clean --project-dir tests/integration/jaffle-shop --profiles-dir tests/integration/jaffle-shop --no-clean-project-files-only dbt-py compile --project-dir tests/integration/jaffle-shop --profiles-dir tests/integration/jaffle-shop rm dbt_packages diff --git a/tests/integration/jaffle-shop/models/example.sql b/tests/integration/jaffle-shop/models/example.sql index ef4f2d4..ac076a7 100644 --- a/tests/integration/jaffle-shop/models/example.sql +++ b/tests/integration/jaffle-shop/models/example.sql @@ -1,5 +1,5 @@ -{{ modules.dbt_py_test.some_module.select_final() }} -{{ modules.dbt_py_test.some_module.salutation("World") }} +{{ modules.custom_py.some_module.select_final() }} +{{ modules.custom_py.some_module.salutation("World") }} -{{ modules.dbt_py_test.select_final() }} -{{ modules.dbt_py_test.salutation("World") }} +{{ modules.custom_py.select_final() }} +{{ modules.custom_py.salutation("World") }} diff --git a/tests/integration/test__integration.py b/tests/integration/test__integration.py index d0126d5..b364be5 100644 --- a/tests/integration/test__integration.py +++ b/tests/integration/test__integration.py @@ -37,7 +37,7 @@ def mock_env(monkeypatch) -> None: monkeypatch.setenv( "DBT_PY_PACKAGE_ROOT", "tests.integration.jaffle-shop.dbt_py_test" ) - monkeypatch.setenv("DBT_PY_PACKAGE_NAME", "dbt_py_test") + monkeypatch.setenv("DBT_PY_PACKAGE_NAME", "custom_py") def test__integration(mock_env) -> None: