Repository: apache/airflow

Issue:
Reduce memory usage in `S3Hook`
### Body

Original stacktrace from the Slack

```console
Error:
 File "/usr/local/airflow/plugins/plugins/others/data_source_monitor.py", line 53, in retrieve_data
get_time_query = s3_hook.read_key(
File "/usr/local/airflow/.local/lib/python3.10/site-packages/airflow/providers/amazon/aws/hooks/s3.py", line 64, in wrapper
return func(*bound_args.args, **bound_args.kwargs)
File "/usr/local/airflow/.local/lib/python3.10/site-packages/airflow/providers/amazon/aws/hooks/s3.py", line 92, in wrapper
return func(*bound_args.args, **bound_args.kwargs)
File "/usr/local/airflow/.local/lib/python3.10/site-packages/airflow/providers/amazon/aws/hooks/s3.py", line 514, in read_key
obj = self.get_key(key, bucket_name)
File "/usr/local/airflow/.local/lib/python3.10/site-packages/airflow/providers/amazon/aws/hooks/s3.py", line 64, in wrapper
return func(*bound_args.args, **bound_args.kwargs)
File "/usr/local/airflow/.local/lib/python3.10/site-packages/airflow/providers/amazon/aws/hooks/s3.py", line 92, in wrapper
return func(*bound_args.args, **bound_args.kwargs)
File "/usr/local/airflow/.local/lib/python3.10/site-packages/airflow/providers/amazon/aws/hooks/s3.py", line 493, in get_key
s3_resource = self.get_session().resource(
File "/usr/local/airflow/.local/lib/python3.10/site-packages/boto3/session.py", line 446, in resource
client = self.client(
File "/usr/local/airflow/.local/lib/python3.10/site-packages/boto3/session.py", line 299, in client
return self._session.create_client(
File "/usr/local/airflow/.local/lib/python3.10/site-packages/botocore/session.py", line 976, in create_client
client = client_creator.create_client(
File "/usr/local/airflow/.local/lib/python3.10/site-packages/botocore/client.py", line 116, in create_client
endpoints_ruleset_data = self._load_service_endpoints_ruleset(
File "/usr/local/airflow/.local/lib/python3.10/site-packages/botocore/client.py", line 220, in _load_service_endpoints_ruleset
return self._loader.load_service_model(
File "/usr/local/airflow/.local/lib/python3.10/site-packages/botocore/loaders.py", line 142, in _wrapper
data = func(self, *args, **kwargs)
File "/usr/local/airflow/.local/lib/python3.10/site-packages/botocore/loaders.py", line 406, in load_service_model
known_services = self.list_available_services(type_name)
File "/usr/local/airflow/.local/lib/python3.10/site-packages/botocore/loaders.py", line 142, in _wrapper
data = func(self, *args, **kwargs)
File "/usr/local/airflow/.local/lib/python3.10/site-packages/botocore/loaders.py", line 311, in list_available_services
api_versions = os.listdir(full_dirname)
OSError: [Errno 12] Cannot allocate memory: '/usr/local/airflow/.local/lib/python3.10/site-packages/botocore/data/efs'
```

