Add paginator type support

dlt-hub · Dec 9, 2024 · 4a6351e · 4a6351e
1 parent 5b8421e
commit 4a6351e
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 2 deletions.
diff --git a/dlt/sources/rest_api/config_setup.py b/dlt/sources/rest_api/config_setup.py
@@ -66,6 +66,7 @@
     Endpoint,
     EndpointResource,
     AuthType,
+    PaginatorType,
 )
 
 
@@ -105,6 +106,7 @@ def register_paginator(
             "Your custom paginator has to be a subclass of BasePaginator"
         )
     PAGINATOR_MAP[paginator_name] = paginator_class
+    add_value_to_literal(PaginatorType, paginator_name)
 
 
 def get_paginator_class(paginator_name: str) -> Type[BasePaginator]:

diff --git a/tests/sources/rest_api/configurations/test_custom_auth_config.py b/tests/sources/rest_api/configurations/test_custom_auth_config.py
@@ -78,7 +78,7 @@ class NotAuthConfigBase:
             )
         assert e.match("Invalid auth: NotAuthConfigBase.")
 
-    def test_validate_config_raises_no_error(self, custom_auth_config: AuthConfig) -> None:
+    def test_valid_config_raises_no_error(self, custom_auth_config: AuthConfig) -> None:
         rest_api.config_setup.register_auth("custom_oauth_2", CustomOAuth2)
 
         valid_config: RESTAPIConfig = {

diff --git a/tests/sources/rest_api/configurations/test_custom_paginator_config.py b/tests/sources/rest_api/configurations/test_custom_paginator_config.py
@@ -4,7 +4,7 @@
 
 from dlt.sources import rest_api
 from dlt.sources.helpers.rest_client.paginators import JSONLinkPaginator
-from dlt.sources.rest_api.typing import PaginatorConfig
+from dlt.sources.rest_api.typing import PaginatorConfig, RESTAPIConfig
 
 
 class CustomPaginator(JSONLinkPaginator):
@@ -67,3 +67,13 @@ class NotAPaginator:
         with pytest.raises(ValueError) as e:
             rest_api.config_setup.register_paginator("not_a_paginator", NotAPaginator)  # type: ignore[arg-type]
         assert e.match("Invalid paginator: NotAPaginator.")
+
+    def test_test_valid_config_raises_no_error(self, custom_paginator_config) -> None:
+        rest_api.config_setup.register_paginator("custom_paginator", CustomPaginator)
+
+        valid_config: RESTAPIConfig = {
+            "client": {"base_url": "https://example.com", "paginator": custom_paginator_config},
+            "resources": ["test"],
+        }
+
+        rest_api.rest_api_source(valid_config)