Chunk download latency

[saishreeeee]

What type of PR is this?

• Refactor
• Feature
• Bug Fix
• Other

Description

Record chunk download latency
Added chunk_id to SqlExecutionEvent and ResultDownloadHandler

How is this tested?

• Unit tests
• E2E Tests
• Manually
• N/A

Ran the query: SELECT * FROM RANGE(20000000)
Latency was recorded for 4 chunks.
The latency log for the 1st chunk:

{
  "frontend_log_event_id": "e684ff49-9026-49af-abb6-687a6a16e565",
  "context": {
    "client_context": {
      "timestamp_millis": 1752467452769,
      "user_agent": "PyDatabricksSqlConnector/4.0.5"
    }
  },
  "entry": {
    "sql_driver_log": {
      "session_id": "01f0606b-4825-118d-abd2-87f5748c7dd3",
      "system_configuration": {
        "driver_version": "4.0.5",
        "os_name": "Darwin",
        "os_version": "24.5.0",
        "os_arch": "arm64",
        "runtime_name": "Python 3.13.3",
        "runtime_version": "3.13.3",
        "runtime_vendor": "CPython",
        "driver_name": "Databricks SQL Python Connector",
        "char_set_encoding": "utf-8",
        "locale_name": "en_US"
      },
      "driver_connection_params": {
        "http_path": <HTTP_PATH>,
        "mode": "THRIFT",
        "host_info": {
          "host_url": <SERVER_HOSTNAME>,
          "port": 443
        },
        "auth_mech": "PAT"
      },
      "sql_statement_id": "01f0606b-4854-1f89-a0cb-d5860332ef81",
      "sql_operation": {
        "statement_type": "QUERY",
        "is_compressed": true,
        "execution_result": "EXTERNAL_LINKS",
        "chunk_id": 0
      },
      "operation_latency_ms": 15882
    }
  }
}

Related Tickets & Documents

[vikrantpuppala]

@jayantsing-db can you take a look, this PR adds latency logs (merges into sea-migration) as there are some refactorings that this PR leverages

[jprakash-db]

I feel why is the code modified to fit telemetry? like the statement type within log latency decorator made sense. Why is it now a part of the statement type in the code ?

[saishreeeee]

The chunk downloader function doesn't know the StatementType, hence had to pass it down as arg.
Now have changed it to always set statement_type to StatementType.QUERY as metadata does not use cloud fetch.

[jprakash-db]

@vikrantpuppala I think we should introduce some variables like the connection context in JDBC .

• Passing these variables in so many functions is just bloating the code. A future TODO
• This will also simply the TelemetryExtractor: