rjgoyln commented on code in PR #73647:
URL: https://github.com/apache/airflow/pull/73647#discussion_r4120882572
##########
providers/sftp/src/airflow/providers/sftp/hooks/sftp.py:
##########
@@ -207,6 +207,45 @@ def get_conn_count(self) -> int:
"""Get the number of open connections."""
return self._conn_count
+ def _build_worker_hook(self) -> SFTPHook:
+ """
+ Build a new SFTPHook for a concurrent-transfer worker.
+
+ Mirrors this hook's effective connection settings -- i.e. the result
of merging
+ this hook's constructor overrides (``remote_host``, ``port``,
``username``, etc.)
+ with the underlying Airflow connection -- so worker hooks used by
+ ``store_directory_concurrently`` and
``retrieve_directory_concurrently`` connect
+ the same way the parent hook does, instead of falling back to the
connection's
+ raw defaults.
+ """
+ worker_hook = SFTPHook(
+ ssh_conn_id=self.ssh_conn_id,
+ remote_host=self.remote_host,
+ username=self.username,
+ password=self.password,
+ key_file=self.key_file,
Review Comment:
Re-passing the parent's resolved `key_file` can make worker construction
raise even though the parent succeeds. When `pkey` comes from `private_key` and
`key_file` is resolved from `~/.ssh/config`, passing both back to `SSHHook`
trips the `key_file`/`private_key` guard. Letting the worker re-resolve
`key_file` when `pkey` is set preserves the parent's behavior.
```suggestion
# Re-resolve key_file when pkey is set to avoid the
key_file/private_key guard.
key_file=None if self.pkey else self.key_file,
```
##########
providers/sftp/tests/unit/sftp/hooks/test_sftp.py:
##########
@@ -634,6 +634,103 @@ def test_store_and_retrieve_directory_concurrently(self):
)
assert retrieved_dir_name in os.listdir(os.path.join(self.temp_dir,
TMP_DIR_FOR_TESTS))
+ @patch("airflow.providers.sftp.hooks.sftp.SFTPHook.get_connection")
+ def test_build_worker_hook_inherits_parent_overrides(self,
mock_get_connection):
+ """
+ Regression test for #73585.
+
+ SFTPHook._build_worker_hook() must copy the parent hook's *effective*
+ connection settings (constructor overrides merged with the connection)
+ onto the worker hook it builds for concurrent transfers, not just
+ ssh_conn_id / no_host_key_check.
+ """
+ mock_connection = MagicMock()
+ mock_connection.login = "conn_user"
+ mock_connection.password = "conn_pass"
+ mock_connection.host = "conn.example.com"
+ mock_connection.port = 2222
+ mock_connection.extra = None
+ mock_get_connection.return_value = mock_connection
+
+ parent_hook = SFTPHook(
+ ssh_conn_id="sftp_default",
+ remote_host="override.example.com",
+ port=2022,
+ username="override_user",
+ password="override_pass",
+ key_file="/tmp/override_key",
+ conn_timeout=42,
+ host_proxy_cmd="ncat --proxy proxy_host:1234 %h %p",
+ )
+ # Simulate values that only ever come from the connection's `extra`
+ # field (no constructor parameter exists for these on SSHHook).
+ parent_hook.no_host_key_check = False
+ parent_hook.allow_host_key_change = True
+ parent_hook.look_for_keys = False
+
+ worker_hook = parent_hook._build_worker_hook()
+
+ assert worker_hook is not parent_hook
+ assert worker_hook.remote_host == "override.example.com"
+ assert worker_hook.port == 2022
+ assert worker_hook.username == "override_user"
+ assert worker_hook.password == "override_pass"
+ assert worker_hook.key_file == "/tmp/override_key"
+ assert worker_hook.conn_timeout == 42
+ assert worker_hook.host_proxy_cmd == "ncat --proxy proxy_host:1234 %h
%p"
+ assert worker_hook.no_host_key_check is False
+ assert worker_hook.allow_host_key_change is True
+ assert worker_hook.look_for_keys is False
+
+ def
test_store_and_retrieve_directory_concurrently_use_parent_overrides(self):
Review Comment:
These assertions don't distinguish the fixed code from the old behavior. The
fixture uses a bare `SFTPHook()`, so `remote_host`/`port`/`username` are the
same values that the old `SFTPHook(ssh_conn_id=...)` would resolve. Either give
the parent hook overrides that differ from the connection, or make this a
wiring test by asserting `mock_build.call_count == workers` and leaving the
value checks to `test_build_worker_hook_inherits_parent_overrides`.
##########
providers/sftp/tests/unit/sftp/hooks/test_sftp.py:
##########
@@ -634,6 +634,103 @@ def test_store_and_retrieve_directory_concurrently(self):
)
assert retrieved_dir_name in os.listdir(os.path.join(self.temp_dir,
TMP_DIR_FOR_TESTS))
+ @patch("airflow.providers.sftp.hooks.sftp.SFTPHook.get_connection")
+ def test_build_worker_hook_inherits_parent_overrides(self,
mock_get_connection):
+ """
+ Regression test for #73585.
+
+ SFTPHook._build_worker_hook() must copy the parent hook's *effective*
+ connection settings (constructor overrides merged with the connection)
+ onto the worker hook it builds for concurrent transfers, not just
+ ssh_conn_id / no_host_key_check.
+ """
+ mock_connection = MagicMock()
Review Comment:
`spec=Connection` keeps the mock honest here — the repo's testing standards
ask for `spec`/`autospec` on mocks, and the test still passes with it.
```suggestion
mock_connection = MagicMock(spec=Connection)
```
--
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]
