rest_api: filter, exclude, transform API responses

sources/rest_api/__init__.py

@@ -1,4 +1,5 @@
 """Generic API Source"""
+
 from copy import deepcopy
 from typing import Type, Any, Dict, List, Optional, Generator, Callable, cast, Union
 import graphlib  # type: ignore[import,unused-ignore]
@@ -33,6 +34,7 @@
     IncrementalParamConfig,
     RESTAPIConfig,
     ParamBindType,
+    ProcessingSteps,
 )
 from .config_setup import (
     IncrementalParam,
@@ -222,6 +224,7 @@ def create_resources(
         request_params = endpoint_config.get("params", {})
         request_json = endpoint_config.get("json", None)
         paginator = create_paginator(endpoint_config.get("paginator"))
+        processing_steps = endpoint_resource.pop("processing_steps", [])
 
         resolved_param: ResolvedParam = resolved_param_map[resource_name]
 
@@ -253,6 +256,14 @@ def create_resources(
             endpoint_resource, {"endpoint", "include_from_parent"}
         )
 
+        def process(resource, processing_steps) -> Any:
+            for step in processing_steps:
+                if "filter" in step:
+                    resource.add_filter(step["filter"])
+                if "map" in step:
+                    resource.add_map(step["map"])
+            return resource
+
         if resolved_param is None:
 
             def paginate_resource(
@@ -300,6 +311,9 @@ def paginate_resource(
                 data_selector=endpoint_config.get("data_selector"),
                 hooks=hooks,
             )
+            resources[resource_name] = process(
+                resources[resource_name], processing_steps
+            )
 
         else:
             predecessor = resources[resolved_param.resolve_config["resource"]]
@@ -347,7 +361,7 @@ def paginate_dependent_resource(
                         if parent_record:
                             for child_record in child_page:
                                 child_record.update(parent_record)
-                        yield child_page
+                        yield from child_page
 
             resources[resource_name] = dlt.resource(  # type: ignore[call-overload]
                 paginate_dependent_resource,
@@ -362,6 +376,10 @@ def paginate_dependent_resource(
                 hooks=hooks,
             )
 
+            resources[resource_name] = process(
+                resources[resource_name], processing_steps
+            )
+
     return resources
 
 

sources/rest_api/typing.py

@@ -240,6 +240,11 @@ class Endpoint(TypedDict, total=False):
     incremental: Optional[IncrementalConfig]
 
 
+class ProcessingSteps(TypedDict):
+    filter: Optional[Callable[[Any], bool]]
+    map: Optional[Callable[[Any], Any]]
+
+
 class ResourceBase(TypedDict, total=False):
     """Defines hints that may be passed to `dlt.resource` decorator"""
 
@@ -254,6 +259,7 @@ class ResourceBase(TypedDict, total=False):
     table_format: Optional[TTableHintTemplate[TTableFormat]]
     selected: Optional[bool]
     parallelized: Optional[bool]
+    processing_steps: Optional[List[ProcessingSteps]]
 
 
 class EndpointResourceBase(ResourceBase, total=False):

tests/rest_api/test_rest_api_source_processed.py

@@ -0,0 +1,188 @@
+import dlt
+import pytest
+from dlt.sources.helpers.rest_client.paginators import SinglePagePaginator
+
+from sources.rest_api import rest_api_source, RESTAPIConfig
+from tests.utils import ALL_DESTINATIONS, assert_load_info, load_table_counts
+
+
+def _make_pipeline(destination_name: str):
+    return dlt.pipeline(
+        pipeline_name="rest_api",
+        destination=destination_name,
+        dataset_name="rest_api_data",
+        full_refresh=True,
+    )
+
+
+def test_rest_api_source_filtered(mock_api_server) -> None:
+    config: RESTAPIConfig = {
+        "client": {
+            "base_url": "https://api.example.com",
+        },
+        "resources": [
+            {
+                "name": "posts",
+                "endpoint": "posts",
+                "processing_steps": [
+                    {"filter": lambda x: x["id"] == 1},
+                ],
+            },
+        ],
+    }
+    mock_source = rest_api_source(config)
+
+    data = list(mock_source.with_resources("posts"))
+    assert len(data) == 1
+    assert data[0]["title"] == "Post 1"
+
+
+def test_rest_api_source_map(mock_api_server) -> None:
+
+    def lower_title(row):
+        row["title"] = row["title"].lower()
+        return row
+
+    config: RESTAPIConfig = {
+        "client": {
+            "base_url": "https://api.example.com",
+        },
+        "resources": [
+            {
+                "name": "posts",
+                "endpoint": "posts",
+                "processing_steps": [
+                    {"map": lower_title},
+                ],
+            },
+        ],
+    }
+    mock_source = rest_api_source(config)
+
+    data = list(mock_source.with_resources("posts"))
+
+    assert all(record["title"].startswith("post ") for record in data)
+
+
+def test_rest_api_source_filter_and_map(mock_api_server) -> None:
+
+    def id_by_10(row):
+        row["id"] = row["id"] * 10
+        return row
+
+    config: RESTAPIConfig = {
+        "client": {
+            "base_url": "https://api.example.com",
+        },
+        "resources": [
+            {
+                "name": "posts",
+                "endpoint": "posts",
+                "processing_steps": [
+                    {"map": id_by_10},
+                    {"filter": lambda x: x["id"] == 10},
+                ],
+            },
+            {
+                "name": "posts_2",
+                "endpoint": "posts",
+                "processing_steps": [
+                    {"filter": lambda x: x["id"] == 10},
+                    {"map": id_by_10},
+                ],
+            },
+        ],
+    }
+    mock_source = rest_api_source(config)
+
+    data = list(mock_source.with_resources("posts"))
+    assert len(data) == 1
+    assert data[0]["title"] == "Post 1"
+
+    data = list(mock_source.with_resources("posts_2"))
+    assert len(data) == 1
+    assert data[0]["id"] == 100
+    assert data[0]["title"] == "Post 10"
+
+
+def test_rest_api_source_filtered_child(mock_api_server) -> None:
+    config: RESTAPIConfig = {
+        "client": {
+            "base_url": "https://api.example.com",
+        },
+        "resources": [
+            {
+                "name": "posts",
+                "endpoint": "posts",
+                "processing_steps": [
+                    {"filter": lambda x: x["id"] in (1, 2)},
+                ],
+            },
+            {
+                "name": "comments",
+                "endpoint": {
+                    "path": "/posts/{post_id}/comments",
+                    "params": {
+                        "post_id": {
+                            "type": "resolve",
+                            "resource": "posts",
+                            "field": "id",
+                        }
+                    },
+                },
+                "processing_steps": [
+                    {"filter": lambda x: x["id"] == 1},
+                ],
+            },
+        ],
+    }
+    mock_source = rest_api_source(config)
+
+    data = list(mock_source.with_resources("comments"))
+    assert len(data) == 2
+    # assert data[0]["title"] == "Post 1"
+
+
+def test_rest_api_source_filtered_and_map_child(mock_api_server) -> None:
+
+    def extend_body(row):
+        row["body"] = f"{row['_posts_title']} - {row['body']}"
+        return row
+
+    config: RESTAPIConfig = {
+        "client": {
+            "base_url": "https://api.example.com",
+        },
+        "resources": [
+            {
+                "name": "posts",
+                "endpoint": "posts",
+                "processing_steps": [
+                    {"filter": lambda x: x["id"] in (1, 2)},
+                ],
+            },
+            {
+                "name": "comments",
+                "endpoint": {
+                    "path": "/posts/{post_id}/comments",
+                    "params": {
+                        "post_id": {
+                            "type": "resolve",
+                            "resource": "posts",
+                            "field": "id",
+                        }
+                    },
+                },
+                "include_from_parent": ["title"],
+                "processing_steps": [
+                    {"map": extend_body},
+                    {"filter": lambda x: x["body"].startswith("Post 2")},
+                ],
+            },
+        ],
+    }
+    mock_source = rest_api_source(config)
+
+    data = list(mock_source.with_resources("comments"))
+    # assert len(data) == 1
+    assert data[0]["body"] == "Post 2 - Comment 0 for post 2"