Lee-W commented on code in PR #52700:
URL: https://github.com/apache/airflow/pull/52700#discussion_r2256490033
##########
providers/sftp/src/airflow/providers/sftp/hooks/sftp.py:
##########
@@ -49,22 +50,17 @@
from airflow.models.connection import Connection
-exclude_methods = [
- "__init__",
- "get_ui_field_behaviour",
- "get_conn",
- "close_conn",
- "get_managed_conn",
- "get_conn_count",
-] # exclude the methods which are not using the connection, so that we won't
create a
-# connection when not needed, i.e init
+def handle_connection_management(func):
Review Comment:
```suggestion
def handle_connection_management(func: Callable) -> Callable:
```
##########
providers/sftp/src/airflow/providers/sftp/hooks/sftp.py:
##########
@@ -49,22 +50,17 @@
from airflow.models.connection import Connection
-exclude_methods = [
- "__init__",
- "get_ui_field_behaviour",
- "get_conn",
- "close_conn",
- "get_managed_conn",
- "get_conn_count",
-] # exclude the methods which are not using the connection, so that we won't
create a
-# connection when not needed, i.e init
+def handle_connection_management(func):
+ CONNECTION_NOT_OPEN_EXCEPTION: ConnectionNotOpenedException =
ConnectionNotOpenedException(
+ "Connection not open, use with hook.get_managed_conn() Managed
Connection in order to create and open the connection"
+ )
Review Comment:
```suggestion
```
##########
providers/sftp/src/airflow/providers/sftp/hooks/sftp.py:
##########
@@ -49,22 +50,17 @@
from airflow.models.connection import Connection
-exclude_methods = [
- "__init__",
- "get_ui_field_behaviour",
- "get_conn",
- "close_conn",
- "get_managed_conn",
- "get_conn_count",
-] # exclude the methods which are not using the connection, so that we won't
create a
-# connection when not needed, i.e init
+def handle_connection_management(func):
+ CONNECTION_NOT_OPEN_EXCEPTION: ConnectionNotOpenedException =
ConnectionNotOpenedException(
+ "Connection not open, use with hook.get_managed_conn() Managed
Connection in order to create and open the connection"
+ )
-def handle_connection_management_decorator(func):
- def handle_connection_management(self, *args, **kwargs):
+ @functools.wraps(func)
+ def handle_connection_management_wrapper(self, *args, **kwargs):
Review Comment:
```suggestion
def handle_connection_management_wrapper(self, *args: Any, **kwargs:
dict[str, Any]) -> Callable:
```
##########
providers/sftp/src/airflow/providers/sftp/hooks/sftp.py:
##########
@@ -49,22 +50,17 @@
from airflow.models.connection import Connection
-exclude_methods = [
- "__init__",
- "get_ui_field_behaviour",
- "get_conn",
- "close_conn",
- "get_managed_conn",
- "get_conn_count",
-] # exclude the methods which are not using the connection, so that we won't
create a
-# connection when not needed, i.e init
+def handle_connection_management(func):
+ CONNECTION_NOT_OPEN_EXCEPTION: ConnectionNotOpenedException =
ConnectionNotOpenedException(
+ "Connection not open, use with hook.get_managed_conn() Managed
Connection in order to create and open the connection"
+ )
-def handle_connection_management_decorator(func):
- def handle_connection_management(self, *args, **kwargs):
+ @functools.wraps(func)
+ def handle_connection_management_wrapper(self, *args, **kwargs):
if not self.managed_conn:
if self.conn is None:
- raise self.CONNECTION_NOT_OPEN_EXCEPTION
+ raise CONNECTION_NOT_OPEN_EXCEPTION
Review Comment:
```suggestion
raise ConnectionNotOpenedException(
"Connection not open, use with hook.get_managed_conn()
Managed Connection in order to create and open the connection"
)
```
##########
providers/sftp/src/airflow/providers/sftp/sensors/sftp.py:
##########
@@ -135,6 +137,20 @@ def poke(self, context: Context) -> PokeReturnValue | bool:
else:
files_found = actual_files_present
+ return files_found
+
+ def poke(self, context: Context) -> PokeReturnValue | bool:
+ self.hook = SFTPHook(self.sftp_conn_id, managed_conn=self.managed_conn)
+
+ self.log.info("Poking for %s, with pattern %s", self.path,
self.file_pattern)
+ files_found = []
+
+ if not self.managed_conn:
+ with self.hook.get_managed_conn():
+ files_found = self._get_files()
Review Comment:
I think we should be able to reuse the decorator
##########
providers/sftp/src/airflow/providers/sftp/hooks/sftp.py:
##########
@@ -86,10 +112,12 @@ def __init__(
self,
ssh_conn_id: str | None = "sftp_default",
host_proxy_cmd: str | None = None,
+ managed_conn: bool = True,
Review Comment:
not sure whether `use_managed_conn` is a better name
##########
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:
would be nice if we could list what are the type issue ignored
##########
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:
same here
##########
providers/sftp/src/airflow/providers/sftp/sensors/sftp.py:
##########
@@ -76,14 +77,15 @@ def __init__(
self.hook: SFTPHook | None = None
self.sftp_conn_id = sftp_conn_id
self.newer_than: datetime | str | None = newer_than
+ self.managed_conn = managed_conn
self.python_callable: Callable | None = python_callable
self.op_args = op_args or []
self.op_kwargs = op_kwargs or {}
self.deferrable = deferrable
- def poke(self, context: Context) -> PokeReturnValue | bool:
- self.hook = SFTPHook(self.sftp_conn_id)
- self.log.info("Poking for %s, with pattern %s", self.path,
self.file_pattern)
+ def _get_files(self):
Review Comment:
```suggestion
def _get_files(self) -> list[Path | str]:
```
kinda not sure the type
##########
providers/sftp/src/airflow/providers/sftp/sensors/sftp.py:
##########
@@ -67,6 +67,7 @@ def __init__(
python_callable: Callable | None = None,
op_args: list | None = None,
op_kwargs: dict[str, Any] | None = None,
+ managed_conn: bool = True,
Review Comment:
not sure whether use_managed_conn is a better name
##########
providers/sftp/src/airflow/providers/sftp/hooks/sftp.py:
##########
@@ -816,7 +831,7 @@ async def get_mod_time(self, path: str) -> str: # type:
ignore[return]
sftp_client = await ssh_conn.start_sftp_client()
ftp_mdtm = await sftp_client.stat(path)
modified_time = ftp_mdtm.mtime
- mod_time =
datetime.datetime.fromtimestamp(modified_time).strftime("%Y%m%d%H%M%S") #
type: ignore[arg-type]
+ mod_time =
datetime.datetime.fromtimestamp(modified_time).strftime("%Y%m%d%H%M%S") #
type: ignore
Review Comment:
Why do we remove this
--
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]
