Skip to content

[ES-1717039] Fix 60 seconds delay in gov cloud connections + Fix PR check failures in the repo #735

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.

Sign up for GitHub

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

Merged
samikshya-db merged 8 commits into from
Feb 5, 2026
+94 −55
Merged

[ES-1717039] Fix 60 seconds delay in gov cloud connections + Fix PR check failures in the repo #735

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
65e8587
Fix 60 seconds delay in gov cloud connections
samikshya-db Feb 4, 2026
ca25c61
keep it simple :)
samikshya-db Feb 4, 2026
616f5b2
Add fix for krb error
samikshya-db Feb 4, 2026
0a6b7b0
pin poetry
samikshya-db Feb 4, 2026
9f8d98f
Pin for publish flow too
samikshya-db Feb 4, 2026
44b09b1
Fix failing tests
samikshya-db Feb 4, 2026
eed311b
Edit order for pypi
samikshya-db Feb 4, 2026
b6a417d
One last fix : pls work
samikshya-db Feb 4, 2026
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
5 changes: 5 additions & 0 deletions .github/workflows/code-coverage.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ jobs:
- name: Install Poetry
uses: snok/install-poetry@v1
with:
version: "2.2.1"
virtualenvs-create: true
virtualenvs-in-project: true
installer-parallel: true
Expand All @@ -58,6 +59,10 @@ jobs:
#----------------------------------------------
# install your root project, if required
#----------------------------------------------
- name: Install Kerberos system dependencies
run: |
sudo apt-get update
sudo apt-get install -y libkrb5-dev
- name: Install library
run: poetry install --no-interaction --all-extras
#----------------------------------------------
Expand Down
8 changes: 8 additions & 0 deletions .github/workflows/code-quality-checks.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ jobs:
- name: Install Poetry
uses: snok/install-poetry@v1
with:
version: "2.2.1"
virtualenvs-create: true
virtualenvs-in-project: true
installer-parallel: true
Expand Down Expand Up @@ -118,6 +119,7 @@ jobs:
- name: Install Poetry
uses: snok/install-poetry@v1
with:
version: "2.2.1"
virtualenvs-create: true
virtualenvs-in-project: true
installer-parallel: true
Expand All @@ -140,6 +142,10 @@ jobs:
#----------------------------------------------
# install your root project, if required
#----------------------------------------------
- name: Install Kerberos system dependencies
run: |
sudo apt-get update
sudo apt-get install -y libkrb5-dev
- name: Install library
run: poetry install --no-interaction --all-extras
#----------------------------------------------
Expand Down Expand Up @@ -191,6 +197,7 @@ jobs:
- name: Install Poetry
uses: snok/install-poetry@v1
with:
version: "2.2.1"
virtualenvs-create: true
virtualenvs-in-project: true
installer-parallel: true
Expand Down Expand Up @@ -243,6 +250,7 @@ jobs:
- name: Install Poetry
uses: snok/install-poetry@v1
with:
version: "2.2.1"
virtualenvs-create: true
virtualenvs-in-project: true
installer-parallel: true
Expand Down
5 changes: 5 additions & 0 deletions .github/workflows/daily-telemetry-e2e.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ jobs:
- name: Install Poetry
uses: snok/install-poetry@v1
with:
version: "2.2.1"
virtualenvs-create: true
virtualenvs-in-project: true
installer-parallel: true
Expand All @@ -60,6 +61,10 @@ jobs:
#----------------------------------------------
# install dependencies if cache does not exist
#----------------------------------------------
- name: Install Kerberos system dependencies
run: |
sudo apt-get update
sudo apt-get install -y libkrb5-dev
- name: Install dependencies
run: poetry install --no-interaction --all-extras

Expand Down
5 changes: 5 additions & 0 deletions .github/workflows/integration.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ jobs:
- name: Install Poetry
uses: snok/install-poetry@v1
with:
version: "2.2.1"
virtualenvs-create: true
virtualenvs-in-project: true
installer-parallel: true
Expand All @@ -49,6 +50,10 @@ jobs:
#----------------------------------------------
# install dependencies if cache does not exist
#----------------------------------------------
- name: Install Kerberos system dependencies
run: |
sudo apt-get update
sudo apt-get install -y libkrb5-dev
- name: Install dependencies
run: poetry install --no-interaction --all-extras
#----------------------------------------------
Expand Down
9 changes: 9 additions & 0 deletions .github/workflows/publish-manual.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,10 +31,19 @@ jobs:
- name: Install Poetry
uses: snok/install-poetry@v1 # Install Poetry, the Python package manager
with:
version: "2.2.1"
virtualenvs-create: true
virtualenvs-in-project: true
installer-parallel: true

