feat: add local Spark session support - #273
Conversation
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
|
I think this PR is the right place to add Spark job-group metadata for the new A concrete example from a current run: dbt is executing The submitted SQL already has dbt metadata including 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._rowsThe 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 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 |
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
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
method: sessionwith a lazy-loaded PySpark connection and cursor implementation.dbt-fabricspark[spark]extra, while pinning ANSI mode off for consistent behavior across both generations.Test
npx nx run dbt-fabricspark:lint --output-style=streamnpx nx run dbt-fabricspark:build --output-style=streamnpx nx run dbt-fabricspark:test:local-e2e --output-style=streamnpx nx run dbt-fabricspark:test --output-style=stream