potiuk commented on code in PR #66746:
URL: https://github.com/apache/airflow/pull/66746#discussion_r3236995672
##########
providers/google/src/airflow/providers/google/cloud/hooks/compute_ssh.py:
##########
@@ -158,8 +159,47 @@ def __init__(
self.cmd_timeout = cmd_timeout
self.max_retries = max_retries
self.impersonation_chain = impersonation_chain
+ self.host_key_policy = host_key_policy
self._conn: Any | None = None
+ def _resolve_host_key_policy(self) -> paramiko.MissingHostKeyPolicy:
+ """
+ Resolve ``self.host_key_policy`` to a concrete paramiko policy
instance.
+
+ Accepts:
+
+ - the string aliases ``"auto_add"``, ``"reject"`` or ``"warning"`` —
+ mapped to the matching ``paramiko`` policy class;
+ - an instance of ``paramiko.MissingHostKeyPolicy`` — used as-is, so
+ callers can plug in a custom policy (e.g. one that loads pinned
+ host keys from GCE guest attributes).
+
+ Any other value raises :class:`AirflowException`.
Review Comment:
Done in `8e6c5c5682` — the code at line 197 actually raises `ValueError`, so
the docstring now matches. Thanks.
---
Drafted-by: Claude Opus 4.7 (1M context); reviewed by @potiuk before posting
##########
providers/google/tests/unit/google/cloud/hooks/test_compute_ssh.py:
##########
@@ -560,3 +560,45 @@ def test__authorize_compute_engine_instance_metadata(
mock_set_instance_metadata.call_args.kwargs["metadata"]["items"].sort(key=lambda
x: x["key"])
expected_metadata["items"].sort(key=lambda x: x["key"])
assert mock_set_instance_metadata.call_args.kwargs["metadata"] ==
expected_metadata
+
+
+class TestHostKeyPolicyResolution:
+ """Tests for the ``host_key_policy`` constructor argument."""
+
+ def test_default_is_auto_add(self):
+ import paramiko
Review Comment:
No reason — moved `import paramiko` to module top in `8e6c5c5682` and
removed the three function-local imports. Was an oversight when this test class
was extracted.
---
Drafted-by: Claude Opus 4.7 (1M context); reviewed by @potiuk before posting
##########
providers/google/src/airflow/providers/google/cloud/hooks/compute_ssh.py:
##########
@@ -158,8 +159,47 @@ def __init__(
self.cmd_timeout = cmd_timeout
self.max_retries = max_retries
self.impersonation_chain = impersonation_chain
+ self.host_key_policy = host_key_policy
self._conn: Any | None = None
+ def _resolve_host_key_policy(self) -> paramiko.MissingHostKeyPolicy:
+ """
+ Resolve ``self.host_key_policy`` to a concrete paramiko policy
instance.
+
+ Accepts:
+
+ - the string aliases ``"auto_add"``, ``"reject"`` or ``"warning"`` —
+ mapped to the matching ``paramiko`` policy class;
+ - an instance of ``paramiko.MissingHostKeyPolicy`` — used as-is, so
+ callers can plug in a custom policy (e.g. one that loads pinned
+ host keys from GCE guest attributes).
+
+ Any other value raises :class:`AirflowException`.
+
+ The default value (``"auto_add"``) preserves the historical behaviour
+ of this hook. Callers that want the remote SSH server authenticated
+ before the session opens should pass ``"reject"`` together with a
+ populated ``known_hosts`` file, or supply a custom policy that
+ looks the remote host's key up from an out-of-band source.
+ """
+ if not isinstance(self.host_key_policy, str):
+ # Trust the caller: an explicit paramiko.MissingHostKeyPolicy
+ # instance, or a subclass instance with custom behaviour.
+ return self.host_key_policy
+ builtins = {
+ "auto_add": paramiko.AutoAddPolicy,
+ "reject": paramiko.RejectPolicy,
+ "warning": paramiko.WarningPolicy,
+ }
+ try:
+ return builtins[self.host_key_policy]()
+ except KeyError:
+ raise ValueError(
+ f"Unknown host_key_policy {self.host_key_policy!r}. "
+ "Expected one of 'auto_add', 'reject', 'warning', "
+ "or an instance of paramiko.MissingHostKeyPolicy."
+ )
Review Comment:
Agreed — `raise ValueError(...) from None` in `8e6c5c5682`. The `KeyError`
from the dict lookup was indeed an implementation detail leaking into the
chained traceback.
---
Drafted-by: Claude Opus 4.7 (1M context); reviewed by @potiuk before posting
##########
providers/google/src/airflow/providers/google/cloud/hooks/compute_ssh.py:
##########
@@ -141,6 +141,7 @@ def __init__(
cmd_timeout: int | ArgNotSet = NOTSET,
max_retries: int = 10,
impersonation_chain: str | None = None,
+ host_key_policy: str | paramiko.MissingHostKeyPolicy = "auto_add",
Review Comment:
Added a `:param host_key_policy:` entry to the class docstring in
`8e6c5c5682` — covers the three string aliases, the
`paramiko.MissingHostKeyPolicy`-instance shape, and the `ValueError` failure
mode.
---
Drafted-by: Claude Opus 4.7 (1M context); reviewed by @potiuk before posting
--
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]
