Jul 3, 2025 · Jun 10, 2025 · Jun 11, 2025 · Jun 11, 2025 · Jun 12, 2025 · Jun 12, 2025
diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py
    DriverConnectionParameters,
    HostDetails,
 )

 from databricks.sql.telemetry.latency_logger import log_latency
 from databricks.sql.telemetry.models.enums import StatementType

 logger = logging.getLogger(__name__)

                session_id_hex=self.connection.get_session_id_hex(),
            )

    @log_latency(StatementType.SQL)
    def _handle_staging_put(
        self, presigned_url: str, local_file: str, headers: Optional[dict] = None
    ):
                + "but not yet applied on the server. It's possible this command may fail later."
            )

    @log_latency(StatementType.SQL)
    def _handle_staging_get(
        self, local_file: str, presigned_url: str, headers: Optional[dict] = None
    ):
        with open(local_file, "wb") as fp:
            fp.write(r.content)

    @log_latency(StatementType.SQL)
    def _handle_staging_remove(
        self, presigned_url: str, headers: Optional[dict] = None
    ):
                session_id_hex=self.connection.get_session_id_hex(),
            )

    @log_latency(StatementType.QUERY)
    def execute(
        self,
        operation: str,

        return self

    @log_latency(StatementType.QUERY)
    def execute_async(
        self,
        operation: str,
            self.execute(operation, parameters)
        return self

    @log_latency(StatementType.METADATA)
    def catalogs(self) -> "Cursor":
        """
        Get all available catalogs.
        )
        return self

    @log_latency(StatementType.METADATA)
    def schemas(
        self, catalog_name: Optional[str] = None, schema_name: Optional[str] = None
    ) -> "Cursor":
        )
        return self

    @log_latency(StatementType.METADATA)
    def tables(
        self,
        catalog_name: Optional[str] = None,
        )
        return self

    @log_latency(StatementType.METADATA)
    def columns(
        self,
        catalog_name: Optional[str] = None,
        self.results = results
        self.has_more_rows = has_more_rows

    @log_latency()
    def _convert_columnar_table(self, table):
        column_names = [c[0] for c in self.description]
        ResultRow = Row(*column_names)

        return result

    @log_latency()
    def _convert_arrow_table(self, table):
        column_names = [c[0] for c in self.description]
        ResultRow = Row(*column_names)
    def rownumber(self):
        return self._next_row_index

    @log_latency()
    def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
        """
        Fetch the next set of rows of a query result, returning a PyArrow table.
        ]
        return ColumnTable(merged_result, result1.column_names)

    @log_latency()
    def fetchmany_columnar(self, size: int):
        """
        Fetch the next set of rows of a query result, returning a Columnar Table.

        return results

    @log_latency()
    def fetchall_arrow(self) -> "pyarrow.Table":
        """Fetch all (remaining) rows of a query result, returning them as a PyArrow table."""
        results = self.results.remaining_rows()
            return pyarrow.Table.from_pydict(data)
        return results

    @log_latency()
    def fetchall_columnar(self):
        """Fetch all (remaining) rows of a query result, returning them as a Columnar table."""
        results = self.results.remaining_rows()
diff --git a/src/databricks/sql/telemetry/latency_logger.py b/src/databricks/sql/telemetry/latency_logger.py
 import time
 import functools
 from typing import Optional
 from databricks.sql.telemetry.telemetry_client import TelemetryClientFactory
 from databricks.sql.telemetry.models.event import (
    SqlExecutionEvent,
 )
 from databricks.sql.telemetry.models.enums import ExecutionResultFormat, StatementType
 from databricks.sql.utils import ColumnQueue, CloudFetchQueue, ArrowQueue
 from uuid import UUID


 class TelemetryExtractor:
    """
    Base class for extracting telemetry information from various object types.

    This class serves as a proxy that delegates attribute access to the wrapped object
    while providing a common interface for extracting telemetry-related data.
    """

    def __init__(self, obj):
        """
        Initialize the extractor with an object to wrap.

        Args:
            obj: The object to extract telemetry information from.
        """
        self._obj = obj

    def __getattr__(self, name):
        """
        Delegate attribute access to the wrapped object.

        Args:
            name (str): The name of the attribute to access.

        Returns:
            The attribute value from the wrapped object.
        """
        return getattr(self._obj, name)

    def get_session_id_hex(self):
        pass

    def get_statement_id(self):
        pass

    def get_is_compressed(self):
        pass

    def get_execution_result(self):
        pass

    def get_retry_count(self):
        pass


 class CursorExtractor(TelemetryExtractor):
    """
    Telemetry extractor specialized for Cursor objects.

    Extracts telemetry information from database cursor objects, including
    statement IDs, session information, compression settings, and result formats.
    """

    def get_statement_id(self) -> Optional[str]:
        return self.query_id

    def get_session_id_hex(self) -> Optional[str]:
        return self.connection.get_session_id_hex()

    def get_is_compressed(self) -> bool:
        return self.connection.lz4_compression

    def get_execution_result(self) -> ExecutionResultFormat:
        if self.active_result_set is None:
            return ExecutionResultFormat.FORMAT_UNSPECIFIED

        if isinstance(self.active_result_set.results, ColumnQueue):
            return ExecutionResultFormat.COLUMNAR_INLINE
        elif isinstance(self.active_result_set.results, CloudFetchQueue):
            return ExecutionResultFormat.EXTERNAL_LINKS
        elif isinstance(self.active_result_set.results, ArrowQueue):
            return ExecutionResultFormat.INLINE_ARROW
        return ExecutionResultFormat.FORMAT_UNSPECIFIED

    def get_retry_count(self) -> int:
        if (
            hasattr(self.thrift_backend, "retry_policy")
            and self.thrift_backend.retry_policy
        ):
            return len(self.thrift_backend.retry_policy.history)
        return 0


 class ResultSetExtractor(TelemetryExtractor):
    """
    Telemetry extractor specialized for ResultSet objects.

    Extracts telemetry information from database result set objects, including
    operation IDs, session information, compression settings, and result formats.
    """

    def get_statement_id(self) -> Optional[str]:
        if self.command_id:
            return str(UUID(bytes=self.command_id.operationId.guid))
        return None

    def get_session_id_hex(self) -> Optional[str]:
        return self.connection.get_session_id_hex()

    def get_is_compressed(self) -> bool:
        return self.lz4_compressed

    def get_execution_result(self) -> ExecutionResultFormat:
        if isinstance(self.results, ColumnQueue):
            return ExecutionResultFormat.COLUMNAR_INLINE
        elif isinstance(self.results, CloudFetchQueue):
            return ExecutionResultFormat.EXTERNAL_LINKS
        elif isinstance(self.results, ArrowQueue):
            return ExecutionResultFormat.INLINE_ARROW
        return ExecutionResultFormat.FORMAT_UNSPECIFIED

    def get_retry_count(self) -> int:
        if (
            hasattr(self.thrift_backend, "retry_policy")
            and self.thrift_backend.retry_policy
        ):
            return len(self.thrift_backend.retry_policy.history)
        return 0


 def get_extractor(obj):
    """
    Factory function to create the appropriate telemetry extractor for an object.

    Determines the object type and returns the corresponding specialized extractor
    that can extract telemetry information from that object type.

    Args:
        obj: The object to create an extractor for. Can be a Cursor, ResultSet,
             or any other object.

    Returns:
        TelemetryExtractor: A specialized extractor instance:
            - CursorExtractor for Cursor objects
            - ResultSetExtractor for ResultSet objects
            - Throws an NotImplementedError for all other objects
    """
    if obj.__class__.__name__ == "Cursor":
        return CursorExtractor(obj)
    elif obj.__class__.__name__ == "ResultSet":
        return ResultSetExtractor(obj)
    else:
        raise NotImplementedError(f"No extractor found for {obj.__class__.__name__}")


 def log_latency(statement_type: StatementType = StatementType.NONE):
    """
    Decorator for logging execution latency and telemetry information.

    This decorator measures the execution time of a method and sends telemetry
    data about the operation, including latency, statement information, and
    execution context.

    The decorator automatically:
    - Measures execution time using high-precision performance counters
    - Extracts telemetry information from the method's object (self)
    - Creates a SqlExecutionEvent with execution details
    - Sends the telemetry data asynchronously via TelemetryClient

    Args:
        statement_type (StatementType): The type of SQL statement being executed.

    Usage:
        @log_latency(StatementType.SQL)
        def execute(self, query):
            # Method implementation
            pass

    Returns:
        function: A decorator that wraps methods to add latency logging.

    Note:
        The wrapped method's object (self) must be compatible with the
        telemetry extractor system (e.g., Cursor or ResultSet objects).
    """

    def decorator(func):
        @functools.wraps(func)
        def wrapper(self, *args, **kwargs):
            start_time = time.perf_counter()
            result = None
            try:
                result = func(self, *args, **kwargs)
                return result
            finally:

                def _safe_call(func_to_call):
                    """Calls a function and returns a default value on any exception."""
                    try:
                        return func_to_call()
                    except Exception:
                        return None

                end_time = time.perf_counter()
                duration_ms = int((end_time - start_time) * 1000)

                extractor = get_extractor(self)
                session_id_hex = _safe_call(extractor.get_session_id_hex)
                statement_id = _safe_call(extractor.get_statement_id)

                sql_exec_event = SqlExecutionEvent(
                    statement_type=statement_type,
                    is_compressed=_safe_call(extractor.get_is_compressed),
                    execution_result=_safe_call(extractor.get_execution_result),
                    retry_count=_safe_call(extractor.get_retry_count),
                )

                telemetry_client = TelemetryClientFactory.get_telemetry_client(
                    session_id_hex
                )
                telemetry_client.export_latency_log(
                    latency_ms=duration_ms,
                    sql_execution_event=sql_exec_event,
                    sql_statement_id=statement_id,
                )

        return wrapper

    return decorator
Original file line number	Diff line number	Diff line change
Expand Up		@@ -61,7 +61,8 @@
		DriverConnectionParameters,
		HostDetails,
		)

		from databricks.sql.telemetry.latency_logger import log_latency
		from databricks.sql.telemetry.models.enums import StatementType

		logger = logging.getLogger(__name__)

Expand DownExpand Up		@@ -745,6 +746,7 @@ def _handle_staging_operation(
		session_id_hex=self.connection.get_session_id_hex(),
		)

		@log_latency(StatementType.SQL)
		def _handle_staging_put(
		self, presigned_url: str, local_file: str, headers: Optional[dict] = None
		):
Expand DownExpand Up		@@ -784,6 +786,7 @@ def _handle_staging_put(
		+ "but not yet applied on the server. It's possible this command may fail later."
		)

		@log_latency(StatementType.SQL)
		def _handle_staging_get(
		self, local_file: str, presigned_url: str, headers: Optional[dict] = None
		):
Expand DownExpand Up		@@ -811,6 +814,7 @@ def _handle_staging_get(
		with open(local_file, "wb") as fp:
		fp.write(r.content)

		@log_latency(StatementType.SQL)
		def _handle_staging_remove(
		self, presigned_url: str, headers: Optional[dict] = None
		):
Expand All		@@ -824,6 +828,7 @@ def _handle_staging_remove(
		session_id_hex=self.connection.get_session_id_hex(),
		)

		@log_latency(StatementType.QUERY)
		def execute(
		self,
		operation: str,
Expand DownExpand Up		@@ -914,6 +919,7 @@ def execute(

		return self

		@log_latency(StatementType.QUERY)
		def execute_async(
		self,
		operation: str,
Expand DownExpand Up		@@ -1039,6 +1045,7 @@ def executemany(self, operation, seq_of_parameters):
		self.execute(operation, parameters)
		return self

		@log_latency(StatementType.METADATA)
		def catalogs(self) -> "Cursor":
		"""
		Get all available catalogs.
Expand All		@@ -1062,6 +1069,7 @@ def catalogs(self) -> "Cursor":
		)
		return self

		@log_latency(StatementType.METADATA)
		def schemas(
		self, catalog_name: Optional[str] = None, schema_name: Optional[str] = None
		) -> "Cursor":
