Skip to content

feat: add local Spark session support - #273

Open
Raki (mdrakiburrahman) wants to merge 3 commits into
mainfrom
dev/mdrrahman/local-spark-cursor
Open

feat: add local Spark session support#273
Raki (mdrakiburrahman) wants to merge 3 commits into
mainfrom
dev/mdrrahman/local-spark-cursor

Conversation

@mdrakiburrahman

Copy link
Copy Markdown
Collaborator

Why this change is needed

Customers need to run dbt-fabricspark directly through an in-process Spark Session, including environments such as Fabric notebooks where PySpark is already available, without requiring Livy or making PySpark a mandatory adapter dependency.

How

  • Add method: session with a lazy-loaded PySpark connection and cursor implementation.
  • Support PySpark 3.5 and 4.x through the optional dbt-fabricspark[spark] extra, while pinning ANSI mode off for consistent behavior across both generations.
  • Add profile-template and README documentation for session configuration.
  • Run the Livy and Spark Session local E2E flows concurrently against separate databases while sharing the devcontainer SQL Server Hive metastore.
  • Keep the Livy wheel-install path free of PySpark and install a PySpark version matching the devcontainer runtime for the session path.

Test

  • npx nx run dbt-fabricspark:lint --output-style=stream
  • npx nx run dbt-fabricspark:build --output-style=stream
  • npx nx run dbt-fabricspark:test:local-e2e --output-style=stream
  • npx nx run dbt-fabricspark:test --output-style=stream
  • Real in-process smoke test with PySpark 4.2.0

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@mdrakiburrahman

Raki (mdrakiburrahman) commented Aug 24, 2026

Copy link
Copy Markdown
Collaborator Author

I think this PR is the right place to add Spark job-group metadata for the new session method. It would make the active dbt model visible in the Spark UI and status APIs instead of showing only the adapter call site.

A concrete example from a current run: dbt is executing model.insights_raw.mon_analytics_db_snapshot_b2bbe3, while Spark reports the job as:

collect at .../dbt/adapters/fabricspark/session.py:93

The submitted SQL already has dbt metadata including node_id. The following is a drop-in implementation that labels both eager work in SparkSession.sql() and lazy work in DataFrame.collect(). Reapplying the group in fetchall() is important because Spark local properties are thread-local and SELECT execution is lazy.

diff --git a/src/dbt/adapters/fabricspark/session.py b/src/dbt/adapters/fabricspark/session.py
index 0000000..0000000 100644
--- a/src/dbt/adapters/fabricspark/session.py
+++ b/src/dbt/adapters/fabricspark/session.py
@@
 import datetime as dt
+import json
+import re
+import uuid
+from contextlib import contextmanager
 from types import TracebackType
-from typing import TYPE_CHECKING, Any, Optional, Sequence, Tuple, Union
+from typing import TYPE_CHECKING, Any, Iterator, Optional, Sequence, Tuple, Union
@@
 logger = AdapterLogger("Microsoft Fabric-Spark")
 NUMBERS = DECIMALS + (int, float)
+DBT_QUERY_COMMENT_PATTERN = re.compile(r"/\*\s*(\{.*?\})\s*\*/", re.DOTALL)
+SPARK_JOB_GROUP_PROPERTIES = (
+    "spark.jobGroup.id",
+    "spark.job.description",
+    "spark.job.interruptOnCancel",
+)
+
+
+def _dbt_job_description(sql: str) -> str:
+    for match in DBT_QUERY_COMMENT_PATTERN.finditer(sql):
+        try:
+            metadata = json.loads(match.group(1))
+        except json.JSONDecodeError:
+            continue
+        if not isinstance(metadata, dict) or metadata.get("app") != "dbt":
+            continue
+        context = metadata.get("node_id") or metadata.get("connection_name")
+        if isinstance(context, str) and context:
+            return context
+    return "dbt query"
@@
 class SessionCursor:
@@
         self._df: Optional[DataFrame] = None
         self._rows: Optional[list[Row]] = None
         self._fetch_index = 0
+        self._job_group_id: Optional[str] = None
+        self._job_description: Optional[str] = None
@@
     def close(self) -> None:
         self._df = None
         self._rows = None
         self._fetch_index = 0
+        self._job_group_id = None
+        self._job_description = None
+
+    @contextmanager
+    def _job_group(self) -> Iterator[None]:
+        if self._job_group_id is None or self._job_description is None:
+            yield
+            return
+
+        spark_context = self._spark_session.sparkContext
+        previous_properties = {
+            name: spark_context.getLocalProperty(name)
+            for name in SPARK_JOB_GROUP_PROPERTIES
+        }
+        spark_context.setJobGroup(
+            self._job_group_id,
+            self._job_description,
+            interruptOnCancel=True,
+        )
+        try:
+            yield
+        finally:
+            for name, value in previous_properties.items():
+                spark_context.setLocalProperty(name, value)
@@
     def execute(self, sql: str, *parameters: Any) -> None:
         if parameters:
             sql = sql % parameters
 
         self._df = None
         self._rows = None
         self._fetch_index = 0
+        self._job_description = _dbt_job_description(sql)
+        self._job_group_id = (
+            f"dbt:{self._job_description}:{uuid.uuid4().hex}"
+        )
         try:
-            self._df = self._spark_session.sql(sql)
+            with self._job_group():
+                self._df = self._spark_session.sql(sql)
         except self._analysis_error as exc:
             raise DbtRuntimeError(str(exc)) from exc
 
     def fetchall(self) -> Optional[list[Row]]:
         if self._rows is None and self._df is not None:
-            self._rows = self._df.collect()
+            with self._job_group():
+                self._rows = self._df.collect()
         return self._rows

The context manager restores prior values rather than blindly clearing them, which is safer if dbt is hosted inside a notebook or another application that already set a job group.

Add these focused tests to tests/unit/test_session.py:

def test_session_cursor_labels_eager_and_lazy_spark_work() -> None:
    spark_context = MagicMock()
    spark_context.getLocalProperty.return_value = None
    dataframe = MagicMock()
    dataframe.collect.return_value = [(1,)]
    spark_session = MagicMock()
    spark_session.sparkContext = spark_context
    spark_session.sql.return_value = dataframe
    cursor = SessionCursor(spark_session, FakeAnalysisException)

    cursor.execute(
        '/* {"app": "dbt", "node_id": "model.example.orders"} */ select 1'
    )
    assert cursor.fetchall() == [(1,)]

    group_id = spark_context.setJobGroup.call_args_list[0].args[0]
    assert group_id.startswith("dbt:model.example.orders:")
    assert spark_context.setJobGroup.call_args_list == [
        call(group_id, "model.example.orders", interruptOnCancel=True),
        call(group_id, "model.example.orders", interruptOnCancel=True),
    ]
    cleanup = [
        call("spark.jobGroup.id", None),
        call("spark.job.description", None),
        call("spark.job.interruptOnCancel", None),
    ]
    assert spark_context.setLocalProperty.call_args_list == cleanup * 2


def test_session_cursor_restores_job_group_after_collect_failure() -> None:
    spark_context = MagicMock()
    spark_context.getLocalProperty.return_value = None
    dataframe = MagicMock()
    dataframe.collect.side_effect = RuntimeError("collect failed")
    spark_session = MagicMock()
    spark_session.sparkContext = spark_context
    spark_session.sql.return_value = dataframe
    cursor = SessionCursor(spark_session, FakeAnalysisException)
    cursor.execute(
        '/* {"app": "dbt", "connection_name": "model.example.orders"} */ '
        "select 1"
    )

    spark_context.reset_mock()
    spark_context.getLocalProperty.side_effect = [
        "outer-group",
        "outer description",
        "false",
    ]

    with pytest.raises(RuntimeError, match="collect failed"):
        cursor.fetchall()

    assert spark_context.setLocalProperty.call_args_list == [
        call("spark.jobGroup.id", "outer-group"),
        call("spark.job.description", "outer description"),
        call("spark.job.interruptOnCancel", "false"),
    ]

This directly addresses the visibility example in #261. The generated group ID is also sufficient to implement SessionConnectionWrapper.cancel() through cancelJobGroup() later; Livy would still need a separate mechanism.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

1 participant