feat(backend/kernel): match Thrift backend's user-visible surface

src/databricks/sql/backend/kernel/_errors.py

@@ -43,9 +43,9 @@
     Error,
     OperationalError,
     ProgrammingError,
+    ServerOperationError,
 )
 
-
 try:
     import databricks_sql_kernel as _kernel  # type: ignore[import-not-found]
 except ImportError as exc:  # pragma: no cover - same hint as client.py
@@ -76,7 +76,12 @@
     "Internal": DatabaseError,
     "InvalidStatementHandle": ProgrammingError,
     "NetworkError": OperationalError,
-    "SqlError": DatabaseError,
+    # `SqlError` is a server-side query failure (syntax error, missing
+    # object, etc.) — exactly what the Thrift backend surfaces as
+    # `ServerOperationError`. Match Thrift's contract so user code that
+    # catches `ServerOperationError` (a subclass of `DatabaseError`)
+    # works equivalently with `use_kernel=True`.
+    "SqlError": ServerOperationError,
     "Unknown": DatabaseError,
 }
 

src/databricks/sql/backend/kernel/client.py

@@ -517,14 +517,13 @@ def get_columns(
     ) -> "ResultSet":
         if self._kernel_session is None:
             raise InterfaceError("get_columns requires an open session.")
-        if not catalog_name:
-            # Kernel's list_columns requires a catalog (SEA `SHOW
-            # COLUMNS` cannot span catalogs). Surface the constraint
-            # explicitly rather than letting the kernel error.
-            raise ProgrammingError(
-                "get_columns requires catalog_name on the kernel backend."
-            )
         try:
