From 52291878d31cc4e1c5b94bd0d601c6cfa7657ed6 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 15 Oct 2026 19:19:45 +0000 Subject: [PATCH 1/3] feat: add opt-in metadata catalog name escaping Signed-off-by: Vu Anh Phung --- CONNECTION_PARAMETERS.md | 1 + src/databricks/sql/backend/kernel/client.py | 6 +- src/databricks/sql/backend/thrift_backend.py | 14 ++- src/databricks/sql/client.py | 5 ++ src/databricks/sql/session.py | 1 + src/databricks/sql/utils.py | 9 ++ tests/unit/test_kernel_client.py | 85 ++++++++++++++++++ tests/unit/test_thrift_backend.py | 91 ++++++++++++++++++++ 8 files changed, 208 insertions(+), 4 deletions(-) diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md index eaf69c57b..c02789af1 100644 --- a/CONNECTION_PARAMETERS.md +++ b/CONNECTION_PARAMETERS.md @@ -166,6 +166,7 @@ to change without notice. | `schema` | `str` | ✅ | ✅ | `None` | Initial schema for the session (DBR 9.0+). | | `query_tags` | `Dict[str, Optional[str]]` | ✅ | ✅ | `None` | Connection-level key/value tags serialized into the reserved `QUERY_TAGS` session configuration on both backends. Cursor-level `query_tags` are applied per statement separately. | | `enable_metric_view_metadata` | `bool` | ✅ | ✅ | `False` | Sets `spark.sql.thriftserver.metadata.metricview.enabled` via session config so metric-view metadata surfaces. | +| `escape_catalog_name` | `bool` | ✅ | ✅ | `False` | Escape unescaped `_` in metadata catalog names, preserving existing `\_` escapes and `%` wildcards. Applies to `schemas()`, `tables()`, and `columns()` on Thrift and `tables()` on Kernel; Kernel schemas/columns already use exact catalog identifiers. | | `use_inline_params` | `bool` \| `"silent"` | ✅ | ✅ | `False` | Render parameters inline (legacy) vs. native bound params (DBR 14.1+). Shared cursor code performs the rendering before dispatch, so both backends honor it identically. | | `ignore_transactions` | `bool` | ✅ | ✅ | `True` | When `True`: `commit()` is a no-op, `rollback()` raises `NotSupportedError`, and setting `autocommit` is a no-op. | | `fetch_autocommit_from_server`| `bool` | ✅ | ✅ | `False` | Query the server (`SET AUTOCOMMIT`) for autocommit state instead of returning the cached value. | diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index ab1ef0dd7..b9aa33faf 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -49,6 +49,7 @@ ProgrammingError, ) from databricks.sql.telemetry.telemetry_client import TelemetryHelper +from databricks.sql.utils import escape_metadata_catalog_name if TYPE_CHECKING: from databricks.sql.client import Cursor @@ -267,6 +268,7 @@ def __init__( self._telemetry_options = kwargs.get("telemetry_options") or {} self._catalog = catalog self._schema = schema + self._escape_catalog_name = kwargs.get("escape_catalog_name", False) # ``_use_arrow_native_complex_types`` is the connector-side # toggle for whether complex columns (ARRAY / MAP / STRUCT) # are surfaced as native Arrow shapes or as compact JSON @@ -1019,7 +1021,9 @@ def get_tables( # do the work — no connector-side drain + refilter. Passing it # through preserves streaming for large schemas. stream = self._kernel_session.metadata().list_tables( - catalog=_catalog_or_none(catalog_name), + catalog=escape_metadata_catalog_name( + _catalog_or_none(catalog_name), self._escape_catalog_name + ), schema_pattern=schema_name, table_pattern=table_name, table_types=table_types if table_types else None, diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py index 8c6cb69e7..c2ec4a8c4 100644 --- a/src/databricks/sql/backend/thrift_backend.py +++ b/src/databricks/sql/backend/thrift_backend.py @@ -52,6 +52,7 @@ convert_arrow_based_set_to_arrow_table, convert_decimals_in_arrow_table, convert_column_based_set_to_arrow_table, + escape_metadata_catalog_name, serialize_query_tags, ) from databricks.sql.types import SSLOptions @@ -177,6 +178,7 @@ def __init__( raise ValueError("No valid connection settings.") self._host = server_hostname + self._escape_catalog_name = kwargs.get("escape_catalog_name", False) self._initialize_retry_args(kwargs) self._use_arrow_native_complex_types = kwargs.get( "_use_arrow_native_complex_types", True @@ -1193,7 +1195,9 @@ def get_schemas( getDirectResults=ttypes.TSparkGetDirectResults( maxRows=max_rows, maxBytes=max_bytes ), - catalogName=catalog_name, + catalogName=escape_metadata_catalog_name( + catalog_name, self._escape_catalog_name + ), schemaName=schema_name, ) resp = self.make_request(self._client.GetSchemas, req) @@ -1239,7 +1243,9 @@ def get_tables( getDirectResults=ttypes.TSparkGetDirectResults( maxRows=max_rows, maxBytes=max_bytes ), - catalogName=catalog_name, + catalogName=escape_metadata_catalog_name( + catalog_name, self._escape_catalog_name + ), schemaName=schema_name, tableName=table_name, tableTypes=table_types, @@ -1287,7 +1293,9 @@ def get_columns( getDirectResults=ttypes.TSparkGetDirectResults( maxRows=max_rows, maxBytes=max_bytes ), - catalogName=catalog_name, + catalogName=escape_metadata_catalog_name( + catalog_name, self._escape_catalog_name + ), schemaName=schema_name, tableName=table_name, columnName=column_name, diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 3e06aa39b..60d044acc 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -315,6 +315,11 @@ def read(self) -> Optional[OAuthToken]: This allows 1. cursor.tables() to return METRIC_VIEW table type 2. cursor.columns() to return "measure" column type + :param escape_catalog_name: `bool`, optional (default is False) + Escape unescaped underscores in catalog names for schemas(), tables(), + and columns() metadata requests. Existing escapes and % wildcards are + preserved. Kernel schemas() and columns() already use exact catalog + identifiers, so their catalog names are passed through unchanged. :param fetch_autocommit_from_server: `bool`, optional (default is False) When True, the connection.autocommit property queries the server for current state using SET AUTOCOMMIT instead of returning cached value. diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index 74d1f4661..a34f733f0 100644 --- a/src/databricks/sql/session.py +++ b/src/databricks/sql/session.py @@ -333,6 +333,7 @@ def _create_backend( request_timeout_secs=kwargs.get("_socket_timeout"), max_connections=kwargs.get("_pool_maxsize") or None, telemetry_options=kernel_telemetry_options, + escape_catalog_name=kwargs.get("escape_catalog_name", False), ) # These reference the lazily-resolved module attributes defined via diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py index 0914bb168..de36e1769 100644 --- a/src/databricks/sql/utils.py +++ b/src/databricks/sql/utils.py @@ -46,6 +46,15 @@ logger = logging.getLogger(__name__) +def escape_metadata_catalog_name( + catalog_name: Optional[str], enabled: bool +) -> Optional[str]: + """Escape unescaped underscores, preserving existing escapes and % wildcards.""" + if not enabled or catalog_name is None: + return catalog_name + return re.sub(r"(? Optional[str]: diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index 3eecb0145..70730bbbc 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -1812,6 +1812,91 @@ def test_sync_execute_leaves_rowcount_default_when_num_modified_rows_none(): # --------------------------------------------------------------------------- +@pytest.mark.parametrize( + "options", + [{}, {"escape_catalog_name": False}, {"escape_catalog_name": True}], + ids=["default", "disabled", "enabled"], +) +@pytest.mark.parametrize( + "method, filters, kernel_filters", + [ + ("schemas", {"schema_name": "schema_%"}, {"schema_pattern": "schema_%"}), + ( + "tables", + { + "schema_name": "schema_%", + "table_name": "table_%", + "table_types": ["TABLE"], + }, + { + "schema_pattern": "schema_%", + "table_pattern": "table_%", + "table_types": ["TABLE"], + }, + ), + ( + "columns", + { + "schema_name": "schema_%", + "table_name": "table_%", + "column_name": "column_%", + }, + { + "schema_pattern": "schema_%", + "table_pattern": "table_%", + "column_pattern": "column_%", + }, + ), + ], +) +@pytest.mark.parametrize( + "catalog_name, escaped", + [ + (None, None), + ("", ""), + ("%", "%"), + ("*", "*"), + ("simple", "simple"), + ("my_catalog", r"my\_catalog"), + ("a_b_c", r"a\_b\_c"), + (r"my\_catalog", r"my\_catalog"), + (r"a\_b_c%", r"a\_b\_c%"), + ], +) +def test_metadata_catalog_escaping( + monkeypatch, options, method, filters, kernel_filters, catalog_name, escaped +): + from databricks import sql + + kernel_session = MagicMock() + kernel_session.session_id = "session-id" + rpc = getattr(kernel_session.metadata.return_value, f"list_{method}") + rpc.return_value = _stream_with_schema() + monkeypatch.setattr( + kernel_client._kernel, "Session", lambda **kwargs: kernel_session + ) + + with sql.connect( + server_hostname="example.cloud.databricks.com", + http_path="/sql/1.0/warehouses/abc", + access_token="test-token", + use_kernel=True, + enable_telemetry=False, + **options, + ) as connection: + cursor = connection.cursor() + assert getattr(cursor, method)(catalog_name=catalog_name, **filters) is cursor + + expected = ( + escaped + if method == "tables" and options.get("escape_catalog_name", False) + else catalog_name + ) + if expected in ("%", "*"): + expected = None + rpc.assert_called_once_with(catalog=expected, **kernel_filters) + + @pytest.mark.parametrize("catalog_wildcard", ["%", "*"]) def test_get_columns_normalizes_all_catalog_wildcard(catalog_wildcard): c = _make_client() diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py index 17a86464d..075cf9502 100644 --- a/tests/unit/test_thrift_backend.py +++ b/tests/unit/test_thrift_backend.py @@ -32,6 +32,97 @@ def retry_policy_factory(): } +@pytest.mark.parametrize( + "options", + [{}, {"escape_catalog_name": False}, {"escape_catalog_name": True}], + ids=["default", "disabled", "enabled"], +) +@pytest.mark.parametrize( + "method, rpc_name, filters", + [ + ("schemas", "GetSchemas", {"schema_name": "schema_%"}), + ( + "tables", + "GetTables", + { + "schema_name": "schema_%", + "table_name": "table_%", + "table_types": ["TABLE"], + }, + ), + ( + "columns", + "GetColumns", + { + "schema_name": "schema_%", + "table_name": "table_%", + "column_name": "column_%", + }, + ), + ], +) +@pytest.mark.parametrize( + "catalog_name, escaped", + [ + (None, None), + ("", ""), + ("%", "%"), + ("*", "*"), + ("simple", "simple"), + ("my_catalog", r"my\_catalog"), + ("a_b_c", r"a\_b\_c"), + (r"my\_catalog", r"my\_catalog"), + (r"a\_b_c%", r"a\_b\_c%"), + ], +) +def test_metadata_catalog_escaping( + options, method, rpc_name, filters, catalog_name, escaped +): + with ( + patch( + "databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True + ) as client_class, + patch("databricks.sql.result_set.ThriftResultSet"), + patch.object( + ThriftDatabricksClient, + "_handle_execute_response", + return_value=(Mock(), False), + ), + ): + rpc_client = client_class.return_value + rpc_client.OpenSession.return_value = ThriftBackendTestSuite.open_session_resp + rpc = getattr(rpc_client, rpc_name) + rpc.return_value = Mock( + status=ThriftBackendTestSuite.okay_status, directResults=None + ) + + with databricks.sql.connect( + server_hostname="example.cloud.databricks.com", + http_path="/sql/1.0/warehouses/abc", + access_token="test-token", + enable_telemetry=False, + **options, + ) as connection: + cursor = connection.cursor() + assert ( + getattr(cursor, method)(catalog_name=catalog_name, **filters) is cursor + ) + + request = rpc.call_args.args[0] + expected = ( + escaped if options.get("escape_catalog_name", False) else catalog_name + ) + assert request.catalogName == expected + for arg, value in filters.items(): + thrift_field = { + "schema_name": "schemaName", + "table_name": "tableName", + "column_name": "columnName", + "table_types": "tableTypes", + }[arg] + assert getattr(request, thrift_field) == value + + @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed") class ThriftBackendTestSuite(unittest.TestCase): okay_status = ttypes.TStatus(statusCode=ttypes.TStatusCode.SUCCESS_STATUS) From 2463ec8433a9e9de1df80c7d4a62cae7e2523511 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 15 Oct 2026 19:22:23 +0000 Subject: [PATCH 2/3] docs: note metadata catalog escaping configuration in changelog Signed-off-by: Vu Anh Phung --- CHANGELOG.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index bce9e16e4..3a5e33347 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,8 @@ # Release History +# Unreleased +- Add the `escape_catalog_name` connection parameter (default `False`) to opt in to escaping unescaped `_` in metadata catalog names on Thrift and kernel connections. Existing escapes and `%` wildcards are preserved. + # 4.6.0 (2026-09-24) - Upgrade Databricks SQL Kernel to 1.1.0; the kernel dependency is now stable and no longer experimental. - Transparently auto-recover Thrift connections to Reyden / Real-Time warehouses: when a warehouse rejects the default Thrift protocol (SQLSTATE `KP001`), the session is re-opened on the kernel backend and the warehouse is remembered so later connections skip Thrift. Applies only when no backend was chosen explicitly. From 78790895e5c5dce6131fbdcc1122a6750eb7668e Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 15 Oct 2026 19:39:02 +0000 Subject: [PATCH 3/3] test: document JDBC-compatible metadata backslash escaping Signed-off-by: Vu Anh Phung --- src/databricks/sql/utils.py | 2 +- tests/unit/test_kernel_client.py | 11 +++++++++++ tests/unit/test_thrift_backend.py | 11 +++++++++++ 3 files changed, 23 insertions(+), 1 deletion(-) diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py index de36e1769..dd3525547 100644 --- a/src/databricks/sql/utils.py +++ b/src/databricks/sql/utils.py @@ -49,7 +49,7 @@ def escape_metadata_catalog_name( catalog_name: Optional[str], enabled: bool ) -> Optional[str]: - """Escape unescaped underscores, preserving existing escapes and % wildcards.""" + r"""Match JDBC: escape _ unless immediately preceded by \; preserve %.""" if not enabled or catalog_name is None: return catalog_name return re.sub(r"(?