Skip to content

Commit cafed60

Browse files
samikshya-dbsamikshya-db
authored
[ES-1717039] Fix 60 seconds delay in gov cloud connections + Fix PR check failures in the repo (#735)
* Fix 60 seconds delay in gov cloud connections * keep it simple :) * Add fix for krb error * pin poetry * Pin for publish flow too * Fix failing tests * Edit order for pypi * One last fix : pls work
1 parent 4b7df5b commit cafed60

File tree

12 files changed

+94
-55
lines changed

12 files changed

+94
-55
lines changed

‎.github/workflows/code-coverage.yml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ jobs:
3636
- name: Install Poetry
3737
uses: snok/install-poetry@v1
3838
with:
39+
version: "2.2.1"
3940
virtualenvs-create: true
4041
virtualenvs-in-project: true
4142
installer-parallel: true
@@ -58,6 +59,10 @@ jobs:
5859
#----------------------------------------------
5960
# install your root project, if required
6061
#----------------------------------------------
62+
- name: Install Kerberos system dependencies
63+
run: |
64+
sudo apt-get update
65+
sudo apt-get install -y libkrb5-dev
6166
- name: Install library
6267
run: poetry install --no-interaction --all-extras
6368
#----------------------------------------------

‎.github/workflows/code-quality-checks.yml‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ jobs:
3535
- name: Install Poetry
3636
uses: snok/install-poetry@v1
3737
with:
38+
version: "2.2.1"
3839
virtualenvs-create: true
3940
virtualenvs-in-project: true
4041
installer-parallel: true
@@ -118,6 +119,7 @@ jobs:
118119
- name: Install Poetry
119120
uses: snok/install-poetry@v1
120121
with:
122+
version: "2.2.1"
121123
virtualenvs-create: true
122124
virtualenvs-in-project: true
123125
installer-parallel: true
@@ -140,6 +142,10 @@ jobs:
140142
#----------------------------------------------
141143
# install your root project, if required
142144
#----------------------------------------------
145+
- name: Install Kerberos system dependencies
146+
run: |
147+
sudo apt-get update
148+
sudo apt-get install -y libkrb5-dev
143149
- name: Install library
144150
run: poetry install --no-interaction --all-extras
145151
#----------------------------------------------
@@ -191,6 +197,7 @@ jobs:
191197
- name: Install Poetry
192198
uses: snok/install-poetry@v1
193199
with:
200+
version: "2.2.1"
194201
virtualenvs-create: true
195202
virtualenvs-in-project: true
196203
installer-parallel: true
@@ -243,6 +250,7 @@ jobs:
243250
- name: Install Poetry
244251
uses: snok/install-poetry@v1
245252
with:
253+
version: "2.2.1"
246254
virtualenvs-create: true
247255
virtualenvs-in-project: true
248256
installer-parallel: true

‎.github/workflows/daily-telemetry-e2e.yml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ jobs:
4343
- name: Install Poetry
4444
uses: snok/install-poetry@v1
4545
with:
46+
version: "2.2.1"
4647
virtualenvs-create: true
4748
virtualenvs-in-project: true
4849
installer-parallel: true
@@ -60,6 +61,10 @@ jobs:
6061
#----------------------------------------------
6162
# install dependencies if cache does not exist
6263
#----------------------------------------------
64+
- name: Install Kerberos system dependencies
65+
run: |
66+
sudo apt-get update
67+
sudo apt-get install -y libkrb5-dev
6368
- name: Install dependencies
6469
run: poetry install --no-interaction --all-extras
6570

‎.github/workflows/integration.yml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ jobs:
3333
- name: Install Poetry
3434
uses: snok/install-poetry@v1
3535
with:
36+
version: "2.2.1"
3637
virtualenvs-create: true
3738
virtualenvs-in-project: true
3839
installer-parallel: true
@@ -49,6 +50,10 @@ jobs:
4950
#----------------------------------------------
5051
# install dependencies if cache does not exist
5152
#----------------------------------------------
53+
- name: Install Kerberos system dependencies
54+
run: |
55+
sudo apt-get update
56+
sudo apt-get install -y libkrb5-dev
5257
- name: Install dependencies
5358
run: poetry install --no-interaction --all-extras
5459
#----------------------------------------------

‎.github/workflows/publish-manual.yml‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,10 +31,19 @@ jobs:
3131
- name: Install Poetry
3232
uses: snok/install-poetry@v1 # Install Poetry, the Python package manager
3333
with:
34+
version: "2.2.1"
3435
virtualenvs-create: true
3536
virtualenvs-in-project: true
3637
installer-parallel: true
3738

39+
#----------------------------------------------
40+
# Step 3.5: Install Kerberos system dependencies
41+
#----------------------------------------------
42+
- name: Install Kerberos system dependencies
43+
run: |
44+
sudo apt-get update
45+
sudo apt-get install -y libkrb5-dev
46+
3847
# #----------------------------------------------
3948
# # Step 4: Load cached virtual environment (if available)
4049
# #----------------------------------------------

‎.github/workflows/publish-test.yml‎

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ jobs:
2121
- name: Install Poetry
2222
uses: snok/install-poetry@v1
2323
with:
24+
version: "2.2.1"
2425
virtualenvs-create: true
2526
virtualenvs-in-project: true
2627
installer-parallel: true
@@ -36,6 +37,10 @@ jobs:
3637
#----------------------------------------------
3738
# install dependencies if cache does not exist
3839
#----------------------------------------------
40+
- name: Install Kerberos system dependencies
41+
run: |
42+
sudo apt-get update
43+
sudo apt-get install -y libkrb5-dev
3944
- name: Install dependencies
4045
if: steps.cached-poetry-dependencies.outputs.cache-hit != 'true'
4146
run: poetry install --no-interaction --no-root
@@ -54,11 +59,17 @@ jobs:
5459
- name: Update pyproject.toml
5560
run: poetry version ${{ steps.version.outputs.major-version }}.${{ steps.version.outputs.minor-version }}.dev$(date +%s)
5661
#----------------------------------------------
62+
# Build the package (before publish action)
63+
#----------------------------------------------
64+
- name: Build package
65+
run: poetry build
66+
#----------------------------------------------
67+
# Configure test-pypi repository
68+
#----------------------------------------------
69+
- name: Configure test-pypi repository
70+
run: poetry config repositories.testpypi https://test.pypi.org/legacy/
71+
#----------------------------------------------
5772
# Attempt push to test-pypi
5873
#----------------------------------------------
59-
- name: Build and publish to pypi
60-
uses: JRubics/poetry-publish@v1.10
61-
with:
62-
pypi_token: ${{ secrets.TEST_PYPI_TOKEN }}
63-
repository_name: "testpypi"
64-
repository_url: "https://test.pypi.org/legacy/"
74+
- name: Publish to test-pypi
75+
run: poetry publish --username __token__ --password ${{ secrets.TEST_PYPI_TOKEN }} --repository testpypi

‎.github/workflows/publish.yml‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ jobs:
2323
- name: Install Poetry
2424
uses: snok/install-poetry@v1
2525
with:
26+
version: "2.2.1"
2627
virtualenvs-create: true
2728
virtualenvs-in-project: true
2829
installer-parallel: true
@@ -38,6 +39,10 @@ jobs:
3839
#----------------------------------------------
3940
# install dependencies if cache does not exist
4041
#----------------------------------------------
42+
- name: Install Kerberos system dependencies
43+
run: |
44+
sudo apt-get update
45+
sudo apt-get install -y libkrb5-dev
4146
- name: Install dependencies
4247
if: steps.cached-poetry-dependencies.outputs.cache-hit != 'true'
4348
run: poetry install --no-interaction --no-root
@@ -56,9 +61,12 @@ jobs:
5661
- name: Update pyproject.toml
5762
run: poetry version ${{ steps.version.outputs.current-version }}
5863
#----------------------------------------------
59-
# Attempt push to test-pypi
64+
# Build the package (before publish)
6065
#----------------------------------------------
61-
- name: Build and publish to pypi
62-
uses: JRubics/poetry-publish@v1.10
63-
with:
64-
pypi_token: ${{ secrets.PROD_PYPI_TOKEN }}
66+
- name: Build package
67+
run: poetry build
68+
#----------------------------------------------
69+
# Publish to pypi
70+
#----------------------------------------------
71+
- name: Publish to pypi
72+
run: poetry publish --username __token__ --password ${{ secrets.PROD_PYPI_TOKEN }}

