shahar1 commented on code in PR #66746:
URL: https://github.com/apache/airflow/pull/66746#discussion_r3223453740
##########
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:
Consider replacing with `raise ValueError(...) from None` as the `KeyError`
is an implementation detail of the dict lookup
##########
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:
Is there a reason for making this import internal? (same goes for the other
tests)
##########
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:
Could you please add a docstring?
##########
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:
```suggestion
Any other value raises :class:`ValueError`.
```
--
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]
