dbt-python-imports 0.1.0__tar.gz

dbt_python_imports-0.1.0/.envrc

@@ -0,0 +1,2 @@
+source_up_secure
+layout_venv

dbt_python_imports-0.1.0/.github/workflows/test.yml

@@ -0,0 +1,57 @@
+
+name: Run tests
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+
+jobs:
+
+  get-hatch-test-envs:
+    name: Get hatch test environments
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+      - name: Install the latest version of uv
+        uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+      - name: List hatch envs
+        id: get-test-envs
+        run: |-
+          uv run hatch env show --json | jq -r '
+            with_entries(select(.key | startswith("hatch-test.")))
+            | "envs=\(@json)"
+          ' >> "$GITHUB_OUTPUT"
+
+    outputs:
+      envs: ${{ steps.get-test-envs.outputs.envs }}
+
+  test:
+    name: Run Hatch tests
+    needs: get-hatch-test-envs
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        env-name: ${{ fromJSON(needs.get-hatch-test-envs.outputs.envs).*.env-vars.ENV_NAME }}
+    steps:
+
+      - name: Checkout
+        uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ fromJSON(needs.get-hatch-test-envs.outputs.envs)[matrix.env-name].python }}
+
+      - name: Install the latest version of uv
+        uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+
+      - name: Run tests for ${{ matrix.env-name }}
+        run: uv run hatch --env=${{ matrix.env-name }} test
+

dbt_python_imports-0.1.0/.gitignore

@@ -0,0 +1,3 @@
+*_cache
+*.pyc
+.vscode

dbt_python_imports-0.1.0/PKG-INFO

