Repository: apache/airflow

Issue:
`GCSToLocalFilesystemOperator` unnecessarily downloads objects when it checks object size
`GCSToLocalFilesystemOperator` in `airflow/providers/google/cloud/transfers/gcs_to_local.py` checks the file size if `store_to_xcom_key` is `True`.

https://github.com/apache/airflow/blob/b40dffa08547b610162f8cacfa75847f3c4ca364/airflow/providers/google/cloud/transfers/gcs_to_local.py#L137-L142

How it checks size is to download the object as `bytes` then check the size. This unnecessarily downloads the objects. `google.cloud.storage.blob.Blob` itself already has a `size` property ([documentation reference](https://googleapis.dev/python/storage/1.30.0/blobs.html#google.cloud.storage.blob.Blob.size)), and it should be used instead.

In extreme cases, if the object is big size, it adds unnecessary burden on the instance resources.

A new method, `object_size()`, can be added to `GCSHook`, then this can be addressed in `GCSToLocalFilesystemOperator`.


Discussion:
CC @potiuk as you were the original author if I'm not wrong :)

Please correct me if anything in my analysis above is incorrect, or maybe there was some reasons for what you implemented in this way? 

Cheers!


Sorry @potiuk I think I was wrong. You were not the original author of this chunk, and this code has been implemented 4+ years ago.

But it would still be appreciated if you have time to help share your inputs.
>In extreme cases, if the object is big size, it adds unnecessary burden on the instance resources.

I'm not sure about that. If user is using `GCSToLocalFilesystemOperator` then the download action is mandatory. This is what the operator suppose to do. There is just an extra step to check if we should/can push the file to xcom.

> A new method, object_size(), can be added to GCSHook, then this can be addressed in GCSToLocalFilesystemOperator.

We already have that method in the hook:

https://github.com/apache/airflow/blob/1a85ba9e93d44601a322546e31814bd9ef11c125/airflow/providers/google/cloud/hooks/gcs.py#L776-L793



Resolving patch:
diff --git a/airflow/providers/google/cloud/transfers/gcs_to_local.py b/airflow/providers/google/cloud/transfers/gcs_to_local.py
--- a/airflow/providers/google/cloud/transfers/gcs_to_local.py
+++ b/airflow/providers/google/cloud/transfers/gcs_to_local.py
@@ -15,7 +15,6 @@
 # specific language governing permissions and limitations
 # under the License.
 
-import sys
 import warnings
 from typing import Optional, Sequence, Union
 
@@ -133,8 +132,9 @@ def execute(self, context):
         )
 
         if self.store_to_xcom_key:
-            file_bytes = hook.download(bucket_name=self.bucket, object_name=self.object_name)
-            if sys.getsizeof(file_bytes) < MAX_XCOM_SIZE:
+            file_size = hook.get_size(bucket_name=self.bucket, object_name=self.object_name)
+            if file_size < MAX_XCOM_SIZE:
+                file_bytes = hook.download(bucket_name=self.bucket, object_name=self.object_name)
                 context['ti'].xcom_push(key=self.store_to_xcom_key, value=str(file_bytes))
             else:
                 raise AirflowException('The size of the downloaded file is too large to push to XCom!')


Test patch:
diff --git a/tests/providers/google/cloud/transfers/test_gcs_to_local.py b/tests/providers/google/cloud/transfers/test_gcs_to_local.py
--- a/tests/providers/google/cloud/transfers/test_gcs_to_local.py
+++ b/tests/providers/google/cloud/transfers/test_gcs_to_local.py
@@ -17,7 +17,12 @@
 
 import unittest
 from unittest import mock
+from unittest.mock import MagicMock
 
+import pytest
+
+from airflow import AirflowException
+from airflow.models.xcom import MAX_XCOM_SIZE
 from airflow.providers.google.cloud.transfers.gcs_to_local import GCSToLocalFilesystemOperator
 
 TASK_ID = "test-gcs-operator"
@@ -28,6 +33,8 @@
 MOCK_FILES = ["TEST1.csv", "TEST2.csv", "TEST3.csv"]
 TEST_OBJECT = "dir1/test-object"
 LOCAL_FILE_PATH = "/home/airflow/gcp/test-object"
+XCOM_KEY = "some_xkom_key"
+FILE_CONTENT = "some file content"
 
 
 class TestGoogleCloudStorageDownloadOperator(unittest.TestCase):
@@ -44,3 +51,39 @@ def test_execute(self, mock_hook):
         mock_hook.return_value.download.assert_called_once_with(
             bucket_name=TEST_BUCKET, object_name=TEST_OBJECT, filename=LOCAL_FILE_PATH
         )
+
+    @mock.patch("airflow.providers.google.cloud.transfers.gcs_to_local.GCSHook")
+    def test_size_lt_max_xcom_size(self, mock_hook):
+        operator = GCSToLocalFilesystemOperator(
+            task_id=TASK_ID,
+            bucket=TEST_BUCKET,
+            object_name=TEST_OBJECT,
+            store_to_xcom_key=XCOM_KEY,
+        )
+        context = {"ti": MagicMock()}
+        mock_hook.return_value.download.return_value = FILE_CONTENT
+        mock_hook.return_value.get_size.return_value = MAX_XCOM_SIZE - 1
+
+        operator.execute(context=context)
+        mock_hook.return_value.get_size.assert_called_once_with(
+            bucket_name=TEST_BUCKET, object_name=TEST_OBJECT
+        )
+        mock_hook.return_value.download.assert_called_once_with(
+            bucket_name=TEST_BUCKET, object_name=TEST_OBJECT
+        )
+        context["ti"].xcom_push.assert_called_once_with(key=XCOM_KEY, value=FILE_CONTENT)
+
+    @mock.patch("airflow.providers.google.cloud.transfers.gcs_to_local.GCSHook")
+    def test_size_gt_max_xcom_size(self, mock_hook):
+        operator = GCSToLocalFilesystemOperator(
+            task_id=TASK_ID,
+            bucket=TEST_BUCKET,
+            object_name=TEST_OBJECT,
+            store_to_xcom_key=XCOM_KEY,
+        )
+        context = {"ti": MagicMock()}
+        mock_hook.return_value.download.return_value = FILE_CONTENT
+        mock_hook.return_value.get_size.return_value = MAX_XCOM_SIZE + 1
+
+        with pytest.raises(AirflowException, match="file is too large"):
+            operator.execute(context=context)