‎src/databricks/sql/auth/retry.py‎

Lines changed: 7 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -373,6 +373,13 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]:
373373
if status_code == 403:
374374
return False, "403 codes are not retried"
375375

376+
# Request failed with 404. Don't retry for any command type.
377+
if status_code == 404:
378+
return (
379+
False,
380+
"Received 404 - NOT_FOUND. The requested resource does not exist.",
381+
)
382+
376383
# Request failed and server said NotImplemented. This isn't recoverable. Don't retry.
377384
if status_code == 501:
378385
return False, "Received code 501 from server."
@@ -381,33 +388,6 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]:
381388
if not self._is_method_retryable(method):
382389
return False, "Only POST requests are retried"
383390

384-
# Request failed with 404 and was a GetOperationStatus. This is not recoverable. Don't retry.
385-
if status_code == 404 and self.command_type == CommandType.GET_OPERATION_STATUS:
386-
return (
387-
False,
388-
"GetOperationStatus received 404 code from Databricks. Operation was canceled.",
389-
)
390-
391-
# Request failed with 404 because CloseSession returns 404 if you repeat the request.
392-
if (
393-
status_code == 404
394-
and self.command_type == CommandType.CLOSE_SESSION
395-
and len(self.history) > 0
396-
):
397-
raise SessionAlreadyClosedError(
398-
"CloseSession received 404 code from Databricks. Session is already closed."
399-
)
400-
401-
# Request failed with 404 because CloseOperation returns 404 if you repeat the request.
402-
if (
403-
status_code == 404
404-
and self.command_type == CommandType.CLOSE_OPERATION
405-
and len(self.history) > 0
406-
):
407-
raise CursorAlreadyClosedError(
408-
"CloseOperation received 404 code from Databricks. Cursor is already closed."
409-
)
410-
411391
# Request failed, was an ExecuteStatement and the command may have reached the server
412392
if (
413393
self.command_type == CommandType.EXECUTE_STATEMENT

‎tests/e2e/common/retry_test_mixins.py‎

Lines changed: 11 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -278,7 +278,7 @@ def test_retry_max_count_not_exceeded(self, mock_send_telemetry, extra_params):
278278
THEN the connector issues six request (original plus five retries)
279279
before raising an exception
280280
"""
281-
with mocked_server_response(status=404) as mock_obj:
281+
with mocked_server_response(status=429, headers={"Retry-After": "0"}) as mock_obj:
282282
with pytest.raises(MaxRetryError) as cm:
283283
extra_params = {**extra_params, **self._retry_policy}
284284
with self.connection(extra_params=extra_params) as conn:
@@ -467,22 +467,21 @@ def test_retry_safe_execute_statement_retry_condition(self, extra_params):
467467
)
468468
def test_retry_abort_close_session_on_404(self, extra_params, caplog):
469469
"""GIVEN the connector sends a CloseSession command
470-
WHEN server sends a 404 (which is normally retried)
471-
THEN nothing is retried because 404 means the session already closed
470+
WHEN server sends a 404 (which is not retried since commit 41b28159)
471+
THEN nothing is retried because 404 is globally non-retryable
472472
"""
473473

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

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

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

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

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

523519
@pytest.mark.parametrize(
524520
"extra_params",

‎tests/e2e/common/staging_ingestion_tests.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,7 @@ def test_staging_ingestion_life_cycle(self, ingestion_user):
8181
# GET after REMOVE should fail
8282

8383
with pytest.raises(
84-
Error, match="too many 404 error responses"
84+
Error, match="Staging operation over HTTP was unsuccessful: 404"
8585
):
8686
cursor = conn.cursor()
8787
query = f"GET 'stage://tmp/{ingestion_user}/tmp/11/16/file1.csv' TO '{new_temp_path}'"

0 commit comments

Comments
 (0)