Skip to content

Commit

Permalink
FIX: rm nested headers functionality
Browse files Browse the repository at this point in the history
  • Loading branch information
@ArneDePeuter
ArneDePeuter committed Nov 25, 2024
1 parent ce40342 commit a070b23
Show file tree
Hide file tree
Showing 3 changed files with 18 additions and 195 deletions.
32 changes: 14 additions & 18 deletions dlt/sources/rest_api/config_setup.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -431,7 +431,13 @@ def _bind_header_params(resource: EndpointResource) -> None:
bind_params[k] = "{" + k + "}"

headers = resource["endpoint"].get("headers", {})
resource["endpoint"]["headers"] = generic_format(headers, bind_params) # type: ignore[typeddict-item]
formatted_headers = {
k.format(**bind_params) if isinstance(k, str) else str(k): (
v.format(**bind_params) if isinstance(v, str) else str(v)
)
for k, v in headers.items()
}
resource["endpoint"]["headers"] = formatted_headers


def _setup_single_entity_endpoint(endpoint: Endpoint) -> Endpoint:
Expand Down Expand Up @@ -584,21 +590,6 @@ def remove_field(response: Response, *args, **kwargs) -> Response:
return None


def generic_format(
to_format: Union[Dict[str, Any], str, List[Any]], param_values: Dict[str, Any]
) -> Union[Dict[str, Any], str, List[Any]]:
if isinstance(to_format, dict):
return {
generic_format(key, param_values): generic_format(val, param_values) # type: ignore[misc]
for key, val in to_format.items()
}
if isinstance(to_format, list):
return [generic_format(item, param_values) for item in to_format]
if isinstance(to_format, str):
return to_format.format(**param_values)
return str(to_format)


def process_parent_data_item(
path: str,
item: Dict[str, Any],
Expand Down Expand Up @@ -639,8 +630,13 @@ def process_parent_data_item(
parent_record[child_key] = item[parent_key]

if headers is not None:
formatted_headers = generic_format(headers, param_values)
return bound_path, formatted_headers, parent_record # type: ignore[return-value]
formatted_headers = {
k.format(**param_values) if isinstance(k, str) else str(k): (
v.format(**param_values) if isinstance(v, str) else str(v)
)
for k, v in headers.items()
}
return bound_path, formatted_headers, parent_record
return bound_path, {}, parent_record


Expand Down
64 changes: 2 additions & 62 deletions tests/sources/rest_api/configurations/test_resolve_config.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -193,52 +193,11 @@ def test_process_parent_data_item_headers() -> None:
)
assert resolved_headers == {"Authorization": "12345", "num": "2"}

# nested params
resolve_params = [
ResolvedParam(
"token", {"field": "auth.token", "resource": "authenticate", "type": "resolve"}
),
ResolvedParam("num", {"field": "auth.num", "resource": "authenticate", "type": "resolve"}),
]
_, resolved_headers, parent_record = process_parent_data_item(
"chicken",
{"auth": {"token": 12345, "num": 2}},
resolve_params,
None,
{"Authorization": "{token}", "num": "{num}"},
)
assert resolved_headers == {"Authorization": "12345", "num": "2"}

# nested header dict
resolve_params = [
ResolvedParam(
"token", {"field": "auth.token", "resource": "authenticate", "type": "resolve"}
)
]
_, resolved_headers, parent_record = process_parent_data_item(
"chicken",
{"auth": {"token": 12345}},
resolve_params,
None,
{"Authorization": {"Bearer": "{token}"}},
)
assert resolved_headers == {"Authorization": {"Bearer": "12345"}}

# nested header list
resolve_params = [
ResolvedParam(
"token", {"field": "auth.token", "resource": "authenticate", "type": "resolve"}
)
]
_, resolved_headers, parent_record = process_parent_data_item(
"chicken",
{"auth": {"token": 12345}},
resolve_params,
None,
{"Authorization": ["Bearer", "{token}"]},
)
assert resolved_headers == {"Authorization": ["Bearer", "12345"]}

# param path not found
with pytest.raises(ValueError) as val_ex:
_, _, parent_record = process_parent_data_item(
Expand Down Expand Up @@ -444,31 +403,12 @@ def test_bind_header_params() -> None:
"name": "test_resource",
"endpoint": {
"path": "test/path",
"headers": {"Authorization": "{token}"},
"headers": {"Authorization": "Bearer {token}"},
"params": {
"token": "test_token",
},
},
}
_bind_header_params(resource_with_headers)
assert resource_with_headers["endpoint"]["headers"]["Authorization"] == "test_token" # type: ignore[index]
assert len(resource_with_headers["endpoint"]["params"]) == 0 # type: ignore[index]


def test_bind_header_params_nested() -> None:
resource_with_headers: EndpointResource = {
"name": "test_resource",
"endpoint": {
"path": "test/path",
"headers": {"{token}": "{token}", "deeper": {"{token}": ["{token}"]}},
"params": {
"token": "test_token",
},
},
}
_bind_header_params(resource_with_headers)
assert resource_with_headers["endpoint"]["headers"] == { # type: ignore[index]
"test_token": "test_token",
"deeper": {"test_token": ["test_token"]},
}
assert resource_with_headers["endpoint"]["headers"]["Authorization"] == "Bearer test_token" # type: ignore[index]
assert len(resource_with_headers["endpoint"]["params"]) == 0 # type: ignore[index]
117 changes: 2 additions & 115 deletions tests/sources/rest_api/test_rest_api_source.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -174,7 +174,7 @@ def authenticate():
"name": "chicken",
"endpoint": {
"path": "chicken",
"headers": {"token": "{token}", "num": "2"},
"headers": {"token": "Bearer {token}", "num": "2"},
"params": {
"token": {
"type": "resolve",
Expand All @@ -200,117 +200,4 @@ def authenticate():
request_param: Request = args[0]

assert request_param.url == f"{base_url}/chicken"
assert request_param.headers == {"foo": "bar", "token": "1", "num": "2"}


@pytest.mark.parametrize("destination_name", ALL_DESTINATIONS)
@pytest.mark.parametrize("invocation_type", ("deco", "factory"))
@patch("dlt.sources.helpers.rest_client.client.RESTClient._send_request")
def test_request_headers_dynamic_key(
mock: MagicMock, destination_name: str, invocation_type: str
) -> None:
mock_resp = Response()
mock_resp.status_code = 200
mock_resp.json = lambda: {"success": "ok"} # type: ignore
mock.return_value = mock_resp

@dlt.resource
def authenticate():
yield [{"token": 1}]

base_url = "https://api.example.com"
config: RESTAPIConfig = {
"client": {"base_url": base_url, "headers": {"foo": "bar"}},
"resources": [
{
"name": "chicken",
"endpoint": {
"path": "chicken",
"headers": {"{token}": "{token}", "num": "2"},
"params": {
"token": {
"type": "resolve",
"field": "token",
"resource": "authenticate",
},
},
},
},
authenticate(),
],
}

if invocation_type == "deco":
data = rest_api(**config)
else:
data = rest_api_source(config)
pipeline = _make_pipeline(destination_name)
pipeline.run(data)

mock.assert_called()
args, kwargs = mock.call_args
request_param: Request = args[0]

assert request_param.url == f"{base_url}/chicken"
assert request_param.headers == {"foo": "bar", "1": "1", "num": "2"}


@pytest.mark.parametrize("destination_name", ALL_DESTINATIONS)
@pytest.mark.parametrize("invocation_type", ("deco", "factory"))
@patch("dlt.sources.helpers.rest_client.client.RESTClient._send_request")
def test_request_headers_nested(
mock: MagicMock, destination_name: str, invocation_type: str
) -> None:
mock_resp = Response()
mock_resp.status_code = 200
mock_resp.json = lambda: {"success": "ok"} # type: ignore
mock.return_value = mock_resp

@dlt.resource
def authenticate():
yield [{"token": 1}]

base_url = "https://api.example.com"
config: RESTAPIConfig = {
"client": {"base_url": base_url, "headers": {"foo": "bar"}},
"resources": [
{
"name": "chicken",
"endpoint": {
"path": "chicken",
"headers": {
"{token}": "{token}",
"num": "2",
"nested": {"nested": "{token}", "{token}": "other"},
},
"params": {
"token": {
"type": "resolve",
"field": "token",
"resource": "authenticate",
},
},
},
},
authenticate(),
],
}

if invocation_type == "deco":
data = rest_api(**config)
else:
data = rest_api_source(config)
pipeline = _make_pipeline(destination_name)
pipeline.run(data)

mock.assert_called()
args, kwargs = mock.call_args
request_param: Request = args[0]

assert request_param.url == f"{base_url}/chicken"
assert request_param.headers == {
"foo": "bar",
"1": "1",
"num": "2",
"nested": {"nested": "1", "1": "other"},
}
assert request_param.headers == {"foo": "bar", "token": "Bearer 1", "num": "2"}

0 comments on commit a070b23

Please sign in to comment.