Skip to content

Add optional telemetry support to the python connector #628

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 252 commits into
base: main
Choose a base branch
from

Conversation

saishreeeee
Copy link
Collaborator

What type of PR is this?

  • Refactor
  • Feature
  • Bug Fix
  • Other

Description

Initial telemetry logs, failure logs, latency logs via telemetry

How is this tested?

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

Related Tickets & Documents

Design doc

Jesse and others added 30 commits June 24, 2022 12:47
* Remove databricks-specific test fixtures
* Introduce pytest, black, and mypy workflows
* Skip test_fetches_bench case
* Remove references to databricks-specific testing infra
The code for version <=1.0 is not contained in this repository. There has been no decision 
about whether or not to open source it.
* Add license and contributing sections to README.
* Add environment setup docs to CONTRIBUTING
* Clarify example of connection details in example
* Add badges from pypi
* Explicitly call out Python 3.7 or above is needed
Signed-off-by: Jesse Whitehouse <jesse@whitehouse.dev>
Advise developers to use Python 3.7, 3.8, or 3.9 until #26 is fixed.
Includes a GitHub Action which checks for a valid sign-off on every proposed commit
* Isolate delay bounding logic
* Move error details scope up one-level.
* Retry GetOperationStatus if an OSError was raised during execution. Add retry_delay_default to use in this case.
* Log when a request is retried due to an OSError. Emit warnings for unexpected OSError codes
* Update docstring for make_request
* Nit: unit tests show the .warn message is deprecated. DeprecationWarning: The 'warn' function is deprecated, use 'warning' instead

Signed-off-by: Jesse Whitehouse <jesse@whitehouse.dev>
Signed-off-by: Jesse Whitehouse <jesse@whitehouse.dev>
Signed-off-by: Jesse Whitehouse <jesse@whitehouse.dev>
* Test with multiple python versions.
* Update pyarrow to version 9.0.0 to address issue in relation to python 3.10 & a specific version of numpy being pulled in by pyarrow.

Closes #26 

Signed-off-by: David Black <dblack@atlassian.com>
* Update changelog and bump to v2.0.4
* Specifically thank @dbaxa for this change.

Signed-off-by: Jesse Whitehouse <jesse@whitehouse.dev>
Signed-off-by: Jesse Whitehouse <jesse@whitehouse.dev>
Signed-off-by: Jesse Whitehouse <jesse@whitehouse.dev>
* Add test: cursors are closed when connection closes

Signed-off-by: Jesse Whitehouse <jesse@whitehouse.dev>
Signed-off-by: Jesse Whitehouse <jesse@whitehouse.dev>
Signed-off-by: Moe Derakhshani <moe.derakhshani@databricks.com>
Signed-off-by: Moe Derakhshani <moe.derakhshani@databricks.com>
Signed-off-by: Moe Derakhshani <moe.derakhshani@databricks.com>