The reason of this error simple, for some operations S3Hook create [resource](https://boto3.amazonaws.com/v1/documentation/api/latest/guide/resources.html) (High Level client) in addition to [`S3.Client`](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/s3.html) and this resource created every time when some method of S3Hook called as result additional memory required, for example if run [`S3Hook.download_file`](https://github.com/apache/airflow/blob/95980a9bc50c1accd34166ba608bbe2b4ebd6d52/airflow/providers/amazon/aws/hooks/s3.py#L1338) into the loop it might be reason for this error

As usual there are at least two solutions:
**Option 1**: use caching into the internal methods of S3Hook
**Option 2**: Get rid of resource usage in S3 hook and replace it by `S3.Client` methods. It might be better solution:
   - Seems like resources do not actively maintained in `boto3`
   - It required for about 30-40 MB of memory for create new resource object, however everything (and even more) could be done by `S3.Client`

### Committer

- [X] I acknowledge that I am a maintainer/committer of the Apache Airflow project.


Discussion:
I'd love to see someone take/make the time to do a full rewrite on the S3Hook.  It's so convoluted and nothing like any of the other hooks.  Knowing what we know now about how to create a hook, I bet we could greatly simplify it.
@ferruzzi I'd like to take a stab at this one if it's available.

@ellisms - That would be great.  You can look at the newer hooks and see lots of examples of how this could be reworked to use boto3.client instead of boto3.resources, and they should really help clean up a LOT of this spaghetti code.  let me know if you need any help.
@ferruzzi Can you assign this one to me? Finally have some time to start looking at it.

I started digging into this, and converting boto3.resource to boto3.client could introduce a breaking change. For example, `get_key` returns a resource.Object (and the object.load is nothing more than a wrapper around client.head_object), but a client.head_object returns a dictionary. It's easy enough to modify the hook and operator, but it would break for anyone that directly calls `hook.get_key()`.  Wanted to throw this out there for discussion. @Taragolis I thought maybe you'd have some thoughts. Are hook methods considered part of the public API?


Resolving patch:
diff --git a/airflow/providers/amazon/aws/hooks/s3.py b/airflow/providers/amazon/aws/hooks/s3.py
--- a/airflow/providers/amazon/aws/hooks/s3.py
+++ b/airflow/providers/amazon/aws/hooks/s3.py
@@ -30,16 +30,18 @@
 from contextlib import suppress
 from copy import deepcopy
 from datetime import datetime
-from functools import wraps
+from functools import cached_property, wraps
 from inspect import signature
 from io import BytesIO
 from pathlib import Path
 from tempfile import NamedTemporaryFile, gettempdir
-from typing import TYPE_CHECKING, Any, Callable, TypeVar, cast
+from typing import TYPE_CHECKING, Any, Callable
 from urllib.parse import urlsplit
 from uuid import uuid4
 
 if TYPE_CHECKING:
+    from mypy_boto3_s3.service_resource import Bucket as S3Bucket, Object as S3ResourceObject
+
     from airflow.utils.types import ArgNotSet
 
     with suppress(ImportError):
@@ -55,22 +57,17 @@
 from airflow.providers.amazon.aws.utils.tags import format_tags
 from airflow.utils.helpers import chunks
 
-if TYPE_CHECKING:
-    from mypy_boto3_s3.service_resource import Bucket as S3Bucket, Object as S3ResourceObject
-
-T = TypeVar("T", bound=Callable)
-
 logger = logging.getLogger(__name__)
 
 
-def provide_bucket_name(func: T) -> T:
+def provide_bucket_name(func: Callable) -> Callable:
     """Provide a bucket name taken from the connection if no bucket name has been passed to the function."""
     if hasattr(func, "_unify_bucket_name_and_key_wrapped"):
         logger.warning("`unify_bucket_name_and_key` should wrap `provide_bucket_name`.")
     function_signature = signature(func)
 
     @wraps(func)
-    def wrapper(*args, **kwargs) -> T:
+    def wrapper(*args, **kwargs) -> Callable:
         bound_args = function_signature.bind(*args, **kwargs)
 
         if "bucket_name" not in bound_args.arguments:
@@ -90,10 +87,10 @@ def wrapper(*args, **kwargs) -> T:
 
         return func(*bound_args.args, **bound_args.kwargs)
 
-    return cast(T, wrapper)
+    return wrapper
 
 
-def provide_bucket_name_async(func: T) -> T:
+def provide_bucket_name_async(func: Callable) -> Callable:
     """Provide a bucket name taken from the connection if no bucket name has been passed to the function."""
     function_signature = signature(func)
 
@@ -110,15 +107,15 @@ async def wrapper(*args: Any, **kwargs: Any) -> Any:
 
         return await func(*bound_args.args, **bound_args.kwargs)
 
-    return cast(T, wrapper)
+    return wrapper
 
 
-def unify_bucket_name_and_key(func: T) -> T:
+def unify_bucket_name_and_key(func: Callable) -> Callable:
     """Unify bucket name and key in case no bucket name and at least a key has been passed to the function."""
     function_signature = signature(func)
 
     @wraps(func)
-    def wrapper(*args, **kwargs) -> T:
+    def wrapper(*args, **kwargs) -> Callable:
         bound_args = function_signature.bind(*args, **kwargs)
 
         if "wildcard_key" in bound_args.arguments:
@@ -141,7 +138,7 @@ def wrapper(*args, **kwargs) -> T:
     # if provide_bucket_name is applied first, and there's a bucket defined in conn
     # then if user supplies full key, bucket in key is not respected
     wrapper._unify_bucket_name_and_key_wrapped = True  # type: ignore[attr-defined]
-    return cast(T, wrapper)
+    return wrapper
 
 
 class S3Hook(AwsBaseHook):
@@ -188,6 +185,15 @@ def __init__(
 
         super().__init__(*args, **kwargs)
 
+    @cached_property
+    def resource(self):
+        return self.get_session().resource(
+            self.service_name,
+            endpoint_url=self.conn_config.get_service_endpoint_url(service_name=self.service_name),
+            config=self.config,
+            verify=self.verify,
+        )
+
     @property
     def extra_args(self):
         """Return hook's extra arguments (immutable)."""
@@ -307,13 +313,7 @@ def get_bucket(self, bucket_name: str | None = None) -> S3Bucket:
         :param bucket_name: the name of the bucket
         :return: the bucket object to the bucket name.
         """
-        s3_resource = self.get_session().resource(
-            "s3",
-            endpoint_url=self.conn_config.endpoint_url,
-            config=self.config,
-            verify=self.verify,
-        )
-        return s3_resource.Bucket(bucket_name)
+        return self.resource.Bucket(bucket_name)
 
     @provide_bucket_name
     def create_bucket(self, bucket_name: str | None = None, region_name: str | None = None) -> None:
@@ -943,14 +943,7 @@ def sanitize_extra_args() -> dict[str, str]:
                 if arg_name in S3Transfer.ALLOWED_DOWNLOAD_ARGS
             }
 
-        s3_resource = self.get_session().resource(
-            "s3",
-            endpoint_url=self.conn_config.endpoint_url,
-            config=self.config,
-            verify=self.verify,
-        )
-        obj = s3_resource.Object(bucket_name, key)
-
+        obj = self.resource.Object(bucket_name, key)
         obj.load(**sanitize_extra_args())
         return obj
 


Test patch:
diff --git a/tests/providers/amazon/aws/hooks/test_s3.py b/tests/providers/amazon/aws/hooks/test_s3.py
--- a/tests/providers/amazon/aws/hooks/test_s3.py
+++ b/tests/providers/amazon/aws/hooks/test_s3.py
@@ -62,6 +62,10 @@ def test_get_conn(self):
         hook = S3Hook()
         assert hook.get_conn() is not None
 
+    def test_resource(self):
+        hook = S3Hook()
+        assert hook.resource is not None
+
     def test_use_threads_default_value(self):
         hook = S3Hook()
         assert hook.transfer_config.use_threads is True