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

.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
       #----------------------------------------------

.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

.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
 

.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
       #----------------------------------------------

.github/workflows/publish-manual.yml

@@ -31,6 +31,7 @@ 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

.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

.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

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

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