fix: can instantiate connectors without callbacks/secrets-keeper (TCTC-10081)

tests/google_sheets/test_google_sheets.py

@@ -337,10 +337,11 @@ def test_can_instantiate_without_retrieve_token_callback():
         name="test_connector",
         auth_id="test_auth_id",
     )
-    assert gsheet_connector
+    assert gsheet_connector.auth_id.get_secret_value() == "test_auth_id"
+    assert gsheet_connector.retrieve_token is None
 
 
-def test_raise_when_trying_to_retrieve_token_if_callable_missing():
+def test_raises_when_trying_to_retrieve_token_if_callable_missing():
     gsheet_connector = GoogleSheetsConnector(
         name="test_connector",
         auth_id="test_auth_id",

tests/oauth2_connector/test_oauth2connector.py

@@ -11,6 +11,7 @@
     NoOAuth2RefreshToken,
     OAuth2Connector,
     OAuth2ConnectorConfig,
+    SecretKeeperMissingError,
 )
 from toucan_connectors.snowflake_oauth2.snowflake_oauth2_connector import SnowflakeoAuth2Connector
 from toucan_connectors.toucan_connector import get_oauth2_configuration
@@ -248,18 +249,18 @@ def test_get_refresh_token(mocker, oauth2_connector):
     assert token == "bla"
 
 
-def test_raise_error_if_secret_keeper_not_set(oauth2_connector_without_secret_keeper: OAuth2Connector):
-    with pytest.raises(ValueError):
+def test_raises_exception_if_secret_keeper_not_set(oauth2_connector_without_secret_keeper: OAuth2Connector):
+    with pytest.raises(SecretKeeperMissingError):
         oauth2_connector_without_secret_keeper.get_access_token()
 
-    with pytest.raises(ValueError):
+    with pytest.raises(SecretKeeperMissingError):
         oauth2_connector_without_secret_keeper.retrieve_tokens(authorization_response="")
 
-    with pytest.raises(ValueError):
+    with pytest.raises(SecretKeeperMissingError):
         oauth2_connector_without_secret_keeper.build_authorization_url()
 
-    with pytest.raises(ValueError):
+    with pytest.raises(SecretKeeperMissingError):
         oauth2_connector_without_secret_keeper.get_access_data()
 
-    with pytest.raises(ValueError):
+    with pytest.raises(SecretKeeperMissingError):
         oauth2_connector_without_secret_keeper.get_refresh_token()

toucan_connectors/google_sheets/google_sheets_connector.py

@@ -17,10 +17,10 @@
     getLogger(__name__).warning(f"Missing dependencies for {__name__}: {exc}")
     CONNECTOR_OK = False
 
-from pydantic import Field, PrivateAttr, create_model
+from pydantic import Field, create_model
 from pydantic.json_schema import DEFAULT_REF_TEMPLATE, GenerateJsonSchema, JsonSchemaMode
 
