Added model fields for chunk/result latency

Signed-off-by: Nikhil Suri <[email protected]>

Author: nikhilsuri-db

src/databricks/sql/common/unified_http_client.py

@@ -50,7 +50,7 @@ def __init__(self, client_context):
         """
         self.config = client_context
         # Since the unified http client is used for all requests, we need to have proxy and direct pool managers
-        # for per-request proxy decisions.
+        # for per-reques ̰ˇt proxy decisions.
         self._direct_pool_manager = None
         self._proxy_pool_manager = None
         self._retry_policy = None

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

@@ -149,6 +149,100 @@ class DriverErrorInfo(JsonSerializableMixin):
     stack_trace: str
 
 
+@dataclass
+class ChunkDetails(JsonSerializableMixin):
+    """
+    Contains detailed metrics about chunk downloads during result fetching.
+    
+    These metrics are accumulated across all chunk downloads for a single statement.
+    In Java, this is populated by the StatementTelemetryDetails tracker as chunks are downloaded.
+    
+    Tracking approach:
+    - Initialize total_chunks_present from result manifest
+    - For each chunk downloaded:
+      * Increment total_chunks_iterated
+      * Add chunk latency to sum_chunks_download_time_millis
+      * Update initial_chunk_latency_millis (first chunk only)
+      * Update slowest_chunk_latency_millis (if current chunk is slower)
+
+    Attributes:
+        initial_chunk_latency_millis (int): Latency of the first chunk download
+        slowest_chunk_latency_millis (int): Latency of the slowest chunk download
+        total_chunks_present (int): Total number of chunks available
+        total_chunks_iterated (int): Number of chunks actually downloaded
+        sum_chunks_download_time_millis (int): Total time spent downloading all chunks
+    """
+
+    initial_chunk_latency_millis: Optional[int] = None
+    slowest_chunk_latency_millis: Optional[int] = None
+    total_chunks_present: Optional[int] = None
+    total_chunks_iterated: Optional[int] = None
+    sum_chunks_download_time_millis: Optional[int] = None
+
+
+@dataclass
+class ResultLatency(JsonSerializableMixin):
+    """
+    Contains latency metrics for different phases of query execution.
+    
+    This tracks two distinct phases:
+    1. result_set_ready_latency_millis: Time from query submission until results are available (execute phase)
+       - Set when execute() completes
+    2. result_set_consumption_latency_millis: Time spent iterating/fetching results (fetch phase)
+       - Measured from first fetch call until no more rows available
+       - In Java: tracked via markResultSetConsumption(hasNext) method
+       - Records start time on first fetch, calculates total on last fetch
+
+    Attributes:
+        result_set_ready_latency_millis (int): Time until query results are ready (execution phase)
+        result_set_consumption_latency_millis (int): Time spent fetching/consuming results (fetch phase)
+        
+    Note:
+        Java implementation includes private field 'startTimeOfResultSetIterationNano' for internal
+        tracking (not serialized to JSON). When implementing tracking in Python, use similar approach:
+        - Record start time on first fetchone/fetchmany/fetchall call
+        - Calculate total consumption latency when iteration completes or cursor closes
+    """
+
+    result_set_ready_latency_millis: Optional[int] = None
+    result_set_consumption_latency_millis: Optional[int] = None
+
+
+@dataclass
+class OperationDetail(JsonSerializableMixin):
+    """
+    Contains detailed information about the operation being performed.
+    
+    This provides more granular operation tracking than statement_type, allowing
+    differentiation between similar operations (e.g., EXECUTE_STATEMENT vs EXECUTE_STATEMENT_ASYNC).
+    
+    Tracking approach:
+    - operation_type: Map method name to operation type enum
+      * Java maps: executeStatement -> EXECUTE_STATEMENT
+      * Java maps: listTables -> LIST_TABLES
+      * Python could use similar mapping from method names
+    
+    - is_internal_call: Track if operation is initiated by driver internally
+      * Set to true for driver-initiated metadata calls
+      * Set to false for user-initiated operations
+    
+    - Status polling: For async operations
+      * Increment n_operation_status_calls for each status check
+      * Accumulate operation_status_latency_millis across all status calls
+
+    Attributes:
+        n_operation_status_calls (int): Number of status polling calls made
+        operation_status_latency_millis (int): Total latency of all status calls
+        operation_type (str): Specific operation type (e.g., EXECUTE_STATEMENT, LIST_TABLES, CANCEL_STATEMENT)
+        is_internal_call (bool): Whether this is an internal driver operation
+    """
+
+    n_operation_status_calls: Optional[int] = None
+    operation_status_latency_millis: Optional[int] = None
+    operation_type: Optional[str] = None
+    is_internal_call: Optional[bool] = None
+
+
 @dataclass
 class SqlExecutionEvent(JsonSerializableMixin):
     """
@@ -160,14 +254,20 @@ class SqlExecutionEvent(JsonSerializableMixin):
         is_compressed (bool): Whether the result is compressed
         execution_result (ExecutionResultFormat): Format of the execution result
         retry_count (int): Number of retry attempts made
-        chunk_id (int): ID of the chunk if applicable
+        chunk_id (int): ID of the chunk if applicable (used for error tracking)
+        chunk_details (ChunkDetails): Aggregated chunk download metrics
+        result_latency (ResultLatency): Latency breakdown by execution phase
+        operation_detail (OperationDetail): Detailed operation information
     """
 
     statement_type: StatementType
     is_compressed: bool
     execution_result: ExecutionResultFormat
     retry_count: Optional[int]
     chunk_id: Optional[int]
+    chunk_details: Optional[ChunkDetails] = None
+    result_latency: Optional[ResultLatency] = None
+    operation_detail: Optional[OperationDetail] = None
 
 
 @dataclass

src/databricks/sql/telemetry/telemetry_client.py

@@ -380,7 +380,7 @@ class TelemetryClientFactory:
     # Shared flush thread for all clients
     _flush_thread = None
     _flush_event = threading.Event()
-    _flush_interval_seconds = 90
+    _flush_interval_seconds = 300  # 5 minutes
 
     DEFAULT_BATCH_SIZE = 100