diff --git a/airbyte_cdk/sources/declarative/auth/oauth.py b/airbyte_cdk/sources/declarative/auth/oauth.py index 9f4826edeb..3b70fe8228 100644 --- a/airbyte_cdk/sources/declarative/auth/oauth.py +++ b/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: diff --git a/airbyte_cdk/sources/streams/http/requests_native_auth/abstract_oauth.py b/airbyte_cdk/sources/streams/http/requests_native_auth/abstract_oauth.py index ff842c9ce4..8fa300b176 100644 --- a/airbyte_cdk/sources/streams/http/requests_native_auth/abstract_oauth.py +++ b/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( diff --git a/airbyte_cdk/sources/streams/http/requests_native_auth/oauth.py b/airbyte_cdk/sources/streams/http/requests_native_auth/oauth.py index 59e46bdfac..06643ba82c 100644 --- a/airbyte_cdk/sources/streams/http/requests_native_auth/oauth.py +++ b/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 @@ -199,7 +204,7 @@ def __init__( access_token_config_path (Sequence[str]): Dpath to the access_token field in the connector configuration. Defaults to ("credentials", "access_token"). refresh_token_config_path (Sequence[str]): Dpath to the refresh_token field in the connector configuration. Defaults to ("credentials", "refresh_token"). token_expiry_date_config_path (Sequence[str]): Dpath to the token_expiry_date field in the connector configuration. Defaults to ("credentials", "token_expiry_date"). - token_expiry_date_format (Optional[str]): Date format of the token expiry date field (set by expires_in_name). If not specified the token expiry date is interpreted as number of seconds until expiration. + token_expiry_date_format (Optional[str]): Date format of the token expiry date field (set by expires_in_name). If specified the token expiry date is interpreted as time of expiration, as with token_expiry_is_time_of_expiration, otherwise as number of seconds until expiration. token_expiry_is_time_of_expiration bool: set True it if expires_in is returned as time of expiration instead of the number seconds until expiration message_repository (MessageRepository): the message repository used to emit logs on HTTP requests and control message on config update """ @@ -238,7 +243,8 @@ def __init__( grant_type_name=self._grant_type_name, grant_type=grant_type, token_expiry_date_format=token_expiry_date_format, - token_expiry_is_time_of_expiration=token_expiry_is_time_of_expiration, + token_expiry_is_time_of_expiration=token_expiry_is_time_of_expiration + or bool(token_expiry_date_format), refresh_token_error_status_codes=refresh_token_error_status_codes, refresh_token_error_key=refresh_token_error_key, refresh_token_error_values=refresh_token_error_values, @@ -316,6 +322,13 @@ 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 + ): + # 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) ) diff --git a/unit_tests/sources/declarative/auth/test_oauth.py b/unit_tests/sources/declarative/auth/test_oauth.py index ddd96a5474..71b8328b33 100644 --- a/unit_tests/sources/declarative/auth/test_oauth.py +++ b/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: diff --git a/unit_tests/sources/streams/http/requests_native_auth/test_requests_native_auth.py b/unit_tests/sources/streams/http/requests_native_auth/test_requests_native_auth.py index bbc6a17086..923292cc25 100644 --- a/unit_tests/sources/streams/http/requests_native_auth/test_requests_native_auth.py +++ b/unit_tests/sources/streams/http/requests_native_auth/test_requests_native_auth.py @@ -88,6 +88,19 @@ def test_multiple_token_authenticator(): assert {"Authorization": "Bearer token1"} == header3 +class _ComparableWithDatetime: + """Like `arrow.Arrow`: not a `datetime`, but comparable with one.""" + + def __init__(self, value): + self.value = value + + def __lt__(self, other): + return self.value < other + + def __gt__(self, other): + return self.value > other + + @freezegun.freeze_time("2022-01-01") class TestOauth2Authenticator: """ @@ -554,6 +567,154 @@ def test_parse_refresh_token_lifespan( assert expires_datetime == expected_token_expiry_date assert token == "access_token" + @pytest.mark.parametrize( + "expires_in, token_expiry_date_format", + [ + (3600, None), + ("3600", None), + (3600.0, None), + (1640998800, "%s"), + (AirbyteDateTime(year=2022, month=1, day=1, hour=1), None), + ], + ids=["seconds", "string_of_seconds", "float_seconds", "epoch", "already_parsed"], + ) + def test_get_access_token_parses_raw_expires_in_from_overridden_refresh_access_token( + self, mocker, expires_in, token_expiry_date_format + ): + """Overrides written before 6.45.5 return the raw `expires_in` from `refresh_access_token`.""" + oauth = Oauth2Authenticator( + token_refresh_endpoint=TestOauth2Authenticator.refresh_endpoint, + client_id=TestOauth2Authenticator.client_id, + client_secret=TestOauth2Authenticator.client_secret, + refresh_token=TestOauth2Authenticator.refresh_token, + token_expiry_date_format=token_expiry_date_format, + token_expiry_is_time_of_expiration=bool(token_expiry_date_format), + ) + mocked_refresh = mocker.patch.object( + Oauth2Authenticator, "refresh_access_token", return_value=("access_token", expires_in) + ) + + assert oauth.get_access_token() == "access_token" + assert oauth.get_access_token() == "access_token" + mocked_refresh.assert_called_once() + assert oauth.get_token_expiry_date() == AirbyteDateTime(year=2022, month=1, day=1, hour=1) + + @pytest.mark.parametrize("override_refresh", [True, False], ids=["raw_override", "cdk_refresh"]) + def test_set_token_expiry_date_override_that_parses_keeps_working( + self, mocker, requests_mock, override_refresh + ): + """Setters that parse, as `set_token_expiry_date` did before 6.45.5, get raw values from + older `refresh_access_token` overrides and datetimes from the CDK one.""" + + class ParsingSetterAuthenticator(Oauth2Authenticator): + def set_token_expiry_date(self, value): + self._token_expiry_date = self._parse_token_expiration_date(value) + + oauth = ParsingSetterAuthenticator( + f"https://{TestOauth2Authenticator.refresh_endpoint}", + TestOauth2Authenticator.client_id, + TestOauth2Authenticator.client_secret, + TestOauth2Authenticator.refresh_token, + ) + requests_mock.post( + f"https://{TestOauth2Authenticator.refresh_endpoint}", + json={"access_token": "access_token", "expires_in": 3600}, + ) + if override_refresh: + mocker.patch.object( + ParsingSetterAuthenticator, + "refresh_access_token", + return_value=("access_token", 3600), + ) + + assert oauth.get_access_token() == "access_token" + assert oauth.get_token_expiry_date() == ab_datetime_now() + timedelta(seconds=3600) + + def test_set_token_expiry_date_override_that_takes_raw_expires_in_keeps_working(self): + class RawSetterAuthenticator(Oauth2Authenticator): + def refresh_access_token(self): + return "access_token", 3600 + + def set_token_expiry_date(self, value): + self._token_expiry_date = ab_datetime_now() + timedelta(seconds=value) + + oauth = RawSetterAuthenticator( + TestOauth2Authenticator.refresh_endpoint, + TestOauth2Authenticator.client_id, + TestOauth2Authenticator.client_secret, + TestOauth2Authenticator.refresh_token, + ) + + 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=3600) + + def test_get_access_token_override_passing_raw_expires_in_to_set_token_expiry_date(self): + class LegacyAuthenticator(Oauth2Authenticator): + def get_access_token(self): + if self.token_has_expired(): + token, expires_in = self.refresh_access_token() + self.access_token = token + self.set_token_expiry_date(expires_in) + return self.access_token + + def refresh_access_token(self): + return "access_token", 3600 + + oauth = LegacyAuthenticator( + TestOauth2Authenticator.refresh_endpoint, + TestOauth2Authenticator.client_id, + TestOauth2Authenticator.client_secret, + TestOauth2Authenticator.refresh_token, + ) + + 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=3600) + + def test_parse_token_expiration_date_override_only_gets_raw_values(self, requests_mock): + parsed = [] + + class SecondsParserAuthenticator(Oauth2Authenticator): + def _parse_token_expiration_date(self, value): + parsed.append(value) + return ab_datetime_now() + timedelta(seconds=int(value)) + + oauth = SecondsParserAuthenticator( + f"https://{TestOauth2Authenticator.refresh_endpoint}", + TestOauth2Authenticator.client_id, + TestOauth2Authenticator.client_secret, + TestOauth2Authenticator.refresh_token, + ) + requests_mock.post( + f"https://{TestOauth2Authenticator.refresh_endpoint}", + json={"access_token": "access_token", "expires_in": 3600}, + ) + + assert oauth.get_access_token() == "access_token" + assert oauth.get_access_token() == "access_token" + assert parsed == [3600] + assert oauth.get_token_expiry_date() == ab_datetime_now() + timedelta(seconds=3600) + + @pytest.mark.parametrize( + "expires_in", + [_ComparableWithDatetime(AirbyteDateTime(year=2022, month=1, day=2)), None, 1641081600000], + ids=["comparable_object", "none", "epoch_millis"], + ) + def test_get_access_token_stores_an_unparseable_expiry_as_is(self, mocker, expires_in): + oauth = Oauth2Authenticator( + TestOauth2Authenticator.refresh_endpoint, + TestOauth2Authenticator.client_id, + TestOauth2Authenticator.client_secret, + TestOauth2Authenticator.refresh_token, + ) + mocker.patch.object( + Oauth2Authenticator, "refresh_access_token", return_value=("access_token", expires_in) + ) + + assert oauth.get_access_token() == "access_token" + assert oauth.get_token_expiry_date() is expires_in + @pytest.mark.usefixtures("mock_sleep") @pytest.mark.parametrize("error_code", (429, 500, 502, 504)) def test_refresh_access_token_retry(self, error_code, requests_mock): @@ -1036,6 +1197,7 @@ def test_init(self, connector_config): ("number_of_seconds", 42, None, "2022-12-31T00:00:42+00:00"), ("string_of_seconds", "42", None, "2022-12-31T00:00:42+00:00"), ("date_format", "2023-04-04", "YYYY-MM-DD", "2023-04-04T00:00:00+00:00"), + ("epoch_format", 1672531200, "%s", "2023-01-01T00:00:00+00:00"), ], ) def test_given_no_message_repository_get_access_token( @@ -1054,7 +1216,6 @@ def test_given_no_message_repository_get_access_token( client_id=connector_config["credentials"]["client_id"], client_secret=connector_config["credentials"]["client_secret"], token_expiry_date_format=expiry_date_format, - token_expiry_is_time_of_expiration=bool(expiry_date_format), ) # Mock the response from the refresh token endpoint @@ -1198,6 +1359,142 @@ def test_refresh_access_token(self, mocker, connector_config): "new_refresh_token", ) + def test_get_access_token_parses_raw_expires_in_from_overridden_refresh_access_token( + self, mocker, connector_config + ): + """A raw `expires_in` must not be stored as-is, where it reads back as an epoch in 1970.""" + connector_config["credentials"]["token_expiry_date"] = "" + authenticator = SingleUseRefreshTokenOauth2Authenticator( + connector_config, token_refresh_endpoint="https://refresh_endpoint.com" + ) + mocked_refresh = mocker.patch.object( + SingleUseRefreshTokenOauth2Authenticator, + "refresh_access_token", + return_value=("new_access_token", 3600, "new_refresh_token"), + ) + mocker.patch.object(SingleUseRefreshTokenOauth2Authenticator, "_emit_control_message") + + for _ in range(3): + assert authenticator.get_access_token() == "new_access_token" + + mocked_refresh.assert_called_once() + assert connector_config["credentials"]["token_expiry_date"] == "2022-12-31T01:00:00+00:00" + + @pytest.mark.parametrize( + "expires_in, token_expiry_date_format", + [ + (1672531200, "%s"), + ("2023-01-01T00:00:00Z", None), + (None, None), + (1672531200000, None), + ], + ids=["epoch_with_format", "date", "none", "epoch_millis"], + ) + def test_refresh_and_set_access_token_stores_other_overridden_expires_in_as_is( + self, capsys, mocker, connector_config, expires_in, token_expiry_date_format + ): + """Values other than a raw `expires_in` in seconds are stored and emitted as they were before.""" + authenticator = SingleUseRefreshTokenOauth2Authenticator( + connector_config, + token_refresh_endpoint="https://refresh_endpoint.com", + token_expiry_date_format=token_expiry_date_format, + ) + mocker.patch.object( + SingleUseRefreshTokenOauth2Authenticator, + "refresh_access_token", + return_value=("new_access_token", expires_in, "new_refresh_token"), + ) + + authenticator.refresh_and_set_access_token() + + airbyte_message = json.loads(capsys.readouterr().out) + emitted_credentials = airbyte_message["control"]["connectorConfig"]["config"]["credentials"] + assert emitted_credentials["refresh_token"] == "new_refresh_token" + assert emitted_credentials["token_expiry_date"] == str(expires_in) + + def test_parse_token_expiration_date_override_only_gets_raw_values( + self, mocker, connector_config + ): + parsed = [] + + class SecondsParserAuthenticator(SingleUseRefreshTokenOauth2Authenticator): + def _parse_token_expiration_date(self, value): + parsed.append(value) + return ab_datetime_now() + timedelta(seconds=int(value)) + + authenticator = SecondsParserAuthenticator( + connector_config, token_refresh_endpoint="https://refresh_endpoint.com" + ) + resp.status_code = 200 + mocker.patch.object( + resp, + "json", + return_value={"access_token": "new_access_token", "expires_in": 3600}, + ) + mocker.patch.object(requests, "request", side_effect=mock_request, autospec=True) + mocker.patch.object(SingleUseRefreshTokenOauth2Authenticator, "_emit_control_message") + + authenticator.refresh_and_set_access_token() + + assert parsed == [3600] + assert connector_config["credentials"]["token_expiry_date"] == "2022-12-31T01:00:00+00:00" + + def test_set_token_expiry_date_override_that_takes_raw_expires_in_keeps_working( + self, capsys, connector_config + ): + class RawSetterAuthenticator(SingleUseRefreshTokenOauth2Authenticator): + refresh_count = 0 + + def refresh_access_token(self): + RawSetterAuthenticator.refresh_count += 1 + return "new_access_token", 3600, "new_refresh_token" + + def set_token_expiry_date(self, value): + super().set_token_expiry_date(ab_datetime_now() + timedelta(seconds=value)) + + connector_config["credentials"]["token_expiry_date"] = "" + authenticator = RawSetterAuthenticator( + connector_config, token_refresh_endpoint="https://refresh_endpoint.com" + ) + + for _ in range(3): + assert authenticator.get_access_token() == "new_access_token" + + assert RawSetterAuthenticator.refresh_count == 1 + airbyte_message = json.loads(capsys.readouterr().out) + emitted_credentials = airbyte_message["control"]["connectorConfig"]["config"]["credentials"] + assert emitted_credentials["refresh_token"] == "new_refresh_token" + assert emitted_credentials["token_expiry_date"] == "2022-12-31T01:00:00+00:00" + + def test_set_token_expiry_date_does_not_parse_a_time_of_expiration(self, connector_config): + parsed = [] + + class RecordingParserAuthenticator(SingleUseRefreshTokenOauth2Authenticator): + def _parse_token_expiration_date(self, value): + parsed.append(value) + return super()._parse_token_expiration_date(value) + + authenticator = RecordingParserAuthenticator( + connector_config, + token_refresh_endpoint="https://refresh_endpoint.com", + token_expiry_is_time_of_expiration=True, + ) + + authenticator.set_token_expiry_date(1672531200) + + assert parsed == [] + assert connector_config["credentials"]["token_expiry_date"] == "1672531200" + + def test_set_token_expiry_date_stores_a_date_string_as_is(self, connector_config): + authenticator = SingleUseRefreshTokenOauth2Authenticator( + connector_config, token_refresh_endpoint="https://refresh_endpoint.com" + ) + + authenticator.set_token_expiry_date("2023-01-01T00:00:00Z") + + assert connector_config["credentials"]["token_expiry_date"] == "2023-01-01T00:00:00Z" + assert authenticator.get_token_expiry_date() == AirbyteDateTime(year=2023, month=1, day=1) + def test_send_refresh_request_as_query_params_picks_up_rotated_refresh_token( self, mocker, connector_config ):