Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 20 additions & 22 deletions .github/workflows/trigger-integration-tests.yml
Original file line number Diff line number Diff line change
Expand Up @@ -194,15 +194,6 @@ jobs:
owner: databricks
repositories: databricks-driver-test

- name: Generate GitHub App Token (public repo)
id: public-token
uses: actions/create-github-app-token@f8d387b68d61c58ab83c6c016672934102569859 # v3.0.0
with:
app-id: ${{ secrets.INTEGRATION_TEST_APP_ID }}
private-key: ${{ secrets.INTEGRATION_TEST_PRIVATE_KEY }}
owner: databricks
repositories: databricks-sql-python

- name: Sanitize PR title
id: sanitize
uses: actions/github-script@f28e40c7f34bde8b3046d885e986cb6290c5673b # v7.1.0
Expand Down Expand Up @@ -235,7 +226,11 @@ jobs:
if: steps.changed.outputs.python != 'true'
uses: actions/github-script@f28e40c7f34bde8b3046d885e986cb6290c5673b # v7.1.0
with:
github-token: ${{ steps.public-token.outputs.token }}
# Default workflow token, not the App token — same rationale
# as the failure handler below. We don't want a missing-secret
# state to silently swallow the green check for path-filtered
# no-op runs.
github-token: ${{ github.token }}
script: |
await github.rest.checks.create({
owner: context.repo.owner,
Expand All @@ -255,7 +250,15 @@ jobs:
if: failure() && steps.changed.outputs.python == 'true'
uses: actions/github-script@f28e40c7f34bde8b3046d885e986cb6290c5673b # v7.1.0
with:
github-token: ${{ steps.public-token.outputs.token }}
# Use the default workflow token, not the App token. The
# App-token-generating step is the *most likely* thing to
# fail (missing/rotated secrets, App uninstalled), and using
# it here means a token-generation failure also kills this
# handler — leaving the gate silently green on the stale
# synthetic-success from skip-integration-tests-pr. The
# default token has checks:write (declared on this job)
# which is all we need.
github-token: ${{ github.token }}
script: |
await github.rest.checks.create({
owner: context.repo.owner,
Expand Down Expand Up @@ -316,20 +319,13 @@ jobs:
echo "No driver files changed — will auto-pass"
fi

- name: Generate GitHub App Token (public repo)
id: public-token
uses: actions/create-github-app-token@f8d387b68d61c58ab83c6c016672934102569859 # v3.0.0
with:
app-id: ${{ secrets.INTEGRATION_TEST_APP_ID }}
private-key: ${{ secrets.INTEGRATION_TEST_PRIVATE_KEY }}
owner: databricks
repositories: databricks-sql-python

- name: Auto-pass (no driver changes)
if: steps.changed.outputs.changed != 'true'
uses: actions/github-script@f28e40c7f34bde8b3046d885e986cb6290c5673b # v7.1.0
with:
github-token: ${{ steps.public-token.outputs.token }}
# Default workflow token — see the trigger-tests-pr job's
# equivalent step above for the rationale.
github-token: ${{ github.token }}
script: |
await github.rest.checks.create({
owner: context.repo.owner,
Expand Down Expand Up @@ -392,7 +388,9 @@ jobs:
if: failure() && steps.changed.outputs.changed == 'true'
uses: actions/github-script@f28e40c7f34bde8b3046d885e986cb6290c5673b # v7.1.0
with:
github-token: ${{ steps.public-token.outputs.token }}
# Use the default workflow token, not the App token — see
# the rationale in the trigger-tests-pr job above.
github-token: ${{ github.token }}
script: |
await github.rest.checks.create({
owner: context.repo.owner,
Expand Down
60 changes: 57 additions & 3 deletions tests/e2e/common/retry_test_mixins.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,58 @@
SessionAlreadyClosedError,
UnsafeToRetryError,
)
from databricks.sql.telemetry.telemetry_client import (
NoopTelemetryClient,
TelemetryClientFactory,
_TelemetryClientHolder,
)


@contextmanager
def _isolated_from_telemetry():
# Tests that mock urllib3 globally (via _get_conn / _validate_conn) also
# intercept background telemetry pushes from the shared
# TelemetryClientFactory executor — inflating mock.call_count and breaking
# assertions like `call_count == 6`. Drain the factory and force any new
# connection to use NoopTelemetryClient for the duration of the test.
with TelemetryClientFactory._lock:
saved_clients = TelemetryClientFactory._clients
saved_executor = TelemetryClientFactory._executor
saved_initialized = TelemetryClientFactory._initialized
for holder in list(saved_clients.values()):
try:
holder.client.close()
except Exception:
pass
if saved_executor is not None:
try:
TelemetryClientFactory._stop_flush_thread()
saved_executor.shutdown(wait=True)
except Exception:
pass
TelemetryClientFactory._clients = {}
TelemetryClientFactory._executor = None
TelemetryClientFactory._initialized = False

def _install_noop(*args, host_url=None, **kwargs):
host_url = TelemetryClientFactory.getHostUrlSafely(host_url)
with TelemetryClientFactory._lock:
TelemetryClientFactory._clients[host_url] = _TelemetryClientHolder(
NoopTelemetryClient()
)

try:
with patch.object(
TelemetryClientFactory,
"initialize_telemetry_client",
staticmethod(_install_noop),
):
yield
finally:
with TelemetryClientFactory._lock:
TelemetryClientFactory._clients = saved_clients
TelemetryClientFactory._executor = saved_executor
TelemetryClientFactory._initialized = saved_initialized


class Client429ResponseMixin:
Expand Down Expand Up @@ -253,7 +305,7 @@ def test_retry_urllib3_settings_are_honored(
@patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry")
def test_oserror_retries(self, mock_send_telemetry, extra_params):
"""If a network error occurs during make_request, the request is retried according to policy"""
with patch(
with _isolated_from_telemetry(), patch(
"urllib3.connectionpool.HTTPSConnectionPool._validate_conn",
) as mock_validate_conn:
mock_validate_conn.side_effect = OSError("Some arbitrary network error")
Expand All @@ -278,7 +330,9 @@ def test_retry_max_count_not_exceeded(self, mock_send_telemetry, extra_params):
THEN the connector issues six request (original plus five retries)
before raising an exception
"""
with mocked_server_response(status=429, headers={"Retry-After": "0"}) as mock_obj:
with _isolated_from_telemetry(), mocked_server_response(
status=429, headers={"Retry-After": "0"}
) as mock_obj:
with pytest.raises(MaxRetryError) as cm:
extra_params = {**extra_params, **self._retry_policy}
with self.connection(extra_params=extra_params) as conn:
Expand All @@ -302,7 +356,7 @@ def test_retry_exponential_backoff(self, mock_send_telemetry, extra_params):
retry_policy["_retry_delay_min"] = 1

time_start = time.time()
with mocked_server_response(
with _isolated_from_telemetry(), mocked_server_response(
status=429, headers={"Retry-After": "8"}
) as mock_obj:
with pytest.raises(RequestError) as cm:
Expand Down
Loading