Fix: ensure custom session can be provided to rest client

[z3z1ma]

Description

A simple fix for a small oversight in the rest client. Additionally a test proving you can pass a custom session and there is no crash.

Related Issues

• Fixes Cannot pass session to rest client #1395

Additional Context

[netlify]

✅ Deploy Preview for dlt-hub-docs canceled.

Name	Link
🔨 Latest commit	f2742b9
🔍 Latest deploy log	https://app.netlify.com/sites/dlt-hub-docs/deploys/66549f81bd7d7d00082178f7

[z3z1ma]

This is also confusing

dlt/dlt/sources/helpers/rest_client/client.py

Lines 186 to 190 in 64c9538

	def raise_for_status(response: Response, *args: Any, **kwargs: Any) -> None:
	response.raise_for_status()
	if "response" not in hooks:
	hooks["response"] = [raise_for_status]

Why is this there is raise for status is intended to be configured at session level

EDIT: Its nice being able to configure these though. Just confusing contrasted with the requests Client which may be the source of my confusion.

[z3z1ma]

Also note, you bypass session.request though this code path here

dlt/dlt/sources/helpers/rest_client/client.py

Lines 92 to 133 in 64c9538

	def _create_request(
	self,
	path: str,
	method: HTTPMethod,
	params: Dict[str, Any],
	json: Optional[Dict[str, Any]] = None,
	auth: Optional[AuthConfigBase] = None,
	hooks: Optional[Hooks] = None,
	) -> Request:
	parsed_url = urlparse(path)
	if parsed_url.scheme in ("http", "https"):
	url = path
	else:
	url = join_url(self.base_url, path)
	return Request(
	method=method,
	url=url,
	headers=self.headers,
	params=params,
	json=json,
	auth=auth or self.auth,
	hooks=hooks,
	)
	def _send_request(self, request: Request) -> Response:
	logger.info(
	f"Making {request.method.upper()} request to {request.url}"
	f" with params={request.params}, json={request.json}"
	)
	prepared_request = self.session.prepare_request(request)
	return self.session.send(prepared_request)
	def request(self, path: str = "", method: HTTPMethod = "GET", **kwargs: Any) -> Response:
	prepared_request = self._create_request(
	path=path,
	method=method,
	**kwargs,
	)
	return self._send_request(prepared_request)

And call session.send directly BUT there is no retry logic on send. Its only monkey-patched into session.request here

dlt/dlt/sources/helpers/requests/retry.py

Line 242 in 64c9538

session.request = retry.wraps(session.request) # type: ignore[method-assign]

I can fix that too. It will specifically fix passing in the dlt requests Cient into RESTClient and benefitting from the retry logic.

[burnash]

Thanks for spotting this and for the fix! I have a couple of comments to the PR below.

[burnash]

you bypass session.request though this code path here

Yes, very good catch! Thanks for the fix again.

This is also confusing

def raise_for_status(response: Response, *args: Any, **kwargs: Any) -> None: 
    response.raise_for_status() 
if "response" not in hooks: 
    hooks["response"] = [raise_for_status]

Agreed, not the most elegant solution, the idea is to enable hooks while having the rase_for_status as a default. This is also related to the session check. I wanted to give the user a way to override the default behaviour. Naturally the suggestions of how to do this better are very welcome!

[rudolfix]

LGTM!

one question: are you trying to use a session that is not from requests library? ie httpx?

[rudolfix]

I think that your change somehow broke one test where I expect the host is mocked.

FAILED tests/sources/helpers/rest_client/test_client.py::TestRESTClient::test_oauth_jwt_auth_success - requests.exceptions.ConnectionError: HTTPSConnectionPool(host='api.example.com', port=443): Max retries exceeded with url: /oauth/token (Caused by NewConnectionError('<urllib3.connection.HTTPSConnection object at 0x7f08d16f1710>: Failed to establish a new connection: [Errno -2] Name or service not known'))

@z3z1ma pls ping us if you have time to investigate that soon, otherwise we'll take over

[z3z1ma]

Test is passing for me 🤔 interesting, it shouldnt be flaky?

Also was doing my PR reply using https://github.com/pwntester/octo.nvim and totally replied via appending/editing your comments by accident (didnt even know that was possible)

[z3z1ma]

LGTM!

one question: are you trying to use a session that is not from requests library? ie httpx?

Actually I was trying to use the dlt requests Client.session (I had done some custom stuff but wanted the default retry logic).

Agreed, not the most elegant solution, the idea is to enable hooks while having the raise_for_status as a default. This is also related to the session check. I wanted to give the user a way to override the default behaviour. Naturally the suggestions of how to do this better are very welcome!

Would it make more sense to inherit from the session by default? This way the behavior is consistent within a source but can still be overridden on a per endpoint basis? I need to double check if/how sessions exposes this.

[burnash]

Would it make more sense to inherit from the session by default? This way the behavior is consistent within a source but can still be overridden on a per endpoint basis? I need to double check if/how sessions exposes this.

Yeah I was also contemplating this. But in that case how do we allow the user to pass their own Session?

[burnash]

Test is passing for me 🤔 interesting, it shouldnt be flaky?

I've checked out the branch locally and also don't see this test_oauth_jwt_auth_success error. 🤔

[rudolfix]

this is due to the OAuthAuth using default dlt session which is created before the mock by some other test... let me take a look and fix it

[burnash]

this is due to the OAuthAuth using default dlt session which is created before the mock by some other test... let me take a look and fix it

Interesting

[burnash]

Looks good, @z3z1ma! Thanks for fixing it!