#----------------------------------------------
# Step 3.5: Install Kerberos system dependencies
#----------------------------------------------
- name: Install Kerberos system dependencies
run: |
sudo apt-get update
sudo apt-get install -y libkrb5-dev

# #----------------------------------------------
# # Step 4: Load cached virtual environment (if available)
# #----------------------------------------------
Expand Down
23 changes: 17 additions & 6 deletions .github/workflows/publish-test.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ jobs:
- name: Install Poetry
uses: snok/install-poetry@v1
with:
version: "2.2.1"
virtualenvs-create: true
virtualenvs-in-project: true
installer-parallel: true
Expand All @@ -36,6 +37,10 @@ jobs:
#----------------------------------------------
# install dependencies if cache does not exist
#----------------------------------------------
- name: Install Kerberos system dependencies
run: |
sudo apt-get update
sudo apt-get install -y libkrb5-dev
- name: Install dependencies
if: steps.cached-poetry-dependencies.outputs.cache-hit != 'true'
run: poetry install --no-interaction --no-root
Expand All @@ -54,11 +59,17 @@ jobs:
- name: Update pyproject.toml
run: poetry version ${{ steps.version.outputs.major-version }}.${{ steps.version.outputs.minor-version }}.dev$(date +%s)
#----------------------------------------------
# Build the package (before publish action)
#----------------------------------------------
- name: Build package
run: poetry build
#----------------------------------------------
# Configure test-pypi repository
#----------------------------------------------
- name: Configure test-pypi repository
run: poetry config repositories.testpypi https://test.pypi.org/legacy/
#----------------------------------------------
# Attempt push to test-pypi
#----------------------------------------------
- name: Build and publish to pypi
uses: JRubics/poetry-publish@v1.10
with:
pypi_token: ${{ secrets.TEST_PYPI_TOKEN }}
repository_name: "testpypi"
repository_url: "https://test.pypi.org/legacy/"
- name: Publish to test-pypi
run: poetry publish --username __token__ --password ${{ secrets.TEST_PYPI_TOKEN }} --repository testpypi
18 changes: 13 additions & 5 deletions .github/workflows/publish.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ jobs:
- name: Install Poetry
uses: snok/install-poetry@v1
with:
version: "2.2.1"
virtualenvs-create: true
virtualenvs-in-project: true
installer-parallel: true
Expand All @@ -38,6 +39,10 @@ jobs:
#----------------------------------------------
# install dependencies if cache does not exist
#----------------------------------------------
- name: Install Kerberos system dependencies
run: |
sudo apt-get update
sudo apt-get install -y libkrb5-dev
- name: Install dependencies
if: steps.cached-poetry-dependencies.outputs.cache-hit != 'true'
run: poetry install --no-interaction --no-root
Expand All @@ -56,9 +61,12 @@ jobs:
- name: Update pyproject.toml
run: poetry version ${{ steps.version.outputs.current-version }}
#----------------------------------------------
# Attempt push to test-pypi
# Build the package (before publish)
#----------------------------------------------
- name: Build and publish to pypi
uses: JRubics/poetry-publish@v1.10
with:
pypi_token: ${{ secrets.PROD_PYPI_TOKEN }}
- name: Build package
run: poetry build
#----------------------------------------------
# Publish to pypi
#----------------------------------------------
- name: Publish to pypi
run: poetry publish --username __token__ --password ${{ secrets.PROD_PYPI_TOKEN }}
34 changes: 7 additions & 27 deletions src/databricks/sql/auth/retry.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -373,6 +373,13 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]:
if status_code == 403:
return False, "403 codes are not retried"

# Request failed with 404. Don't retry for any command type.
if status_code == 404:
return (
False,
"Received 404 - NOT_FOUND. The requested resource does not exist.",
)

# Request failed and server said NotImplemented. This isn't recoverable. Don't retry.
if status_code == 501:
return False, "Received code 501 from server."
Expand All @@ -381,33 +388,6 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]:
if not self._is_method_retryable(method):
return False, "Only POST requests are retried"

# Request failed with 404 and was a GetOperationStatus. This is not recoverable. Don't retry.
if status_code == 404 and self.command_type == CommandType.GET_OPERATION_STATUS:
return (
False,
"GetOperationStatus received 404 code from Databricks. Operation was canceled.",
)

# Request failed with 404 because CloseSession returns 404 if you repeat the request.
if (
status_code == 404
and self.command_type == CommandType.CLOSE_SESSION
and len(self.history) > 0
):
raise SessionAlreadyClosedError(
"CloseSession received 404 code from Databricks. Session is already closed."
)

# Request failed with 404 because CloseOperation returns 404 if you repeat the request.
if (
status_code == 404
and self.command_type == CommandType.CLOSE_OPERATION
and len(self.history) > 0
):
raise CursorAlreadyClosedError(
"CloseOperation received 404 code from Databricks. Cursor is already closed."
)

