changed enums to follow proto, get_extractor returns None if not Cursor/ResultSet, shifted log_latency decorator to fetchall

Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>

Author: saishreeeee

src/databricks/sql/__init__.py

@@ -9,7 +9,7 @@
 paramstyle = "named"
 
 import re
-
+import logging
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:

src/databricks/sql/client.py

@@ -1389,7 +1389,6 @@ 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)
@@ -1402,7 +1401,6 @@ 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)
@@ -1445,7 +1443,6 @@ 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.
@@ -1488,7 +1485,6 @@ 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.
@@ -1514,7 +1510,6 @@ 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()
@@ -1541,7 +1536,6 @@ 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()
@@ -1555,6 +1549,7 @@ def fetchall_columnar(self):
 
         return results
 
+    @log_latency()
     def fetchone(self) -> Optional[Row]:
         """
         Fetch the next row of a query result set, returning a single sequence,
@@ -1571,6 +1566,7 @@ def fetchone(self) -> Optional[Row]:
         else:
             return None
 
+    @log_latency()
     def fetchall(self) -> List[Row]:
         """
         Fetch all (remaining) rows of a query result, returning them as a list of rows.
@@ -1580,6 +1576,7 @@ def fetchall(self) -> List[Row]:
         else:
             return self._convert_arrow_table(self.fetchall_arrow())
 
+    @log_latency()
     def fetchmany(self, size: int) -> List[Row]:
         """
         Fetch the next set of rows of a query result, returning a list of rows.

src/databricks/sql/exc.py

@@ -1,9 +1,6 @@
 import json
-import logging
-
 from databricks.sql.telemetry.telemetry_client import TelemetryClientFactory
 
-logger = logging.getLogger(__name__)
 
 ### PEP-249 Mandated ###
 # https://peps.python.org/pep-0249/#exceptions

src/databricks/sql/telemetry/latency_logger.py

@@ -1,6 +1,7 @@
 import time
 import functools
 from typing import Optional
+import logging
 from databricks.sql.telemetry.telemetry_client import TelemetryClientFactory
 from databricks.sql.telemetry.models.event import (
     SqlExecutionEvent,
@@ -9,6 +10,8 @@
 from databricks.sql.utils import ColumnQueue, CloudFetchQueue, ArrowQueue
 from uuid import UUID
 
+logger = logging.getLogger(__name__)
+
 
 class TelemetryExtractor:
     """
@@ -145,14 +148,15 @@ def get_extractor(obj):
         TelemetryExtractor: A specialized extractor instance:
             - CursorExtractor for Cursor objects
             - ResultSetExtractor for ResultSet objects
-            - Throws an NotImplementedError for all other objects
+            - None 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__}")
+        logger.error(f"No extractor found for {obj.__class__.__name__}")
+        return None
 
 
 def log_latency(statement_type: StatementType = StatementType.NONE):
@@ -207,24 +211,26 @@ def _safe_call(func_to_call):
                 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,
-                )
+
+                if extractor is not None:
+                    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
 

src/databricks/sql/telemetry/models/enums.py

@@ -2,15 +2,17 @@
 
 
 class AuthFlow(Enum):
-    TOKEN_PASSTHROUGH = "token_passthrough"
-    BROWSER_BASED_AUTHENTICATION = "browser_based_authentication"
+    TYPE_UNSPECIFIED = "TYPE_UNSPECIFIED"
+    TOKEN_PASSTHROUGH = "TOKEN_PASSTHROUGH"
+    CLIENT_CREDENTIALS = "CLIENT_CREDENTIALS"
+    BROWSER_BASED_AUTHENTICATION = "BROWSER_BASED_AUTHENTICATION"
 
 
 class AuthMech(Enum):
-    CLIENT_CERT = "CLIENT_CERT"  # ssl certificate authentication
-    PAT = "PAT"  # Personal Access Token authentication
-    DATABRICKS_OAUTH = "DATABRICKS_OAUTH"  # Databricks-managed OAuth flow
-    EXTERNAL_AUTH = "EXTERNAL_AUTH"  # External identity provider (AWS, Azure, etc.)
+    TYPE_UNSPECIFIED = "TYPE_UNSPECIFIED"
+    OTHER = "OTHER"
+    PAT = "PAT"
+    OAUTH = "OAUTH"
 
 
 class DatabricksClientType(Enum):
@@ -19,24 +21,24 @@ class DatabricksClientType(Enum):
 
 
 class DriverVolumeOperationType(Enum):
-    TYPE_UNSPECIFIED = "type_unspecified"
-    PUT = "put"
-    GET = "get"
-    DELETE = "delete"
-    LIST = "list"
-    QUERY = "query"
+    TYPE_UNSPECIFIED = "TYPE_UNSPECIFIED"
+    PUT = "PUT"
+    GET = "GET"
+    DELETE = "DELETE"
+    LIST = "LIST"
+    QUERY = "QUERY"
 
 
 class ExecutionResultFormat(Enum):
-    FORMAT_UNSPECIFIED = "format_unspecified"
-    INLINE_ARROW = "inline_arrow"
-    EXTERNAL_LINKS = "external_links"
-    COLUMNAR_INLINE = "columnar_inline"
+    FORMAT_UNSPECIFIED = "FORMAT_UNSPECIFIED"
+    INLINE_ARROW = "INLINE_ARROW"
+    EXTERNAL_LINKS = "EXTERNAL_LINKS"
+    COLUMNAR_INLINE = "COLUMNAR_INLINE"
 
 
 class StatementType(Enum):
-    NONE = "none"
-    QUERY = "query"
-    SQL = "sql"
-    UPDATE = "update"
-    METADATA = "metadata"
+    NONE = "NONE"
+    QUERY = "QUERY"
+    SQL = "SQL"
+    UPDATE = "UPDATE"
+    METADATA = "METADATA"

src/databricks/sql/telemetry/telemetry_client.py

@@ -60,40 +60,34 @@ def get_driver_system_configuration(cls) -> DriverSystemConfiguration:
     def get_auth_mechanism(auth_provider):
         """Get the auth mechanism for the auth provider."""
         # AuthMech is an enum with the following values:
-        # PAT, DATABRICKS_OAUTH, EXTERNAL_AUTH, CLIENT_CERT
+        # TYPE_UNSPECIFIED, OTHER, PAT, OAUTH
 
         if not auth_provider:
             return None
         if isinstance(auth_provider, AccessTokenAuthProvider):
-            return AuthMech.PAT  # Personal Access Token authentication
+            return AuthMech.PAT
         elif isinstance(auth_provider, DatabricksOAuthProvider):
-            return AuthMech.DATABRICKS_OAUTH  # Databricks-managed OAuth flow
-        elif isinstance(auth_provider, ExternalAuthProvider):
-            return (
-                AuthMech.EXTERNAL_AUTH
-            )  # External identity provider (AWS, Azure, etc.)
-        return AuthMech.CLIENT_CERT  # Client certificate (ssl)
+            return AuthMech.OAUTH
+        else:
+            return AuthMech.OTHER
 
     @staticmethod
     def get_auth_flow(auth_provider):
         """Get the auth flow for the auth provider."""
         # AuthFlow is an enum with the following values:
-        # TOKEN_PASSTHROUGH, BROWSER_BASED_AUTHENTICATION
+        # TYPE_UNSPECIFIED, TOKEN_PASSTHROUGH, CLIENT_CREDENTIALS, BROWSER_BASED_AUTHENTICATION
 
         if not auth_provider:
             return None
-
         if isinstance(auth_provider, DatabricksOAuthProvider):
             if auth_provider._access_token and auth_provider._refresh_token:
-                return (
-                    AuthFlow.TOKEN_PASSTHROUGH
-                )  # Has existing tokens, no user interaction needed
-            if hasattr(auth_provider, "oauth_manager"):
-                return (
-                    AuthFlow.BROWSER_BASED_AUTHENTICATION
-                )  # Will initiate OAuth flow requiring browser
-
-        return None
+                return AuthFlow.TOKEN_PASSTHROUGH
+            else:
+                return AuthFlow.BROWSER_BASED_AUTHENTICATION
+        elif isinstance(auth_provider, ExternalAuthProvider):
+            return AuthFlow.CLIENT_CREDENTIALS
+        else:
+            return None
 
 
 class BaseTelemetryClient(ABC):
@@ -104,21 +98,23 @@ class BaseTelemetryClient(ABC):
 
     @abstractmethod
     def export_initial_telemetry_log(self, driver_connection_params, user_agent):
-        raise NotImplementedError(
-            "Subclasses must implement export_initial_telemetry_log"
-        )
+        logger.debug("subclass must implement export_initial_telemetry_log")
+        pass
 
     @abstractmethod
     def export_failure_log(self, error_name, error_message):
-        raise NotImplementedError("Subclasses must implement export_failure_log")
+        logger.debug("subclass must implement export_failure_log")
+        pass
 
     @abstractmethod
     def export_latency_log(self, latency_ms, sql_execution_event, sql_statement_id):
-        raise NotImplementedError("Subclasses must implement export_latency_log")
+        logger.debug("subclass must implement export_latency_log")
+        pass
 
     @abstractmethod
     def close(self):
-        raise NotImplementedError("Subclasses must implement close")
+        logger.debug("subclass must implement close")
+        pass
 
 
 class NoopTelemetryClient(BaseTelemetryClient):
@@ -157,6 +153,8 @@ class TelemetryClient(BaseTelemetryClient):
     TELEMETRY_AUTHENTICATED_PATH = "/telemetry-ext"
     TELEMETRY_UNAUTHENTICATED_PATH = "/telemetry-unauth"
 
+    DEFAULT_BATCH_SIZE = 100
+
     def __init__(
         self,
         telemetry_enabled,
@@ -167,12 +165,12 @@ def __init__(
     ):
         logger.debug("Initializing TelemetryClient for connection: %s", session_id_hex)
         self._telemetry_enabled = telemetry_enabled
-        self._batch_size = 10  # TODO: Decide on batch size
+        self._batch_size = self.DEFAULT_BATCH_SIZE
         self._session_id_hex = session_id_hex
         self._auth_provider = auth_provider
         self._user_agent = None
         self._events_batch = []
-        self._lock = threading.Lock()
+        self._lock = threading.RLock()
         self._driver_connection_params = None
         self._host_url = host_url
         self._executor = executor