Expand DownExpand Up		@@ -1090,6 +1098,7 @@ def schemas(
		)
		return self

		@log_latency(StatementType.METADATA)
		def tables(
		self,
		catalog_name: Optional[str] = None,
Expand DownExpand Up		@@ -1125,6 +1134,7 @@ def tables(
		)
		return self

		@log_latency(StatementType.METADATA)
		def columns(
		self,
		catalog_name: Optional[str] = None,
Expand DownExpand Up		@@ -1379,6 +1389,7 @@ def _fill_results_buffer(self):
		self.results = results
		self.has_more_rows = has_more_rows

		@log_latency()
		def _convert_columnar_table(self, table):
		column_names = [c[0] for c in self.description]
		ResultRow = Row(*column_names)
Expand All		@@ -1391,6 +1402,7 @@ def _convert_columnar_table(self, table):

		return result

		@log_latency()
		def _convert_arrow_table(self, table):
		column_names = [c[0] for c in self.description]
		ResultRow = Row(*column_names)
Expand DownExpand Up		@@ -1433,6 +1445,7 @@ def _convert_arrow_table(self, table):
		def rownumber(self):
		return self._next_row_index

		@log_latency()
		def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
		"""
		Fetch the next set of rows of a query result, returning a PyArrow table.
Expand DownExpand Up		@@ -1475,6 +1488,7 @@ def merge_columnar(self, result1, result2):
		]
		return ColumnTable(merged_result, result1.column_names)

		@log_latency()
		def fetchmany_columnar(self, size: int):
		"""
		Fetch the next set of rows of a query result, returning a Columnar Table.
Expand All		@@ -1500,6 +1514,7 @@ def fetchmany_columnar(self, size: int):

		return results

		@log_latency()
		def fetchall_arrow(self) -> "pyarrow.Table":
		"""Fetch all (remaining) rows of a query result, returning them as a PyArrow table."""
		results = self.results.remaining_rows()
Expand All		@@ -1526,6 +1541,7 @@ def fetchall_arrow(self) -> "pyarrow.Table":
		return pyarrow.Table.from_pydict(data)
		return results

		@log_latency()
		def fetchall_columnar(self):
		"""Fetch all (remaining) rows of a query result, returning them as a Columnar table."""
		results = self.results.remaining_rows()
Expand Down
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,231 @@
		import time
		import functools
		from typing import Optional
		from databricks.sql.telemetry.telemetry_client import TelemetryClientFactory
		from databricks.sql.telemetry.models.event import (
		SqlExecutionEvent,
		)
		from databricks.sql.telemetry.models.enums import ExecutionResultFormat, StatementType
		from databricks.sql.utils import ColumnQueue, CloudFetchQueue, ArrowQueue
		from uuid import UUID


		class TelemetryExtractor:
saishreeeee marked this conversation as resolved. Show resolvedHide resolved
		"""
		Base class for extracting telemetry information from various object types.

		This class serves as a proxy that delegates attribute access to the wrapped object
		while providing a common interface for extracting telemetry-related data.
		"""

		def __init__(self, obj):
		"""
		Initialize the extractor with an object to wrap.

		Args:
		obj: The object to extract telemetry information from.
		"""
		self._obj = obj

		def __getattr__(self, name):
		"""
		Delegate attribute access to the wrapped object.

		Args:
		name (str): The name of the attribute to access.

		Returns:
		The attribute value from the wrapped object.
		"""
		return getattr(self._obj, name)

		def get_session_id_hex(self):
		pass

		def get_statement_id(self):
		pass

		def get_is_compressed(self):
		pass

		def get_execution_result(self):
		pass

		def get_retry_count(self):
		pass


		class CursorExtractor(TelemetryExtractor):
		"""
		Telemetry extractor specialized for Cursor objects.

		Extracts telemetry information from database cursor objects, including
		statement IDs, session information, compression settings, and result formats.
		"""

		def get_statement_id(self) -> Optional[str]:
		return self.query_id

		def get_session_id_hex(self) -> Optional[str]:
		return self.connection.get_session_id_hex()

		def get_is_compressed(self) -> bool:
		return self.connection.lz4_compression

		def get_execution_result(self) -> ExecutionResultFormat:
		if self.active_result_set is None:
		return ExecutionResultFormat.FORMAT_UNSPECIFIED

saishreeeee marked this conversation as resolved. Show resolvedHide resolved
		if isinstance(self.active_result_set.results, ColumnQueue):
		return ExecutionResultFormat.COLUMNAR_INLINE
		elif isinstance(self.active_result_set.results, CloudFetchQueue):
		return ExecutionResultFormat.EXTERNAL_LINKS
		elif isinstance(self.active_result_set.results, ArrowQueue):
		return ExecutionResultFormat.INLINE_ARROW
		return ExecutionResultFormat.FORMAT_UNSPECIFIED

		def get_retry_count(self) -> int:
		if (
		hasattr(self.thrift_backend, "retry_policy")
		and self.thrift_backend.retry_policy
		):
		return len(self.thrift_backend.retry_policy.history)
		return 0


		class ResultSetExtractor(TelemetryExtractor):
		"""
		Telemetry extractor specialized for ResultSet objects.

		Extracts telemetry information from database result set objects, including
		operation IDs, session information, compression settings, and result formats.
		"""

		def get_statement_id(self) -> Optional[str]:
		if self.command_id:
		return str(UUID(bytes=self.command_id.operationId.guid))
		return None

		def get_session_id_hex(self) -> Optional[str]:
		return self.connection.get_session_id_hex()

		def get_is_compressed(self) -> bool:
		return self.lz4_compressed

		def get_execution_result(self) -> ExecutionResultFormat:
		if isinstance(self.results, ColumnQueue):
		return ExecutionResultFormat.COLUMNAR_INLINE
		elif isinstance(self.results, CloudFetchQueue):
		return ExecutionResultFormat.EXTERNAL_LINKS
		elif isinstance(self.results, ArrowQueue):
		return ExecutionResultFormat.INLINE_ARROW
		return ExecutionResultFormat.FORMAT_UNSPECIFIED
saishreeeee marked this conversation as resolved. Show resolvedHide resolved

		def get_retry_count(self) -> int:
		if (
		hasattr(self.thrift_backend, "retry_policy")
		and self.thrift_backend.retry_policy
		):
		return len(self.thrift_backend.retry_policy.history)
		return 0


		def get_extractor(obj):
		"""
		Factory function to create the appropriate telemetry extractor for an object.

		Determines the object type and returns the corresponding specialized extractor
		that can extract telemetry information from that object type.

		Args:
		obj: The object to create an extractor for. Can be a Cursor, ResultSet,
		or any other object.

		Returns:
		TelemetryExtractor: A specialized extractor instance:
		- CursorExtractor for Cursor objects
		- ResultSetExtractor for ResultSet objects
		- Throws an NotImplementedError for all other objects
		"""
		if obj.__class__.__name__ == "Cursor":
saishreeeee marked this conversation as resolved. Show resolvedHide resolved
		return CursorExtractor(obj)
		elif obj.__class__.__name__ == "ResultSet":
		return ResultSetExtractor(obj)
		else:
		raise NotImplementedError(f"No extractor found for {obj.__class__.__name__}")


		def log_latency(statement_type: StatementType = StatementType.NONE):
		"""
		Decorator for logging execution latency and telemetry information.

		This decorator measures the execution time of a method and sends telemetry
		data about the operation, including latency, statement information, and
		execution context.

		The decorator automatically:
		- Measures execution time using high-precision performance counters
		- Extracts telemetry information from the method's object (self)
		- Creates a SqlExecutionEvent with execution details
		- Sends the telemetry data asynchronously via TelemetryClient

		Args:
		statement_type (StatementType): The type of SQL statement being executed.

		Usage:
		@log_latency(StatementType.SQL)
		def execute(self, query):
		# Method implementation
		pass

		Returns:
		function: A decorator that wraps methods to add latency logging.

		Note:
		The wrapped method's object (self) must be compatible with the
		telemetry extractor system (e.g., Cursor or ResultSet objects).
		"""

		def decorator(func):
		@functools.wraps(func)
		def wrapper(self, args, *kwargs):
saishreeeee marked this conversation as resolved. Show resolvedHide resolved
		start_time = time.perf_counter()
		result = None
		try:
		result = func(self, args, *kwargs)
		return result
		finally:

		def _safe_call(func_to_call):
		"""Calls a function and returns a default value on any exception."""
		try:
		return func_to_call()
		except Exception:
		return None

		end_time = time.perf_counter()
		duration_ms = int((end_time - start_time) * 1000)

		extractor = get_extractor(self)
		session_id_hex = _safe_call(extractor.get_session_id_hex)
		statement_id = _safe_call(extractor.get_statement_id)

		sql_exec_event = SqlExecutionEvent(
		statement_type=statement_type,
saishreeeee marked this conversation as resolved. Show resolvedHide resolved
		is_compressed=_safe_call(extractor.get_is_compressed),
		execution_result=_safe_call(extractor.get_execution_result),
		retry_count=_safe_call(extractor.get_retry_count),
		)

		telemetry_client = TelemetryClientFactory.get_telemetry_client(
		session_id_hex
		)
		telemetry_client.export_latency_log(
		latency_ms=duration_ms,
		sql_execution_event=sql_exec_event,
		sql_statement_id=statement_id,
		)

		return wrapper

		return decorator