# Request failed, was an ExecuteStatement and the command may have reached the server
if (
self.command_type == CommandType.EXECUTE_STATEMENT
Expand Down
26 changes: 11 additions & 15 deletions tests/e2e/common/retry_test_mixins.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -278,7 +278,7 @@ 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=404) as mock_obj:
with 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 Down Expand Up @@ -467,22 +467,21 @@ def test_retry_safe_execute_statement_retry_condition(self, extra_params):
)
def test_retry_abort_close_session_on_404(self, extra_params, caplog):
"""GIVEN the connector sends a CloseSession command
WHEN server sends a 404 (which is normally retried)
THEN nothing is retried because 404 means the session already closed
WHEN server sends a 404 (which is not retried since commit 41b28159)
THEN nothing is retried because 404 is globally non-retryable
"""

# First response is a Bad Gateway -> Result is the command actually goes through
# Second response is a 404 because the session is no longer found
# With the idempotency-based retry refactor, 404 is now globally non-retryable
# regardless of command type. The close() method catches RequestError and proceeds.
responses = [
{"status": 502, "headers": {"Retry-After": "1"}, "redirect_location": None},
{"status": 404, "headers": {}, "redirect_location": None},
]

extra_params = {**extra_params, **self._retry_policy}
with self.connection(extra_params=extra_params) as conn:
with mock_sequential_server_responses(responses):
# Should not raise an exception, the error is caught internally
conn.close()
assert "Session was closed by a prior request" in caplog.text

@pytest.mark.parametrize(
"extra_params",
Expand All @@ -493,14 +492,13 @@ def test_retry_abort_close_session_on_404(self, extra_params, caplog):
)
def test_retry_abort_close_operation_on_404(self, extra_params, caplog):
"""GIVEN the connector sends a CancelOperation command
WHEN server sends a 404 (which is normally retried)
THEN nothing is retried because 404 means the operation was already canceled
WHEN server sends a 404 (which is not retried since commit 41b28159)
THEN nothing is retried because 404 is globally non-retryable
"""

# First response is a Bad Gateway -> Result is the command actually goes through
# Second response is a 404 because the session is no longer found
# With the idempotency-based retry refactor, 404 is now globally non-retryable
# regardless of command type. The close() method catches RequestError and proceeds.
responses = [
{"status": 502, "headers": {"Retry-After": "1"}, "redirect_location": None},
{"status": 404, "headers": {}, "redirect_location": None},
]

Expand All @@ -515,10 +513,8 @@ def test_retry_abort_close_operation_on_404(self, extra_params, caplog):
# This call guarantees we have an open cursor at the server
curs.execute("SELECT 1")
with mock_sequential_server_responses(responses):
# Should not raise an exception, the error is caught internally
curs.close()
assert (
"Operation was canceled by a prior request" in caplog.text
)

@pytest.mark.parametrize(
"extra_params",
Expand Down
2 changes: 1 addition & 1 deletion tests/e2e/common/staging_ingestion_tests.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@ def test_staging_ingestion_life_cycle(self, ingestion_user):
# GET after REMOVE should fail

with pytest.raises(
Error, match="too many 404 error responses"
Error, match="Staging operation over HTTP was unsuccessful: 404"
):
cursor = conn.cursor()
query = f"GET 'stage://tmp/{ingestion_user}/tmp/11/16/file1.csv' TO '{new_temp_path}'"
Expand Down
2 changes: 1 addition & 1 deletion tests/e2e/common/uc_volume_tests.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@ def test_uc_volume_life_cycle(self, catalog, schema):
# GET after REMOVE should fail

with pytest.raises(
Error, match="too many 404 error responses"
Error, match="Staging operation over HTTP was unsuccessful: 404"
):
cursor = conn.cursor()
query = f"GET '/Volumes/{catalog}/{schema}/e2etests/file1.csv' TO '{new_temp_path}'"
Expand Down
12 changes: 12 additions & 0 deletions tests/unit/test_retry.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,3 +83,15 @@ def test_excessive_retry_attempts_error(self, t_mock, retry_policy):
retry_policy.sleep(HTTPResponse(status=503))
# Internally urllib3 calls the increment function generating a new instance for every retry
retry_policy = retry_policy.increment()

def test_404_does_not_retry_for_any_command_type(self, retry_policy):
"""Test that 404 never retries for any CommandType"""
retry_policy._retry_start_time = time.time()

# Test for each CommandType
for command_type in CommandType:
retry_policy.command_type = command_type
should_retry, msg = retry_policy.should_retry("POST", 404)

assert should_retry is False, f"404 should not retry for {command_type}"
assert "404" in msg or "NOT_FOUND" in msg
Loading