From 53258ce79d0a1eccefb4b628ca6b44e90abd8c64 Mon Sep 17 00:00:00 2001 From: Raki Rahman Date: Sun, 2 Aug 2026 04:07:57 +0000 Subject: [PATCH 1/8] 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/8] 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/8] 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/8] 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/8] 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/8] 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/8] 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 From e27086e0d206c706c99581a9f90c043306f5b5fd Mon Sep 17 00:00:00 2001 From: Raki Rahman Date: Sat, 12 Sep 2026 14:29:30 +0000 Subject: [PATCH 8/8] fix: isolate Privy requests and pin reviewed source Keep generated Spark request variables local while preserving synchronous client calls and automatic async transport. Cover interleaved rows, schemas, job groups, probes, DDL, errors, and timeouts using the pinned Privy executor. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pyproject.toml | 2 +- src/dbt/adapters/fabricspark/privysession.py | 6 +- tests/unit/test_privy_exec_snippet.py | 318 ++++++++++++++++++- uv.lock | 17 +- 4 files changed, 321 insertions(+), 22 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index e6bf7273..f859b550 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,7 +39,7 @@ 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", + "privy @ git+https://github.com/mdrakiburrahman/privy.git@2f39703cc4165c91958efbb0c5e1ab3ec27ac9ab", ] [project.optional-dependencies] diff --git a/src/dbt/adapters/fabricspark/privysession.py b/src/dbt/adapters/fabricspark/privysession.py index 070bffbd..5fc523b7 100644 --- a/src/dbt/adapters/fabricspark/privysession.py +++ b/src/dbt/adapters/fabricspark/privysession.py @@ -418,6 +418,9 @@ def _build_exec_snippet(sql: str, marker: str) -> str: 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. + Each request gets its own globals dictionary containing only ``spark`` + from the notebook, so concurrent requests cannot overwrite result variables. + ``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 @@ -433,7 +436,7 @@ def _build_exec_snippet(sql: str, marker: str) -> str: marker_literal = json.dumps(marker) group_literal = json.dumps(_job_group_for(sql)) description_literal = json.dumps(" ".join(sql.split())[:400]) - return ( + body = ( "import json as __privy_json\n" f"spark.sparkContext.setJobGroup({group_literal}, {description_literal}, True)\n" "try:\n" @@ -457,6 +460,7 @@ def _build_exec_snippet(sql: str, marker: str) -> str: "{'data': __privy_rows, 'schema': {'fields': __privy_fields}}, default=str))\n" f"print({marker_literal})\n" ) + return f"exec({body!r}, {{'spark': spark}})\n" def _extract_marked_json(stdout: str, marker: str) -> Dict[str, Any]: diff --git a/tests/unit/test_privy_exec_snippet.py b/tests/unit/test_privy_exec_snippet.py index d6b42760..9e4dbafc 100644 --- a/tests/unit/test_privy_exec_snippet.py +++ b/tests/unit/test_privy_exec_snippet.py @@ -1,15 +1,37 @@ +import builtins +import datetime as dt +import io import json +import sys +import threading +from concurrent.futures import ThreadPoolExecutor +from contextlib import redirect_stdout +from types import SimpleNamespace +from unittest.mock import Mock, call +import pytest +from dbt_common.exceptions import DbtDatabaseError +from privy import executor +from privy.client import RELAY_RESPONSE_LIMIT_S, ExecResult +from privy.protocol import ExecRequest, ExecResponse + +from dbt.adapters.fabricspark.credentials import FabricSparkCredentials from dbt.adapters.fabricspark.privysession import ( + _PROBE_TIMEOUT_S, + _UNBOUNDED_TIMEOUT_S, + PrivyConnectionWrapper, _build_exec_snippet, + _build_relay_client, _extract_marked_json, _job_group_for, + _probe, ) CTAS = ( '/* {"app": "dbt", "node_id": "model.insights.fact_machine"} */ ' "create or replace table dbo.fact_machine as select 1 as a" ) +JOB_PROPERTIES = ("spark.jobGroup.id", "spark.job.description", "spark.job.interruptOnCancel") def test_job_group_uses_node_id_from_query_comment(): @@ -26,7 +48,7 @@ def test_snippet_sets_and_clears_job_group(): 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"): + for prop in JOB_PROPERTIES: assert prop in snippet @@ -74,10 +96,11 @@ 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") + out = io.StringIO() + env = {"spark": spark} + with redirect_stdout(out): + exec(snippet, env) # noqa: S102 - exercising generated code is the point + payload = _extract_marked_json(out.getvalue(), "MARKER") return payload, collected, spark.sparkContext.props @@ -95,3 +118,288 @@ def test_query_with_output_schema_collects_rows(): assert payload["schema"]["fields"][0]["name"] == "id" assert collected == [True] assert props["spark.jobGroup.id"] is None + + +def _credentials(statement_timeout): + return FabricSparkCredentials( + method="privy", + privy_relay_namespace="test-relay", + privy_relay_path="test", + privy_relay_keyrule="test", + privy_relay_key="unused-test-key", + privy_notebook_url="https://app.fabric.microsoft.com/groups/" + "00000000-0000-0000-0000-000000000001/synapsenotebooks/" + "00000000-0000-0000-0000-000000000002", + spark_config={"name": "test"}, + statement_timeout=statement_timeout, + ) + + +def _field(name, data_type, nullable): + return SimpleNamespace( + name=name, + dataType=SimpleNamespace(simpleString=lambda: data_type), + nullable=nullable, + ) + + +@pytest.fixture +def inprocess_relay(monkeypatch): + spark = Mock() + shared_globals = { + "__builtins__": builtins.__dict__, + "spark": spark, + "sc": spark.sparkContext, + "notebook_state": object(), + **{ + name: object() + for name in ( + "__privy_json", + "__privy_df", + "__privy_fields", + "__privy_rows", + "__privy_prop", + ) + }, + } + monkeypatch.setattr(executor, "_INPROCESS_GLOBALS", shared_globals) + monkeypatch.setattr(executor, "_SERIALIZE_INPROCESS", False) + monkeypatch.setattr(executor, "_JOBS", {}) + monkeypatch.setattr(executor, "_ROUTERS", {}) + # Privy installs per-thread stream routers; restore pytest's streams afterwards. + monkeypatch.setattr(sys, "stdout", sys.stdout) + monkeypatch.setattr(sys, "stderr", sys.stderr) + + client = _build_relay_client(_credentials(10), http_timeout_s=40) + requests = [] + responses = [] + + def post_json(payload, *, http_timeout_s=None): + request = ExecRequest.from_json(payload) + requests.append(request) + response = executor.execute(request) + responses.append(response) + return json.loads(response.to_json()) + + monkeypatch.setattr(client, "_post_json", post_json) + return SimpleNamespace( + client=client, + spark=spark, + globals=shared_globals, + requests=requests, + responses=responses, + ) + + +@pytest.mark.parametrize("statement_timeout", [10, 55, 56, 120, 0]) +def test_wrapper_preserves_sync_api_and_automatic_long_request_protocol( + inprocess_relay, statement_timeout +): + relay = inprocess_relay + before = relay.globals.copy() + rows = [[7, dt.datetime(2026, 9, 12, 12, 0), dt.date(2026, 9, 12)]] + frame = Mock() + frame.schema.fields = [ + _field("id", "int", False), + _field("created_at", "timestamp", True), + _field("day", "date", True), + ] + frame.collect.return_value = rows + relay.spark.sql.return_value = frame + wrapper = PrivyConnectionWrapper(relay.client, _credentials(statement_timeout)) + + wrapper.execute(" select 'quoted \\' value\\n雪' as id; ") + + relay.spark.sql.assert_called_once_with("select 'quoted \\' value\\n雪' as id") + assert wrapper.fetchall() == rows + assert wrapper.description == [ + ("id", "int", None, None, None, None, False), + ("created_at", "timestamp", None, None, None, None, True), + ("day", "date", None, None, None, None, True), + ] + timeout_s = statement_timeout or _UNBOUNDED_TIMEOUT_S + actions = [request.action for request in relay.requests] + if timeout_s > RELAY_RESPONSE_LIMIT_S: + assert actions == ["submit", "poll"] + else: + assert actions == ["exec"] + assert all(request.mode == "inprocess" for request in relay.requests) + assert all(request.timeout_s == timeout_s for request in relay.requests) + assert relay.globals == before + assert relay.spark.sparkContext.setLocalProperty.call_args_list == [ + call(prop, None) for prop in JOB_PROPERTIES + ] + + +@pytest.mark.parametrize("statement_timeout", [10, 120]) +@pytest.mark.parametrize("pause_at", ["schema", "collect", "cleanup"]) +def test_interleaved_requests_keep_distinct_results_and_job_groups( + inprocess_relay, statement_timeout, pause_at +): + relay = inprocess_relay + before = relay.globals.copy() + paused = threading.Event() + second_finished = threading.Event() + active_group = threading.local() + cleared = {"model.first": [], "model.second": []} + first_fields = [_field("first_id", "int", False)] + second_fields = [_field("second_name", "string", True), _field("value", "double", False)] + first_rows = [[1], [2]] + second_rows = [["second", 3.5]] + + def pause(): + paused.set() + assert second_finished.wait(5), "second request did not execute concurrently" + + class FirstFrame: + @property + def schema(self): + if pause_at == "schema": + pause() + return SimpleNamespace(fields=first_fields) + + def collect(self): + if pause_at == "collect": + pause() + return first_rows + + second_frame = Mock() + second_frame.schema.fields = second_fields + second_frame.collect.return_value = second_rows + + def sql(query): + print("Spark diagnostic output") + return FirstFrame() if "model.first" in query else second_frame + + def set_job_group(group, description, interrupt): + active_group.name = group + assert interrupt is True + + def clear_property(prop, value): + group = active_group.name + if group == "model.first" and prop == JOB_PROPERTIES[0] and pause_at == "cleanup": + pause() + cleared[group].append((prop, value)) + + relay.spark.sql.side_effect = sql + relay.spark.sparkContext.setJobGroup.side_effect = set_job_group + relay.spark.sparkContext.setLocalProperty.side_effect = clear_property + first = PrivyConnectionWrapper(relay.client, _credentials(statement_timeout)) + second = PrivyConnectionWrapper(relay.client, _credentials(statement_timeout)) + + def run_second(): + try: + second.execute('/* {"node_id": "model.second"} */ select second_name, value') + finally: + second_finished.set() + + with ThreadPoolExecutor(max_workers=2) as pool: + first_result = pool.submit( + first.execute, '/* {"node_id": "model.first"} */ select first_id' + ) + assert paused.wait(5), "first request did not reach the interleaving point" + second_result = pool.submit(run_second) + second_result.result(timeout=10) + first_result.result(timeout=10) + + assert first.fetchall() == first_rows + assert first.description == [("first_id", "int", None, None, None, None, False)] + assert second.fetchall() == second_rows + assert second.description == [ + ("second_name", "string", None, None, None, None, True), + ("value", "double", None, None, None, None, False), + ] + for props in cleared.values(): + assert props == [(prop, None) for prop in JOB_PROPERTIES] + assert relay.globals == before + + +@pytest.mark.parametrize("statement_timeout", [10, 120]) +def test_wrapper_ddl_skips_collect(inprocess_relay, statement_timeout): + relay = inprocess_relay + frame = relay.spark.sql.return_value + frame.schema.fields = [] + wrapper = PrivyConnectionWrapper(relay.client, _credentials(statement_timeout)) + + wrapper.execute(CTAS) + + assert wrapper.fetchall() == [] + assert wrapper.description == [] + frame.collect.assert_not_called() + assert relay.spark.sparkContext.setLocalProperty.call_count == len(JOB_PROPERTIES) + + +@pytest.mark.parametrize("statement_timeout", [10, 120]) +@pytest.mark.parametrize("failure_at", ["sql", "collect"]) +def test_wrapper_query_errors_clear_job_group(inprocess_relay, statement_timeout, failure_at): + relay = inprocess_relay + frame = relay.spark.sql.return_value + frame.schema.fields = [_field("id", "int", False)] + failing_call = relay.spark.sql if failure_at == "sql" else frame.collect + failing_call.side_effect = RuntimeError("query failed") + wrapper = PrivyConnectionWrapper(relay.client, _credentials(statement_timeout)) + + with pytest.raises(DbtDatabaseError, match="query failed"): + wrapper.execute("select 1") + + assert relay.responses[-1].stdout == b"" + assert relay.spark.sparkContext.setLocalProperty.call_args_list == [ + call(prop, None) for prop in JOB_PROPERTIES + ] + + +def test_probe_uses_short_inprocess_request(inprocess_relay): + relay = inprocess_relay + + assert _probe(relay.client) is True + + assert len(relay.requests) == 1 + request = relay.requests[0] + assert (request.code, request.mode, request.action, request.timeout_s) == ( + "1", + "inprocess", + "exec", + _PROBE_TIMEOUT_S, + ) + relay.spark.sql.assert_not_called() + + +@pytest.mark.parametrize("statement_timeout", [10, 120, 0]) +def test_wrapper_propagates_timeout_without_retry(statement_timeout): + client = Mock() + client.run_python.return_value = ExecResult.from_response( + ExecResponse.from_output( + exit_code=1, stdout=b"", stderr=b"deadline exceeded", duration_ms=0, timed_out=True + ) + ) + wrapper = PrivyConnectionWrapper(client, _credentials(statement_timeout)) + + with pytest.raises(DbtDatabaseError, match=r"\(timed out\).*deadline exceeded"): + wrapper.execute("select 1") + + client.run_python.assert_called_once() + assert client.run_python.call_args.kwargs == { + "mode": "inprocess", + "timeout_s": statement_timeout or _UNBOUNDED_TIMEOUT_S, + } + + +def test_probe_returns_false_on_transport_error(): + client = Mock() + client.run_python.side_effect = RuntimeError("relay unavailable") + + assert _probe(client) is False + client.run_python.assert_called_once_with("1", mode="inprocess", timeout_s=_PROBE_TIMEOUT_S) + + +@pytest.mark.parametrize( + ("stdout", "error"), + [ + ("Spark output only", "missing the result marker"), + ("MARKER\n{}", "missing the closing result marker"), + ("MARKER\nnot-json\nMARKER", "Could not parse Privy result JSON"), + ], +) +def test_invalid_result_markers_raise_database_error(stdout, error): + with pytest.raises(DbtDatabaseError, match=error): + _extract_marked_json(stdout, "MARKER") diff --git a/uv.lock b/uv.lock index cd7f0280..bb6ce185 100644 --- a/uv.lock +++ b/uv.lock @@ -1915,7 +1915,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", url = "https://rakirahman.blob.core.windows.net/public/whls/privy-0.1.0-py3-none-any.whl" }, + { name = "privy", git = "https://github.com/mdrakiburrahman/privy.git?rev=2f39703cc4165c91958efbb0c5e1ab3ec27ac9ab" }, { name = "requests", specifier = ">=2.32.0" }, ] provides-extras = ["cli"] @@ -2649,24 +2649,11 @@ wheels = [ [[package]] name = "privy" version = "0.1.0" -source = { url = "https://rakirahman.blob.core.windows.net/public/whls/privy-0.1.0-py3-none-any.whl" } +source = { git = "https://github.com/mdrakiburrahman/privy.git?rev=2f39703cc4165c91958efbb0c5e1ab3ec27ac9ab#2f39703cc4165c91958efbb0c5e1ab3ec27ac9ab" } 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:777cfa492e6cc008cf5138364a44e9113cde6fecbe095cb4325263b38f8cbb70" }, -] - -[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"