diff --git a/.github/workflows/code-coverage.yml b/.github/workflows/code-coverage.yml index 3c76be728..9cb68dbc9 100644 --- a/.github/workflows/code-coverage.yml +++ b/.github/workflows/code-coverage.yml @@ -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 @@ -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 #---------------------------------------------- diff --git a/.github/workflows/code-quality-checks.yml b/.github/workflows/code-quality-checks.yml index 3c368abef..cc3952920 100644 --- a/.github/workflows/code-quality-checks.yml +++ b/.github/workflows/code-quality-checks.yml @@ -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 @@ -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 @@ -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 #---------------------------------------------- @@ -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 @@ -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 diff --git a/.github/workflows/daily-telemetry-e2e.yml b/.github/workflows/daily-telemetry-e2e.yml index 3d61cf177..d60b7f5a9 100644 --- a/.github/workflows/daily-telemetry-e2e.yml +++ b/.github/workflows/daily-telemetry-e2e.yml @@ -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 @@ -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 diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index ad5369997..7fd6d98f1 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -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 @@ -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 #---------------------------------------------- diff --git a/.github/workflows/publish-manual.yml b/.github/workflows/publish-manual.yml index ecad71a29..2f2a7a4dd 100644 --- a/.github/workflows/publish-manual.yml +++ b/.github/workflows/publish-manual.yml @@ -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) # #---------------------------------------------- diff --git a/.github/workflows/publish-test.yml b/.github/workflows/publish-test.yml index 2e6359a78..ea4158934 100644 --- a/.github/workflows/publish-test.yml +++ b/.github/workflows/publish-test.yml @@ -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 @@ -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 @@ -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 diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index dde6cc2dc..b101f421c 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -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 @@ -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 @@ -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 }} diff --git a/src/databricks/sql/auth/retry.py b/src/databricks/sql/auth/retry.py index 4281883da..b0c2f497d 100755 --- a/src/databricks/sql/auth/retry.py +++ b/src/databricks/sql/auth/retry.py @@ -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." @@ -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 diff --git a/tests/e2e/common/retry_test_mixins.py b/tests/e2e/common/retry_test_mixins.py index b2350bd98..80822ba47 100755 --- a/tests/e2e/common/retry_test_mixins.py +++ b/tests/e2e/common/retry_test_mixins.py @@ -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: @@ -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", @@ -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}, ] @@ -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", diff --git a/tests/e2e/common/staging_ingestion_tests.py b/tests/e2e/common/staging_ingestion_tests.py index 73aa0a113..a88f55238 100644 --- a/tests/e2e/common/staging_ingestion_tests.py +++ b/tests/e2e/common/staging_ingestion_tests.py @@ -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}'" diff --git a/tests/e2e/common/uc_volume_tests.py b/tests/e2e/common/uc_volume_tests.py index 93e63bd28..5b4086f91 100644 --- a/tests/e2e/common/uc_volume_tests.py +++ b/tests/e2e/common/uc_volume_tests.py @@ -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}'" diff --git a/tests/unit/test_retry.py b/tests/unit/test_retry.py index 897a1d111..0d01d8675 100644 --- a/tests/unit/test_retry.py +++ b/tests/unit/test_retry.py @@ -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