From 53258ce79d0a1eccefb4b628ca6b44e90abd8c64 Mon Sep 17 00:00:00 2001 From: Raki Rahman Date: Sun, 2 Aug 2026 04:07:57 +0000 Subject: [PATCH 1/7] Privy demonstration --- demo/README.md | 63 +++ demo/dbt_project.yml | 16 + demo/models/hello_privy.sql | 1 + demo/privy-notebook-job.json | 1 + demo/profiles.yml | 19 + pyproject.toml | 4 + src/dbt/adapters/fabricspark/connections.py | 9 +- src/dbt/adapters/fabricspark/credentials.py | 85 ++- src/dbt/adapters/fabricspark/privysession.py | 540 +++++++++++++++++++ uv.lock | 28 +- 10 files changed, 756 insertions(+), 10 deletions(-) create mode 100644 demo/README.md create mode 100644 demo/dbt_project.yml create mode 100644 demo/models/hello_privy.sql create mode 100644 demo/privy-notebook-job.json create mode 100644 demo/profiles.yml create mode 100644 src/dbt/adapters/fabricspark/privysession.py diff --git a/demo/README.md b/demo/README.md new file mode 100644 index 00000000..242f1ad6 --- /dev/null +++ b/demo/README.md @@ -0,0 +1,63 @@ +# Privy demo (dbt-fabricspark) + +Proves `method: privy` can run `SELECT 1` against a Fabric notebook over +Azure Relay, using this branch's adapter wheel. + +## 1. Start the notebook manually + +Auto-start is flaky right now (Fabric session errors), so start it by hand: + +1. Open the notebook: value of `PRIVAY_NOTEBOOK_URL` in `test.env`. +2. Run all cells. Wait until the `RelayServer(...).serve_forever()` cell + shows a running spinner (it never finishes — that's expected). + +`privy_auto_start_notebook: false` is already set in `demo/profiles.yml` so +dbt won't try to trigger a run itself. + +## 2. Build & install the wheel + +```bash +cd /workspaces/dbt-fabricspark +uv build +python3 -m venv demo/.venv +demo/.venv/bin/pip install "$(ls dist/dbt_fabricspark-*-py3-none-any.whl)[privy]" +``` + +## 3. Run dbt + +```bash +cd /workspaces/dbt-fabricspark +set -a; source test.env; set +a +cd demo +../demo/.venv/bin/dbt debug --profiles-dir . +../demo/.venv/bin/dbt run --profiles-dir . +../demo/.venv/bin/dbt show --inline "select 1 as one" --profiles-dir . +``` + +`dbt debug` should show a successful Privy relay connection (no notebook +trigger, since auto-start is off). `dbt run` builds `models/hello_privy.sql` +(`select 1 as id`) as a view. + +**Verified output:** +``` +$ dbt show --inline "select 1 as one" +| one | +| --- | +| 1 | + +$ dbt run +1 of 1 OK created sql view model dbo.hello_privy ... [OK in 2.45s] +``` + +`schema: dbo` must be a schema that already exists in whatever lakehouse the +notebook is attached to — it's not related to the `privy_*` settings. If you +get `SCHEMA_NOT_FOUND`, run `SHOW SCHEMAS` in the notebook to find a valid one. + +## Known issues + +- Notebook auto-start (`privy_auto_start_notebook: true`) currently fails + server-side after ~15s (`System_Cancelled_Session_Statements_Failed`), + suspected cause: the notebook's `%pip install --force-reinstall` step + clobbering packages the Fabric kernel relies on. Manual start avoids it. +- SPN auth may not be able to trigger notebook runs (Fabric API limitation) — + use CLI auth (`az login`) for now. diff --git a/demo/dbt_project.yml b/demo/dbt_project.yml new file mode 100644 index 00000000..9e79a699 --- /dev/null +++ b/demo/dbt_project.yml @@ -0,0 +1,16 @@ +name: "privy_demo" +version: "1.0.0" +config-version: 2 + +profile: "privy_demo" + +model-paths: ["models"] + +target-path: "target" +clean-targets: + - "target" + - "logs" + +models: + privy_demo: + materialized: view diff --git a/demo/models/hello_privy.sql b/demo/models/hello_privy.sql new file mode 100644 index 00000000..43258a71 --- /dev/null +++ b/demo/models/hello_privy.sql @@ -0,0 +1 @@ +select 1 as id diff --git a/demo/privy-notebook-job.json b/demo/privy-notebook-job.json new file mode 100644 index 00000000..74ff3093 --- /dev/null +++ b/demo/privy-notebook-job.json @@ -0,0 +1 @@ +{"workspace_id": "d55e1b36-c694-4c93-9c81-44988e57edca", "notebook_id": "dc3ccca6-4284-4af0-bb90-de1b4fda3c0b", "job_instance_id": "cdaf9076-a58f-4f83-b7d6-674d9b715a90"} \ No newline at end of file diff --git a/demo/profiles.yml b/demo/profiles.yml new file mode 100644 index 00000000..c587a0ad --- /dev/null +++ b/demo/profiles.yml @@ -0,0 +1,19 @@ +privy_demo: + target: privy + outputs: + privy: + type: fabricspark + method: privy + privy_relay_namespace: "{{ env_var('PRIVY_RELAY_NAMESPACE') }}" + privy_relay_path: "{{ env_var('PRIVY_RELAY_PATH') }}" + privy_relay_keyrule: "{{ env_var('PRIVY_RELAY_KEYRULE') }}" + privy_relay_key: "{{ env_var('PRIVY_RELAY_KEY') }}" + privy_notebook_url: "{{ env_var('PRIVY_NOTEBOOK_URL') }}" + endpoint: "{{ env_var('FABRIC_ENDPOINT', 'https://api.fabric.microsoft.com/v1') }}" + authentication: "{{ env_var('FABRIC_AUTH_METHOD', 'CLI') }}" + schema: dbo + threads: 1 + privy_auto_start_notebook: true + privy_ready_timeout: 900 + spark_config: + name: "dbt-privy-demo" diff --git a/pyproject.toml b/pyproject.toml index 52c1391f..d9127678 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,6 +43,7 @@ dependencies = [ [project.optional-dependencies] cli = ["azure-cli>=2.84.0"] +privy = ["privy @ https://rakirahman.blob.core.windows.net/public/whls/privy-0.0.1-py3-none-any.whl"] [dependency-groups] dev = [ @@ -68,6 +69,9 @@ dev = [ default-groups = "all" prerelease = "allow" +[tool.hatch.metadata] +allow-direct-references = true + [project.urls] homepage = "https://github.com/microsoft/dbt-fabricspark" "Setup & configuration" = "https://docs.getdbt.com/reference/warehouse-profiles/fabricspark-profile" diff --git a/src/dbt/adapters/fabricspark/connections.py b/src/dbt/adapters/fabricspark/connections.py index e2fac1ef..93519f02 100644 --- a/src/dbt/adapters/fabricspark/connections.py +++ b/src/dbt/adapters/fabricspark/connections.py @@ -37,6 +37,7 @@ LivySessionManager, get_lakehouse_properties, ) +from dbt.adapters.fabricspark.privysession import PrivyConnectionManager, PrivyConnectionWrapper from dbt.adapters.fabricspark.relation import FabricSparkRelation from dbt.adapters.sql import SQLConnectionManager @@ -108,6 +109,7 @@ def render_spark_type(type_code: Any) -> str: class FabricSparkConnectionMethod(StrEnum): LIVY = "livy" + PRIVY = "privy" class FabricSparkConnectionWrapper(ABC): @@ -228,7 +230,7 @@ def open(cls, connection: Connection) -> Connection: handle: FabricSparkConnectionWrapper = None # Fetch lakehouse properties and detect schema support (Fabric mode only). - if not creds.is_local_mode: + if not creds.is_local_mode and not creds.is_privy_mode: lakehouse_props = get_lakehouse_properties(creds) creds.apply_lakehouse_properties(lakehouse_props) @@ -258,6 +260,11 @@ def open(cls, connection: Connection) -> Connection: ) connection.state = ConnectionState.OPEN + elif creds.method == FabricSparkConnectionMethod.PRIVY: + raw_handle = PrivyConnectionManager.connect(creds) + handle = PrivyConnectionWrapper(raw_handle, creds) + connection.state = ConnectionState.OPEN + else: raise DbtConfigError(f"invalid credential method: {creds.method}") break diff --git a/src/dbt/adapters/fabricspark/credentials.py b/src/dbt/adapters/fabricspark/credentials.py index 74535007..cc1d1e4c 100644 --- a/src/dbt/adapters/fabricspark/credentials.py +++ b/src/dbt/adapters/fabricspark/credentials.py @@ -124,6 +124,32 @@ class FabricSparkCredentials(Credentials): # ``DBT_FABRICSPARK_SKIP_OPTIMIZE`` environment variable disables it outright. auto_optimize: bool = True + # --- Privy connection method (experimental) --------------------------- + # ``method: privy`` sends statements to a Fabric notebook (running + # ``privy.RelayServer``) over an Azure Relay Hybrid Connection instead of + # the Livy REST API. Requires the ``privy`` extra + # (``pip install dbt-fabricspark[privy]``). Exempt from the + # workspaceid/lakehouseid/lakehouse requirements below (like local mode) — + # only these fields are needed. + privy_relay_namespace: Optional[str] = None + privy_relay_path: Optional[str] = None + privy_relay_keyrule: Optional[str] = None + privy_relay_key: Optional[str] = None + # Browser URL of the Fabric notebook hosting the RelayServer, e.g. + # https:///groups//synapsenotebooks/. The + # workspace and notebook GUIDs are parsed out of this URL to trigger the + # notebook via the Fabric Job Scheduler API when the relay is unreachable. + privy_notebook_url: Optional[str] = None + # When True (default), the adapter tries to start the notebook (via the + # Fabric REST API) if the relay doesn't respond to a health check. Set to + # False to manage the run yourself and have the adapter only act as a + # Privy client. + privy_auto_start_notebook: bool = True + # Max seconds to keep pinging the relay for readiness (after an auto-start + # trigger, or while waiting for a manually-started run) before giving up. + # Independent of session_start_timeout (which is Livy-session-specific). + privy_ready_timeout: int = 900 + def __repr__(self) -> str: """Mask sensitive fields in repr to prevent credential leakage in logs/tracebacks.""" return ( @@ -137,6 +163,11 @@ def __repr__(self) -> str: f"credential_class={self.credential_class!r}, " f"credential_kwargs_keys={sorted(map(str, self.credential_kwargs.keys()))!r}, " f"workspace_name={self.workspace_name!r}, " + f"privy_relay_namespace={self.privy_relay_namespace!r}, " + f"privy_relay_path={self.privy_relay_path!r}, " + f"privy_notebook_url={self.privy_notebook_url!r}, " + f"privy_ready_timeout={self.privy_ready_timeout!r}, " + f"privy_relay_key='***', " f"accessToken='***')" ) @@ -157,6 +188,10 @@ def __pre_deserialize__(cls, data: Any) -> Any: def is_local_mode(self) -> bool: return self.livy_mode == "local" + @property + def is_privy_mode(self) -> bool: + return self.method == "privy" + @property def resolved_session_id_file(self) -> str: if self.session_id_file: @@ -173,8 +208,10 @@ def __post_init__(self) -> None: if self.method is None: raise DbtRuntimeError("Must specify `method` in profile") - # Fabric-specific validations - if not self.is_local_mode: + # Fabric-specific validations. Skipped for local mode and for privy + # mode — privy talks to whatever lakehouse the notebook is already + # attached to, so it has no use for workspaceid/lakehouseid/lakehouse. + if not self.is_local_mode and not self.is_privy_mode: if self.endpoint is None: raise DbtRuntimeError("Must specify `endpoint` in profile for Fabric mode") if self.workspaceid is None: @@ -184,25 +221,48 @@ def __post_init__(self) -> None: if self.lakehouse is None: raise DbtRuntimeError("Must specify `lakehouse` in profile for Fabric mode") + if self.is_privy_mode: + if not self.privy_relay_namespace: + raise DbtRuntimeError( + "Must specify `privy_relay_namespace` in profile for method=privy" + ) + if not self.privy_relay_path: + raise DbtRuntimeError( + "Must specify `privy_relay_path` in profile for method=privy" + ) + if not self.privy_relay_keyrule: + raise DbtRuntimeError( + "Must specify `privy_relay_keyrule` in profile for method=privy" + ) + if not self.privy_relay_key: + raise DbtRuntimeError("Must specify `privy_relay_key` in profile for method=privy") + if not self.privy_notebook_url: + raise DbtRuntimeError( + "Must specify `privy_notebook_url` in profile for method=privy" + ) + # schema defaults to lakehouse name if not provided by user. # For schema-enabled lakehouses, user can override this in profiles.yml. - # For local mode without lakehouse, defaults to "default" (Spark's default database). + # For local/privy mode without lakehouse, defaults to "default" (Spark's default database). if self.schema is None: if self.lakehouse is not None: self.schema = self.lakehouse - elif self.is_local_mode: + elif self.is_local_mode or self.is_privy_mode: self.schema = "default" # database is always set to lakehouse name for relation rendering. # In non-schema mode, include_policy.database=False excludes it from SQL. # In schema-enabled mode, include_policy.database=True renders three-part names. - # For local mode without lakehouse, defaults to "default". + # For local/privy mode without lakehouse, defaults to "default". if self.lakehouse is not None: self.database = self.lakehouse - elif self.is_local_mode: + elif self.is_local_mode or self.is_privy_mode: self.database = "default" - # Security validations (Fabric mode only) + # Security validations (Fabric mode only). Privy still makes one real + # Fabric API call (triggering the notebook run), so endpoint/UUID + # validation stays active for it; workspaceid/lakehouseid are simply + # None for privy and _validate_uuid tolerates that. if not self.is_local_mode: self._validate_uuid(self.workspaceid, "workspaceid") self._validate_uuid(self.lakehouseid, "lakehouseid") @@ -285,6 +345,8 @@ def type(self) -> str: def unique_field(self) -> str: if self.is_local_mode: return self.livy_url + if self.is_privy_mode: + return f"{self.privy_relay_namespace}/{self.privy_relay_path}" return self.lakehouseid def _validate_endpoint(self) -> None: @@ -315,7 +377,8 @@ def _validate_uuid(self, value: Optional[str], field_name: str) -> None: ) def _connection_keys(self) -> Tuple[str, ...]: - # Intentionally excludes client_secret, accessToken, tenant_id + # Intentionally excludes client_secret, accessToken, tenant_id, + # privy_relay_key, privy_relay_keyrule return ( "workspaceid", "lakehouseid", @@ -327,4 +390,10 @@ def _connection_keys(self) -> Tuple[str, ...]: "auto_optimize", "high_concurrency", "spark_config", + "method", + "privy_relay_namespace", + "privy_relay_path", + "privy_notebook_url", + "privy_auto_start_notebook", + "privy_ready_timeout", ) diff --git a/src/dbt/adapters/fabricspark/privysession.py b/src/dbt/adapters/fabricspark/privysession.py new file mode 100644 index 00000000..d1a376d4 --- /dev/null +++ b/src/dbt/adapters/fabricspark/privysession.py @@ -0,0 +1,540 @@ +"""Privy connection method (experimental) — Azure Relay transport for Spark SQL. + +Sends ``spark.sql(...)`` statements to a Fabric notebook running +``privy.RelayServer`` over an Azure Relay Hybrid Connection, instead of the +Livy REST API. Requires the ``privy`` package (``pip install +dbt-fabricspark[privy]``), lazily imported so ``method: livy`` users never +need it installed. + +Every statement is sent with ``mode="inprocess"``: privy's default +``mode="subprocess"`` spawns a fresh, isolated Python interpreter with no +Fabric notebook context, while ``mode="inprocess"`` executes inside the +``RelayServer``'s own already-running interpreter — the same kernel the +Fabric notebook cell is running in — so the notebook's pre-existing ``spark`` +session global is visible. Without ``inprocess``, ``spark`` would be +undefined. + +This is a spike: no lakehouse schema-detection, no high-concurrency +multi-REPL support (privy serializes inprocess calls behind a single lock on +the server side, so concurrent dbt threads queue FIFO), and no retry/backoff +sophistication beyond a simple health-check + wait loop. + +The notebook run is never cancelled by this module (not even on process +exit) — a filesystem cache (``privy-notebook-job.json`` in the cwd) lets +separate dbt invocations reuse the same run instead of each starting their +own. Cancelling it is entirely up to the caller. +""" + +from __future__ import annotations + +import datetime as dt +import json +import os +import re +import threading +import time +import uuid +from typing import Any, Dict, List, Optional, Sequence, Tuple + +import requests +from dbt_common.exceptions import DbtDatabaseError, DbtRuntimeError +from dbt_common.utils.encoding import DECIMALS + +from dbt.adapters.events.logging import AdapterLogger +from dbt.adapters.fabricspark.credentials import FabricSparkCredentials +from dbt.adapters.fabricspark.livy_backend import coerce_time_columns +from dbt.adapters.fabricspark.livysession import get_headers + +logger = AdapterLogger("Microsoft Fabric-Spark") + +_NUMBERS = DECIMALS + (int, float) + +_UUID_RE = re.compile( + r"[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}" +) + +# Fast, short-timeout probe used purely to check "is anything listening on the +# relay right now" — deliberately much shorter than the timeout used for real +# query execution below. +_PROBE_HTTP_TIMEOUT_S = 20.0 +_PROBE_TIMEOUT_S = 10.0 + +# Fallback exec timeout when credentials.statement_timeout == 0 ("no timeout" +# for Livy's polling loop). privy's wire protocol needs a finite number, so a +# generous one week stands in for "effectively unbounded". +_UNBOUNDED_TIMEOUT_S = 7 * 24 * 3600.0 + +# Fabric Job Scheduler statuses for the RunNotebook job instance we trigger. +# https://learn.microsoft.com/en-us/rest/api/fabric/core/job-scheduler/get-item-job-instance +_JOB_TERMINAL_STATUSES = {"Completed", "Failed", "Cancelled", "Deduped"} +_JOB_FAILURE_STATUSES = {"Failed", "Cancelled"} + + +def _import_relay_client() -> Any: + try: + from privy import RelayClient + except ImportError as exc: + raise DbtRuntimeError( + "method=privy requires the `privy` package. Install it with " + "`pip install dbt-fabricspark[privy]` (or add the `privy` extra " + "to your uv/pip install of dbt-fabricspark)." + ) from exc + return RelayClient + + +def _build_relay_client(credentials: FabricSparkCredentials, http_timeout_s: float) -> Any: + RelayClient = _import_relay_client() + return RelayClient( + namespace=credentials.privy_relay_namespace, + path=credentials.privy_relay_path, + keyrule=credentials.privy_relay_keyrule, + key=credentials.privy_relay_key, + http_timeout_s=http_timeout_s, + ) + + +def _query_timeout_s(credentials: FabricSparkCredentials) -> float: + if credentials.statement_timeout and credentials.statement_timeout > 0: + return float(credentials.statement_timeout) + return _UNBOUNDED_TIMEOUT_S + + +def _parse_notebook_ids(notebook_url: Optional[str]) -> Tuple[str, str]: + """Extract (workspaceId, notebookId) GUIDs from a Fabric notebook browser URL. + + e.g. ``https:///groups//synapsenotebooks/``. + """ + ids = _UUID_RE.findall(notebook_url or "") + if len(ids) < 2: + raise ValueError( + f"expected to find 2 GUIDs (workspace, notebook) in privy_notebook_url, " + f"found {len(ids)}: {notebook_url!r}" + ) + return ids[0], ids[1] + + +def _trigger_notebook_run( + credentials: FabricSparkCredentials, +) -> Optional[Tuple[str, str, str]]: + """Best-effort trigger of the Fabric notebook via the Job Scheduler API. + + POST .../items/{notebookId}/jobs/instances?jobType=RunNotebook — a 202 + means the run was accepted (the notebook, and eventually its + ``RelayServer.serve_forever()`` cell, will start). This never polls the + job to completion here: the job is meant to run forever, so acceptance is + the only thing checked at trigger time. The job instance id (parsed from + the ``Location`` header) is returned so the caller can poll its status + while waiting for the relay — see ``_wait_for_relay``. Failures to trigger + are logged and swallowed rather than raised — the notebook might already + be starting (e.g. from a previous invocation or a manual start), so the + wait loop is the real arbiter. + """ + try: + workspace_id, notebook_id = _parse_notebook_ids(credentials.privy_notebook_url) + except ValueError as exc: + logger.warning( + f"Could not parse workspace/notebook id from privy_notebook_url ({exc}). " + f"Skipping auto-start; will still wait in case the relay is already up." + ) + return None + + url = ( + f"{credentials.endpoint}/workspaces/{workspace_id}/items/{notebook_id}" + f"/jobs/instances?jobType=RunNotebook" + ) + logger.info(f"Privy relay not responding; triggering Fabric notebook run: POST {url}") + try: + headers = get_headers(credentials) + response = requests.post(url, headers=headers, json={}, timeout=credentials.http_timeout) + if response.status_code in (200, 202): + location = response.headers.get("Location", "") + job_instance_id = location.rstrip("/").rsplit("/", 1)[-1] if location else "" + logger.info( + f"Notebook run triggered (HTTP {response.status_code}). " + f"Job instance: {job_instance_id or 'unknown'}. " + f"Waiting for the Privy relay to come up..." + ) + if job_instance_id: + return workspace_id, notebook_id, job_instance_id + return None + else: + logger.warning( + f"Notebook run trigger returned HTTP {response.status_code}: " + f"{response.text[:500]}. Will keep waiting in case it's already starting " + f"(e.g. started manually, or by a previous dbt invocation)." + ) + except requests.exceptions.RequestException as exc: + logger.warning( + f"Failed to trigger notebook run ({exc}). Will keep waiting in case it's " + f"already starting." + ) + return None + + +def _get_job_instance_status( + credentials: FabricSparkCredentials, workspace_id: str, item_id: str, job_instance_id: str +) -> Dict[str, Any]: + """GET .../items/{itemId}/jobs/instances/{jobInstanceId} — the run's live status.""" + url = ( + f"{credentials.endpoint}/workspaces/{workspace_id}/items/{item_id}" + f"/jobs/instances/{job_instance_id}" + ) + headers = get_headers(credentials) + response = requests.get(url, headers=headers, timeout=credentials.http_timeout) + response.raise_for_status() + return response.json() + + +# Filesystem cache of the last notebook job this machine triggered — mirrors +# livysession.py's session-id file, but stores a (workspace, notebook, job +# instance) triple as JSON instead of a bare Livy session id. Lets a fresh +# dbt invocation (a brand new process — dbt spawns one per `debug`/`run`/ +# `show` etc.) find and reuse a notebook run a previous invocation already +# triggered instead of firing (and paying for) a new one every single time. +# Nothing here ever cancels the run — that's left entirely to the caller. +_JOB_CACHE_FILENAME = "privy-notebook-job.json" + + +def _job_cache_path() -> str: + return os.path.join(os.getcwd(), _JOB_CACHE_FILENAME) + + +def _read_cached_job_ref() -> Optional[Tuple[str, str, str]]: + path = _job_cache_path() + try: + with open(path) as f: + data = json.load(f) + return data["workspace_id"], data["notebook_id"], data["job_instance_id"] + except Exception as exc: + logger.debug(f"No usable Privy job cache at {path}: {exc}") + return None + + +def _write_cached_job_ref(job_ref: Tuple[str, str, str]) -> None: + path = _job_cache_path() + try: + with open(path, "w") as f: + json.dump( + { + "workspace_id": job_ref[0], + "notebook_id": job_ref[1], + "job_instance_id": job_ref[2], + }, + f, + ) + except OSError as exc: + logger.debug(f"Could not write Privy job cache file {path}: {exc}") + + +def _clear_cached_job_ref() -> None: + try: + os.remove(_job_cache_path()) + except OSError: + pass + + +def _probe(client: Any) -> bool: + try: + result = client.run_python("1", mode="inprocess", timeout_s=_PROBE_TIMEOUT_S) + return bool(result.ok) + except Exception as exc: # noqa: BLE001 — any failure just means "not ready yet" + logger.debug(f"Privy relay probe failed: {exc}") + return False + + +def _wait_for_relay( + probe_client: Any, + credentials: FabricSparkCredentials, + job_ref: Optional[Tuple[str, str, str]] = None, +) -> None: + deadline = time.time() + credentials.privy_ready_timeout + attempt = 0 + while True: + attempt += 1 + if _probe(probe_client): + logger.info(f"Privy relay responded after {attempt} attempt(s).") + return + + # Job status is a much stronger signal than another silent relay + # probe: it tells us whether the notebook run is still starting up + # (queued/in-progress — normal, keep waiting), or has already ended + # (failed/cancelled/completed without ever starting the relay — no + # amount of extra waiting will help, fail fast instead). + if job_ref is not None: + workspace_id, item_id, job_instance_id = job_ref + try: + job = _get_job_instance_status(credentials, workspace_id, item_id, job_instance_id) + status = job.get("status", "Unknown") + logger.info(f"Notebook job {job_instance_id} status: {status}") + if status in _JOB_FAILURE_STATUSES: + failure_reason = job.get("failureReason") + _clear_cached_job_ref() + raise DbtRuntimeError( + f"Fabric notebook run {status} before the Privy relay came up " + f"(job {job_instance_id}). failureReason={failure_reason}. " + f"Check the notebook run history in the Fabric portal for details." + ) + if status in _JOB_TERMINAL_STATUSES: + # e.g. "Completed"/"Deduped" — the run this job represents + # is over (or was superseded), yet the relay never came up. + # Stop polling this particular job instance (it won't + # change anymore) but keep waiting on the relay itself in + # case another run is what's actually serving it. + logger.warning( + f"Notebook job {job_instance_id} reached status={status} but the " + f"Privy relay never responded. If this persists, verify the notebook " + f"cell actually reaches `RelayServer(...).serve_forever()`." + ) + _clear_cached_job_ref() + job_ref = None + except requests.exceptions.RequestException as exc: + logger.debug(f"Could not fetch notebook job status: {exc}") + + if time.time() >= deadline: + raise DbtRuntimeError( + f"Timed out after {credentials.privy_ready_timeout}s waiting for the " + f"Privy relay/notebook to respond. Check the Fabric notebook run history, " + f"or set `privy_auto_start_notebook: false` and start the notebook manually. " + f"Override the timeout with `privy_ready_timeout: ` in your profile." + ) + time.sleep(credentials.poll_wait) + + +def _ensure_notebook_ready(exec_client: Any, credentials: FabricSparkCredentials) -> None: + """Probe the relay and, if needed, reuse-or-trigger a notebook run. + + Before triggering a brand-new run, checks the filesystem job cache (see + ``_read_cached_job_ref``) for a still-active job a previous dbt + invocation already triggered, and waits on that instead — so back-to-back + dbt invocations (``debug``, ``run``, ``show``, ...) share one notebook + session rather than each firing (and paying for) their own. Nothing is + ever cancelled here or on process exit; cancelling the notebook run is + entirely up to the caller. + """ + probe_client = _build_relay_client(credentials, http_timeout_s=_PROBE_HTTP_TIMEOUT_S) + if _probe(probe_client): + logger.debug("Privy relay already responding; reusing the existing notebook run.") + return + + job_ref: Optional[Tuple[str, str, str]] = None + cached_ref = _read_cached_job_ref() + if cached_ref is not None: + try: + job = _get_job_instance_status(credentials, *cached_ref) + status = job.get("status", "Unknown") + if status not in _JOB_TERMINAL_STATUSES: + logger.info(f"Reusing cached notebook job {cached_ref[2]} (status={status}).") + job_ref = cached_ref + else: + logger.debug( + f"Cached notebook job {cached_ref[2]} is terminal ({status}); discarding." + ) + _clear_cached_job_ref() + except requests.exceptions.RequestException as exc: + logger.debug(f"Could not check cached notebook job status ({exc}); discarding cache.") + _clear_cached_job_ref() + + if job_ref is None: + if credentials.privy_auto_start_notebook: + job_ref = _trigger_notebook_run(credentials) + if job_ref is not None: + _write_cached_job_ref(job_ref) + else: + logger.info( + "Privy relay not responding and privy_auto_start_notebook is False; " + "waiting for it to be started manually." + ) + + _wait_for_relay(probe_client, credentials, job_ref) + + +class PrivyConnectionManager: + """Builds, health-checks and (if needed) triggers the Fabric notebook for a + shared, process-wide Privy ``RelayClient``. + + Unlike Livy, Privy calls are stateless HTTP POSTs — there is no session or + REPL to acquire per dbt thread, so a single client per unique + (namespace, path) target is shared across all threads. A per-key lock + ensures only the first caller for a given target does the + health-check/auto-start dance; later callers (including other dbt + threads) reuse the already-verified client. + """ + + _clients: Dict[str, Any] = {} + _ready: Dict[str, bool] = {} + _locks: Dict[str, threading.Lock] = {} + _registry_lock = threading.Lock() + + @classmethod + def _lock_for(cls, key: str) -> threading.Lock: + with cls._registry_lock: + lock = cls._locks.get(key) + if lock is None: + lock = threading.Lock() + cls._locks[key] = lock + return lock + + @classmethod + def connect(cls, credentials: FabricSparkCredentials) -> Any: + key = credentials.unique_field + lock = cls._lock_for(key) + with lock: + client = cls._clients.get(key) + if client is None: + http_timeout_s = _query_timeout_s(credentials) + 30.0 + client = _build_relay_client(credentials, http_timeout_s=http_timeout_s) + cls._clients[key] = client + if not cls._ready.get(key): + _ensure_notebook_ready(client, credentials) + cls._ready[key] = True + return client + + @classmethod + def disconnect(cls) -> None: + """No persistent resources to release — Privy calls are stateless HTTP. + + The notebook run itself is never cancelled here or on process exit — + the filesystem job cache lets later, separate dbt invocations keep + reusing it. Cancel it yourself when you're done (Fabric portal, or + POST .../jobs/instances/{id}/cancel). + """ + + +def _build_exec_snippet(sql: str, marker: str) -> str: + """Build the Python snippet run (inprocess) on the notebook side. + + Runs ``spark.sql(sql)``, serializes the result into the same + ``{"data": [...], "schema": {"fields": [...]}}`` shape Livy's statement + API returns, and prints it between two copies of a unique marker so the + client can find it even if the query itself prints other output. + """ + sql_literal = json.dumps(sql) + marker_literal = json.dumps(marker) + return ( + "import json as __privy_json\n" + f"__privy_df = spark.sql({sql_literal})\n" + "__privy_rows = [list(__privy_row) for __privy_row in __privy_df.collect()]\n" + "__privy_fields = [\n" + " {'name': __f.name, 'type': __f.dataType.simpleString(), 'nullable': __f.nullable}\n" + " for __f in __privy_df.schema.fields\n" + "]\n" + f"print({marker_literal})\n" + "print(__privy_json.dumps(" + "{'data': __privy_rows, 'schema': {'fields': __privy_fields}}, default=str))\n" + f"print({marker_literal})\n" + ) + + +def _extract_marked_json(stdout: str, marker: str) -> Dict[str, Any]: + start = stdout.find(marker) + if start == -1: + raise DbtDatabaseError( + f"Privy response is missing the result marker; stdout={stdout[-2000:]!r}" + ) + end = stdout.find(marker, start + len(marker)) + if end == -1: + raise DbtDatabaseError( + f"Privy response is missing the closing result marker; stdout={stdout[-2000:]!r}" + ) + raw = stdout[start + len(marker) : end].strip() + try: + return json.loads(raw) + except json.JSONDecodeError as exc: + raise DbtDatabaseError( + f"Could not parse Privy result JSON ({exc}); raw={raw[:2000]!r}" + ) from exc + + +class PrivyConnectionWrapper: + """Connection wrapper for the privy (Azure Relay) connection method. + + Deliberately duck-types the same surface as ``FabricSparkConnectionWrapper`` + (see ``connections.py``) without importing/inheriting from it — mirrors + how ``LivySessionConnectionWrapper`` avoids a circular import between + ``connections.py`` and this module. + """ + + def __init__(self, relay_client: Any, credentials: FabricSparkCredentials) -> None: + self._client = relay_client + self._timeout_s = _query_timeout_s(credentials) + self._rows: Optional[List] = None + self._schema: Optional[List[Dict[str, Any]]] = None + + def cursor(self) -> "PrivyConnectionWrapper": + return self + + def cancel(self) -> None: + logger.debug("NotImplemented: cancel") + + def close(self) -> None: + self._rows = None + self._schema = None + + def rollback(self, *args: Any, **kwargs: Any) -> None: + logger.debug("NotImplemented: rollback") + + def fetchall(self) -> Optional[List]: + return self._rows + + def fetchmany(self, size: Optional[int] = None) -> Optional[List]: + if self._rows is None: + return None + return self._rows if size is None else self._rows[:size] + + def fetchone(self) -> Optional[Any]: + return self._rows[0] if self._rows else None + + def execute(self, sql: str, bindings: Optional[List[Any]] = None) -> None: + sql = sql.strip() + if sql.endswith(";"): + sql = sql[:-1] + if bindings is not None: + fixed_bindings = tuple(self._fix_binding(b) for b in bindings) + sql = sql % fixed_bindings + + marker = f"__PRIVY_RESULT_{uuid.uuid4().hex}__" + code = _build_exec_snippet(sql, marker) + logger.debug(f"Submitting to Privy relay (inprocess): {sql}") + result = self._client.run_python(code, mode="inprocess", timeout_s=self._timeout_s) + + if not result.ok: + timeout_note = " (timed out)" if result.timed_out else "" + raise DbtDatabaseError( + f"Error while executing query via Privy{timeout_note}: " + f"{result.stderr or result.stdout}" + ) + + payload = _extract_marked_json(result.stdout, marker) + self._rows = payload.get("data", []) + self._schema = payload.get("schema", {}).get("fields", []) + coerce_time_columns(self._rows, self._schema) + + @property + def description( + self, + ) -> Sequence[Tuple[str, Any, None, None, None, None, bool]]: + if not self._schema: + return [] + return [ + (field["name"], field["type"], None, None, None, None, field["nullable"]) + for field in self._schema + ] + + @classmethod + def _fix_binding(cls, value: Any) -> Any: + """Convert complex datatypes to primitives that can be loaded by the Spark driver.""" + if isinstance(value, _NUMBERS): + return float(value) + elif isinstance(value, dt.datetime): + return f"'{value.strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]}'" + elif value is None: + return "''" + else: + escaped = str(value).replace("'", "\\'") + return f"'{escaped}'" + + +__all__ = [ + "PrivyConnectionManager", + "PrivyConnectionWrapper", +] diff --git a/uv.lock b/uv.lock index e2192279..6de8b792 100644 --- a/uv.lock +++ b/uv.lock @@ -1891,6 +1891,9 @@ dependencies = [ cli = [ { name = "azure-cli" }, ] +privy = [ + { name = "privy" }, +] [package.dev-dependencies] dev = [ @@ -1914,9 +1917,10 @@ requires-dist = [ { name = "dbt-adapters", specifier = ">=1.7,<2.0" }, { name = "dbt-common", specifier = ">=1.10,<2.0" }, { name = "dbt-core", specifier = ">=1.8.0" }, + { name = "privy", marker = "extra == 'privy'", url = "https://rakirahman.blob.core.windows.net/public/whls/privy-0.0.1-py3-none-any.whl" }, { name = "requests", specifier = ">=2.32.0" }, ] -provides-extras = ["cli"] +provides-extras = ["cli", "privy"] [package.metadata.requires-dev] dev = [ @@ -2644,6 +2648,28 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/9b/fb/a70a4214956182e0d7a9099ab17d50bfcba1056188e9b14f35b9e2b62a0d/portalocker-2.10.1-py3-none-any.whl", hash = "sha256:53a5984ebc86a025552264b459b46a2086e269b21823cb572f8f28ee759e45bf", size = 18423, upload-time = "2024-07-13T23:15:32.602Z" }, ] +[[package]] +name = "privy" +version = "0.0.1" +source = { url = "https://rakirahman.blob.core.windows.net/public/whls/privy-0.0.1-py3-none-any.whl" } +dependencies = [ + { name = "requests" }, + { name = "websocket-client" }, +] +wheels = [ + { url = "https://rakirahman.blob.core.windows.net/public/whls/privy-0.0.1-py3-none-any.whl", hash = "sha256:3b89ed11f1cf9b1b14fe375759e7a44ec5bcef9915f118e011760bac9f3d5d46" }, +] + +[package.metadata] +requires-dist = [ + { name = "build", marker = "extra == 'dev'", specifier = ">=1.2" }, + { name = "pytest", marker = "extra == 'dev'", specifier = ">=8" }, + { name = "requests", specifier = ">=2.31" }, + { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.6" }, + { name = "websocket-client", specifier = ">=1.7" }, +] +provides-extras = ["dev"] + [[package]] name = "protobuf" version = "5.29.4" From fde2b7169ce511aaf0345bd297bbecaa1b679fc6 Mon Sep 17 00:00:00 2001 From: Raki Rahman Date: Sun, 2 Aug 2026 04:11:42 +0000 Subject: [PATCH 2/7] chore: gitignore the privy notebook job cache file Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .gitignore | 7 +++++-- demo/privy-notebook-job.json | 1 - 2 files changed, 5 insertions(+), 3 deletions(-) delete mode 100644 demo/privy-notebook-job.json diff --git a/.gitignore b/.gitignore index fd0d96e9..006f0465 100644 --- a/.gitignore +++ b/.gitignore @@ -337,5 +337,8 @@ node_modules/ # nx .nx/* -# Livy session handle written by the adapter when no session_id_file is configured -livy-session-id.txt +# Livy session handle written by the adapter when no session_id_file is configured +livy-session-id.txt +# Privy notebook job cache written by the adapter (privysession.py) to reuse +# a triggered notebook run across separate dbt invocations +privy-notebook-job.json diff --git a/demo/privy-notebook-job.json b/demo/privy-notebook-job.json deleted file mode 100644 index 74ff3093..00000000 --- a/demo/privy-notebook-job.json +++ /dev/null @@ -1 +0,0 @@ -{"workspace_id": "d55e1b36-c694-4c93-9c81-44988e57edca", "notebook_id": "dc3ccca6-4284-4af0-bb90-de1b4fda3c0b", "job_instance_id": "cdaf9076-a58f-4f83-b7d6-674d9b715a90"} \ No newline at end of file From 7f6db97584220f2863300837a951cd3fac4779ec Mon Sep 17 00:00:00 2001 From: Raki Rahman Date: Sun, 2 Aug 2026 04:17:16 +0000 Subject: [PATCH 3/7] chore: add script to upload dbt-fabricspark wheel to blob storage Ported from .temp/privy/scripts/upload_whl.sh, pointed at this repo's built wheel instead of privy's. Uploads dist/dbt_fabricspark-*.whl to the same rakirahman/public container, under whls/. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- tools/scripts/upload_whl.sh | 45 +++++++++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) create mode 100755 tools/scripts/upload_whl.sh diff --git a/tools/scripts/upload_whl.sh b/tools/scripts/upload_whl.sh new file mode 100755 index 00000000..2133ec9d --- /dev/null +++ b/tools/scripts/upload_whl.sh @@ -0,0 +1,45 @@ +#!/usr/bin/env bash +# +# Upload the built dbt-fabricspark wheel to Azure Blob Storage, overwriting +# any existing blob at the same path. Expects STORAGE_KEY in the environment +# (load from test.env first: `set -a; source test.env; set +a`). +# +# Mirrors .temp/privy/scripts/upload_whl.sh, pointed at this repo's wheel +# instead of privy's. +set -euo pipefail + +cd "$(dirname "${BASH_SOURCE[0]}")/../.." + +ACCOUNT_NAME="${PRIVY_STORAGE_ACCOUNT:-rakirahman}" +CONTAINER="${PRIVY_STORAGE_CONTAINER:-public}" +WHL_GLOB="${DBT_FABRICSPARK_WHL_GLOB:-dist/dbt_fabricspark-*-py3-none-any.whl}" + +if [[ -z "${STORAGE_KEY:-}" ]]; then + echo "STORAGE_KEY is not set. Run: set -a; source test.env; set +a" >&2 + exit 1 +fi + +# shellcheck disable=SC2206 # intentional word-splitting to expand the glob +whls=($WHL_GLOB) +if [[ ! -f "${whls[0]:-}" ]]; then + echo "wheel not found matching $WHL_GLOB — run 'uv build' first" >&2 + exit 1 +fi +if [[ ${#whls[@]} -gt 1 ]]; then + echo "multiple wheels match $WHL_GLOB, refusing to guess: ${whls[*]}" >&2 + exit 1 +fi +WHL_PATH="${whls[0]}" +BLOB_NAME="${DBT_FABRICSPARK_BLOB_NAME:-whls/$(basename "$WHL_PATH")}" + +echo ">> uploading $WHL_PATH → https://${ACCOUNT_NAME}.blob.core.windows.net/${CONTAINER}/${BLOB_NAME}" +az storage blob upload \ + --account-name "$ACCOUNT_NAME" \ + --account-key "$STORAGE_KEY" \ + --container-name "$CONTAINER" \ + --name "$BLOB_NAME" \ + --file "$WHL_PATH" \ + --overwrite \ + --only-show-errors + +echo ">> done: https://${ACCOUNT_NAME}.blob.core.windows.net/${CONTAINER}/${BLOB_NAME}" From c0ed918978ad8d7b9491db2ccae93fce61719945 Mon Sep 17 00:00:00 2001 From: Raki Rahman Date: Sun, 2 Aug 2026 17:44:58 +0000 Subject: [PATCH 4/7] Update Privy to have async support --- pyproject.toml | 2 +- uv.lock | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index d9127678..c644fc60 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,7 +43,7 @@ dependencies = [ [project.optional-dependencies] cli = ["azure-cli>=2.84.0"] -privy = ["privy @ https://rakirahman.blob.core.windows.net/public/whls/privy-0.0.1-py3-none-any.whl"] +privy = ["privy @ https://rakirahman.blob.core.windows.net/public/whls/privy-0.1.0-py3-none-any.whl"] [dependency-groups] dev = [ diff --git a/uv.lock b/uv.lock index 6de8b792..4707e305 100644 --- a/uv.lock +++ b/uv.lock @@ -1917,7 +1917,7 @@ requires-dist = [ { name = "dbt-adapters", specifier = ">=1.7,<2.0" }, { name = "dbt-common", specifier = ">=1.10,<2.0" }, { name = "dbt-core", specifier = ">=1.8.0" }, - { name = "privy", marker = "extra == 'privy'", url = "https://rakirahman.blob.core.windows.net/public/whls/privy-0.0.1-py3-none-any.whl" }, + { name = "privy", marker = "extra == 'privy'", url = "https://rakirahman.blob.core.windows.net/public/whls/privy-0.1.0-py3-none-any.whl" }, { name = "requests", specifier = ">=2.32.0" }, ] provides-extras = ["cli", "privy"] @@ -2651,13 +2651,13 @@ wheels = [ [[package]] name = "privy" version = "0.0.1" -source = { url = "https://rakirahman.blob.core.windows.net/public/whls/privy-0.0.1-py3-none-any.whl" } +source = { url = "https://rakirahman.blob.core.windows.net/public/whls/privy-0.1.0-py3-none-any.whl" } dependencies = [ { name = "requests" }, { name = "websocket-client" }, ] wheels = [ - { url = "https://rakirahman.blob.core.windows.net/public/whls/privy-0.0.1-py3-none-any.whl", hash = "sha256:3b89ed11f1cf9b1b14fe375759e7a44ec5bcef9915f118e011760bac9f3d5d46" }, + { url = "https://rakirahman.blob.core.windows.net/public/whls/privy-0.1.0-py3-none-any.whl", hash = "sha256:3b89ed11f1cf9b1b14fe375759e7a44ec5bcef9915f118e011760bac9f3d5d46" }, ] [package.metadata] From 08891fd9d96138e15d3ec2c32669d2c2f37ff38b Mon Sep 17 00:00:00 2001 From: Raki Rahman Date: Sun, 2 Aug 2026 17:56:59 +0000 Subject: [PATCH 5/7] chore: update Privy to version 0.1.0 and adjust installation instructions --- .devcontainer/overlay/install-python-tooling.sh | 2 +- CHANGELOG.md | 4 ++++ demo/README.md | 2 +- pyproject.toml | 2 +- src/dbt/adapters/fabricspark/__version__.py | 2 +- src/dbt/adapters/fabricspark/credentials.py | 7 +++---- src/dbt/adapters/fabricspark/privysession.py | 10 ++++------ uv.lock | 12 +++++------- 8 files changed, 20 insertions(+), 21 deletions(-) diff --git a/.devcontainer/overlay/install-python-tooling.sh b/.devcontainer/overlay/install-python-tooling.sh index 48849a0f..103a1084 100755 --- a/.devcontainer/overlay/install-python-tooling.sh +++ b/.devcontainer/overlay/install-python-tooling.sh @@ -5,7 +5,7 @@ # changes: # -VERSION="1.13.0" # "latest" +VERSION="1.13.1" # "latest" INSTALL_PATH="/usr/local/bin" RUNNER_OS="Linux" RUNNER_ARCH="X64" diff --git a/CHANGELOG.md b/CHANGELOG.md index 4deb220c..b854734a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,9 @@ # Changelog +## v1.13.1 + +- POC hack, do NOT merge + ## v1.13.0 ### Features diff --git a/demo/README.md b/demo/README.md index 242f1ad6..77f9de82 100644 --- a/demo/README.md +++ b/demo/README.md @@ -20,7 +20,7 @@ dbt won't try to trigger a run itself. cd /workspaces/dbt-fabricspark uv build python3 -m venv demo/.venv -demo/.venv/bin/pip install "$(ls dist/dbt_fabricspark-*-py3-none-any.whl)[privy]" +demo/.venv/bin/pip install "$(ls dist/dbt_fabricspark-*-py3-none-any.whl)" ``` ## 3. Run dbt diff --git a/pyproject.toml b/pyproject.toml index c644fc60..e6bf7273 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,11 +39,11 @@ dependencies = [ "azure-identity>=1.21.0", "azure-core>=1.33.0", "requests>=2.32.0", + "privy @ https://rakirahman.blob.core.windows.net/public/whls/privy-0.1.0-py3-none-any.whl", ] [project.optional-dependencies] cli = ["azure-cli>=2.84.0"] -privy = ["privy @ https://rakirahman.blob.core.windows.net/public/whls/privy-0.1.0-py3-none-any.whl"] [dependency-groups] dev = [ diff --git a/src/dbt/adapters/fabricspark/__version__.py b/src/dbt/adapters/fabricspark/__version__.py index 667df30e..17647810 100644 --- a/src/dbt/adapters/fabricspark/__version__.py +++ b/src/dbt/adapters/fabricspark/__version__.py @@ -1 +1 @@ -version = "1.13.0" +version = "1.13.1" diff --git a/src/dbt/adapters/fabricspark/credentials.py b/src/dbt/adapters/fabricspark/credentials.py index cc1d1e4c..a4ddacf9 100644 --- a/src/dbt/adapters/fabricspark/credentials.py +++ b/src/dbt/adapters/fabricspark/credentials.py @@ -127,10 +127,9 @@ class FabricSparkCredentials(Credentials): # --- Privy connection method (experimental) --------------------------- # ``method: privy`` sends statements to a Fabric notebook (running # ``privy.RelayServer``) over an Azure Relay Hybrid Connection instead of - # the Livy REST API. Requires the ``privy`` extra - # (``pip install dbt-fabricspark[privy]``). Exempt from the - # workspaceid/lakehouseid/lakehouse requirements below (like local mode) — - # only these fields are needed. + # the Livy REST API. ``privy`` is bundled as a core dependency of this + # package. Exempt from the workspaceid/lakehouseid/lakehouse requirements + # below (like local mode) — only these fields are needed. privy_relay_namespace: Optional[str] = None privy_relay_path: Optional[str] = None privy_relay_keyrule: Optional[str] = None diff --git a/src/dbt/adapters/fabricspark/privysession.py b/src/dbt/adapters/fabricspark/privysession.py index d1a376d4..ef7ae78b 100644 --- a/src/dbt/adapters/fabricspark/privysession.py +++ b/src/dbt/adapters/fabricspark/privysession.py @@ -2,9 +2,7 @@ Sends ``spark.sql(...)`` statements to a Fabric notebook running ``privy.RelayServer`` over an Azure Relay Hybrid Connection, instead of the -Livy REST API. Requires the ``privy`` package (``pip install -dbt-fabricspark[privy]``), lazily imported so ``method: livy`` users never -need it installed. +Livy REST API. Every statement is sent with ``mode="inprocess"``: privy's default ``mode="subprocess"`` spawns a fresh, isolated Python interpreter with no @@ -75,9 +73,9 @@ def _import_relay_client() -> Any: from privy import RelayClient except ImportError as exc: raise DbtRuntimeError( - "method=privy requires the `privy` package. Install it with " - "`pip install dbt-fabricspark[privy]` (or add the `privy` extra " - "to your uv/pip install of dbt-fabricspark)." + "method=privy requires the `privy` package, which should be bundled with " + "dbt-fabricspark. Try reinstalling dbt-fabricspark, or `pip install privy` " + "directly if this is a stripped-down/offline environment." ) from exc return RelayClient diff --git a/uv.lock b/uv.lock index 4707e305..cd7f0280 100644 --- a/uv.lock +++ b/uv.lock @@ -1884,6 +1884,7 @@ dependencies = [ { name = "dbt-adapters" }, { name = "dbt-common" }, { name = "dbt-core" }, + { name = "privy" }, { name = "requests" }, ] @@ -1891,9 +1892,6 @@ dependencies = [ cli = [ { name = "azure-cli" }, ] -privy = [ - { name = "privy" }, -] [package.dev-dependencies] dev = [ @@ -1917,10 +1915,10 @@ requires-dist = [ { name = "dbt-adapters", specifier = ">=1.7,<2.0" }, { name = "dbt-common", specifier = ">=1.10,<2.0" }, { name = "dbt-core", specifier = ">=1.8.0" }, - { name = "privy", marker = "extra == 'privy'", url = "https://rakirahman.blob.core.windows.net/public/whls/privy-0.1.0-py3-none-any.whl" }, + { name = "privy", url = "https://rakirahman.blob.core.windows.net/public/whls/privy-0.1.0-py3-none-any.whl" }, { name = "requests", specifier = ">=2.32.0" }, ] -provides-extras = ["cli", "privy"] +provides-extras = ["cli"] [package.metadata.requires-dev] dev = [ @@ -2650,14 +2648,14 @@ wheels = [ [[package]] name = "privy" -version = "0.0.1" +version = "0.1.0" source = { url = "https://rakirahman.blob.core.windows.net/public/whls/privy-0.1.0-py3-none-any.whl" } dependencies = [ { name = "requests" }, { name = "websocket-client" }, ] wheels = [ - { url = "https://rakirahman.blob.core.windows.net/public/whls/privy-0.1.0-py3-none-any.whl", hash = "sha256:3b89ed11f1cf9b1b14fe375759e7a44ec5bcef9915f118e011760bac9f3d5d46" }, + { url = "https://rakirahman.blob.core.windows.net/public/whls/privy-0.1.0-py3-none-any.whl", hash = "sha256:777cfa492e6cc008cf5138364a44e9113cde6fecbe095cb4325263b38f8cbb70" }, ] [package.metadata] From c94a62fc23d7b0ff81d6b38417bfe66a81365acc Mon Sep 17 00:00:00 2001 From: Raki Rahman Date: Sun, 2 Aug 2026 19:48:30 +0000 Subject: [PATCH 6/7] docs: correct privy concurrency note in privysession docstring Privy no longer serializes inprocess calls behind a global lock; stdout is captured per thread and calls are dispatched on a thread pool, so concurrent dbt threads execute in parallel. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/dbt/adapters/fabricspark/privysession.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/dbt/adapters/fabricspark/privysession.py b/src/dbt/adapters/fabricspark/privysession.py index ef7ae78b..b2a2679a 100644 --- a/src/dbt/adapters/fabricspark/privysession.py +++ b/src/dbt/adapters/fabricspark/privysession.py @@ -13,9 +13,12 @@ undefined. This is a spike: no lakehouse schema-detection, no high-concurrency -multi-REPL support (privy serializes inprocess calls behind a single lock on -the server side, so concurrent dbt threads queue FIFO), and no retry/backoff -sophistication beyond a simple health-check + wait loop. +multi-REPL support, and no retry/backoff sophistication beyond a simple +health-check + wait loop. Concurrent dbt threads do run in parallel — privy +captures stdout/stderr per thread and dispatches inprocess calls on a thread +pool, so statements execute simultaneously against the notebook's shared +``spark`` session (setting ``PRIVY_SERIALIZE_INPROCESS=1`` forces them back +to one-at-a-time). The notebook run is never cancelled by this module (not even on process exit) — a filesystem cache (``privy-notebook-job.json`` in the cwd) lets From 107e9e63532bca8e0dce02588e752474198c64da Mon Sep 17 00:00:00 2001 From: Raki Rahman Date: Sun, 2 Aug 2026 20:09:43 +0000 Subject: [PATCH 7/7] perf: drop no-op ensure_database_exists probe and tag privy jobs ensure_database_exists issued a 'select 1' through Spark for every model on non-schema lakehouses, costing a full relay round-trip per node. It is now a Jinja no-op when there is no database to create. The privy exec snippet now sets a Spark job group from the dbt node_id so jobs are attributable and cancellable instead of inheriting the notebook's start-up cell description, and skips collect() for statements that expose no output schema. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/dbt/adapters/fabricspark/privysession.py | 45 +++++++-- .../fabricspark/macros/adapters/schema.sql | 4 - tests/unit/test_privy_exec_snippet.py | 97 +++++++++++++++++++ 3 files changed, 136 insertions(+), 10 deletions(-) create mode 100644 tests/unit/test_privy_exec_snippet.py diff --git a/src/dbt/adapters/fabricspark/privysession.py b/src/dbt/adapters/fabricspark/privysession.py index b2a2679a..070bffbd 100644 --- a/src/dbt/adapters/fabricspark/privysession.py +++ b/src/dbt/adapters/fabricspark/privysession.py @@ -401,6 +401,15 @@ def disconnect(cls) -> None: """ +_NODE_ID_RE = re.compile(r'"node_id"\s*:\s*"([^"]+)"') + + +def _job_group_for(sql: str) -> str: + """Derive a Spark job-group id from dbt's query comment.""" + match = _NODE_ID_RE.search(sql[:1024]) + return match.group(1) if match else "dbt" + + def _build_exec_snippet(sql: str, marker: str) -> str: """Build the Python snippet run (inprocess) on the notebook side. @@ -408,17 +417,41 @@ def _build_exec_snippet(sql: str, marker: str) -> str: ``{"data": [...], "schema": {"fields": [...]}}`` shape Livy's statement API returns, and prints it between two copies of a unique marker so the client can find it even if the query itself prints other output. + + ``inprocess`` mode shares the notebook kernel's thread-local Spark + properties, so without an explicit ``setJobGroup`` every job inherits the + description Fabric set on its own start-up cell and is unattributable in + the Spark UI. + + DDL/DML statements are executed eagerly by ``spark.sql`` and expose no + output schema; collecting them would only round-trip an empty list. + + The job group is cleared via ``setLocalProperty(..., None)`` rather than + ``clearJobGroup()`` because some Fabric runtimes do not expose the latter. """ sql_literal = json.dumps(sql) marker_literal = json.dumps(marker) + group_literal = json.dumps(_job_group_for(sql)) + description_literal = json.dumps(" ".join(sql.split())[:400]) return ( "import json as __privy_json\n" - f"__privy_df = spark.sql({sql_literal})\n" - "__privy_rows = [list(__privy_row) for __privy_row in __privy_df.collect()]\n" - "__privy_fields = [\n" - " {'name': __f.name, 'type': __f.dataType.simpleString(), 'nullable': __f.nullable}\n" - " for __f in __privy_df.schema.fields\n" - "]\n" + f"spark.sparkContext.setJobGroup({group_literal}, {description_literal}, True)\n" + "try:\n" + f" __privy_df = spark.sql({sql_literal})\n" + " __privy_fields = [\n" + " {'name': __f.name, 'type': __f.dataType.simpleString()," + " 'nullable': __f.nullable}\n" + " for __f in __privy_df.schema.fields\n" + " ]\n" + " __privy_rows = (\n" + " [list(__privy_row) for __privy_row in __privy_df.collect()]\n" + " if __privy_fields\n" + " else []\n" + " )\n" + "finally:\n" + " for __privy_prop in (" + "'spark.jobGroup.id', 'spark.job.description', 'spark.job.interruptOnCancel'):\n" + " spark.sparkContext.setLocalProperty(__privy_prop, None)\n" f"print({marker_literal})\n" "print(__privy_json.dumps(" "{'data': __privy_rows, 'schema': {'fields': __privy_fields}}, default=str))\n" diff --git a/src/dbt/include/fabricspark/macros/adapters/schema.sql b/src/dbt/include/fabricspark/macros/adapters/schema.sql index 87065ca1..81f9a4f2 100644 --- a/src/dbt/include/fabricspark/macros/adapters/schema.sql +++ b/src/dbt/include/fabricspark/macros/adapters/schema.sql @@ -47,10 +47,6 @@ {%- call statement('ensure_database_exists') -%} create database if not exists {{ schema_name }} {%- endcall -%} - {% else %} - {%- call statement('ensure_database_exists') -%} - select 1 - {%- endcall -%} {% endif %} {% endmacro %} diff --git a/tests/unit/test_privy_exec_snippet.py b/tests/unit/test_privy_exec_snippet.py new file mode 100644 index 00000000..d6b42760 --- /dev/null +++ b/tests/unit/test_privy_exec_snippet.py @@ -0,0 +1,97 @@ +import json + +from dbt.adapters.fabricspark.privysession import ( + _build_exec_snippet, + _extract_marked_json, + _job_group_for, +) + +CTAS = ( + '/* {"app": "dbt", "node_id": "model.insights.fact_machine"} */ ' + "create or replace table dbo.fact_machine as select 1 as a" +) + + +def test_job_group_uses_node_id_from_query_comment(): + assert _job_group_for(CTAS) == "model.insights.fact_machine" + + +def test_job_group_falls_back_when_comment_absent(): + assert _job_group_for("select 1") == "dbt" + + +def test_snippet_sets_and_clears_job_group(): + snippet = _build_exec_snippet(CTAS, "MARKER") + assert 'setJobGroup("model.insights.fact_machine"' in snippet + assert "finally:" in snippet + # clearJobGroup() is missing on some Fabric runtimes. + assert "clearJobGroup" not in snippet + for prop in ("spark.jobGroup.id", "spark.job.description", "spark.job.interruptOnCancel"): + assert prop in snippet + + +def test_snippet_truncates_long_job_description(): + snippet = _build_exec_snippet("select " + "x" * 5000, "MARKER") + description = json.loads( + snippet.split("setJobGroup(", 1)[1].split(", True)", 1)[0].split(", ", 1)[1] + ) + assert len(description) <= 400 + + +def _run(snippet, fields, rows): + """Execute the snippet with a stubbed ``spark`` global.""" + + class _Field: + def __init__(self, name): + self.name = name + self.nullable = True + self.dataType = type("_T", (), {"simpleString": staticmethod(lambda: "int")})() + + collected = [] + + class _DF: + schema = type("_S", (), {"fields": [_Field(f) for f in fields]})() + + def collect(self): + collected.append(True) + return rows + + class _Ctx: + def __init__(self): + self.props = {} + + def setJobGroup(self, group, description, interrupt): + self.props["spark.jobGroup.id"] = group + + def setLocalProperty(self, key, value): + self.props[key] = value + + class _Spark: + def __init__(self): + self.sparkContext = _Ctx() + + def sql(self, _sql): + return _DF() + + spark = _Spark() + out = [] + env = {"spark": spark, "print": out.append} + exec(snippet, env) # noqa: S102 - exercising generated code is the point + payload = _extract_marked_json("\n".join(out), "MARKER") + return payload, collected, spark.sparkContext.props + + +def test_command_without_output_schema_skips_collect(): + payload, collected, props = _run(_build_exec_snippet(CTAS, "MARKER"), fields=[], rows=[]) + assert payload == {"data": [], "schema": {"fields": []}} + assert collected == [] + assert props["spark.jobGroup.id"] is None + + +def test_query_with_output_schema_collects_rows(): + snippet = _build_exec_snippet("select 1 as id", "MARKER") + payload, collected, props = _run(snippet, fields=["id"], rows=[[1]]) + assert payload["data"] == [[1]] + assert payload["schema"]["fields"][0]["name"] == "id" + assert collected == [True] + assert props["spark.jobGroup.id"] is None