+            # `catalog_name=None` is supported: the kernel issues
+            # `SHOW COLUMNS IN ALL CATALOGS` server-side and the
+            # response carries `catalogName` per row, so each result
+            # row's `TABLE_CAT` is correctly attributed. Matches the
+            # Thrift backend's `getColumns(null, …)` behaviour from
+            # the user's perspective.
             stream = self._kernel_session.metadata().list_columns(
                 catalog=catalog_name,
                 schema_pattern=schema_name,

src/databricks/sql/backend/kernel/type_mapping.py

@@ -85,24 +85,56 @@ def description_from_arrow_schema(schema: pyarrow.Schema) -> List[Tuple]:
     """Build a PEP 249 ``description`` list from a pyarrow Schema.
 
     Each tuple is ``(name, type_code, display_size, internal_size,
-    precision, scale, null_ok)``. ``null_ok`` is taken from
-    ``field.nullable``; the other four are not reported by the
-    kernel today.
+    precision, scale, null_ok)``. PEP 249 allows ``null_ok`` to be
+    either a bool or ``None``; the Thrift backend always reports
+    ``None``, so we match that here for drop-in parity. The actual
+    nullability bit is still available via ``schema.field(i).nullable``
+    for callers that want it from the Arrow schema directly.
+
+    ``type_code`` normally comes from the Arrow ``DataType`` via
+    ``_arrow_type_to_dbapi_string``, which collapses
+    Databricks-specific types into their nearest Arrow shape (e.g.
+    ``VARIANT`` → ``Utf8``). To recover the precise Databricks type
+    name, we consult the field's metadata first — the kernel writes
+    the server-reported type into ``databricks.type_name`` (see
+    ``databricks_sql_kernel::reader::metadata_keys``). Today only
+    ``VARIANT`` is special-cased here for parity with the Thrift
+    backend's behaviour; other precise types (``INTERVAL_*``,
+    ``GEOMETRY``, ``GEOGRAPHY``) collapse to their Arrow shape on
+    both backends and don't need a remap.
     """
     return [
         (
             field.name,
-            _arrow_type_to_dbapi_string(field.type),
+            _databricks_type_for_field(field),
+            None,
             None,
             None,
             None,
             None,
-            field.nullable,
         )
         for field in schema
     ]
 
 
+def _databricks_type_for_field(field: pyarrow.Field) -> str:
+    """Pick the PEP 249 type code for a single field.
+
+    Consults the field's Arrow metadata under
+    ``databricks.type_name`` (written by the kernel from the SEA
+    response's column type) so types that collapse onto a generic
+    Arrow shape can still be distinguished. Today only ``VARIANT``
+    is mapped; everything else delegates to
+    ``_arrow_type_to_dbapi_string``.
+    """
+    md = field.metadata or {}
+    # `databricks.type_name` is bytes (Arrow metadata is always
+    # bytes); compare against bytes to avoid one encode per field.
+    if md.get(b"databricks.type_name") == b"VARIANT":
+        return "variant"
+    return _arrow_type_to_dbapi_string(field.type)
+
+
 def _tspark_param_value_str(param: ttypes.TSparkParameter) -> Any:
     """Extract the string-encoded value from a ``TSparkParameter``,
     or ``None`` for SQL NULL.

tests/unit/test_kernel_client.py

@@ -70,9 +70,9 @@ def __init__(
     NotSupportedError,
     OperationalError,
     ProgrammingError,
+    ServerOperationError,
 )
 
-
 # ---------------------------------------------------------------------------
 # Error mapping
 # ---------------------------------------------------------------------------
@@ -93,7 +93,14 @@ def __init__(
         ("Internal", DatabaseError),
         ("InvalidStatementHandle", ProgrammingError),
         ("NetworkError", OperationalError),
-        ("SqlError", DatabaseError),
+        # `SqlError` is the kernel's slug for server-side query
+        # failures (syntax error, missing object, etc.) — exactly the
+        # case Thrift's backend surfaces as ``ServerOperationError``.
+        # Match Thrift so user code that catches the specific class
+        # works equivalently. ``ServerOperationError`` is itself a
+        # ``DatabaseError`` subclass, so existing catches of the base
+        # class are unaffected.
+        ("SqlError", ServerOperationError),
         ("Unknown", DatabaseError),
     ],
 )
@@ -309,20 +316,37 @@ def test_execute_command_rejects_query_tags():
         )
 
 
-def test_get_columns_requires_catalog():
+def test_get_columns_accepts_none_catalog():
+    """The kernel's `list_columns` honours `catalog=None` by issuing
+    `SHOW COLUMNS IN ALL CATALOGS` server-side. The connector should
+    pass `None` through rather than rejecting it, matching the Thrift
+    backend's `getColumns(null, …)` behaviour."""
     c = _make_client()
     c._kernel_session = MagicMock()
     cursor = MagicMock()
     cursor.arraysize = 100
     cursor.buffer_size_bytes = 1024
-    with pytest.raises(ProgrammingError, match="catalog_name"):
-        c.get_columns(
-            session_id=MagicMock(),
-            max_rows=1,
-            max_bytes=1,
-            cursor=cursor,
-            catalog_name=None,
-        )
+    cursor.connection = MagicMock()
+    # `list_columns` returns a stream the result-set wrapper will try
+    # to call `arrow_schema()` on; give it a minimal fake.
+    fake_stream = MagicMock()
+    fake_stream.arrow_schema.return_value = MagicMock(__iter__=lambda self: iter([]))
+    c._kernel_session.metadata.return_value.list_columns.return_value = fake_stream
+
+    c.get_columns(
+        session_id=MagicMock(),
+        max_rows=1,
+        max_bytes=1,
+        cursor=cursor,
+        catalog_name=None,
+    )
+    # `list_columns` should be called with catalog=None, not rejected.
+    c._kernel_session.metadata.return_value.list_columns.assert_called_once_with(
+        catalog=None,
+        schema_pattern=None,
+        table_pattern=None,
+        column_pattern=None,
+    )
 
 
 # ---------------------------------------------------------------------------
@@ -521,7 +545,9 @@ def test_pyo3_native_exception_wrapped_as_operational_error():
     stmt = MagicMock()
     stmt.execute.side_effect = TypeError("argument 'foo' must be str, not int")
     c._kernel_session.statement.return_value = stmt
-    with pytest.raises(OperationalError, match="Unexpected error from databricks_sql_kernel"):
+    with pytest.raises(
+        OperationalError, match="Unexpected error from databricks_sql_kernel"
+    ):
         c.execute_command(
             operation="SELECT 1",
             session_id=MagicMock(),
@@ -546,9 +572,7 @@ def test_pyo3_native_exception_wrapped_for_metadata_calls():
     cursor.arraysize = 100
     cursor.buffer_size_bytes = 1024
     with pytest.raises(OperationalError):
-        c.get_catalogs(
-            session_id=MagicMock(), max_rows=1, max_bytes=1, cursor=cursor
-        )
+        c.get_catalogs(session_id=MagicMock(), max_rows=1, max_bytes=1, cursor=cursor)
 
 
 # ---------------------------------------------------------------------------
@@ -574,9 +598,7 @@ def test_kernel_error_during_result_set_construction_is_mapped():
     cursor.arraysize = 100
     cursor.buffer_size_bytes = 1024
     with pytest.raises(DatabaseError, match="schema unavailable"):
-        c.get_catalogs(
-            session_id=MagicMock(), max_rows=1, max_bytes=1, cursor=cursor
-        )
+        c.get_catalogs(session_id=MagicMock(), max_rows=1, max_bytes=1, cursor=cursor)
 
 
 # ---------------------------------------------------------------------------

tests/unit/test_kernel_result_set.py

@@ -77,8 +77,9 @@ def int_schema():
 def test_description_built_from_kernel_schema(int_schema):
     handle = _FakeKernelHandle(int_schema, [])
     rs = _make_rs(handle)
-    # null_ok slot reflects pyarrow.Field.nullable (True by default).
-    assert rs.description == [("n", "bigint", None, None, None, None, True)]
+    # `null_ok` slot is always `None` to match the Thrift backend's
+    # PEP-249 description shape.
+    assert rs.description == [("n", "bigint", None, None, None, None, None)]
 
 
 def test_fetchall_arrow_drains_all_batches(int_schema):
@@ -107,7 +108,11 @@ def test_fetchmany_arrow_slices_within_batch(int_schema):
 def test_fetchmany_arrow_spans_batch_boundary(int_schema):
     handle = _FakeKernelHandle(
         int_schema,
-        [_batch(int_schema, [1, 2]), _batch(int_schema, [3, 4]), _batch(int_schema, [5, 6])],
+        [
+            _batch(int_schema, [1, 2]),
+            _batch(int_schema, [3, 4]),
+            _batch(int_schema, [5, 6]),
+        ],
     )
     rs = _make_rs(handle)
     t = rs.fetchmany_arrow(5)

tests/unit/test_kernel_type_mapping.py

@@ -65,25 +65,53 @@ def test_description_from_schema_preserves_field_names_and_order():
         ("name", "string"),
         ("created_at", "timestamp"),
     ]
-    # PEP 249 says 7-tuples. We don't report display_size /
-    # internal_size / precision / scale (all None); ``null_ok`` is
-    # taken from ``pyarrow.Field.nullable`` — True by default for
-    # schemas built from (name, type) pairs.
+    # PEP 249 7-tuples; we report column name and type only. PEP 249
+    # allows ``null_ok`` to be ``None`` and that's what the Thrift
+    # backend has always returned; match it so kernel-backed cursors
+    # are drop-in compatible. The Arrow ``field.nullable`` bit is still
+    # available via ``schema.field(i).nullable`` for callers that need
+    # the real value.
     for d in desc:
         assert len(d) == 7
-        assert d[2:] == (None, None, None, None, True)
+        assert d[2:] == (None, None, None, None, None)
 
 
-def test_description_from_schema_reports_non_nullable_fields():
+def test_description_null_ok_always_none_regardless_of_field_nullable():
+    # Match Thrift backend's behaviour: ``null_ok`` is ``None`` for
+    # every column even when the Arrow ``field.nullable`` bit is
+    # meaningful.
     schema = pa.schema(
         [
             pa.field("id", pa.int64(), nullable=False),
             pa.field("name", pa.string(), nullable=True),
         ]
     )
     desc = description_from_arrow_schema(schema)
-    assert desc[0][6] is False
-    assert desc[1][6] is True
+    assert desc[0][6] is None
+    assert desc[1][6] is None
+
+
+def test_description_uses_databricks_type_name_for_variant():
+    """VARIANT columns arrive over SEA as Arrow ``Utf8``; the kernel
+    annotates them with ``databricks.type_name=VARIANT`` so the
+    connector can recover the precise type for PEP-249 description.
+    Matches the Thrift backend, which exposes the same column as
+    ``variant``."""
+    schema = pa.schema(
+        [
+            pa.field(
+                "v",
+                pa.string(),
+                metadata={b"databricks.type_name": b"VARIANT"},
+            ),
+            # Plain Utf8 column without the metadata stays ``string``
+            # so we don't claim "variant" for everything.
+            pa.field("s", pa.string()),
+        ]
+    )
+    desc = description_from_arrow_schema(schema)
+    assert desc[0][1] == "variant"
+    assert desc[1][1] == "string"
 
 
 # ─── bind_tspark_params ──────────────────────────────────────────────────
@@ -94,7 +122,9 @@ def _mk_param(*, type, value, ordinal=True, name=None):
     from databricks.sql.thrift_api.TCLIService import ttypes
 
     p = ttypes.TSparkParameter(ordinal=ordinal, name=name, type=type)
-    p.value = ttypes.TSparkParameterValue(stringValue=value) if value is not None else None
+    p.value = (
+        ttypes.TSparkParameterValue(stringValue=value) if value is not None else None
+    )
     return p