feat: expose dbt_task job_id, job_run_id, task_run_id in adapter_response

CHANGELOG.md

@@ -1,3 +1,9 @@
+## dbt-databricks 1.12.1 (TBD)
+
+### Features
+
+- Expose `job_id`, `job_run_id`, and `task_run_id` from the Databricks Jobs `dbt_task` runtime in `adapter_response`, enabling correlation between dbt runs and Databricks workflow executions via `run_results.json` ([#1451](https://github.com/databricks/dbt-databricks/pull/1451) closes [#722](https://github.com/databricks/dbt-databricks/issues/722))
+
 ## dbt-databricks 1.12.0 (May 18, 2026)
 
 ### Features

dbt/adapters/databricks/connections.py

@@ -46,7 +46,12 @@
     ConnectionCreateError,
 )
 from dbt.adapters.databricks.events.other_events import QueryError
-from dbt.adapters.databricks.handle import CursorWrapper, DatabricksHandle, SqlUtils
+from dbt.adapters.databricks.handle import (
+    CursorWrapper,
+    DatabricksAdapterResponse,
+    DatabricksHandle,
+    SqlUtils,
+)
 from dbt.adapters.databricks.logging import logger
 from dbt.adapters.databricks.python_models.run_tracking import PythonRunTracker
 from dbt.adapters.databricks.utils import QueryTagsUtils, is_cluster_http_path, redact_credentials
@@ -540,7 +545,7 @@ def get_response(cls, cursor: Any) -> AdapterResponse:
         if isinstance(cursor, CursorWrapper):
             return cursor.get_response()
         else:
-            return AdapterResponse("OK")
+            return DatabricksAdapterResponse.from_cursor(cursor)
 
     def clear_transaction(self) -> None:
         """Noop."""

dbt/adapters/databricks/handle.py

@@ -1,7 +1,11 @@
+import base64
 import decimal
+import json
+import os
 import re
 import sys
 from collections.abc import Callable, Sequence
+from dataclasses import dataclass
 from types import TracebackType
 from typing import TYPE_CHECKING, Any, Optional, TypeVar
 
@@ -28,6 +32,71 @@
 FailLogOp = Callable[[Exception], str]
 
 
+# Set by the Databricks Jobs `dbt_task` runtime on the dbt CLI subprocess.
+_DBT_TASK_HEADERS_ENV = "DBT_DATABRICKS_HTTP_SESSION_HEADERS"
+
+
+def _get_job_run_context() -> dict[str, Optional[str]]:
+    """Return Databricks `dbt_task` job/run identifiers from the runtime headers.
+
+    All three values are ``None`` when dbt is not running inside a `dbt_task`
+    or when the header env var is missing or malformed.
+    """
+    empty: dict[str, Optional[str]] = {
+        "job_id": None,
+        "job_run_id": None,
+        "task_run_id": None,
+    }
+    raw = os.environ.get(_DBT_TASK_HEADERS_ENV)
+    if not raw:
+        return empty
+    try:
+        headers = json.loads(raw)
+    except (json.JSONDecodeError, TypeError):
+        return empty
+    if not isinstance(headers, dict):
+        return empty
+
+    def _str_or_none(value: Any) -> Optional[str]:
+        return str(value) if value is not None else None
+
+    job_run_id: Optional[str] = None
+    encoded = headers.get("X-Databricks-Sql-Query-Source")
+    if isinstance(encoded, str):
+        try:
+            decoded = json.loads(base64.b64decode(encoded))
+        except (ValueError, json.JSONDecodeError, TypeError):
+            decoded = None
+        if isinstance(decoded, dict):
+            job_run_id = _str_or_none(decoded.get("job_run_id"))
+
+    return {
+        "job_id": _str_or_none(headers.get("X-Databricks-Dbsql-Job-Id")),
+        "job_run_id": job_run_id,
+        "task_run_id": _str_or_none(headers.get("X-Databricks-Dbsql-Run-Id")),
+    }
+
+
+@dataclass
+class DatabricksAdapterResponse(AdapterResponse):
+    """Extends ``AdapterResponse`` with Databricks `dbt_task` run identifiers,
+    so dbt runs can be linked to Databricks `system.lakeflow.*` tables for
+    additional execution info. Populated only when dbt runs as a `dbt_task`.
+    """
+
+    job_id: Optional[str] = None
+    job_run_id: Optional[str] = None
+    task_run_id: Optional[str] = None
+
+    @classmethod
+    def from_cursor(cls, cursor: Any) -> "DatabricksAdapterResponse":
+        return cls(
+            _message="OK",
+            query_id=getattr(cursor, "query_id", None) or "N/A",
+            **_get_job_run_context(),
+        )
+
+
 class CursorWrapper:
     """
     Wrap the DBSQL cursor to abstract the details from DatabricksConnectionManager.
@@ -79,7 +148,7 @@ def fetchmany(self, size: int) -> Sequence[tuple]:
         return self._safe_execute(lambda cursor: cursor.fetchmany(size))
 
     def get_response(self) -> AdapterResponse:
-        return AdapterResponse(_message="OK", query_id=self._cursor.query_id or "N/A")
+        return DatabricksAdapterResponse.from_cursor(self._cursor)
 
     T = TypeVar("T")
 

tests/unit/test_handle.py

@@ -1,13 +1,57 @@
+import base64
+import json
+import os
 import sys
 from decimal import Decimal
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
 
 import pytest
 from databricks.sql.client import Cursor
-from dbt.adapters.contracts.connection import AdapterResponse
 from dbt_common.exceptions import DbtRuntimeError
 
-from dbt.adapters.databricks.handle import CursorWrapper, DatabricksHandle, SqlUtils
+from dbt.adapters.databricks.handle import (
+    CursorWrapper,
+    DatabricksAdapterResponse,
+    DatabricksHandle,
+    SqlUtils,
+    _get_job_run_context,
+)
+
+
+def _build_session_headers(
+    job_id: str = "222",
+    task_run_id: str = "333",
+    job_run_id: str = "111",
+    include_source: bool = True,
+    source_override: object = None,
+) -> str:
+    """Build a realistic DBT_DATABRICKS_HTTP_SESSION_HEADERS JSON string."""
+    payload: dict[str, object] = {
+        "X-Databricks-Dbsql-Attribution-Flags": {
+            "dbsqlTriggerSource": "jobsScheduler",
+            "jobId": job_id,
+            "dbsqlTriggerExecutionType": "manual",
+            "dbsqlTriggerAssetType": "dbt",
+            "runId": task_run_id,
+        },
+        "X-Databricks-Dbsql-Job-Id": job_id,
+        "X-Databricks-Dbsql-Run-Id": task_run_id,
+    }
+    if include_source:
+        if source_override is not None:
+            payload["X-Databricks-Sql-Query-Source"] = source_override
+        else:
+            inner = json.dumps(
+                {
+                    "job_run_id": job_run_id,
+                    "job_id": job_id,
+                    "run_id": task_run_id,
+                    "scheduled": False,
+                    "job_type": "EPHEMERAL",
+                }
+            )
+            payload["X-Databricks-Sql-Query-Source"] = base64.b64encode(inner.encode()).decode()
+    return json.dumps(payload)
 
 
 class TestSqlUtils:
@@ -213,15 +257,37 @@ def test_fetchmany(self, cursor):
         wrapper = CursorWrapper(cursor)
         assert wrapper.fetchmany(1) == [("foo", "bar")]
 
+    @patch.dict(os.environ, {}, clear=True)
     def test_get_response__no_query_id(self, cursor):
         cursor.query_id = None
         wrapper = CursorWrapper(cursor)
-        assert wrapper.get_response() == AdapterResponse("OK", query_id="N/A")
-
+        response = wrapper.get_response()
+        assert isinstance(response, DatabricksAdapterResponse)
+        assert response.query_id == "N/A"
+        assert response.job_id is None
+        assert response.job_run_id is None
+        assert response.task_run_id is None
+
+    @patch.dict(os.environ, {}, clear=True)
     def test_get_response__with_query_id(self, cursor):
         cursor.query_id = "id"
         wrapper = CursorWrapper(cursor)
-        assert wrapper.get_response() == AdapterResponse("OK", query_id="id")
+        response = wrapper.get_response()
+        assert isinstance(response, DatabricksAdapterResponse)
+        assert response.query_id == "id"
+
+    def test_get_response__with_job_context(self, cursor):
+        cursor.query_id = "qid"
+        wrapper = CursorWrapper(cursor)
+        with patch.dict(
+            os.environ,
+            {"DBT_DATABRICKS_HTTP_SESSION_HEADERS": _build_session_headers()},
+        ):
+            response = wrapper.get_response()
+        assert response.job_id == "222"
+        assert response.job_run_id == "111"
+        assert response.task_run_id == "333"
+        assert response.query_id == "qid"
 
     def test_with__no_exception(self, cursor):
         with CursorWrapper(cursor) as c:
@@ -324,3 +390,123 @@ def test_close__open_raising_exception(self, conn, cursor):
         handle.close()
         cursor.close.assert_called_once()
         conn.close.assert_called_once()
+
+
+class TestGetJobRunContext:
+    @patch.dict(os.environ, {}, clear=True)
+    def test_env_absent(self):
+        assert _get_job_run_context() == {
+            "job_id": None,
+            "job_run_id": None,
+            "task_run_id": None,
+        }
+
+    @patch.dict(os.environ, {"DBT_DATABRICKS_HTTP_SESSION_HEADERS": ""})
+    def test_env_empty_string(self):
+        assert _get_job_run_context() == {
+            "job_id": None,
+            "job_run_id": None,
+            "task_run_id": None,
+        }
+
+    def test_real_headers_full(self):
+        with patch.dict(
+            os.environ,
+            {"DBT_DATABRICKS_HTTP_SESSION_HEADERS": _build_session_headers()},
+        ):
+            assert _get_job_run_context() == {
+                "job_id": "222",
+                "job_run_id": "111",
+                "task_run_id": "333",
+            }
+
+    def test_headers_without_query_source(self):
+        """job_run_id is None when X-Databricks-Sql-Query-Source is absent."""
+        with patch.dict(
+            os.environ,
+            {"DBT_DATABRICKS_HTTP_SESSION_HEADERS": _build_session_headers(include_source=False)},
+        ):
+            ctx = _get_job_run_context()
+        assert ctx == {"job_id": "222", "job_run_id": None, "task_run_id": "333"}
+
+    @patch.dict(os.environ, {"DBT_DATABRICKS_HTTP_SESSION_HEADERS": "not-json{"}, clear=False)
+    def test_malformed_json(self):
+        assert _get_job_run_context() == {
+            "job_id": None,
+            "job_run_id": None,
+            "task_run_id": None,
+        }
+
+    def test_malformed_base64_source(self):
+        """Bad base64 in X-Databricks-Sql-Query-Source leaves job_run_id None
+        but the directly-readable job_id and task_run_id still resolve."""
+        with patch.dict(
+            os.environ,
+            {
+                "DBT_DATABRICKS_HTTP_SESSION_HEADERS": _build_session_headers(
+                    source_override="!!!not-base64!!!"
+                )
+            },
+        ):
+            ctx = _get_job_run_context()
+        assert ctx == {"job_id": "222", "job_run_id": None, "task_run_id": "333"}
+
+    def test_base64_decodes_to_non_json(self):
+        garbage_b64 = base64.b64encode(b"not json").decode()
+        with patch.dict(
+            os.environ,
+            {
+                "DBT_DATABRICKS_HTTP_SESSION_HEADERS": _build_session_headers(
+                    source_override=garbage_b64
+                )
+            },
+        ):
+            ctx = _get_job_run_context()
+        assert ctx == {"job_id": "222", "job_run_id": None, "task_run_id": "333"}
+
+    @pytest.mark.parametrize("raw", ["null", "[]", "123", '"foo"', "true"])
+    def test_outer_json_non_object(self, raw):
+        """Outer JSON is valid but not an object — must not raise."""
+        with patch.dict(os.environ, {"DBT_DATABRICKS_HTTP_SESSION_HEADERS": raw}):
+            ctx = _get_job_run_context()
+        assert ctx == {"job_id": None, "job_run_id": None, "task_run_id": None}
+
+    @pytest.mark.parametrize("inner_payload", [b"[1,2,3]", b"null", b"123", b'"x"', b"true"])
+    def test_inner_base64_json_non_object(self, inner_payload):
+        """Inner base64 decodes to valid but non-object JSON — job_run_id stays None,
+        outer fields still resolve."""
+        bad_source = base64.b64encode(inner_payload).decode()
+        with patch.dict(
+            os.environ,
+            {
+                "DBT_DATABRICKS_HTTP_SESSION_HEADERS": _build_session_headers(
+                    source_override=bad_source
+                )
+            },
+        ):
+            ctx = _get_job_run_context()
+        assert ctx == {"job_id": "222", "job_run_id": None, "task_run_id": "333"}
+
+
+class TestDatabricksAdapterResponse:
+    @patch.dict(os.environ, {}, clear=True)
+    def test_from_cursor__no_context(self):
+        cursor = Mock()
+        cursor.query_id = "q1"
+        resp = DatabricksAdapterResponse.from_cursor(cursor)
+        assert resp.query_id == "q1"
+        assert resp.job_id is None
+        assert resp.job_run_id is None
+        assert resp.task_run_id is None
+
+    def test_from_cursor__with_context(self):
+        cursor = Mock()
+        cursor.query_id = "qid"
+        with patch.dict(
+            os.environ,
+            {"DBT_DATABRICKS_HTTP_SESSION_HEADERS": _build_session_headers()},
+        ):
+            resp = DatabricksAdapterResponse.from_cursor(cursor)
+        assert resp.job_id == "222"
+        assert resp.job_run_id == "111"
+        assert resp.task_run_id == "333"