@@ -0,0 +1,52 @@
+Metadata-Version: 2.4
+Name: dbt-python-imports
+Version: 0.1.0
+Summary: dbt plugin to add `modules.import` method to Jinja context to import arbitrary Python modules
+Author-email: Tim Vergenz <vergenzt@gmail.com>
+Requires-Python: >=3.8
+Requires-Dist: dbt-core>=1.6
+Description-Content-Type: text/markdown
+
+# dbt-python-imports
+
+This is a [dbt plugin][dbt-plugin] to allow importing arbitrary Python modules from Jinja templates.
+
+[dbt-plugin]: https://github.com/dbt-labs/dbt-core/blob/fa96acb15f79ae4f10b1d78f311f5ef2f4ed645e/core/dbt/plugins/manager.py
+
+## Usage
+
+1. Install `dbt-python-imports` via PyPI into the same Python environment as dbt itself.
+2. Call `modules.import(module_name)` to import the given module from a dbt Jinja expression.
+
+> [!warning]
+> **This allows arbitrary code execution.** (That's kindof the point.)
+>
+> Jinja's sandboxing is not foolproof (e.g. some adapters allow arbitrary reads/writes to disk), so you already shouldn't incorporate any untrusted dbt code/packages already. However the limited Jinja context available by default does currently make it *harder* to run arbitrary code.
+>
+> Make sure you're not installing any packages (dbt or Python) whose source you haven't inspected.
+
+## Example
+
+```console
+# download the artifacts from your latest Databricks dbt job run
+
+$ echo '
+{%- macro fetch_dbt_artifacts(job_name='dbt build', extract_to=flags.TARGET_PATH ~ '/remote-state') %}
+  {#- https://databricks-sdk-py.readthedocs.io/en/latest/workspace/jobs/jobs.html #}
+  {%- set jobs_api = adapter.config.credentials.authenticate().api_client.jobs %}
+  {%- set job = jobs_api.list(name=job_name, limit=1) | first %}
+  {%- set job_run = jobs_api.list_runs(job_id=job["job_id"], limit=1, completed_only=true, expand_tasks=true) | first %}
+  {%- set task_run = job_run["tasks"] | first %}
+  {%- set dbt_output = jobs_api.get_run_output(run_id=task_run["run_id"]).dbt_output %}
+  {%- set download_url = dbt_output.artifacts_link %}
+
+  {%- set system = modules.import("dbt_common.clients.system") %}
+  {%- set tar_path = extract_to ~ '/dbt-artifacts.tar.gz' %}
+  {%- do system.download(download_url, tar_path) %}
+  {%- do system.untar_package(tar_path, extract_to) %}
+{%- endmacro %}
+' > macros/fetch_dbt_artifacts.sql
+
+$ dbt run-operation fetch_dbt_artifacts
+
+```

dbt_python_imports-0.1.0/README.md

@@ -0,0 +1,43 @@
+# dbt-python-imports
+
+This is a [dbt plugin][dbt-plugin] to allow importing arbitrary Python modules from Jinja templates.
+
+[dbt-plugin]: https://github.com/dbt-labs/dbt-core/blob/fa96acb15f79ae4f10b1d78f311f5ef2f4ed645e/core/dbt/plugins/manager.py
+
+## Usage
+
+1. Install `dbt-python-imports` via PyPI into the same Python environment as dbt itself.
+2. Call `modules.import(module_name)` to import the given module from a dbt Jinja expression.
+
+> [!warning]
+> **This allows arbitrary code execution.** (That's kindof the point.)
+>
+> Jinja's sandboxing is not foolproof (e.g. some adapters allow arbitrary reads/writes to disk), so you already shouldn't incorporate any untrusted dbt code/packages already. However the limited Jinja context available by default does currently make it *harder* to run arbitrary code.
+>
+> Make sure you're not installing any packages (dbt or Python) whose source you haven't inspected.
+
+## Example
+
+```console
+# download the artifacts from your latest Databricks dbt job run
+
+$ echo '
+{%- macro fetch_dbt_artifacts(job_name='dbt build', extract_to=flags.TARGET_PATH ~ '/remote-state') %}
+  {#- https://databricks-sdk-py.readthedocs.io/en/latest/workspace/jobs/jobs.html #}
+  {%- set jobs_api = adapter.config.credentials.authenticate().api_client.jobs %}
+  {%- set job = jobs_api.list(name=job_name, limit=1) | first %}
+  {%- set job_run = jobs_api.list_runs(job_id=job["job_id"], limit=1, completed_only=true, expand_tasks=true) | first %}
+  {%- set task_run = job_run["tasks"] | first %}
+  {%- set dbt_output = jobs_api.get_run_output(run_id=task_run["run_id"]).dbt_output %}
+  {%- set download_url = dbt_output.artifacts_link %}
+
+  {%- set system = modules.import("dbt_common.clients.system") %}
+  {%- set tar_path = extract_to ~ '/dbt-artifacts.tar.gz' %}
+  {%- do system.download(download_url, tar_path) %}
+  {%- do system.untar_package(tar_path, extract_to) %}
+{%- endmacro %}
+' > macros/fetch_dbt_artifacts.sql
+
+$ dbt run-operation fetch_dbt_artifacts
+
+```

dbt_python_imports-0.1.0/hatch.toml

@@ -0,0 +1,34 @@
+[dirs.env]
+virtual = ".venvs"
+
+[envs.hatch-test]
+installer = "uv"
+dependency-groups = ["dev"]
+dependencies = [
+    "dbt-core~={matrix:dbt-core-version}",
+]
+
+[envs.hatch-test.env-vars]
+ENV_NAME = "{env_name}"
+DBT_SEND_ANONYMOUS_USAGE_STATS = "False"  # https://docs.getdbt.com/reference/global-configs/usage-stats
+
+[[envs.hatch-test.matrix]]
+python = ["3.8"]
+dbt-core-version = [
+    "1.6",
+    "1.7",
+    "1.8",
+]
+
+[[envs.hatch-test.matrix]]
+python = ["3.9"]
+dbt-core-version = [
+    "1.9",
+    "1.10",
+]
+
+[[envs.hatch-test.matrix]]
+python = ["3.10"]
+dbt-core-version = [
+    "1.11",
+]

dbt_python_imports-0.1.0/pyproject.toml

@@ -0,0 +1,24 @@
+[project]
+name = "dbt-python-imports"
+version = "0.1.0"
+description = "dbt plugin to add `modules.import` method to Jinja context to import arbitrary Python modules"
+readme = "README.md"
+authors = [
+    { name = "Tim Vergenz", email = "vergenzt@gmail.com" }
+]
+requires-python = ">=3.8"
+dependencies = [
+    "dbt-core>=1.6",
+]
+
+[dependency-groups]
+dev = [
+    "hatch>=1.14.2",
+    "pytest>=8.3.5",
+    "dbt-postgres",
+    "tinypg>=0.3.0",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"

dbt_python_imports-0.1.0/src/dbt_python_imports/__init__.py

@@ -0,0 +1,30 @@
+from functools import wraps
+from importlib import import_module
+
+import dbt.context.base
+from dbt.plugins.manager import dbtPlugin
+
+
+class DbtPythonImportsPlugin(dbtPlugin):
+    """dbt plugin to add `modules.import` method to Jinja context to import arbitrary Python modules"""
+
+    def initialize(self) -> None:
+        self._get_context_modules_orig = dbt.context.base.get_context_modules
+
+        @wraps(self._get_context_modules_orig)
+        def _wrapper():
+            return {
+                **self._get_context_modules_orig(),
+                "import": self.import_module,
+            }
+
+        dbt.context.base.get_context_modules = _wrapper
+
+    @staticmethod
+    def import_module(module: str):
+        if module.startswith("."):
+            raise ValueError("Relative imports not supported by `modules.import`!")
+        return import_module(module)
+
+
+plugins = [DbtPythonImportsPlugin]

dbt_python_imports-0.1.0/tests/profiles.yml

@@ -0,0 +1,10 @@
+test_profile:
+  outputs:
+    default:
+      type: postgres
+      host:     '{{ env_var("DB_HOST") }}'
+      user:     '{{ env_var("DB_USERNAME") }}'
+      password: '{{ env_var("DB_PASSWORD") }}'
+      port:     '{{ env_var("DB_PORT") | as_native }}'
+      dbname:   '{{ env_var("DB_NAME") }}'
+      schema:   '{{ env_var("DB_SCHEMA") }}'

dbt_python_imports-0.1.0/tests/test_plugin.py

@@ -0,0 +1,206 @@
+import json
+import os
+import re
+from pathlib import Path
+from textwrap import dedent
+from urllib.parse import urlparse
+
+import pytest
+import tinypg
+from pytest import fixture
+
+DBT_PROFILES_DIR = Path(__file__).parent
+DBT_PROFILE = "test_profile"
+DBT_PROJECT_NAME = "test_" + re.sub(
+    r"[^\w]", "_", os.getenv("ENV_NAME", DBT_PROFILES_DIR.parent.name).lower()
+)
+DBT_PROJECT_YML: str = json.dumps(
+    {
+        "name": DBT_PROJECT_NAME,
+        "flags": {
+            "send_anonymous_usage_stats": False,
+        },
+        "version": "0.0.1",
+        "profile": DBT_PROFILE,
+        "model-paths": ["models"],
+        "target-path": "target",
+    }
+)
+
+
+@fixture
+def temp_postgres_db_url():
+    with tinypg.database() as dburl_str:
+        yield dburl_str
+
+
+@fixture
+def dbt_project(tmp_path, temp_postgres_db_url, monkeypatch):
+    dbt_project_path = tmp_path / DBT_PROJECT_NAME
+    dbt_project_path.mkdir()
+    dbt_project_path.joinpath("dbt_project.yml").write_text(DBT_PROJECT_YML)
+    dbt_project_path.joinpath("models").mkdir()
+    dbt_project_path.joinpath("target").mkdir()
+
+    monkeypatch.setenv("DBT_PROFILES_DIR", str(DBT_PROFILES_DIR))
+    monkeypatch.setenv("DBT_PROJECT_DIR", str(dbt_project_path))
+
+    dburl = urlparse(temp_postgres_db_url)
+    monkeypatch.setenv("DB_HOST", dburl.hostname or "")
+    monkeypatch.setenv("DB_USERNAME", dburl.username or "")
+    monkeypatch.setenv("DB_PASSWORD", dburl.password or "")
+    monkeypatch.setenv("DB_PORT", str(dburl.port))
+    monkeypatch.setenv("DB_NAME", dburl.path[1:] or "")
+    monkeypatch.setenv("DB_SCHEMA", "public")
+    monkeypatch.chdir(dbt_project_path)
+
+    return dbt_project_path
+
+
+def test_model_context(dbt_project: Path):
+    model_path = dbt_project / "models" / "my_model.sql"
+    model_path.write_text("""
+    {%- set os_path = modules.import("os.path") %}
+    select '{{ os_path.dirname("foo/bar/baz") }}'
+    """)
+
+    from dbt.cli.main import cli
+
+    cli(["compile"], standalone_mode=False)
+
+    model_compiled_path = (
+        dbt_project / "target/compiled" / model_path.relative_to(dbt_project.parent)
+    )
+    assert model_compiled_path.read_text().strip() == "select 'foo/bar'"
+
+
+def test_properties_context(dbt_project: Path):
+    model_path = dbt_project / "models" / "my_model.sql"
+    model_path.write_text("select 1 as my_column")
+
+    yml_path = dbt_project / "models" / "my_model.yml"
+    yml_path.write_text(
+        dedent(
+            """
+            version: 2
+            models:
+            - name: my_model
+              description: |-
+                {{ modules.import("calendar").TextCalendar().formatyear(2026) }}
+            """
+        )
+    )
+
+    from dbt.cli.main import cli
+
+    cli(["compile"], standalone_mode=False)
+
+    manifest_path = dbt_project / "target" / "manifest.json"
+    manifest = json.loads(manifest_path.read_text())
+    assert manifest["nodes"][f"model.{DBT_PROJECT_NAME}.my_model"]["description"].strip() == """
+                                  2026
+
+      January                   February                   March
+Mo Tu We Th Fr Sa Su      Mo Tu We Th Fr Sa Su      Mo Tu We Th Fr Sa Su
+          1  2  3  4                         1                         1
+ 5  6  7  8  9 10 11       2  3  4  5  6  7  8       2  3  4  5  6  7  8
+12 13 14 15 16 17 18       9 10 11 12 13 14 15       9 10 11 12 13 14 15
+19 20 21 22 23 24 25      16 17 18 19 20 21 22      16 17 18 19 20 21 22
+26 27 28 29 30 31         23 24 25 26 27 28         23 24 25 26 27 28 29
+                                                    30 31
+
+       April                      May                       June
+Mo Tu We Th Fr Sa Su      Mo Tu We Th Fr Sa Su      Mo Tu We Th Fr Sa Su
+       1  2  3  4  5                   1  2  3       1  2  3  4  5  6  7
+ 6  7  8  9 10 11 12       4  5  6  7  8  9 10       8  9 10 11 12 13 14
+13 14 15 16 17 18 19      11 12 13 14 15 16 17      15 16 17 18 19 20 21
+20 21 22 23 24 25 26      18 19 20 21 22 23 24      22 23 24 25 26 27 28
+27 28 29 30               25 26 27 28 29 30 31      29 30
+
+        July                     August                  September
+Mo Tu We Th Fr Sa Su      Mo Tu We Th Fr Sa Su      Mo Tu We Th Fr Sa Su
+       1  2  3  4  5                      1  2          1  2  3  4  5  6
+ 6  7  8  9 10 11 12       3  4  5  6  7  8  9       7  8  9 10 11 12 13
+13 14 15 16 17 18 19      10 11 12 13 14 15 16      14 15 16 17 18 19 20
+20 21 22 23 24 25 26      17 18 19 20 21 22 23      21 22 23 24 25 26 27
+27 28 29 30 31            24 25 26 27 28 29 30      28 29 30
+                          31
+
+      October                   November                  December
+Mo Tu We Th Fr Sa Su      Mo Tu We Th Fr Sa Su      Mo Tu We Th Fr Sa Su
+          1  2  3  4                         1          1  2  3  4  5  6
+ 5  6  7  8  9 10 11       2  3  4  5  6  7  8       7  8  9 10 11 12 13
+12 13 14 15 16 17 18       9 10 11 12 13 14 15      14 15 16 17 18 19 20
+19 20 21 22 23 24 25      16 17 18 19 20 21 22      21 22 23 24 25 26 27
+26 27 28 29 30 31         23 24 25 26 27 28 29      28 29 30 31
+                          30
+""".strip()
+
+
+
+def test_project_yml_context(dbt_project: Path, monkeypatch):
+    monkeypatch.setenv("EXAMPLE_URL", "https://example.org:8080/foo/bar?baz")
+
+    yml_path = dbt_project / "dbt_project.yml"
+    yml_path.write_text(
+        json.dumps(
+            {
+                **json.loads(yml_path.read_text()),
+                "vars": {
+                    DBT_PROJECT_NAME: {
+                        "example_domain": "{{- modules.import('urllib.parse').urlparse(env_var('EXAMPLE_URL')).hostname }}"
+                    },
+                },
+            }
+        )
+    )
+
+    model_path = dbt_project / "models" / "my_model.sql"
+    model_path.write_text("select '{{ var('example_domain') }}'")
+
+    from dbt.cli.main import cli
+
+    cli(["compile"], standalone_mode=False)
+
+    model_compiled_path = (
+        dbt_project / "target/compiled" / model_path.relative_to(dbt_project.parent)
+    )
+    assert model_compiled_path.read_text().strip() == "select 'example.org'"
+
+
+@pytest.mark.xfail(reason="dbt doesn't call set_up_plugin_manager until after profiles.yml is rendered")
+def test_profiles_yml_context(dbt_project: Path, temp_postgres_db_url: str, monkeypatch, capsys):
+    monkeypatch.undo()
+
+    monkeypatch.setenv("DBT_PROJECT_DIR", str(dbt_project))
+    monkeypatch.setenv("DBT_PROFILES_DIR", str(dbt_project))
+    monkeypatch.setenv("DB_URI", temp_postgres_db_url + "?schema=my_schema")
+
+    yml_path = dbt_project / "profiles.yml"
+    yml_path.write_text(
+        dedent(
+            """
+            test_profile:
+              outputs:
+                default:
+                  type: postgres
+                  host:     '{%- set parse = modules.import("urllib.parse") %}{{ parse.urlparse(env_var("DB_URI")).hostname }}'
+                  user:     '{%- set parse = modules.import("urllib.parse") %}{{ parse.urlparse(env_var("DB_URI")).username }}'
+                  password: '{%- set parse = modules.import("urllib.parse") %}{{ parse.urlparse(env_var("DB_URI")).password }}'
+                  port:     '{%- set parse = modules.import("urllib.parse") %}{{ parse.urlparse(env_var("DB_URI")).port }}'
+                  dbname:   '{%- set parse = modules.import("urllib.parse") %}{{ parse.urlparse(env_var("DB_URI")).path.split("/")[1] }}'
+                  schema:   '{%- set parse = modules.import("urllib.parse") %}{{ parse.parse_qs(parse.urlparse(env_var("DB_URI")).query["schema"] }}'
+            """
+        )
+    )
+
+    model_path = dbt_project / "models" / "my_model.sql"
+    model_path.write_text("select 1 as my_column")
+
+    from dbt.cli.main import cli
+
+    cli(["run"], standalone_mode=False)
+    capsys.readouterr()  # reset
+
+    cli(["show", "-q", "--output=json", "--inline", "select * from my_schema.my_model"])
+    assert json.loads(capsys.readouterr().out) == {"show": [{"my_column": 1}]}