-from toucan_connectors.common import ConnectorStatus
+from toucan_connectors.common import UI_HIDDEN, ConnectorStatus
 from toucan_connectors.toucan_connector import (
     PlainJsonSecretStr,
     ToucanConnector,
@@ -100,26 +100,22 @@ class GoogleSheetsConnector(ToucanConnector, data_source_model=GoogleSheetsDataS
     _auth_flow = "managed_oauth2"
     _managed_oauth_service_id = "google-sheets"
     _oauth_trigger = "retrieve_token"
-    _retrieve_token: Callable[[str, str], str] | None = PrivateAttr()
 
+    retrieve_token: Callable[[str, str], str] | None = Field(None, **UI_HIDDEN)
     auth_id: PlainJsonSecretStr = None
 
-    def __init__(self, retrieve_token: Callable[[str, str], str] | None = None, *args, **kwargs):
-        super().__init__(**kwargs)
-        self._retrieve_token = retrieve_token  # Could be async
-
     def _call_retrieve_token(self) -> str:
-        """Retrieve the access token for Google Sheets
+        """Retrieves the access token for Google Sheets
 
-        Raise an GoogleSheetsInvalidConfiguration if retrieve_token callback is not set
-        Raise an GoogleSheetsRetrieveTokenError if an error is encountered while retrieving the token
+        Raises a GoogleSheetsInvalidConfiguration if retrieve_token callback is not set
+        Raises a GoogleSheetsRetrieveTokenError if an error is encountered while retrieving the token
         """
-        if self._retrieve_token is None:
+        if self.retrieve_token is None:
             raise GoogleSheetsInvalidConfiguration(
                 "Retrieve token callback function is not configured. Please provide it at instantiation."
             )
         try:
-            return self._retrieve_token(self._managed_oauth_service_id, self.auth_id.get_secret_value())
+            return self.retrieve_token(self._managed_oauth_service_id, self.auth_id.get_secret_value())
         except Exception as exc:
             raise GoogleSheetsRetrieveTokenError(str(exc)) from exc
 

toucan_connectors/oauth2_connector/oauth2connector.py

@@ -40,6 +40,10 @@ def load(self, key: str, **kwargs) -> Any:
         """
 
 
+class SecretKeeperMissingError(Exception):
+    """Raised when secret_keeper is not set on oauth2 connector"""
+
+
 class OAuth2ConnectorConfig(BaseModel):
     client_id: str
     client_secret: SecretStr
@@ -66,13 +70,15 @@ def __init__(
         self.token_url = token_url
         self.redirect_uri = redirect_uri
 
+    def _check_secret_keeper_exists(self) -> None:
+        if self.secrets_keeper is None:
+            raise SecretKeeperMissingError("Secret keeper is not set on oauth2 connector.")
+
     def build_authorization_url(self, **kwargs) -> str:
         """Build an authorization request that will be sent to the client."""
         from authlib.common.security import generate_token
 
-        if self.secrets_keeper is None:
-            raise ValueError("Secret Keeper not initialized.")
-
+        self._check_secret_keeper_exists()
         client = oauth_client(
             client_id=self.config.client_id,
             client_secret=self.config.client_secret.get_secret_value(),
@@ -86,9 +92,7 @@ def build_authorization_url(self, **kwargs) -> str:
         return uri
 
     def retrieve_tokens(self, authorization_response: str, **kwargs):
-        if self.secrets_keeper is None:
-            raise ValueError("Secret Keeper not initialized.")
-
+        self._check_secret_keeper_exists()
         url = url_parse.urlparse(authorization_response)
         url_params = url_parse.parse_qs(url.query)
         client = oauth_client(
@@ -117,9 +121,7 @@ def get_access_token(self) -> str:
         instance_url parameters are return by service, better to use it
         new method get_access_data return all information to connect (secret and instance_url)
         """
-        if self.secrets_keeper is None:
-            raise ValueError("Secret Keeper not initialized.")
-
+        self._check_secret_keeper_exists()
         token = self.secrets_keeper.load(self.auth_flow_id)
 
         if "expires_at" in token:
@@ -148,9 +150,7 @@ def get_access_data(self):
         Returns the access_token to use to access resources
         If necessary, this token will be refreshed
         """
-        if self.secrets_keeper is None:
-            raise ValueError("Secret Keeper not initialized.")
-
+        self._check_secret_keeper_exists()
         access_data = self.secrets_keeper.load(self.auth_flow_id)
 
         logging.getLogger(__name__).debug("Refresh and get access data")
@@ -177,9 +177,7 @@ def get_refresh_token(self) -> str:
         """
         Return the refresh token, used to obtain an access token
         """
-        if self.secrets_keeper is None:
-            raise ValueError("Secret Keeper not initialized.")
-
+        self._check_secret_keeper_exists()
         return self.secrets_keeper.load(self.auth_flow_id)["refresh_token"]