Repository navigation
feat(volume): enable native SQL file transfers #1025
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
2cd0562
9540245
175efa3
a82dba0
c09ff98
3305d43
67efcd9
b47f25e
37fa359
964533b
4bfd4b9
e1fc615
317bb1a
0f391f7
865f6a6
83fc674
224a044
dc9d19e
9591f93
934cf35
9ddfb80
b07f679
1598dc0
6015540
91ad80a
f7efb41
706c9c4
fb2ef00
a42d021
5e78ed6
dbf9701
3efbcc9
a8664be
5df354a
04a52c5
e9da44a
8ab0fa6
70f6695
43fedf4
73990ee
6fa0c01
31e1e05
e9a48d7
993fe73
e3806ee
8e10949
6794035
1171efc
aa92717
9b37b13
9e0f8ab
a90f623
d119a79
85fe971
a197d58
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1 +1 @@ | ||
| 6f4b6d607a61bdde61dcda2cf82bc8e80fd32cb9 | ||
| 37c7a3ec1b54f4690d9325d65f2220e0bb64b83e | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -16,14 +16,14 @@ | |
|
|
||
| - ``query_tags`` on execute is not supported (kernel exposes | ||
| ``statement_conf`` but PyO3 doesn't surface it). | ||
| - Volume PUT/GET (staging operations): kernel has no Volume API | ||
| yet. Users on Thrift-only paths. | ||
| """ | ||
|
|
||
| from __future__ import annotations | ||
|
|
||
| import inspect | ||
| import logging | ||
| import os | ||
| import re | ||
| import threading | ||
| import uuid | ||
| from typing import Any, Dict, List, Optional, Set, TYPE_CHECKING, Union | ||
|
|
@@ -73,24 +73,19 @@ | |
| # http_path). Forwarding either is redundant and trips the kernel's | ||
| # per-request skip-and-warn. | ||
| _KERNEL_MANAGED_HEADERS = frozenset({"authorization", "x-databricks-org-id"}) | ||
|
|
||
| # Leading verbs of SQL volume/staging statements. Detected by the | ||
| # leading token (case-insensitive) so the kernel backend can fail loud | ||
| # on staging ops it can't service — see ``execute_command``. | ||
| _STAGING_VERBS = ("PUT", "GET", "REMOVE") | ||
| _STAGING_PATHS_TYPE_ERROR = ( | ||
| "staging_allowed_local_path must be a string or a list of " | ||
| "strings or os.PathLike objects" | ||
| ) | ||
|
|
||
|
|
||
| def _strip_leading_sql_comments(sql: str) -> str: | ||
| """Strip leading whitespace and SQL comments (``-- …`` line and | ||
| ``/* … */`` block, possibly several) from ``sql``, returning the | ||
| remainder. | ||
|
|
||
| Needed so staging detection sees the real leading verb: a | ||
| comment-prefixed staging op (``-- upload\\nPUT …`` or | ||
| ``/* c */ PUT …``, common in ETL scripts) must still be classified | ||
| as staging, or it would slip past the guard into the silent-no-op | ||
| bug. Block comments do not nest in Databricks SQL, so a simple | ||
| scan-to-``*/`` is correct. | ||
| Comment-prefixed PUT/GET statements need the permitted local paths | ||
| just as statements that begin with the verb do. | ||
| """ | ||
| i = 0 | ||
| n = len(sql) | ||
|
|
@@ -102,9 +97,15 @@ def _strip_leading_sql_comments(sql: str) -> str: | |
| nl = sql.find("\n", i) | ||
| i = n if nl == -1 else nl + 1 | ||
| elif sql.startswith("/*", i): | ||
| # Block comment: skip to closing */ (or end if unterminated). | ||
| close = sql.find("*/", i + 2) | ||
| i = n if close == -1 else close + 2 | ||
| depth = 1 | ||
| i += 2 | ||
| for delimiter in re.finditer(r"/\*|\*/", sql[i:]): | ||
| depth += 1 if delimiter.group() == "/*" else -1 | ||
| if depth == 0: | ||
| i += delimiter.end() | ||
| break | ||
| else: | ||
| i = n | ||
| else: | ||
| break | ||
| return sql[i:] | ||
|
|
@@ -145,19 +146,10 @@ def _catalog_or_none(value: Optional[str]) -> Optional[str]: | |
| return value | ||
|
|
||
|
|
||
| def _is_staging_statement(operation: str) -> bool: | ||
| """True iff ``operation`` is a volume/staging statement (PUT / GET / | ||
| REMOVE). | ||
|
|
||
| Strips leading whitespace + SQL comments first (so a comment- | ||
| prefixed staging op is still caught), then matches the leading token | ||
| only — so a normal query that merely *contains* the word (e.g. | ||
| ``SELECT 'GET' AS x``) isn't misflagged. | ||
| """ | ||
| stripped = _strip_leading_sql_comments(operation) | ||
| # First whitespace-delimited token, uppercased. | ||
| verb = stripped.split(None, 1)[0].upper() if stripped.strip() else "" | ||
| return verb in _STAGING_VERBS | ||
| def _staging_verb(stripped_sql: str) -> Optional[str]: | ||
|
peco-review-bot[bot] marked this conversation as resolved.
Resolved
|
||
| """Match a leading verb so SELECT 'GET' cannot enable local file access.""" | ||
| match = re.match(r"(PUT|GET|REMOVE)\b", stripped_sql, re.IGNORECASE) | ||
| return match.group(1).upper() if match else None | ||
|
|
||
|
|
||
| def _kernel_session_accepts_kwarg(name: str) -> bool: | ||
|
|
@@ -575,28 +567,32 @@ def execute_command( | |
| # ``_async_statements`` and closed by ``close_command``); the sync | ||
| # path drops it in finally. ``close_stmt`` is the post-success | ||
| # decision flag — it stays True on sync, flips to False on async. | ||
| # Volume/staging (PUT/GET/REMOVE) is not supported on the kernel | ||
| # path: the kernel returns the staging control row as a normal | ||
| # result set (``KernelResultSet.is_staging_operation`` is always | ||
| # False), so the connector's ``_handle_staging_operation`` never | ||
| # fires and NO file is transferred. Rather than silently no-op | ||
| # (the Thrift path performs the presigned-URL upload/download), | ||
| # fail loud at the call site so ETL scripts don't ingest | ||
| # stale/missing data. Detected by the leading SQL verb — the | ||
| # only signal available pre-execute, since the kernel exposes no | ||
| # staging marker today. | ||
| if _is_staging_statement(operation): | ||
| raise NotSupportedError( | ||
| "Volume / staging operations (PUT / GET / REMOVE) are not " | ||
| "supported on the kernel backend (use_kernel=True); the file " | ||
| "transfer would silently not happen. Use the Thrift backend " | ||
| "for staging operations." | ||
| ) | ||
|
|
||
| close_stmt = True | ||
| try: | ||
| try: | ||
| stripped_sql = _strip_leading_sql_comments(operation) | ||
|
peco-review-bot[bot] marked this conversation as resolved.
Resolved
|
||
| staging_verb = _staging_verb(stripped_sql) | ||
| stmt.set_sql(operation) | ||
| allowed_paths = ( | ||
| cursor.connection.staging_allowed_local_path | ||
| if staging_verb in ("PUT", "GET") | ||
| else None | ||
| ) | ||
| if allowed_paths is not None and not isinstance( | ||
| allowed_paths, (str, list) | ||
| ): | ||
| raise ProgrammingError(_STAGING_PATHS_TYPE_ERROR) | ||
| if isinstance(allowed_paths, str): | ||
| allowed_paths = [allowed_paths] | ||
| try: | ||
| normalized_paths = [os.fspath(path) for path in allowed_paths or []] | ||
| if any(not isinstance(path, str) for path in normalized_paths): | ||
| raise TypeError("Local paths must resolve to strings") | ||
| except (TypeError, ValueError) as exc: | ||
| raise ProgrammingError(_STAGING_PATHS_TYPE_ERROR) from exc | ||
| stmt.set_staging_allowed_local_paths( | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🟠 High — The PR description mentions the needed floor bump, but nothing in the diff enforces it, so merging early (or a partial cherry-pick) breaks things for users. Either bump the floor in this PR to the kernel release that ships #439, or gate the call: e.g. only call it when
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Agreed: the native dependency floor is a merge and release gate. Published 1.4.0 is incompatible with this change, including ordinary queries. This PR remains draft and must not merge until the first published kernel release containing the setter is assigned, the dependency and lockfiles are updated, and installed-artifact validation passes. The prerequisite is tracked in databricks/databricks-sql-kernel#443. The intended design is a coordinated native minimum. A second compatibility path and old-wheel tests would add code for an unsupported driver/native pairing. No version is invented and no clean-review claim is made while this gate remains open. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Thanks for confirming. Dropping the compatibility shim is fine if the native floor moves in lockstep, and you don't need old-wheel tests for a pairing you won't support. The current code still makes the call unconditionally, though, and |
||
| normalized_paths, input_stream_token="__input_stream__" | ||
| ) | ||
| stmt.set_row_limit(row_limit) | ||
| if query_tags: | ||
| # Per-statement query tags. The kernel serialises the | ||
|
|
||
Uh oh!
There was an error while loading. Please reload this page.