You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by we...@apache.org on 2025/08/25 01:31:02 UTC
(airflow) branch main updated: Feature: add optional managed connection (#52700)
This is an automated email from the ASF dual-hosted git repository.
weilee pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new bc35370c492 Feature: add optional managed connection (#52700)
bc35370c492 is described below
commit bc35370c49251b2349191db3d8f061b3f714aa3d
Author: Nataneljpwd <10...@users.noreply.github.com>
AuthorDate: Mon Aug 25 04:30:56 2025 +0300
Feature: add optional managed connection (#52700)
Co-authored-by: Natanel Rudyuklakir <na...@gmail.com>
---
.../tests/unit/always/test_project_structure.py | 1 +
.../sftp/src/airflow/providers/sftp/exceptions.py | 23 +++
.../sftp/src/airflow/providers/sftp/hooks/sftp.py | 211 +++++++++++----------
.../src/airflow/providers/sftp/sensors/sftp.py | 30 ++-
.../sftp/tests/unit/sftp/sensors/test_sftp.py | 19 +-
5 files changed, 176 insertions(+), 108 deletions(-)
diff --git a/airflow-core/tests/unit/always/test_project_structure.py b/airflow-core/tests/unit/always/test_project_structure.py
index ec776b98044..48b9de6fc5c 100644
--- a/airflow-core/tests/unit/always/test_project_structure.py
+++ b/airflow-core/tests/unit/always/test_project_structure.py
@@ -197,6 +197,7 @@ class TestProjectStructure:
"providers/standard/tests/unit/standard/sensors/test_external_task.py",
"providers/standard/tests/unit/standard/sensors/test_filesystem.py",
"providers/standard/tests/unit/standard/utils/test_sensor_helper.py",
+ "providers/sftp/tests/unit/sftp/test_exceptions.py",
]
modules_files: list[pathlib.Path] = list(
AIRFLOW_PROVIDERS_ROOT_PATH.glob("**/src/airflow/providers/**/*.py")
diff --git a/providers/sftp/src/airflow/providers/sftp/exceptions.py b/providers/sftp/src/airflow/providers/sftp/exceptions.py
new file mode 100644
index 00000000000..161cf648e3b
--- /dev/null
+++ b/providers/sftp/src/airflow/providers/sftp/exceptions.py
@@ -0,0 +1,23 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from airflow.exceptions import AirflowException
+
+
+class ConnectionNotOpenedException(AirflowException):
+ """Thrown when a connection has not been opened and has been tried to be used."""
diff --git a/providers/sftp/src/airflow/providers/sftp/hooks/sftp.py b/providers/sftp/src/airflow/providers/sftp/hooks/sftp.py
index d6a31752d13..4ea7e3f46f4 100644
--- a/providers/sftp/src/airflow/providers/sftp/hooks/sftp.py
+++ b/providers/sftp/src/airflow/providers/sftp/hooks/sftp.py
@@ -21,6 +21,7 @@ from __future__ import annotations
import concurrent.futures
import datetime
+import functools
import os
import stat
import warnings
@@ -34,7 +35,11 @@ from typing import IO, TYPE_CHECKING, Any, cast
import asyncssh
from asgiref.sync import sync_to_async
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import (
+ AirflowException,
+ AirflowProviderDeprecationWarning,
+)
+from airflow.providers.sftp.exceptions import ConnectionNotOpenedException
from airflow.providers.sftp.version_compat import BaseHook
from airflow.providers.ssh.hooks.ssh import SSHHook
@@ -46,6 +51,25 @@ if TYPE_CHECKING:
from airflow.models.connection import Connection
+def handle_connection_management(func: Callable) -> Callable:
+ @functools.wraps(func)
+ def handle_connection_management_wrapper(self, *args: Any, **kwargs: dict[str, Any]) -> Any:
+ if not self.use_managed_conn:
+ if self.conn is None:
+ raise ConnectionNotOpenedException(
+ "Connection not open, use with hook.get_managed_conn() Managed Connection in order to create and open the connection"
+ )
+
+ return func(self, *args, **kwargs)
+
+ with self.get_managed_conn() as conn:
+ self.conn = conn
+ result = func(self, *args, **kwargs)
+ return result
+
+ return handle_connection_management_wrapper
+
+
class SFTPHook(SSHHook):
"""
Interact with SFTP.
@@ -86,10 +110,12 @@ class SFTPHook(SSHHook):
self,
ssh_conn_id: str | None = "sftp_default",
host_proxy_cmd: str | None = None,
+ use_managed_conn: bool = True,
*args,
**kwargs,
) -> None:
self.conn: SFTPClient | None = None
+ self.use_managed_conn = use_managed_conn
# TODO: remove support for ssh_hook when it is removed from SFTPOperator
if kwargs.get("ssh_hook") is not None:
@@ -155,6 +181,7 @@ class SFTPHook(SSHHook):
"""Get the number of open connections."""
return self._conn_count
+ @handle_connection_management
def describe_directory(self, path: str) -> dict[str, dict[str, str | int | None]]:
"""
Get file information in a directory on the remote system.
@@ -164,36 +191,34 @@ class SFTPHook(SSHHook):
:param path: full path to the remote directory
"""
- with self.get_managed_conn() as conn: # type: SFTPClient
- flist = sorted(conn.listdir_attr(path), key=lambda x: x.filename)
- files = {}
- for f in flist:
- modify = datetime.datetime.fromtimestamp(f.st_mtime).strftime("%Y%m%d%H%M%S") # type: ignore
- files[f.filename] = {
- "size": f.st_size,
- "type": "dir" if stat.S_ISDIR(f.st_mode) else "file", # type: ignore
- "modify": modify,
- }
- return files
+ return {
+ f.filename: {
+ "size": f.st_size,
+ "type": "dir" if stat.S_ISDIR(f.st_mode) else "file", # type: ignore[union-attr]
+ "modify": datetime.datetime.fromtimestamp(f.st_mtime or 0).strftime("%Y%m%d%H%M%S"),
+ }
+ for f in sorted(self.conn.listdir_attr(path), key=lambda f: f.filename) # type: ignore[union-attr]
+ }
+ @handle_connection_management
def list_directory(self, path: str) -> list[str]:
"""
List files in a directory on the remote system.
:param path: full path to the remote directory to list
"""
- with self.get_managed_conn() as conn:
- return sorted(conn.listdir(path))
+ return sorted(self.conn.listdir(path)) # type: ignore[union-attr]
+ @handle_connection_management
def list_directory_with_attr(self, path: str) -> list[SFTPAttributes]:
"""
List files in a directory on the remote system including their SFTPAttributes.
:param path: full path to the remote directory to list
"""
- with self.get_managed_conn() as conn:
- return [file for file in conn.listdir_attr(path)]
+ return [file for file in self.conn.listdir_attr(path)] # type: ignore[union-attr]
+ @handle_connection_management
def mkdir(self, path: str, mode: int = 0o777) -> None:
"""
Create a directory on the remote system.
@@ -204,33 +229,33 @@ class SFTPHook(SSHHook):
:param path: full path to the remote directory to create
:param mode: int permissions of octal mode for directory
"""
- with self.get_managed_conn() as conn:
- conn.mkdir(path, mode=mode)
+ return self.conn.mkdir(path, mode) # type: ignore[union-attr,return-value]
+ @handle_connection_management
def isdir(self, path: str) -> bool:
"""
Check if the path provided is a directory.
:param path: full path to the remote directory to check
"""
- with self.get_managed_conn() as conn:
- try:
- return stat.S_ISDIR(conn.stat(path).st_mode) # type: ignore
- except OSError:
- return False
+ try:
+ return stat.S_ISDIR(self.conn.stat(path).st_mode) # type: ignore[union-attr,arg-type]
+ except OSError:
+ return False
+ @handle_connection_management
def isfile(self, path: str) -> bool:
"""
Check if the path provided is a file.
:param path: full path to the remote file to check
"""
- with self.get_managed_conn() as conn:
- try:
- return stat.S_ISREG(conn.stat(path).st_mode) # type: ignore
- except OSError:
- return False
+ try:
+ return stat.S_ISREG(self.conn.stat(path).st_mode) # type: ignore[arg-type,union-attr]
+ except OSError:
+ return False
+ @handle_connection_management
def create_directory(self, path: str, mode: int = 0o777) -> None:
"""
Create a directory on the remote system.
@@ -253,9 +278,9 @@ class SFTPHook(SSHHook):
self.create_directory(dirname, mode)
if basename:
self.log.info("Creating %s", path)
- with self.get_managed_conn() as conn:
- conn.mkdir(path, mode=mode)
+ self.conn.mkdir(path, mode=mode) # type: ignore
+ @handle_connection_management
def delete_directory(self, path: str, include_files: bool = False) -> None:
"""
Delete a directory on the remote system.
@@ -269,13 +294,13 @@ class SFTPHook(SSHHook):
files, dirs, _ = self.get_tree_map(path)
dirs = dirs[::-1] # reverse the order for deleting deepest directories first
- with self.get_managed_conn() as conn:
- for file_path in files:
- conn.remove(file_path)
- for dir_path in dirs:
- conn.rmdir(dir_path)
- conn.rmdir(path)
+ for file_path in files:
+ self.conn.remove(file_path) # type: ignore
+ for dir_path in dirs:
+ self.conn.rmdir(dir_path) # type: ignore
+ self.conn.rmdir(path) # type: ignore
+ @handle_connection_management
def retrieve_file(self, remote_full_path: str, local_full_path: str, prefetch: bool = True) -> None:
"""
Transfer the remote file to a local location.
@@ -287,28 +312,28 @@ class SFTPHook(SSHHook):
: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 prek hook 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):
+ # It's a file-like object ( BytesIO), so use getfo().
+ self.log.info("Using streaming download for %s", remote_full_path)
+ self.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 prek hook checks
+ stream_full_path = cast("IO[bytes]", local_full_path)
+ self.conn.getfo(remote_full_path, stream_full_path, prefetch=prefetch) # type: ignore[union-attr]
+ 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)
+ self.conn.get(remote_full_path, local_full_path, prefetch=prefetch) # type: ignore[union-attr]
+ # 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."
+ )
+ @handle_connection_management
def store_file(self, remote_full_path: str, local_full_path: str, confirm: bool = True) -> None:
"""
Transfer a local file to the remote location.
@@ -319,20 +344,19 @@ class SFTPHook(SSHHook):
: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
+ else:
+ self.conn.put(local_full_path, remote_full_path, confirm=confirm) # type: ignore
+ @handle_connection_management
def delete_file(self, path: str) -> None:
"""
Remove a file on the server.
:param path: full path to the remote file
"""
- with self.get_managed_conn() as conn:
- conn.remove(path)
+ self.conn.remove(path) # type: ignore[arg-type, union-attr]
def retrieve_directory(self, remote_full_path: str, local_full_path: str, prefetch: bool = True) -> None:
"""
@@ -348,14 +372,13 @@ class SFTPHook(SSHHook):
if Path(local_full_path).exists():
raise AirflowException(f"{local_full_path} already exists")
Path(local_full_path).mkdir(parents=True)
- with self.get_managed_conn():
- files, dirs, _ = self.get_tree_map(remote_full_path)
- for dir_path in dirs:
- new_local_path = os.path.join(local_full_path, os.path.relpath(dir_path, remote_full_path))
- Path(new_local_path).mkdir(parents=True, exist_ok=True)
- for file_path in files:
- new_local_path = os.path.join(local_full_path, os.path.relpath(file_path, remote_full_path))
- self.retrieve_file(file_path, new_local_path, prefetch)
+ files, dirs, _ = self.get_tree_map(remote_full_path)
+ for dir_path in dirs:
+ new_local_path = os.path.join(local_full_path, os.path.relpath(dir_path, remote_full_path))
+ Path(new_local_path).mkdir(parents=True, exist_ok=True)
+ for file_path in files:
+ new_local_path = os.path.join(local_full_path, os.path.relpath(file_path, remote_full_path))
+ self.retrieve_file(file_path, new_local_path, prefetch)
def retrieve_directory_concurrently(
self,
@@ -419,6 +442,7 @@ class SFTPHook(SSHHook):
for conn in conns:
conn.close()
+ @handle_connection_management
def store_directory(self, remote_full_path: str, local_full_path: str, confirm: bool = True) -> None:
"""
Transfer a local directory to the remote location.
@@ -431,21 +455,16 @@ class SFTPHook(SSHHook):
"""
if self.path_exists(remote_full_path):
raise AirflowException(f"{remote_full_path} already exists")
- with self.get_managed_conn():
- self.create_directory(remote_full_path)
- for root, dirs, files in os.walk(local_full_path):
- for dir_name in dirs:
- dir_path = os.path.join(root, dir_name)
- new_remote_path = os.path.join(
- remote_full_path, os.path.relpath(dir_path, local_full_path)
- )
- self.create_directory(new_remote_path)
- for file_name in files:
- file_path = os.path.join(root, file_name)
- new_remote_path = os.path.join(
- remote_full_path, os.path.relpath(file_path, local_full_path)
- )
- self.store_file(new_remote_path, file_path, confirm)
+ self.create_directory(remote_full_path)
+ for root, dirs, files in os.walk(local_full_path):
+ for dir_name in dirs:
+ dir_path = os.path.join(root, dir_name)
+ new_remote_path = os.path.join(remote_full_path, os.path.relpath(dir_path, local_full_path))
+ self.create_directory(new_remote_path)
+ for file_name in files:
+ file_path = os.path.join(root, file_name)
+ new_remote_path = os.path.join(remote_full_path, os.path.relpath(file_path, local_full_path))
+ self.store_file(new_remote_path, file_path, confirm)
def store_directory_concurrently(
self,
@@ -512,28 +531,28 @@ class SFTPHook(SSHHook):
for conn in conns:
conn.close()
+ @handle_connection_management
def get_mod_time(self, path: str) -> str:
"""
Get an entry's modification time.
:param path: full path to the remote file
"""
- with self.get_managed_conn() as conn:
- ftp_mdtm = conn.stat(path).st_mtime
- return datetime.datetime.fromtimestamp(ftp_mdtm).strftime("%Y%m%d%H%M%S") # type: ignore
+ ftp_mdtm = self.conn.stat(path).st_mtime # type: ignore[union-attr]
+ return datetime.datetime.fromtimestamp(ftp_mdtm).strftime("%Y%m%d%H%M%S") # type: ignore
+ @handle_connection_management
def path_exists(self, path: str) -> bool:
"""
Whether a remote entity exists.
:param path: full path to the remote file or directory
"""
- with self.get_managed_conn() as conn:
- try:
- conn.stat(path)
- except OSError:
- return False
- return True
+ try:
+ self.conn.stat(path) # type: ignore[union-attr]
+ except OSError:
+ return False
+ return True
@staticmethod
def _is_path_match(path: str, prefix: str | None = None, delimiter: str | None = None) -> bool:
diff --git a/providers/sftp/src/airflow/providers/sftp/sensors/sftp.py b/providers/sftp/src/airflow/providers/sftp/sensors/sftp.py
index a68d927aea9..79ab593b808 100644
--- a/providers/sftp/src/airflow/providers/sftp/sensors/sftp.py
+++ b/providers/sftp/src/airflow/providers/sftp/sensors/sftp.py
@@ -68,6 +68,7 @@ class SFTPSensor(BaseSensorOperator):
python_callable: Callable | None = None,
op_args: list | None = None,
op_kwargs: dict[str, Any] | None = None,
+ use_managed_conn: bool = True,
deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
**kwargs,
) -> None:
@@ -77,37 +78,37 @@ class SFTPSensor(BaseSensorOperator):
self.hook: SFTPHook | None = None
self.sftp_conn_id = sftp_conn_id
self.newer_than: datetime | str | None = newer_than
+ self.use_managed_conn = use_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)
- files_found = []
+ def _get_files(self) -> list[str]:
+ files_from_pattern: list[str] = []
+ files_found: list[str] = []
if self.file_pattern:
- files_from_pattern = self.hook.get_files_by_pattern(self.path, self.file_pattern)
+ files_from_pattern = self.hook.get_files_by_pattern(self.path, self.file_pattern) # type: ignore[union-attr]
if files_from_pattern:
actual_files_present = [
os.path.join(self.path, file_from_pattern) for file_from_pattern in files_from_pattern
]
else:
- return False
+ return files_found
else:
try:
# If a file is present, it is the single element added to the actual_files_present list to be
# processed. If the file is a directory, actual_file_present will be assigned an empty list,
# since SFTPHook.isfile(...) returns False
- actual_files_present = [self.path] if self.hook.isfile(self.path) else []
+ actual_files_present = [self.path] if self.hook.isfile(self.path) else [] # type: ignore[union-attr]
except Exception as e:
raise AirflowException from e
if self.newer_than:
for actual_file_present in actual_files_present:
try:
- mod_time = self.hook.get_mod_time(actual_file_present)
+ mod_time = self.hook.get_mod_time(actual_file_present) # type: ignore[union-attr]
self.log.info("Found File %s last modified: %s", actual_file_present, mod_time)
except OSError as e:
if e.errno != SFTP_NO_SUCH_FILE:
@@ -136,6 +137,19 @@ class SFTPSensor(BaseSensorOperator):
else:
files_found = actual_files_present
+ return files_found
+
+ def poke(self, context: Context) -> PokeReturnValue | bool:
+ self.hook = SFTPHook(self.sftp_conn_id, use_managed_conn=self.use_managed_conn)
+
+ self.log.info("Poking for %s, with pattern %s", self.path, self.file_pattern)
+
+ if self.use_managed_conn:
+ files_found = self._get_files()
+ else:
+ with self.hook.get_managed_conn():
+ files_found = self._get_files()
+
if not len(files_found):
return False
diff --git a/providers/sftp/tests/unit/sftp/sensors/test_sftp.py b/providers/sftp/tests/unit/sftp/sensors/test_sftp.py
index 2409355ead2..c7f84ee1794 100644
--- a/providers/sftp/tests/unit/sftp/sensors/test_sftp.py
+++ b/providers/sftp/tests/unit/sftp/sensors/test_sftp.py
@@ -97,6 +97,14 @@ class TestSFTPSensor:
sftp_hook_mock.return_value.close_conn.assert_not_called()
assert output
+ @patch("airflow.providers.sftp.sensors.sftp.SFTPHook")
+ def test_only_creating_one_connection_with_unmanaged_conn(self, sftp_hook_mock):
+ sftp_hook_mock.return_value.isfile.return_value = True
+ sftp_hook_mock.return_value.get_mod_time.return_value = "19700101000000"
+ sftp_sensor = SFTPSensor(task_id="test", path="path/to/whatever/test", use_managed_conn=False)
+ sftp_sensor.poke({})
+ assert sftp_hook_mock.return_value.get_managed_conn.called is True
+
@patch("airflow.providers.sftp.sensors.sftp.SFTPHook")
def test_file_not_new_enough(self, sftp_hook_mock):
sftp_hook_mock.return_value.isfile.return_value = True
@@ -193,9 +201,12 @@ class TestSFTPSensor:
)
context = {"ds": "1970-01-00"}
output = sftp_sensor.poke(context)
- sftp_hook_mock.return_value.get_mod_time.assert_has_calls(
- [mock.call("/path/to/file/text_file1.txt"), mock.call("/path/to/file/text_file2.txt")]
- )
+ assert (
+ [
+ mock.call("/path/to/file/text_file1.txt"),
+ mock.call("/path/to/file/text_file2.txt"),
+ ]
+ ) in sftp_hook_mock.return_value.get_mod_time.mock_calls
sftp_hook_mock.return_value.close_conn.assert_not_called()
assert output
@@ -220,7 +231,7 @@ class TestSFTPSensor:
)
context = {"ds": "1970-01-00"}
output = sftp_sensor.poke(context)
- sftp_hook_mock.return_value.get_mod_time.assert_has_calls(
+ sftp_hook_mock.return_value.get_mod_time.mock_call(
[
mock.call("/path/to/file/text_file1.txt"),
mock.call("/path/to/file/text_file2.txt"),