Nataneljpwd commented on code in PR #52700:
URL: https://github.com/apache/airflow/pull/52700#discussion_r2270882254
##########
providers/sftp/src/airflow/providers/sftp/hooks/sftp.py:
##########
@@ -287,28 +314,22 @@ def retrieve_file(self, remote_full_path: str,
local_full_path: str, prefetch: b
:param local_full_path: full path to the local file or a file-like
buffer
:param prefetch: controls whether prefetch is performed (default: True)
"""
- with self.get_managed_conn() as conn:
- if isinstance(local_full_path, BytesIO):
- # It's a file-like object ( BytesIO), so use getfo().
- self.log.info("Using streaming download for %s",
remote_full_path)
- conn.getfo(remote_full_path, local_full_path,
prefetch=prefetch)
- # We use hasattr checking for 'write' for cases like
google.cloud.storage.fileio.BlobWriter
- elif hasattr(local_full_path, "write"):
- self.log.info("Using streaming download for %s",
remote_full_path)
- # We need to cast to pass pre-commit checks
- stream_full_path = cast("IO[bytes]", local_full_path)
- conn.getfo(remote_full_path, stream_full_path,
prefetch=prefetch)
- elif isinstance(local_full_path, (str, bytes, os.PathLike)):
- # It's a string path, so use get().
- self.log.info("Using standard file download for %s",
remote_full_path)
- conn.get(remote_full_path, local_full_path, prefetch=prefetch)
- # If it's neither, it's an unsupported type.
- else:
- raise TypeError(
- f"Unsupported type for local_full_path:
{type(local_full_path)}. "
- "Expected a stream-like object or a path-like object."
- )
-
+ if isinstance(local_full_path, BytesIO):
+ self.conn.getfo(remote_full_path, local_full_path,
prefetch=prefetch) # type: ignore
Review Comment:
done
##########
providers/sftp/src/airflow/providers/sftp/hooks/sftp.py:
##########
@@ -319,20 +340,19 @@ def store_file(self, remote_full_path: str,
local_full_path: str, confirm: bool
:param remote_full_path: full path to the remote file
:param local_full_path: full path to the local file or a file-like
buffer
"""
- with self.get_managed_conn() as conn:
- if isinstance(local_full_path, BytesIO):
- conn.putfo(local_full_path, remote_full_path, confirm=confirm)
- else:
- conn.put(local_full_path, remote_full_path, confirm=confirm)
+ if isinstance(local_full_path, BytesIO):
+ self.conn.putfo(local_full_path, remote_full_path,
confirm=confirm) # type: ignore
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]