my [OAuth PR](https://github.com/databricks/databricks-sql-python/runs/8005844758?check_suite_focus=true) is blocked due to dco validation (following error):
<img width="1202" alt="Screen Shot 2022-08-25 at 12 05 40 PM" src="https://user-images.githubusercontent.com/22279672/186747897-c9d57586-366f-41f9-aa66-609f2bf3911f.png">



We should try to avoid running dco for internal databricks employees:
I am trying to relax the validation based on this guideline:
https://github.com/dcoapp/app/blob/main/README.md#skipping-sign-off-for-organization-members

and here:
https://stackoverflow.com/questions/62969381/is-it-in-line-with-the-dco-that-a-github-sign-off-needs-and-publishes-full-name
Signed-off-by: Moe Derakhshani <moe.derakhshani@databricks.com>
Signed-off-by: Moe Derakhshani <moe.derakhshani@databricks.com>
Signed-off-by: Moe Derakhshani <moe.derakhshani@databricks.com>

this is undo of #42 till we figure out how to fix dco
Signed-off-by: Jesse Whitehouse <jesse.whitehouse@databricks.com>
This PR:
* Adds the foundation for OAuth against Databricks account on AWS with BYOIDP.
* It copies one internal module that Steve Weis @sweisdb wrote for Databricks CLI (oauth.py). Once ecosystem-dev team (Serge, Pieter) build a python sdk core we will move this code to their repo as a dependency. 
* the PR provides authenticators with visitor pattern format for stamping auth-token which later is intended to be moved to the repo owned by Serge @nfx and and Pieter @pietern
Signed-off-by: Jesse Whitehouse <jesse.whitehouse@databricks.com>
Signed-off-by: Jesse Whitehouse <jesse.whitehouse@databricks.com>
Bump to v2.1.0 and update changelog

Signed-off-by: Jesse Whitehouse <jesse.whitehouse@databricks.com>
Signed-off-by: Mohit Singla <mohit.singla@databricks.com>
Co-authored-by: Moe Derakhshani <moe.derakhshani@databricks.com>
madhav-db and others added 4 commits June 25, 2025 12:32
#613)

* Revert "Enhance Cursor close handling and context manager exception management to prevent server side resource leaks (#554)"

This reverts commit edfb283.

* revert e2e
* Release version 4.0.5: Reverted cursor close handling changes to fix user errors. Updated version numbers in pyproject.toml and __init__.py.

* Update CHANGELOG.md to include reference to issue #613 for cursor close handling fix.
* added functionality for export of failure logs

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

* changed logger.error to logger.debug in exc.py

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

* Fix telemetry loss during Python shutdown

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

* unit tests for export_failure_log

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

* try-catch blocks to make telemetry failures non-blocking for connector operations

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

* removed redundant try/catch blocks, added try/catch block to initialize and get telemetry client

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

* skip null fields in telemetry request

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

* removed dup import, renamed func, changed a filter_null_values to lamda

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

* removed unnecassary class variable and a redundant try/except block

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

* public functions defined at interface level

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

* changed export_event and flush to private functions

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

* formatting

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

* changed connection_uuid to thread local in thrift backend

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

* made errors more specific

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

* revert change to connection_uuid

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

* reverting change in close in telemetry client

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

* JsonSerializableMixin

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

* isdataclass check in JsonSerializableMixin

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

* convert TelemetryClientFactory to module-level functions, replace NoopTelemetryClient class with NOOP_TELEMETRY_CLIENT singleton, updated tests accordingly

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

* renamed connection_uuid as session_id_hex

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

* added NotImplementedError to abstract class, added unit tests

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

* formatting

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

* added PEP-249 link, changed NoopTelemetryClient implementation

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

* removed unused import

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

* made telemetry client close a module-level function

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

* unit tests verbose

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

* debug logs in unit tests

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

* debug logs in unit tests

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

* removed ABC from mixin, added try/catch block around executor shutdown

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

* checking stuff

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

* finding out

* finding out more

* more more finding out more nice

* locks are useless anyways

* haha

* normal

* := looks like walrus horizontally

* one more

* walrus again

* old stuff without walrus seems to fail

* manually do the walrussing

* change 3.13t, v2

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

* formatting, added walrus

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

* formatting

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

* removed walrus, removed test before stalling test

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

* changed order of stalling test

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

* removed debugging, added TelemetryClientFactory

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

* remove more debugging

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

* latency logs funcitionality

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

* fixed type of return value in get_session_id_hex() in thrift backend

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

* debug on TelemetryClientFactory lock

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

* formatting

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

* type notation for _waiters

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

* called connection.close() in test_arraysize_buffer_size_passthrough

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

* run all unit tests

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

* more debugging

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

* removed the connection.close() from that test, put debug statement before and after TelemetryClientFactory lock

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

* more debug

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

* more more more

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

* why

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

* whywhy

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

* thread name

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

* added teardown to all tests except finalizer test (gc collect)

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

* added the get_attribute functions to the classes

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

* removed tearDown, added connection.close() to first test

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

* finally

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

* remove debugging

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

* added test for export_latency_log, made mock of thrift backend with retry policy

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

* added multi threaded tests

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

* formatting

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

* added TelemetryExtractor, removed multithreaded tests

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

* formatting

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

* fixes in test

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

* fix in telemetry extractor

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

* added doc strings to latency_logger, abstracted export_telemetry_log

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

* statement type, unit test fix

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

* unit test fix

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

* statement type changes

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

* test_fetches fix

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

* added mocks to resolve the errors caused by log_latency decorator in tests

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

* removed function in test_fetches cuz it is only used once

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

* added _safe_call which returns None in case of errors in the get functions

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

* removed the changes in test_client and test_fetches

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

* removed the changes in test_fetches

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

* test_telemetry

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

* removed test

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

---------

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

github-actions bot commented Jul 7, 2025

Thanks for your contribution! To satisfy the DCO policy in our contributing guide every commit message must include a sign-off message. One or more of your commits is missing this message. You can reword previous commit messages with an interactive rebase (git rebase -i main).

Copy link

github-actions bot commented Jul 7, 2025

Thanks for your contribution! To satisfy the DCO policy in our contributing guide every commit message must include a sign-off message. One or more of your commits is missing this message. You can reword previous commit messages with an interactive rebase (git rebase -i main).

This reverts commit 10375a8, reversing
changes made to 0dfe0f4.

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

github-actions bot commented Jul 7, 2025

Thanks for your contribution! To satisfy the DCO policy in our contributing guide every commit message must include a sign-off message. One or more of your commits is missing this message. You can reword previous commit messages with an interactive rebase (git rebase -i main).

This reverts commit 8c0f474.

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

github-actions bot commented Jul 7, 2025

Thanks for your contribution! To satisfy the DCO policy in our contributing guide every commit message must include a sign-off message. One or more of your commits is missing this message. You can reword previous commit messages with an interactive rebase (git rebase -i main).

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

github-actions bot commented Jul 7, 2025

Thanks for your contribution! To satisfy the DCO policy in our contributing guide every commit message must include a sign-off message. One or more of your commits is missing this message. You can reword previous commit messages with an interactive rebase (git rebase -i main).

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

github-actions bot commented Jul 7, 2025

Thanks for your contribution! To satisfy the DCO policy in our contributing guide every commit message must include a sign-off message. One or more of your commits is missing this message. You can reword previous commit messages with an interactive rebase (git rebase -i main).

@saishreeeee saishreeeee marked this pull request as ready for review July 7, 2025 05:40
@vikrantpuppala vikrantpuppala changed the title Telemetry Add optional telemetry support to the python connector Jul 7, 2025

logger = logging.getLogger(__name__)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

don't think we're using this?

elif obj.__class__.__name__ == "ResultSet":
return ResultSetExtractor(obj)
else:
raise NotImplementedError(f"No extractor found for {obj.__class__.__name__}")
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

will this error break anything or just be printed in debug mode?

try:
result = func(self, *args, **kwargs)
return result
finally:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what if anything in this finally block fails?

class StatementType(Enum):
NONE = "none"
QUERY = "query"
SQL = "sql"
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

query type is used when sql command is expected to return results, in jdbc it makes sense due to the nature of interface, but do we need both query and sql in python connector?

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.)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think @jprakash-db is adding some enums in his latest PR #621, is there scope to re-use?

self._auth_provider = auth_provider
self._user_agent = None
self._events_batch = []
self._lock = threading.Lock()
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do we not need RLock here?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.