Re: [PR] Add disk space check before downloading GCS file in GCSHook [airflow]
potiuk commented on PR #61949: URL: https://github.com/apache/airflow/pull/61949#issuecomment-4658957379 @MonsterChenzhuo — A reviewer (@shahar1) has requested changes on this PR, so I've removed the `ready for maintainer review` label — the next step is on your side. Could you address the review comments (push a fix, or reply in-thread explaining why the feedback doesn't apply)? Once addressed, re-request review from @shahar1 or re-mark the PR ready and it returns to the maintainer queue. Thank you. _Note: This comment was drafted by an AI-assisted triage tool and may contain mistakes. Once you have addressed the points above, an Apache Airflow maintainer — a real person — will take the next look at your PR. We use this [two-stage triage process](https://github.com/apache/airflow/blob/main/contributing-docs/25_maintainer_pr_triage.md#why-the-first-pass-is-automated) so that our maintainers' limited time is spent where it matters most: the conversation with you._ -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add disk space check before downloading GCS file in GCSHook [airflow]
shahar1 commented on code in PR #61949:
URL: https://github.com/apache/airflow/pull/61949#discussion_r3223655275
##
providers/google/src/airflow/providers/google/cloud/hooks/gcs.py:
##
@@ -347,6 +344,16 @@ def download(
blob = bucket.blob(blob_name=object_name,
chunk_size=chunk_size)
if filename:
+blob.reload(timeout=timeout)
+if blob.size is not None:
+directory = os.path.dirname(filename) or os.getcwd()
+free_space = shutil.disk_usage(directory).free
+if free_space < blob.size:
+raise AirflowException(
+f"Insufficient disk space to download file. "
+f"Required: {blob.size} bytes, Available:
{free_space} bytes."
+)
Review Comment:
Could you please replace `AirflowException` with a native Python exception?
(static checks should now fail because of that upon rebase/merge from `main`)
##
providers/google/src/airflow/providers/google/cloud/hooks/gcs.py:
##
@@ -347,6 +344,16 @@ def download(
blob = bucket.blob(blob_name=object_name,
chunk_size=chunk_size)
if filename:
+blob.reload(timeout=timeout)
Review Comment:
1. Can we extract the `blob.reload(...)` outside the retry loop?
2. Could you please add a flag (+docstring) for checking disk space
(default: `False`), as it is an extra API call which might not be relevant when
there's no storage limitations.
##
providers/google/src/airflow/providers/google/cloud/hooks/gcs.py:
##
@@ -347,6 +344,16 @@ def download(
blob = bucket.blob(blob_name=object_name,
chunk_size=chunk_size)
if filename:
+blob.reload(timeout=timeout)
+if blob.size is not None:
+directory = os.path.dirname(filename) or os.getcwd()
+free_space = shutil.disk_usage(directory).free
Review Comment:
It will fail with Python's `FileNotFoundError` if directory doesn't exist,
rather than what GCS client raises.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add disk space check before downloading GCS file in GCSHook [airflow]
shahar1 commented on code in PR #61949: URL: https://github.com/apache/airflow/pull/61949#discussion_r3223660793 ## providers/google/src/airflow/providers/google/cloud/hooks/gcs.py: ## @@ -347,6 +344,16 @@ def download( blob = bucket.blob(blob_name=object_name, chunk_size=chunk_size) if filename: +blob.reload(timeout=timeout) Review Comment: 1. Can we extract the `blob.reload(...)` outside the retry loop? 2. Could you please add a flag (+docstring) for checking disk space (default: `False`), as it is an extra API call which might not be relevant when there are no storage limitations. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add disk space check before downloading GCS file in GCSHook [airflow]
github-actions[bot] commented on PR #61949: URL: https://github.com/apache/airflow/pull/61949#issuecomment-4272146942 This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 5 days if no further activity occurs. Thank you for your contributions. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add disk space check before downloading GCS file in GCSHook [airflow]
shahar1 commented on PR #61949: URL: https://github.com/apache/airflow/pull/61949#issuecomment-3993494936 > @shahar1 PTAL, thanks. It will take me a long while until I'll be able to review it properly due to low availability. Feel free to work on other issues in the meanwhile. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add disk space check before downloading GCS file in GCSHook [airflow]
MonsterChenzhuo commented on code in PR #61949:
URL: https://github.com/apache/airflow/pull/61949#discussion_r2844732210
##
providers/google/tests/unit/google/cloud/hooks/test_gcs.py:
##
@@ -1843,3 +1846,88 @@ def get_logs_string(call_args_list):
assert "GCS object size (15) and local file size (9) differ." in
logs_string
assert f"Downloading dag_03.py to {sync_local_dir}/dag_03.py" in
logs_string
self.gcs_hook.download.assert_called_once()
+
+
+class TestGCSDownloadDiskCheck:
[email protected](GCS_STRING.format("GoogleBaseHook.__init__"),
new=mock_base_gcp_hook_default_project_id)
[email protected](GCS_STRING.format("GCSHook.get_conn"))
[email protected]("shutil.disk_usage")
+def test_download_raises_error_if_not_enough_space(self, mock_disk_usage,
mock_get_conn):
+hook = gcs.GCSHook(gcp_conn_id="test")
+
+# Mock GCS connection and blob
+mock_client = mock.Mock()
+mock_get_conn.return_value = mock_client
+mock_bucket = mock.Mock()
+mock_client.bucket.return_value = mock_bucket
+mock_blob = mock.Mock()
+mock_bucket.blob.return_value = mock_blob
+
+# Blob size 1000 bytes
+mock_blob.size = 1000
+
+# Free space 500 bytes
+# shutil.disk_usage returns a named tuple (total, used, free)
+mock_disk_usage.return_value = mock.Mock(total=2000, used=1500,
free=500)
+
+with pytest.raises(AirflowException, match="Insufficient disk space"):
+hook.download(bucket_name="bucket", object_name="object",
filename="/tmp/file")
+
+# Verify blob.reload() was called to get the size
+mock_blob.reload.assert_called_once()
+
[email protected](GCS_STRING.format("GoogleBaseHook.__init__"),
new=mock_base_gcp_hook_default_project_id)
[email protected](GCS_STRING.format("GCSHook.get_conn"))
[email protected]("shutil.disk_usage")
+def test_download_works_if_enough_space(self, mock_disk_usage,
mock_get_conn):
+hook = gcs.GCSHook(gcp_conn_id="test")
+
+# Mock GCS connection and blob
+mock_client = mock.Mock()
+mock_get_conn.return_value = mock_client
+mock_bucket = mock.Mock()
+mock_client.bucket.return_value = mock_bucket
+mock_blob = mock.Mock()
+mock_bucket.blob.return_value = mock_blob
+
+# Blob size 1000 bytes
+mock_blob.size = 1000
+mock_blob.name = "test_blob"
+mock_bucket.name = "test_bucket"
+
+# Free space 2000 bytes
+mock_disk_usage.return_value = mock.Mock(total=4000, used=2000,
free=2000)
+
+filename = "/tmp/file"
+hook.download(bucket_name="bucket", object_name="object",
filename=filename)
+
+# Verify blob.reload() was called
+mock_blob.reload.assert_called_once()
+# Verify download was called
+mock_blob.download_to_filename.assert_called_once_with(filename,
timeout=mock.ANY)
Review Comment:
done.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add disk space check before downloading GCS file in GCSHook [airflow]
MonsterChenzhuo commented on code in PR #61949:
URL: https://github.com/apache/airflow/pull/61949#discussion_r2844731272
##
providers/google/tests/unit/google/cloud/hooks/test_gcs.py:
##
@@ -1843,3 +1846,88 @@ def get_logs_string(call_args_list):
assert "GCS object size (15) and local file size (9) differ." in
logs_string
assert f"Downloading dag_03.py to {sync_local_dir}/dag_03.py" in
logs_string
self.gcs_hook.download.assert_called_once()
+
+
+class TestGCSDownloadDiskCheck:
[email protected](GCS_STRING.format("GoogleBaseHook.__init__"),
new=mock_base_gcp_hook_default_project_id)
[email protected](GCS_STRING.format("GCSHook.get_conn"))
[email protected]("shutil.disk_usage")
+def test_download_raises_error_if_not_enough_space(self, mock_disk_usage,
mock_get_conn):
+hook = gcs.GCSHook(gcp_conn_id="test")
+
+# Mock GCS connection and blob
+mock_client = mock.Mock()
+mock_get_conn.return_value = mock_client
+mock_bucket = mock.Mock()
+mock_client.bucket.return_value = mock_bucket
+mock_blob = mock.Mock()
+mock_bucket.blob.return_value = mock_blob
+
+# Blob size 1000 bytes
+mock_blob.size = 1000
+
+# Free space 500 bytes
+# shutil.disk_usage returns a named tuple (total, used, free)
+mock_disk_usage.return_value = mock.Mock(total=2000, used=1500,
free=500)
+
+with pytest.raises(AirflowException, match="Insufficient disk space"):
+hook.download(bucket_name="bucket", object_name="object",
filename="/tmp/file")
+
+# Verify blob.reload() was called to get the size
+mock_blob.reload.assert_called_once()
Review Comment:
done.
##
providers/google/src/airflow/providers/google/cloud/hooks/gcs.py:
##
@@ -347,6 +344,16 @@ def download(
blob = bucket.blob(blob_name=object_name,
chunk_size=chunk_size)
if filename:
+blob.reload(timeout=timeout)
+if blob.size:
Review Comment:
done.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add disk space check before downloading GCS file in GCSHook [airflow]
MonsterChenzhuo commented on code in PR #61949:
URL: https://github.com/apache/airflow/pull/61949#discussion_r2844732210
##
providers/google/tests/unit/google/cloud/hooks/test_gcs.py:
##
@@ -1843,3 +1846,88 @@ def get_logs_string(call_args_list):
assert "GCS object size (15) and local file size (9) differ." in
logs_string
assert f"Downloading dag_03.py to {sync_local_dir}/dag_03.py" in
logs_string
self.gcs_hook.download.assert_called_once()
+
+
+class TestGCSDownloadDiskCheck:
[email protected](GCS_STRING.format("GoogleBaseHook.__init__"),
new=mock_base_gcp_hook_default_project_id)
[email protected](GCS_STRING.format("GCSHook.get_conn"))
[email protected]("shutil.disk_usage")
+def test_download_raises_error_if_not_enough_space(self, mock_disk_usage,
mock_get_conn):
+hook = gcs.GCSHook(gcp_conn_id="test")
+
+# Mock GCS connection and blob
+mock_client = mock.Mock()
+mock_get_conn.return_value = mock_client
+mock_bucket = mock.Mock()
+mock_client.bucket.return_value = mock_bucket
+mock_blob = mock.Mock()
+mock_bucket.blob.return_value = mock_blob
+
+# Blob size 1000 bytes
+mock_blob.size = 1000
+
+# Free space 500 bytes
+# shutil.disk_usage returns a named tuple (total, used, free)
+mock_disk_usage.return_value = mock.Mock(total=2000, used=1500,
free=500)
+
+with pytest.raises(AirflowException, match="Insufficient disk space"):
+hook.download(bucket_name="bucket", object_name="object",
filename="/tmp/file")
+
+# Verify blob.reload() was called to get the size
+mock_blob.reload.assert_called_once()
+
[email protected](GCS_STRING.format("GoogleBaseHook.__init__"),
new=mock_base_gcp_hook_default_project_id)
[email protected](GCS_STRING.format("GCSHook.get_conn"))
[email protected]("shutil.disk_usage")
+def test_download_works_if_enough_space(self, mock_disk_usage,
mock_get_conn):
+hook = gcs.GCSHook(gcp_conn_id="test")
+
+# Mock GCS connection and blob
+mock_client = mock.Mock()
+mock_get_conn.return_value = mock_client
+mock_bucket = mock.Mock()
+mock_client.bucket.return_value = mock_bucket
+mock_blob = mock.Mock()
+mock_bucket.blob.return_value = mock_blob
+
+# Blob size 1000 bytes
+mock_blob.size = 1000
+mock_blob.name = "test_blob"
+mock_bucket.name = "test_bucket"
+
+# Free space 2000 bytes
+mock_disk_usage.return_value = mock.Mock(total=4000, used=2000,
free=2000)
+
+filename = "/tmp/file"
+hook.download(bucket_name="bucket", object_name="object",
filename=filename)
+
+# Verify blob.reload() was called
+mock_blob.reload.assert_called_once()
+# Verify download was called
+mock_blob.download_to_filename.assert_called_once_with(filename,
timeout=mock.ANY)
Review Comment:
done.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add disk space check before downloading GCS file in GCSHook [airflow]
MonsterChenzhuo commented on code in PR #61949:
URL: https://github.com/apache/airflow/pull/61949#discussion_r2844731272
##
providers/google/tests/unit/google/cloud/hooks/test_gcs.py:
##
@@ -1843,3 +1846,88 @@ def get_logs_string(call_args_list):
assert "GCS object size (15) and local file size (9) differ." in
logs_string
assert f"Downloading dag_03.py to {sync_local_dir}/dag_03.py" in
logs_string
self.gcs_hook.download.assert_called_once()
+
+
+class TestGCSDownloadDiskCheck:
[email protected](GCS_STRING.format("GoogleBaseHook.__init__"),
new=mock_base_gcp_hook_default_project_id)
[email protected](GCS_STRING.format("GCSHook.get_conn"))
[email protected]("shutil.disk_usage")
+def test_download_raises_error_if_not_enough_space(self, mock_disk_usage,
mock_get_conn):
+hook = gcs.GCSHook(gcp_conn_id="test")
+
+# Mock GCS connection and blob
+mock_client = mock.Mock()
+mock_get_conn.return_value = mock_client
+mock_bucket = mock.Mock()
+mock_client.bucket.return_value = mock_bucket
+mock_blob = mock.Mock()
+mock_bucket.blob.return_value = mock_blob
+
+# Blob size 1000 bytes
+mock_blob.size = 1000
+
+# Free space 500 bytes
+# shutil.disk_usage returns a named tuple (total, used, free)
+mock_disk_usage.return_value = mock.Mock(total=2000, used=1500,
free=500)
+
+with pytest.raises(AirflowException, match="Insufficient disk space"):
+hook.download(bucket_name="bucket", object_name="object",
filename="/tmp/file")
+
+# Verify blob.reload() was called to get the size
+mock_blob.reload.assert_called_once()
Review Comment:
done.
##
providers/google/src/airflow/providers/google/cloud/hooks/gcs.py:
##
@@ -347,6 +344,16 @@ def download(
blob = bucket.blob(blob_name=object_name,
chunk_size=chunk_size)
if filename:
+blob.reload(timeout=timeout)
+if blob.size:
Review Comment:
done.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add disk space check before downloading GCS file in GCSHook [airflow]
MonsterChenzhuo commented on PR #61949: URL: https://github.com/apache/airflow/pull/61949#issuecomment-3949364694 @SameerMesiah97 @Srabasti PTAL, thanks. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add disk space check before downloading GCS file in GCSHook [airflow]
MonsterChenzhuo commented on code in PR #61949:
URL: https://github.com/apache/airflow/pull/61949#discussion_r2844703947
##
providers/google/src/airflow/providers/google/cloud/hooks/gcs.py:
##
@@ -347,6 +344,16 @@ def download(
blob = bucket.blob(blob_name=object_name,
chunk_size=chunk_size)
if filename:
+blob.reload(timeout=timeout)
+if blob.size:
Review Comment:
done.
##
providers/google/tests/unit/google/cloud/hooks/test_gcs.py:
##
@@ -1843,3 +1846,88 @@ def get_logs_string(call_args_list):
assert "GCS object size (15) and local file size (9) differ." in
logs_string
assert f"Downloading dag_03.py to {sync_local_dir}/dag_03.py" in
logs_string
self.gcs_hook.download.assert_called_once()
+
+
+class TestGCSDownloadDiskCheck:
[email protected](GCS_STRING.format("GoogleBaseHook.__init__"),
new=mock_base_gcp_hook_default_project_id)
[email protected](GCS_STRING.format("GCSHook.get_conn"))
[email protected]("shutil.disk_usage")
+def test_download_raises_error_if_not_enough_space(self, mock_disk_usage,
mock_get_conn):
+hook = gcs.GCSHook(gcp_conn_id="test")
+
+# Mock GCS connection and blob
+mock_client = mock.Mock()
+mock_get_conn.return_value = mock_client
+mock_bucket = mock.Mock()
+mock_client.bucket.return_value = mock_bucket
+mock_blob = mock.Mock()
+mock_bucket.blob.return_value = mock_blob
+
+# Blob size 1000 bytes
+mock_blob.size = 1000
+
+# Free space 500 bytes
+# shutil.disk_usage returns a named tuple (total, used, free)
+mock_disk_usage.return_value = mock.Mock(total=2000, used=1500,
free=500)
+
+with pytest.raises(AirflowException, match="Insufficient disk space"):
+hook.download(bucket_name="bucket", object_name="object",
filename="/tmp/file")
+
+# Verify blob.reload() was called to get the size
+mock_blob.reload.assert_called_once()
Review Comment:
done.
##
providers/google/tests/unit/google/cloud/hooks/test_gcs.py:
##
@@ -1843,3 +1846,88 @@ def get_logs_string(call_args_list):
assert "GCS object size (15) and local file size (9) differ." in
logs_string
assert f"Downloading dag_03.py to {sync_local_dir}/dag_03.py" in
logs_string
self.gcs_hook.download.assert_called_once()
+
+
+class TestGCSDownloadDiskCheck:
[email protected](GCS_STRING.format("GoogleBaseHook.__init__"),
new=mock_base_gcp_hook_default_project_id)
[email protected](GCS_STRING.format("GCSHook.get_conn"))
[email protected]("shutil.disk_usage")
+def test_download_raises_error_if_not_enough_space(self, mock_disk_usage,
mock_get_conn):
+hook = gcs.GCSHook(gcp_conn_id="test")
+
+# Mock GCS connection and blob
+mock_client = mock.Mock()
+mock_get_conn.return_value = mock_client
+mock_bucket = mock.Mock()
+mock_client.bucket.return_value = mock_bucket
+mock_blob = mock.Mock()
+mock_bucket.blob.return_value = mock_blob
+
+# Blob size 1000 bytes
+mock_blob.size = 1000
+
+# Free space 500 bytes
+# shutil.disk_usage returns a named tuple (total, used, free)
+mock_disk_usage.return_value = mock.Mock(total=2000, used=1500,
free=500)
+
+with pytest.raises(AirflowException, match="Insufficient disk space"):
+hook.download(bucket_name="bucket", object_name="object",
filename="/tmp/file")
+
+# Verify blob.reload() was called to get the size
+mock_blob.reload.assert_called_once()
+
[email protected](GCS_STRING.format("GoogleBaseHook.__init__"),
new=mock_base_gcp_hook_default_project_id)
[email protected](GCS_STRING.format("GCSHook.get_conn"))
[email protected]("shutil.disk_usage")
+def test_download_works_if_enough_space(self, mock_disk_usage,
mock_get_conn):
+hook = gcs.GCSHook(gcp_conn_id="test")
+
+# Mock GCS connection and blob
+mock_client = mock.Mock()
+mock_get_conn.return_value = mock_client
+mock_bucket = mock.Mock()
+mock_client.bucket.return_value = mock_bucket
+mock_blob = mock.Mock()
+mock_bucket.blob.return_value = mock_blob
+
+# Blob size 1000 bytes
+mock_blob.size = 1000
+mock_blob.name = "test_blob"
+mock_bucket.name = "test_bucket"
+
+# Free space 2000 bytes
+mock_disk_usage.return_value = mock.Mock(total=4000, used=2000,
free=2000)
+
+filename = "/tmp/file"
+hook.download(bucket_name="bucket", object_name="object",
filename=filename)
+
+# Verify blob.reload() was called
+mock_blob.reload.assert_called_once()
+# Verify download was called
+mock_blob.download_to_filename.assert_called_once_with(filename,
timeout=mock.ANY)
Review Comment:
done.
--
This is an automated message f
Re: [PR] Add disk space check before downloading GCS file in GCSHook [airflow]
MonsterChenzhuo commented on code in PR #61949:
URL: https://github.com/apache/airflow/pull/61949#discussion_r2844695794
##
providers/google/tests/unit/google/cloud/hooks/test_gcs.py:
##
@@ -1843,3 +1846,88 @@ def get_logs_string(call_args_list):
assert "GCS object size (15) and local file size (9) differ." in
logs_string
assert f"Downloading dag_03.py to {sync_local_dir}/dag_03.py" in
logs_string
self.gcs_hook.download.assert_called_once()
+
+
+class TestGCSDownloadDiskCheck:
[email protected](GCS_STRING.format("GoogleBaseHook.__init__"),
new=mock_base_gcp_hook_default_project_id)
[email protected](GCS_STRING.format("GCSHook.get_conn"))
[email protected]("shutil.disk_usage")
+def test_download_raises_error_if_not_enough_space(self, mock_disk_usage,
mock_get_conn):
+hook = gcs.GCSHook(gcp_conn_id="test")
+
+# Mock GCS connection and blob
+mock_client = mock.Mock()
+mock_get_conn.return_value = mock_client
+mock_bucket = mock.Mock()
+mock_client.bucket.return_value = mock_bucket
+mock_blob = mock.Mock()
+mock_bucket.blob.return_value = mock_blob
+
+# Blob size 1000 bytes
+mock_blob.size = 1000
+
+# Free space 500 bytes
+# shutil.disk_usage returns a named tuple (total, used, free)
+mock_disk_usage.return_value = mock.Mock(total=2000, used=1500,
free=500)
+
+with pytest.raises(AirflowException, match="Insufficient disk space"):
+hook.download(bucket_name="bucket", object_name="object",
filename="/tmp/file")
+
+# Verify blob.reload() was called to get the size
+mock_blob.reload.assert_called_once()
+
[email protected](GCS_STRING.format("GoogleBaseHook.__init__"),
new=mock_base_gcp_hook_default_project_id)
[email protected](GCS_STRING.format("GCSHook.get_conn"))
[email protected]("shutil.disk_usage")
+def test_download_works_if_enough_space(self, mock_disk_usage,
mock_get_conn):
Review Comment:
done.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add disk space check before downloading GCS file in GCSHook [airflow]
Srabasti commented on PR #61949: URL: https://github.com/apache/airflow/pull/61949#issuecomment-3906279097 Looks like trailing whitespaces at end of line are causing static checks to fail in multiple places @MonsterChenzhuo ! Can you run prek locally first and then commit from your branch? Prek fixes any formatting errors, automatically. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] Add disk space check before downloading GCS file in GCSHook [airflow]
SameerMesiah97 commented on code in PR #61949:
URL: https://github.com/apache/airflow/pull/61949#discussion_r2809743010
##
providers/google/tests/unit/google/cloud/hooks/test_gcs.py:
##
@@ -1843,3 +1846,88 @@ def get_logs_string(call_args_list):
assert "GCS object size (15) and local file size (9) differ." in
logs_string
assert f"Downloading dag_03.py to {sync_local_dir}/dag_03.py" in
logs_string
self.gcs_hook.download.assert_called_once()
+
+
+class TestGCSDownloadDiskCheck:
[email protected](GCS_STRING.format("GoogleBaseHook.__init__"),
new=mock_base_gcp_hook_default_project_id)
[email protected](GCS_STRING.format("GCSHook.get_conn"))
[email protected]("shutil.disk_usage")
+def test_download_raises_error_if_not_enough_space(self, mock_disk_usage,
mock_get_conn):
+hook = gcs.GCSHook(gcp_conn_id="test")
+
+# Mock GCS connection and blob
+mock_client = mock.Mock()
+mock_get_conn.return_value = mock_client
+mock_bucket = mock.Mock()
+mock_client.bucket.return_value = mock_bucket
+mock_blob = mock.Mock()
+mock_bucket.blob.return_value = mock_blob
+
+# Blob size 1000 bytes
+mock_blob.size = 1000
+
+# Free space 500 bytes
+# shutil.disk_usage returns a named tuple (total, used, free)
+mock_disk_usage.return_value = mock.Mock(total=2000, used=1500,
free=500)
+
+with pytest.raises(AirflowException, match="Insufficient disk space"):
+hook.download(bucket_name="bucket", object_name="object",
filename="/tmp/file")
+
+# Verify blob.reload() was called to get the size
+mock_blob.reload.assert_called_once()
+
[email protected](GCS_STRING.format("GoogleBaseHook.__init__"),
new=mock_base_gcp_hook_default_project_id)
[email protected](GCS_STRING.format("GCSHook.get_conn"))
[email protected]("shutil.disk_usage")
+def test_download_works_if_enough_space(self, mock_disk_usage,
mock_get_conn):
+hook = gcs.GCSHook(gcp_conn_id="test")
+
+# Mock GCS connection and blob
+mock_client = mock.Mock()
+mock_get_conn.return_value = mock_client
+mock_bucket = mock.Mock()
+mock_client.bucket.return_value = mock_bucket
+mock_blob = mock.Mock()
+mock_bucket.blob.return_value = mock_blob
+
+# Blob size 1000 bytes
+mock_blob.size = 1000
+mock_blob.name = "test_blob"
+mock_bucket.name = "test_bucket"
+
+# Free space 2000 bytes
+mock_disk_usage.return_value = mock.Mock(total=4000, used=2000,
free=2000)
+
+filename = "/tmp/file"
+hook.download(bucket_name="bucket", object_name="object",
filename=filename)
+
+# Verify blob.reload() was called
+mock_blob.reload.assert_called_once()
+# Verify download was called
+mock_blob.download_to_filename.assert_called_once_with(filename,
timeout=mock.ANY)
Review Comment:
Same here. See comment above (ensure both asserts are present).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] Add disk space check before downloading GCS file in GCSHook [airflow]
SameerMesiah97 commented on code in PR #61949:
URL: https://github.com/apache/airflow/pull/61949#discussion_r2809742430
##
providers/google/tests/unit/google/cloud/hooks/test_gcs.py:
##
@@ -1843,3 +1846,88 @@ def get_logs_string(call_args_list):
assert "GCS object size (15) and local file size (9) differ." in
logs_string
assert f"Downloading dag_03.py to {sync_local_dir}/dag_03.py" in
logs_string
self.gcs_hook.download.assert_called_once()
+
+
+class TestGCSDownloadDiskCheck:
[email protected](GCS_STRING.format("GoogleBaseHook.__init__"),
new=mock_base_gcp_hook_default_project_id)
[email protected](GCS_STRING.format("GCSHook.get_conn"))
[email protected]("shutil.disk_usage")
+def test_download_raises_error_if_not_enough_space(self, mock_disk_usage,
mock_get_conn):
+hook = gcs.GCSHook(gcp_conn_id="test")
+
+# Mock GCS connection and blob
+mock_client = mock.Mock()
+mock_get_conn.return_value = mock_client
+mock_bucket = mock.Mock()
+mock_client.bucket.return_value = mock_bucket
+mock_blob = mock.Mock()
+mock_bucket.blob.return_value = mock_blob
+
+# Blob size 1000 bytes
+mock_blob.size = 1000
+
+# Free space 500 bytes
+# shutil.disk_usage returns a named tuple (total, used, free)
+mock_disk_usage.return_value = mock.Mock(total=2000, used=1500,
free=500)
+
+with pytest.raises(AirflowException, match="Insufficient disk space"):
+hook.download(bucket_name="bucket", object_name="object",
filename="/tmp/file")
+
+# Verify blob.reload() was called to get the size
+mock_blob.reload.assert_called_once()
Review Comment:
I would add assertions for `mock_disk_usage` and `download_to_filename` as
well. Like this:
`mock_disk_usage.assert_called_once()`
`mock_blob.download_to_filename.assert_not_called()`
##
providers/google/src/airflow/providers/google/cloud/hooks/gcs.py:
##
@@ -347,6 +344,16 @@ def download(
blob = bucket.blob(blob_name=object_name,
chunk_size=chunk_size)
if filename:
+blob.reload(timeout=timeout)
+if blob.size:
Review Comment:
What about situations where blob.size is 0? Since 0 is falsy, won't this
result in the disk size check being skipped? I understand that this will not
have a functional impact, but we can be more explicit here. I think this would
be more precise:
`if blob.size is not None:`
##
providers/google/tests/unit/google/cloud/hooks/test_gcs.py:
##
@@ -1843,3 +1846,88 @@ def get_logs_string(call_args_list):
assert "GCS object size (15) and local file size (9) differ." in
logs_string
assert f"Downloading dag_03.py to {sync_local_dir}/dag_03.py" in
logs_string
self.gcs_hook.download.assert_called_once()
+
+
+class TestGCSDownloadDiskCheck:
[email protected](GCS_STRING.format("GoogleBaseHook.__init__"),
new=mock_base_gcp_hook_default_project_id)
[email protected](GCS_STRING.format("GCSHook.get_conn"))
[email protected]("shutil.disk_usage")
+def test_download_raises_error_if_not_enough_space(self, mock_disk_usage,
mock_get_conn):
+hook = gcs.GCSHook(gcp_conn_id="test")
+
+# Mock GCS connection and blob
+mock_client = mock.Mock()
+mock_get_conn.return_value = mock_client
+mock_bucket = mock.Mock()
+mock_client.bucket.return_value = mock_bucket
+mock_blob = mock.Mock()
+mock_bucket.blob.return_value = mock_blob
+
+# Blob size 1000 bytes
+mock_blob.size = 1000
+
+# Free space 500 bytes
+# shutil.disk_usage returns a named tuple (total, used, free)
+mock_disk_usage.return_value = mock.Mock(total=2000, used=1500,
free=500)
+
+with pytest.raises(AirflowException, match="Insufficient disk space"):
+hook.download(bucket_name="bucket", object_name="object",
filename="/tmp/file")
+
+# Verify blob.reload() was called to get the size
+mock_blob.reload.assert_called_once()
+
[email protected](GCS_STRING.format("GoogleBaseHook.__init__"),
new=mock_base_gcp_hook_default_project_id)
[email protected](GCS_STRING.format("GCSHook.get_conn"))
[email protected]("shutil.disk_usage")
+def test_download_works_if_enough_space(self, mock_disk_usage,
mock_get_conn):
Review Comment:
I think this test should handle boundary conditions i.e. `blob.size ==
free_space` and optionally 0. You could parameterize it by setting
`mock_blob.size` to 1000, 2000 and 0.
##
providers/google/tests/unit/google/cloud/hooks/test_gcs.py:
##
@@ -1843,3 +1846,88 @@ def get_logs_string(call_args_list):
assert "GCS object size (15) and local file size (9) differ." in
logs_string
assert f"Downloading dag_03.py to {sync_local_dir}/dag_03.py" in
logs_string
self.gcs_hook.download.assert_ca
