From ffc5568022e3a2b8e5ea128da4b6910b5102c64f Mon Sep 17 00:00:00 2001 From: Shivam Raj Date: Tue, 26 Aug 2025 18:14:16 +0530 Subject: [PATCH 1/3] added support for measure columns --- src/databricks/sql/backend/thrift_backend.py | 6 ++++-- tests/unit/test_thrift_backend.py | 11 +++++++++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py index d2b10e718..16efc297e 100644 --- a/src/databricks/sql/backend/thrift_backend.py +++ b/src/databricks/sql/backend/thrift_backend.py @@ -764,14 +764,16 @@ def _col_to_description(col, field=None, session_id_hex=None): else: precision, scale = None, None - # Extract variant type from field if available + # Extract variant/measure type from field if available if field is not None: try: - # Check for variant type in metadata + # Check for variant/measure type in metadata if field.metadata and b"Spark:DataType:SqlName" in field.metadata: sql_type = field.metadata.get(b"Spark:DataType:SqlName") if sql_type == b"VARIANT": cleaned_type = "variant" + if sql_type and b"measure" in sql_type: + cleaned_type += " measure" except Exception as e: logger.debug(f"Could not extract variant type from field: {e}") diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py index 7254b66cb..7cc54b088 100644 --- a/tests/unit/test_thrift_backend.py +++ b/tests/unit/test_thrift_backend.py @@ -2402,6 +2402,17 @@ def test_hive_schema_to_description(self): ], [("regular_col", "string"), ("variant_col", "variant")], ), + ( + [ + ("measure_col", ttypes.TTypeId.DOUBLE_TYPE), + ("int_measure_col", ttypes.TTypeId.INT_TYPE), + ], + [ + ("measure_col", {b"Spark:DataType:SqlName": b"double measure"}), + ("int_measure_col", {b"Spark:DataType:SqlName": b"int measure"}), + ], + [("measure_col", "double measure"), ("int_measure_col", "int measure")], + ), ( [("regular_col", ttypes.TTypeId.STRING_TYPE)], None, # No arrow schema From 5494793d1c25663d4b11e66eae3b8abbcc5f4481 Mon Sep 17 00:00:00 2001 From: Shivam Raj Date: Thu, 25 Sep 2025 03:27:08 +0530 Subject: [PATCH 2/3] add new connection param --- src/databricks/sql/backend/thrift_backend.py | 6 ++---- src/databricks/sql/client.py | 15 +++++++++++++++ tests/unit/test_session.py | 11 +++++++++++ tests/unit/test_thrift_backend.py | 11 ----------- 4 files changed, 28 insertions(+), 15 deletions(-) diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py index 16efc297e..d2b10e718 100644 --- a/src/databricks/sql/backend/thrift_backend.py +++ b/src/databricks/sql/backend/thrift_backend.py @@ -764,16 +764,14 @@ def _col_to_description(col, field=None, session_id_hex=None): else: precision, scale = None, None - # Extract variant/measure type from field if available + # Extract variant type from field if available if field is not None: try: - # Check for variant/measure type in metadata + # Check for variant type in metadata if field.metadata and b"Spark:DataType:SqlName" in field.metadata: sql_type = field.metadata.get(b"Spark:DataType:SqlName") if sql_type == b"VARIANT": cleaned_type = "variant" - if sql_type and b"measure" in sql_type: - cleaned_type += " measure" except Exception as e: logger.debug(f"Could not extract variant type from field: {e}") diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 78a011421..20a396b34 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -200,6 +200,12 @@ def read(self) -> Optional[OAuthToken]: STRUCT is returned as Dict[str, Any] ARRAY is returned as numpy.ndarray When False, complex types are returned as a strings. These are generally deserializable as JSON. + :param enable_metric_view_metadata: `bool`, optional (default is False) + When True, enables metric view metadata support by setting the + spark.sql.thriftserver.metadata.metricview.enabled session configuration. + This allows + 1. cursor.tables() to return METRIC_VIEW table type + 2. cursor.columns() to return "measure" column type """ # Internal arguments in **kwargs: @@ -248,6 +254,15 @@ def read(self) -> Optional[OAuthToken]: access_token_kv = {"access_token": access_token} kwargs = {**kwargs, **access_token_kv} + # Handle enable_metric_view_metadata parameter + enable_metric_view_metadata = kwargs.get("enable_metric_view_metadata", False) + if enable_metric_view_metadata: + if session_configuration is None: + session_configuration = {} + session_configuration[ + "spark.sql.thriftserver.metadata.metricview.enabled" + ] = "true" + self.disable_pandas = kwargs.get("_disable_pandas", False) self.lz4_compression = kwargs.get("enable_query_result_lz4_compression", True) self.use_cloud_fetch = kwargs.get("use_cloud_fetch", True) diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index c135a846b..1d70ec4c4 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -163,6 +163,17 @@ def test_configuration_passthrough(self, mock_client_class): call_kwargs = mock_client_class.return_value.open_session.call_args[1] assert call_kwargs["session_configuration"] == mock_session_config + @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME) + def test_enable_metric_view_metadata_parameter(self, mock_client_class): + """Test that enable_metric_view_metadata parameter sets the correct session configuration.""" + databricks.sql.connect( + enable_metric_view_metadata=True, **self.DUMMY_CONNECTION_ARGS + ) + + call_kwargs = mock_client_class.return_value.open_session.call_args[1] + expected_config = {"spark.sql.thriftserver.metadata.metricview.enabled": "true"} + assert call_kwargs["session_configuration"] == expected_config + @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME) def test_initial_namespace_passthrough(self, mock_client_class): mock_cat = Mock() diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py index 7cc54b088..7254b66cb 100644 --- a/tests/unit/test_thrift_backend.py +++ b/tests/unit/test_thrift_backend.py @@ -2402,17 +2402,6 @@ def test_hive_schema_to_description(self): ], [("regular_col", "string"), ("variant_col", "variant")], ), - ( - [ - ("measure_col", ttypes.TTypeId.DOUBLE_TYPE), - ("int_measure_col", ttypes.TTypeId.INT_TYPE), - ], - [ - ("measure_col", {b"Spark:DataType:SqlName": b"double measure"}), - ("int_measure_col", {b"Spark:DataType:SqlName": b"int measure"}), - ], - [("measure_col", "double measure"), ("int_measure_col", "int measure")], - ), ( [("regular_col", ttypes.TTypeId.STRING_TYPE)], None, # No arrow schema From 3c89ac1e3795a305adaf61e7173c820c4354ab65 Mon Sep 17 00:00:00 2001 From: Shivam Raj Date: Thu, 25 Sep 2025 09:59:18 +0530 Subject: [PATCH 3/3] remove obvious comment --- src/databricks/sql/client.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 20a396b34..5bb191ca2 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -254,7 +254,6 @@ def read(self) -> Optional[OAuthToken]: access_token_kv = {"access_token": access_token} kwargs = {**kwargs, **access_token_kv} - # Handle enable_metric_view_metadata parameter enable_metric_view_metadata = kwargs.get("enable_metric_view_metadata", False) if enable_metric_view_metadata: if session_configuration is None: