-
Notifications
You must be signed in to change notification settings - Fork 40
Handle boto3 clients more efficiently with lru_cache #361
Handle boto3 clients more efficiently with lru_cache #361
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for the contribution @mattiamatrix! This looks like a great optimization, but I'd like to suggest some tweaks:
- I think we should move the changes into the
get_client
method so all the client types can benefit from this optimization. - We should maintain a cache of clients with their parameters as the cache key. In the current implementation, if any attributes are changed on the class, the client returned would not be updated to reflect those changes. By introducing a LRU cache, we can improve performance while maintaining configurability on the fly.
Let me know if you'd like me to go into more detail for either of those points!
Thank you for the review @desertaxle. I made the changes as you described. Let me know what you think now. |
Hello @desertaxle @zzstoatzz, is the PR ok now? |
hi @mattiamatrix - Alex is currently out-of-office for a while, so I can take over reviewing this. Let me take a look here / get caught up and do some QA and I will get back to you on this - appreciate the contribution! 🙂 |
hey @mattiamatrix - this looks pretty good to me. Can you add a test for the new functionality, maybe something like this? from prefect_aws.credentials import _get_client_cached, AwsCredentials, ClientType
from unittest.mock import patch
@patch('prefect_aws.credentials._get_client_cached')
def test_get_client_cached(mock_get_client_cached):
"""
Test to ensure that _get_client_cached function returns the same instance
for multiple calls with the same parameters and properly utilizes lru_cache.
"""
# Create a mock AwsCredentials instance
aws_credentials_block = AwsCredentials()
# Call _get_client_cached multiple times with the same parameters
_get_client_cached(aws_credentials_block, ClientType.S3)
_get_client_cached(aws_credentials_block, ClientType.S3)
# Verify that _get_client_cached is called only once due to caching
mock_get_client_cached.assert_called_once_with(aws_credentials_block, ClientType.S3)
# Optionally, you can test with different parameters to ensure they are cached separately
_get_client_cached(aws_credentials_block, ClientType.ECS)
assert mock_get_client_cached.call_count == 2, "Should be called twice with different parameters" |
Hi @zzstoatzz, thank you for the snippet. I added the code that you shared but the test is failing with error @zzstoatzz it would be very helpful if you could directly update the code and push the correct test. Thank you. |
sorry @mattiamatrix - i should have clarified that that was pseudocode based on my memory of the mocks / fixtures available i do not have time to do this now but if you're unable then I can come back to this later |
Ok, I pushed a version that is using |
hi @mattiamatrix - after discussing this internally, I think we should change how we're getting that as is:
instead of just the |
Close in favour of #369 |
Closes #332
Example
Screenshots
Checklist
pre-commit
checks.pre-commit install && pre-commit run --all
locally for formatting and linting.mkdocs serve
view documentation locally.