diff --git a/providers/sftp/src/airflow/providers/sftp/hooks/sftp.py b/providers/sftp/src/airflow/providers/sftp/hooks/sftp.py index 80a3a12226372..540150f453046 100644 --- a/providers/sftp/src/airflow/providers/sftp/hooks/sftp.py +++ b/providers/sftp/src/airflow/providers/sftp/hooks/sftp.py @@ -29,7 +29,7 @@ from fnmatch import fnmatch from io import BytesIO from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import IO, TYPE_CHECKING, Any, cast import asyncssh from asgiref.sync import sync_to_async @@ -293,9 +293,25 @@ def retrieve_file(self, remote_full_path: str, local_full_path: str, prefetch: b """ 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) - else: + # 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." + ) def store_file(self, remote_full_path: str, local_full_path: str, confirm: bool = True) -> None: """