[PR-verification] oauth raw expires_in from overrides

airbyte_cdk/sources/declarative/auth/oauth.py

@@ -3,6 +3,7 @@
 #
 
 import logging
+from contextlib import suppress
 from dataclasses import InitVar, dataclass, field
 from datetime import datetime, timedelta
 from typing import Any, List, Mapping, Optional, Tuple, Union
@@ -262,6 +263,10 @@ def _has_access_token_been_initialized(self) -> bool:
         return self._access_token is not None
 
     def set_token_expiry_date(self, value: AirbyteDateTime) -> None:
+        # Overrides of `refresh_access_token` written before 6.45.5 return the raw `expires_in`.
+        if not isinstance(value, datetime):
+            with suppress(ValueError, OverflowError):
+                value = self._parse_token_expiration_date(value)
         self._token_expiry_date = value
 
     def get_assertion_name(self) -> str:

airbyte_cdk/sources/streams/http/requests_native_auth/abstract_oauth.py

@@ -6,7 +6,7 @@
 import re
 import threading
 from abc import abstractmethod
-from datetime import timedelta
+from datetime import datetime, timedelta
 from json import JSONDecodeError
 from typing import Any, List, Mapping, MutableMapping, Optional, Tuple, Union
 
@@ -518,10 +518,13 @@ def _ensure_access_token_in_response(self, response_data: Mapping[str, Any]) ->
 
     def _parse_token_expiration_date(self, value: Union[str, int]) -> AirbyteDateTime:
         """
-        Parse a string or integer token expiration date into a datetime object
+        Parse a string or integer token expiration date into a datetime object. A datetime, which
+        `refresh_access_token` returns since 6.45.5, is passed through for setters written earlier.
 
         :return: expiration datetime
         """
+        if isinstance(value, datetime):
+            return AirbyteDateTime.from_datetime(value)
         if self.token_expiry_is_time_of_expiration:
             if not self.token_expiry_date_format:
                 raise ValueError(

airbyte_cdk/sources/streams/http/requests_native_auth/oauth.py

@@ -2,7 +2,8 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
-from datetime import timedelta
+from contextlib import suppress
+from datetime import datetime, timedelta
 from typing import Any, List, Mapping, Optional, Sequence, Tuple, Union
 
 import dpath
@@ -126,6 +127,10 @@ def get_token_expiry_date(self) -> AirbyteDateTime:
         return self._token_expiry_date
 
     def set_token_expiry_date(self, value: AirbyteDateTime) -> None:
+        # Overrides of `refresh_access_token` written before 6.45.5 return the raw `expires_in`.
+        if not isinstance(value, datetime):
+            with suppress(ValueError, OverflowError):
+                value = self._parse_token_expiration_date(value)
         self._token_expiry_date = value
 
     @property
@@ -316,6 +321,15 @@ def set_token_expiry_date(self, new_token_expiry_date: AirbyteDateTime) -> None:
         Args:
             new_token_expiry_date (AirbyteDateTime): The new expiry date for the token.
         """
+        if not (
+            isinstance(new_token_expiry_date, datetime)
+            or self.token_expiry_is_time_of_expiration
+            or self.token_expiry_date_format
+        ):
+            # Overrides of `refresh_access_token` written before 6.45.5 return the raw `expires_in`
+            # in seconds. Other values are stored as-is and parsed on read.
+            with suppress(ValueError, OverflowError):
+                new_token_expiry_date = self._parse_token_expiration_date(new_token_expiry_date)
         self._set_config_value_by_path(
             self._token_expiry_date_config_path, str(new_token_expiry_date)
         )

unit_tests/sources/declarative/auth/test_oauth.py

@@ -623,6 +623,74 @@ def test_set_token_expiry_date_no_format(self, mocker, expires_in_response, next
             assert "access_token" == token
             assert oauth.get_token_expiry_date() == ab_datetime_parse(next_day)
 
+    @freezegun.freeze_time("2022-01-01")
+    def test_custom_component_returning_raw_expires_in(self):
+        """Custom components written before 6.45.5 (e.g. source-okta) return the raw `expires_in`."""
+
+        class CustomOauth2Authenticator(DeclarativeOauth2Authenticator):
+            def refresh_access_token(self):
+                return "access_token", 948
+
+        oauth = CustomOauth2Authenticator(
+            token_refresh_endpoint="{{ config['refresh_endpoint'] }}",
+            client_id="{{ config['client_id'] }}",
+            client_secret="{{ config['client_secret'] }}",
+            refresh_token="{{ parameters['refresh_token'] }}",
+            config=config,
+            parameters=parameters,
+        )
+
+        assert oauth.get_access_token() == "access_token"
+        assert oauth.get_access_token() == "access_token"
+        assert oauth.get_token_expiry_date() == ab_datetime_now() + timedelta(seconds=948)
+
+    def test_custom_component_returning_unparseable_expires_in_stores_it_as_is(self):
+        class CustomOauth2Authenticator(DeclarativeOauth2Authenticator):
+            def refresh_access_token(self):
+                return "access_token", None
+
+        oauth = CustomOauth2Authenticator(
+            token_refresh_endpoint="{{ config['refresh_endpoint'] }}",
+            client_id="{{ config['client_id'] }}",
+            client_secret="{{ config['client_secret'] }}",
+            refresh_token="{{ parameters['refresh_token'] }}",
+            config=config,
+            parameters=parameters,
+        )
+
+        assert oauth.get_access_token() == "access_token"
+        assert oauth.get_token_expiry_date() is None
+
+    @freezegun.freeze_time("2022-01-01")
+    def test_custom_component_parse_token_expiration_date_override_only_gets_raw_values(
+        self, mocker
+    ):
+        parsed = []
+
+        class CustomOauth2Authenticator(DeclarativeOauth2Authenticator):
+            def _parse_token_expiration_date(self, value):
+                parsed.append(value)
+                return ab_datetime_now() + timedelta(seconds=int(value))
+
+        oauth = CustomOauth2Authenticator(
+            token_refresh_endpoint="{{ config['refresh_endpoint'] }}",
+            client_id="{{ config['client_id'] }}",
+            client_secret="{{ config['client_secret'] }}",
+            refresh_token="{{ parameters['refresh_token'] }}",
+            config=config,
+            parameters=parameters,
+        )
+        resp.status_code = 200
+        mocker.patch.object(
+            resp, "json", return_value={"access_token": "access_token", "expires_in": 948}
+        )
+        mocker.patch.object(requests, "request", side_effect=mock_request, autospec=True)
+
+        assert oauth.get_access_token() == "access_token"
+        assert oauth.get_access_token() == "access_token"
+        assert parsed == [948]
+        assert oauth.get_token_expiry_date() == ab_datetime_now() + timedelta(seconds=948)
+
     @freezegun.freeze_time("2022-01-01")
     def test_profile_assertion(self, mocker):
         with HttpMocker() as http_mocker: