fix(GSheets OAuth2): Re-add UnauthenticatedError (#39785)

superset/db_engine_specs/base.py

@@ -590,7 +590,9 @@ class BaseEngineSpec:  # pylint: disable=too-many-public-methods
     # Driver-specific params to be included in the `get_oauth2_token` request body
     oauth2_additional_token_request_params: dict[str, Any] = {}
     # Driver-specific exception that should be mapped to OAuth2RedirectError
-    oauth2_exception = OAuth2RedirectError
+    oauth2_exception: type[Exception] | tuple[type[Exception], ...] = (
+        OAuth2RedirectError
+    )
 
     # Does the query id related to the connection?
     # The default value is True, which means that the query id is determined when

superset/db_engine_specs/gsheets.py

@@ -31,6 +31,7 @@ from marshmallow import fields, Schema
 from marshmallow.exceptions import ValidationError
 from requests import Session
 from shillelagh.adapters.api.gsheets.lib import SCOPES
+from shillelagh.exceptions import UnauthenticatedError
 from sqlalchemy.engine import create_engine
 from sqlalchemy.engine.reflection import Inspector
 from sqlalchemy.engine.url import URL
@@ -40,7 +41,7 @@ from superset.databases.schemas import encrypted_field_properties, EncryptedStri
 from superset.db_engine_specs.base import DatabaseCategory
 from superset.db_engine_specs.shillelagh import ShillelaghEngineSpec
 from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
-from superset.exceptions import SupersetException
+from superset.exceptions import OAuth2TokenRefreshError, SupersetException
 from superset.utils import json
 from superset.utils.oauth2 import get_oauth2_access_token
 
@@ -151,6 +152,7 @@ class GSheetsEngineSpec(ShillelaghEngineSpec):
         "https://accounts.google.com/o/oauth2/v2/auth"
     )
     oauth2_token_request_uri = "https://oauth2.googleapis.com/token"  # noqa: S105
+    oauth2_exception = (UnauthenticatedError, OAuth2TokenRefreshError)
 
     @classmethod
     def get_oauth2_authorization_uri(

tests/unit_tests/db_engine_specs/test_gsheets.py

@@ -24,6 +24,7 @@ import pandas as pd
 import pytest
 from pytest_mock import MockerFixture
 from requests.exceptions import HTTPError
+from shillelagh.exceptions import UnauthenticatedError
 from sqlalchemy.engine.url import make_url
 
 from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
@@ -789,6 +790,36 @@ def test_needs_oauth2_with_other_error(mocker: MockerFixture) -> None:
     assert GSheetsEngineSpec.needs_oauth2(ex) is False
 
 
+def test_needs_oauth2_with_shillelagh_unauthenticated_error(
+    mocker: MockerFixture,
+) -> None:
+    """
+    Test that needs_oauth2 returns True when UnauthenticatedError is raised.
+    """
+    from superset.db_engine_specs.gsheets import GSheetsEngineSpec
+
+    g = mocker.patch("superset.db_engine_specs.gsheets.g")
+    g.user = mocker.MagicMock()
+
+    ex = UnauthenticatedError("Token has been revoked")
+    assert GSheetsEngineSpec.needs_oauth2(ex) is True
+
+
+def test_needs_oauth2_with_unrelated_exception_type(
+    mocker: MockerFixture,
+) -> None:
+    """
+    Test that an unrelated exception type (with no matching message) returns
+    False.
+    """
+    from superset.db_engine_specs.gsheets import GSheetsEngineSpec
+
+    g = mocker.patch("superset.db_engine_specs.gsheets.g")
+    g.user = mocker.MagicMock()
+
+    assert GSheetsEngineSpec.needs_oauth2(ValueError("unrelated")) is False
+
+
 def test_get_oauth2_fresh_token_success(
     mocker: MockerFixture,
     oauth2_config: OAuth2ClientConfig,