From 96c0fe851375f397922b9274f487c3fd2e052e55 Mon Sep 17 00:00:00 2001 From: Anton Burnashev Date: Wed, 22 May 2024 18:02:55 +0200 Subject: [PATCH 1/3] Fix AuthConfigBase so its instances always evaluate to True in bool context; change docs to suggest direct inheritance from AuthBase --- dlt/sources/helpers/rest_client/auth.py | 6 +++++- docs/website/docs/general-usage/http/rest-client.md | 8 ++++---- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/dlt/sources/helpers/rest_client/auth.py b/dlt/sources/helpers/rest_client/auth.py index 37c0de3db1..020c63a195 100644 --- a/dlt/sources/helpers/rest_client/auth.py +++ b/dlt/sources/helpers/rest_client/auth.py @@ -38,7 +38,11 @@ class AuthConfigBase(AuthBase, CredentialsConfiguration): configurable via env variables or toml files """ - pass + def __bool__(self) -> bool: + # This is needed to avoid AuthConfigBase-derived classes + # which do not implement CredentialsConfiguration interface + # to be evaluated as False in requests.sessions.Session.prepare_request() + return True @configspec diff --git a/docs/website/docs/general-usage/http/rest-client.md b/docs/website/docs/general-usage/http/rest-client.md index 19cc95bf78..1093428b0f 100644 --- a/docs/website/docs/general-usage/http/rest-client.md +++ b/docs/website/docs/general-usage/http/rest-client.md @@ -407,7 +407,7 @@ The available authentication methods are defined in the `dlt.sources.helpers.res - [APIKeyAuth](#api-key-authentication) - [HttpBasicAuth](#http-basic-authentication) -For specific use cases, you can [implement custom authentication](#implementing-custom-authentication) by subclassing the `AuthConfigBase` class. +For specific use cases, you can [implement custom authentication](#implementing-custom-authentication) by subclassing the `AuthBase` class from the Requests library. ### Bearer token authentication @@ -479,12 +479,12 @@ response = client.get("/protected/resource") ### Implementing custom authentication -You can implement custom authentication by subclassing the `AuthConfigBase` class and implementing the `__call__` method: +You can implement custom authentication by subclassing the `AuthBase` class and implementing the `__call__` method: ```py -from dlt.sources.helpers.rest_client.auth import AuthConfigBase +from requests.auth import AuthBase -class CustomAuth(AuthConfigBase): +class CustomAuth(AuthBase): def __init__(self, token): self.token = token From bc20383e0e886f7aaa8df8a86f518094fb79d2d5 Mon Sep 17 00:00:00 2001 From: Anton Burnashev Date: Wed, 22 May 2024 18:31:04 +0200 Subject: [PATCH 2/3] Add tests --- .../helpers/rest_client/test_client.py | 44 ++++++++++++++++++- 1 file changed, 43 insertions(+), 1 deletion(-) diff --git a/tests/sources/helpers/rest_client/test_client.py b/tests/sources/helpers/rest_client/test_client.py index 50defa8edb..1cc6506dc1 100644 --- a/tests/sources/helpers/rest_client/test_client.py +++ b/tests/sources/helpers/rest_client/test_client.py @@ -1,6 +1,7 @@ import os import pytest from typing import Any, cast +from requests.auth import AuthBase from dlt.common.typing import TSecretStrValue from dlt.sources.helpers.requests import Response, Request from dlt.sources.helpers.rest_client import RESTClient @@ -57,7 +58,6 @@ def test_page_context(self, rest_client: RESTClient) -> None: for page in rest_client.paginate( "/posts", paginator=JSONResponsePaginator(next_url_path="next_page"), - auth=AuthConfigBase(), ): # response that produced data assert isinstance(page.response, Response) @@ -183,3 +183,45 @@ def test_oauth_jwt_auth_success(self, rest_client: RESTClient): ) assert_pagination(list(pages_iter)) + + def test_custom_auth_success(self, rest_client: RESTClient): + class CustomAuthConfigBase(AuthConfigBase): + def __init__(self, token: str): + self.token = token + + def __call__(self, request: Request) -> Request: + request.headers["Authorization"] = f"Bearer {self.token}" + return request + + class CustomAuthAuthBase(AuthBase): + def __init__(self, token: str): + self.token = token + + def __call__(self, request: Request) -> Request: + request.headers["Authorization"] = f"Bearer {self.token}" + return request + + + auth_list = [ + CustomAuthConfigBase("test-token"), + CustomAuthAuthBase("test-token"), + ] + + for auth in auth_list: + response = rest_client.get( + "/protected/posts/bearer-token", + auth=auth, + ) + + assert response.status_code == 200 + assert response.json()["data"][0] == {"id": 0, "title": "Post 0"} + + pages_iter = rest_client.paginate( + "/protected/posts/bearer-token", + auth=auth, + ) + + pages_list = list(pages_iter) + assert_pagination(pages_list) + + assert pages_list[0].response.request.headers["Authorization"] == "Bearer test-token" \ No newline at end of file From b4e417e4e2ffe46221506db14d4c8b59734c283d Mon Sep 17 00:00:00 2001 From: Anton Burnashev Date: Wed, 22 May 2024 18:41:01 +0200 Subject: [PATCH 3/3] Fix formatting --- tests/sources/helpers/rest_client/test_client.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/sources/helpers/rest_client/test_client.py b/tests/sources/helpers/rest_client/test_client.py index 1cc6506dc1..e03879b417 100644 --- a/tests/sources/helpers/rest_client/test_client.py +++ b/tests/sources/helpers/rest_client/test_client.py @@ -201,7 +201,6 @@ def __call__(self, request: Request) -> Request: request.headers["Authorization"] = f"Bearer {self.token}" return request - auth_list = [ CustomAuthConfigBase("test-token"), CustomAuthAuthBase("test-token"), @@ -224,4 +223,4 @@ def __call__(self, request: Request) -> Request: pages_list = list(pages_iter) assert_pagination(pages_list) - assert pages_list[0].response.request.headers["Authorization"] == "Bearer test-token" \ No newline at end of file + assert pages_list[0].response.request.headers["Authorization"] == "Bearer test-token"