From 09cda3c07cd5545cebbefbe30166fa8136fbf01f Mon Sep 17 00:00:00 2001 From: gtxu Date: Mon, 3 Aug 2026 23:04:48 -0400 Subject: [PATCH] Add DNS SRV record support to HTTP Operator --- providers/http/docs/connections/http.rst | 23 +- providers/http/docs/index.rst | 17 ++ providers/http/provider.yaml | 21 +- providers/http/pyproject.toml | 7 + .../src/airflow/providers/http/exceptions.py | 4 + .../providers/http/get_provider_info.py | 13 +- .../src/airflow/providers/http/hooks/http.py | 150 ++++++++++- .../http/tests/unit/http/hooks/test_http.py | 250 ++++++++++++++++++ uv.lock | 7 + 9 files changed, 484 insertions(+), 8 deletions(-) diff --git a/providers/http/docs/connections/http.rst b/providers/http/docs/connections/http.rst index ccaf6b550698e..c8e2b3e9517a6 100644 --- a/providers/http/docs/connections/http.rst +++ b/providers/http/docs/connections/http.rst @@ -47,12 +47,22 @@ Password (optional) Host (optional) Specify the entire url or the base of the url for the service. + If "Use DNS SRV Lookup" is enabled, specify the DNS SRV record name instead + (e.g. ``_http._tcp.example.com``) - Note the actual host and port are resolved from DNS at + request time and any value set in the Port field is ignored. + Port (optional) - Specify a port number if applicable. + Specify a port number if applicable. Ignored when SRV lookup is enabled. Schema (optional) Specify the service type etc: http/https. +Use DNS SRV Lookup (optional) + Treat the Host field as a DNS SRV record name and resolve the target host/port at request time. + +SRV Cache TTL (seconds) (optional) + Specify the time to cache a resolved SRV target before re-resolving. (default 60 seconds) + Extra (optional) Specify headers and default requests parameters in json format. Following default requests parameters are taken into account: @@ -64,6 +74,10 @@ Extra (optional) * ``allow_redirects`` * ``max_redirects`` + "Use DNS SRV Lookup" and "SRV Cache TTL" above are stored as the ``srv_lookup`` and + ``srv_cache_ttl`` keys in this same Extra field, so they can also be set directly in json + here, e.g. when configuring the connection via an environment variable. + When specifying the connection in environment variable you should specify it using URI syntax. @@ -75,3 +89,10 @@ For example: .. code-block:: bash export AIRFLOW_CONN_HTTP_DEFAULT='http://username:password@service.com:80/https?headers=header' + +To enable SRV lookup via an environment variable, set ``srv_lookup`` in the Extra query +parameter: + +.. code-block:: bash + + export AIRFLOW_CONN_HTTP_DEFAULT='https://_http._tcp.example.com/https?srv_lookup=true' diff --git a/providers/http/docs/index.rst b/providers/http/docs/index.rst index 2a6a1cf272b44..e7ccea3e84622 100644 --- a/providers/http/docs/index.rst +++ b/providers/http/docs/index.rst @@ -111,6 +111,23 @@ PIP package Version required ``pydantic`` ``>=2.11.0`` ========================================== ====================================== +Optional dependencies +--------------------- + +These extras install optional third-party libraries that enable additional features of the provider. +Install them when installing from PyPI. For example: + +.. code-block:: bash + + pip install apache-airflow-providers-http[srv] + + +======= ==================== +Extra Dependencies +======= ==================== +``srv`` ``dnspython>=2.0.0`` +======= ==================== + Downloading official packages ----------------------------- diff --git a/providers/http/provider.yaml b/providers/http/provider.yaml index f3814465559ee..cbcc7ce1087fe 100644 --- a/providers/http/provider.yaml +++ b/providers/http/provider.yaml @@ -126,4 +126,23 @@ connection-types: hidden-fields: [] relabeling: {} placeholders: {} - conn-fields: {} + conn-fields: + srv_lookup: + label: Use DNS SRV Lookup + description: >- + Whether to treat the Host field as a DNS SRV record name and resolve the target + host/port at request time. + schema: + type: + - boolean + - "null" + default: false + srv_cache_ttl: + label: SRV Cache TTL (seconds) + description: Time to cache a resolved SRV target before re-resolving. + schema: + type: + - number + - "null" + minimum: 0 + default: 60 diff --git a/providers/http/pyproject.toml b/providers/http/pyproject.toml index 73b38de0d0a2e..32bec45189f2a 100644 --- a/providers/http/pyproject.toml +++ b/providers/http/pyproject.toml @@ -71,6 +71,13 @@ dependencies = [ "pydantic>=2.11.0", ] +# The optional dependencies should be modified in place in the generated file +# Any change in the dependencies is preserved when the file is regenerated +[project.optional-dependencies] +"srv" = [ + "dnspython>=2.0.0", +] + [dependency-groups] dev = [ "apache-airflow", diff --git a/providers/http/src/airflow/providers/http/exceptions.py b/providers/http/src/airflow/providers/http/exceptions.py index 3c5f52cf655aa..43ea5c5a23a16 100644 --- a/providers/http/src/airflow/providers/http/exceptions.py +++ b/providers/http/src/airflow/providers/http/exceptions.py @@ -25,3 +25,7 @@ class HttpErrorException(AirflowException): class HttpMethodException(AirflowException): """Exception raised for invalid HTTP methods in Http hook.""" + + +class HttpSrvLookupException(AirflowException): + """Exception raised when DNS SRV record resolution fails or is misconfigured in Http hook.""" diff --git a/providers/http/src/airflow/providers/http/get_provider_info.py b/providers/http/src/airflow/providers/http/get_provider_info.py index 93d137842dea8..99454b61f0d98 100644 --- a/providers/http/src/airflow/providers/http/get_provider_info.py +++ b/providers/http/src/airflow/providers/http/get_provider_info.py @@ -66,7 +66,18 @@ def get_provider_info(): "hook-name": "HTTP", "connection-type": "http", "ui-field-behaviour": {"hidden-fields": [], "relabeling": {}, "placeholders": {}}, - "conn-fields": {}, + "conn-fields": { + "srv_lookup": { + "label": "Use DNS SRV Lookup", + "description": "Whether to treat the Host field as a DNS SRV record name and resolve the target host/port at request time.", + "schema": {"type": ["boolean", "null"], "default": False}, + }, + "srv_cache_ttl": { + "label": "SRV Cache TTL (seconds)", + "description": "Time to cache a resolved SRV target before re-resolving.", + "schema": {"type": ["number", "null"], "minimum": 0, "default": 60}, + }, + }, } ], } diff --git a/providers/http/src/airflow/providers/http/hooks/http.py b/providers/http/src/airflow/providers/http/hooks/http.py index 3401c589c812a..cc53ce50540b3 100644 --- a/providers/http/src/airflow/providers/http/hooks/http.py +++ b/providers/http/src/airflow/providers/http/hooks/http.py @@ -17,8 +17,11 @@ # under the License. from __future__ import annotations +import asyncio import copy -from collections.abc import AsyncGenerator, Awaitable, Callable +import random +import time +from collections.abc import AsyncGenerator, Awaitable, Callable, Iterable from contextlib import asynccontextmanager from typing import TYPE_CHECKING, Any, cast from urllib.parse import urlparse @@ -35,8 +38,9 @@ from tenacity import retry_if_exception from airflow.providers.common.compat.sdk import AirflowException, BaseHook -from airflow.providers.http.exceptions import HttpErrorException, HttpMethodException +from airflow.providers.http.exceptions import HttpErrorException, HttpMethodException, HttpSrvLookupException from airflow.utils.log.logging_mixin import LoggingMixin +from airflow.utils.strings import to_boolean if TYPE_CHECKING: from aiohttp.client_reqrep import ClientResponse @@ -52,6 +56,19 @@ def _url_from_endpoint(base_url: str | None, endpoint: str | None) -> str: return (base_url or "") + (endpoint or "") +def _select_srv_target(answers: Iterable[Any]) -> tuple[str, int]: + """Select a target host and port from resolved DNS SRV records.""" + candidates_by_priority: dict[int, list[Any]] = {} + for record in answers: + candidates_by_priority.setdefault(record.priority, []).append(record) + # Failover mechanism: RFC 2782 + candidates = candidates_by_priority[min(candidates_by_priority)] + chosen = random.choice(candidates) + + target_host = str(chosen.target).rstrip(".") + return target_host, chosen.port + + def _process_extra_options_from_connection( conn, extra_options: dict[str, Any] ) -> tuple[dict[str, Any], dict[str, Any]]: @@ -76,6 +93,9 @@ def _process_extra_options_from_connection( trust_env = conn_extra_options.pop("trust_env", None) check_response = conn_extra_options.pop("check_response", None) + conn_extra_options.pop("srv_lookup", None) + conn_extra_options.pop("srv_cache_ttl", None) + if stream is not None and "stream" not in passed_extra_options: passed_extra_options["stream"] = stream if cert is not None and "cert" not in passed_extra_options: @@ -135,6 +155,11 @@ class HttpHook(BaseHook): :param tcp_keep_alive_count: The TCP Keep Alive count parameter (corresponds to ``socket.TCP_KEEPCNT``) :param tcp_keep_alive_interval: The TCP Keep Alive interval parameter (corresponds to ``socket.TCP_KEEPINTVL``) + + Extra also supports resolving ``host`` via a DNS SRV record: + + * ``srv_lookup`` (bool): treat ``host`` as an SRV record name, e.g. ``_http._tcp.example.com``. + * ``srv_cache_ttl`` (float): SRV cache TTL in seconds (default 60). """ conn_name_attr = "http_conn_id" @@ -162,6 +187,12 @@ def __init__( self._base_url_initialized: bool = False self._retry_obj: Callable[..., Any] self._auth_type: Any = auth_type + self._srv_lookup_enabled: bool = False + self._srv_name: str | None = None + self._srv_scheme: str = "http" + self._srv_cache: tuple[str, int] | None = None + self._srv_cache_time: float = 0.0 + self._srv_cache_ttl: float = 60.0 # If no adapter is provided, use TCPKeepAliveAdapter (default behavior) self.adapter = adapter @@ -218,6 +249,9 @@ def get_conn( def _set_base_url(self, connection) -> None: host = connection.host or self.default_host schema = connection.schema or "http" + extra = connection.extra_dejson + self._srv_lookup_enabled = to_boolean(str(extra.get("srv_lookup", False))) + self._srv_cache_ttl = float(extra.get("srv_cache_ttl", self._srv_cache_ttl)) # RFC 3986 (https://www.rfc-editor.org/rfc/rfc3986.html#page-16) if "://" in host: self.base_url = host @@ -228,8 +262,48 @@ def _set_base_url(self, connection) -> None: parsed = urlparse(self.base_url) if not parsed.scheme: raise ValueError(f"Invalid base URL: Missing scheme in {self.base_url}") + if self._srv_lookup_enabled: + # When SRV lookup is enabled, ``host`` is the SRV record name (e.g. + # ``_http._tcp.example.com``), not a directly connectable hostname. + self._srv_name = parsed.hostname + self._srv_scheme = parsed.scheme + self._srv_cache = None + self._srv_cache_time = 0.0 self._base_url_initialized = True + def _get_dynamic_base_url(self) -> str: + """Return the base URL for the current request, resolving SRV records when enabled.""" + if not self._srv_lookup_enabled: + return self.base_url + now = time.monotonic() + if self._srv_cache is None or (now - self._srv_cache_time) >= self._srv_cache_ttl: + self._srv_cache = self._resolve_srv_record(cast("str", self._srv_name)) + self._srv_cache_time = now + target_host, target_port = self._srv_cache + return f"{self._srv_scheme}://{target_host}:{target_port}" + + def _resolve_srv_record(self, host: str) -> tuple[str, int]: + """ + Resolve a DNS SRV record to a target host and port. + + Requires the optional ``dnspython`` dependency. + """ + try: + import dns.exception + import dns.resolver + except ImportError as e: + raise HttpSrvLookupException( + "To use SRV DNS resolution in HttpHook, the 'dnspython' library must be installed. " + "Install it via the 'srv' extra: pip install apache-airflow-providers-http[srv]" + ) from e + + try: + answers = dns.resolver.resolve(host, "SRV") + except dns.exception.DNSException as e: + self.log.error("Failed to resolve SRV record for %s: %s", host, e) + raise HttpSrvLookupException(f"Failed to resolve SRV record for {host}: {e}") from e + return _select_srv_target(answers) + def _configure_session_from_auth(self, session: Session, connection: Connection) -> Session: session.auth = self._extract_auth(connection) return session @@ -407,12 +481,17 @@ def run_with_advanced_retry(self, _retry_args: dict[Any, Any], *args: Any, **kwa return self._retry_obj(self.run, *args, **kwargs) def url_from_endpoint(self, endpoint: str | None) -> str: - """Combine base url with endpoint.""" + """ + Combine base url with endpoint. + + If SRV lookup is enabled on the connection, the base URL is re-resolved (subject to + caching) before combining it with the endpoint. + """ # Ensure base_url is set by initializing it if it hasn't been initialized yet if not self._base_url_initialized and not self.base_url: connection = self.get_connection(self.http_conn_id) self._set_base_url(connection) - return _url_from_endpoint(base_url=self.base_url, endpoint=endpoint) + return _url_from_endpoint(base_url=self._get_dynamic_base_url(), endpoint=endpoint) def test_connection(self): """Test HTTP Connection.""" @@ -509,7 +588,7 @@ async def run( """ from tenacity import AsyncRetrying, stop_after_attempt, wait_fixed - url = _url_from_endpoint(self.base_url, endpoint) + url = _url_from_endpoint(await self._hook._get_dynamic_base_url_async(), endpoint) merged_headers = {**(self.headers or {}), **(headers or {})} extra_options = {**(self.extra_options or {}), **(extra_options or {})} @@ -558,6 +637,11 @@ class HttpAsyncHook(BaseHook): :param auth_type: The auth type for the service :param retry_limit: Maximum number of times to retry this job if it fails (default is 3) :param retry_delay: Delay between retry attempts (default is 1.0) + + Extra also supports resolving ``host`` via a DNS SRV record: + + * ``srv_lookup`` (bool): treat ``host`` as an SRV record name, e.g. ``_http._tcp.example.com``. + * ``srv_cache_ttl`` (float): SRV cache TTL in seconds (default 60). """ conn_name_attr = "http_conn_id" @@ -583,6 +667,13 @@ def __init__( self.retry_limit = retry_limit self.retry_delay = retry_delay self._config: SessionConfig | None = None + self._srv_lookup_enabled: bool = False + self._srv_name: str | None = None + self._srv_scheme: str = "http" + self._srv_cache: tuple[str, int] | None = None + self._srv_cache_time: float = 0.0 + self._srv_cache_ttl: float = 60.0 + self._srv_lock = asyncio.Lock() def _get_request_func( self, session: aiohttp.ClientSession, method: str | None = None @@ -634,6 +725,16 @@ async def config(self) -> SessionConfig: ) headers.update(conn_extra_options) + extra = conn.extra_dejson + self._srv_lookup_enabled = to_boolean(str(extra.get("srv_lookup", False))) + self._srv_cache_ttl = float(extra.get("srv_cache_ttl", self._srv_cache_ttl)) + if self._srv_lookup_enabled: + # When SRV lookup is enabled, ``host`` is the SRV record name (e.g. + # ``_http._tcp.example.com``), not a directly connectable hostname. + parsed = urlparse(base_url) + self._srv_name = parsed.hostname + self._srv_scheme = parsed.scheme + self._config = SessionConfig( base_url=base_url, headers=headers, @@ -642,6 +743,45 @@ async def config(self) -> SessionConfig: ) return self._config + async def _get_dynamic_base_url_async(self) -> str: + """Return the base URL for the current request, resolving SRV records when enabled.""" + config = await self.config() + if not self._srv_lookup_enabled: + return config.base_url + now = time.monotonic() + if self._srv_cache is None or (now - self._srv_cache_time) >= self._srv_cache_ttl: + async with self._srv_lock: + # Re-check after acquiring the lock: another concurrent request may have + # already refreshed the cache while this one was waiting. + now = time.monotonic() + if self._srv_cache is None or (now - self._srv_cache_time) >= self._srv_cache_ttl: + self._srv_cache = await self._resolve_srv_record_async(cast("str", self._srv_name)) + self._srv_cache_time = now + target_host, target_port = self._srv_cache + return f"{self._srv_scheme}://{target_host}:{target_port}" + + async def _resolve_srv_record_async(self, host: str) -> tuple[str, int]: + """ + Resolve a DNS SRV record to a target host and port without blocking the event loop. + + Requires the optional ``dnspython`` dependency. + """ + try: + import dns.asyncresolver + import dns.exception + except ImportError as e: + raise HttpSrvLookupException( + "To use SRV DNS resolution in HttpAsyncHook, the 'dnspython' library must be installed. " + "Install it via the 'srv' extra: pip install apache-airflow-providers-http[srv]" + ) from e + + try: + answers = await dns.asyncresolver.resolve(host, "SRV") + except dns.exception.DNSException as e: + self.log.error("Failed to resolve SRV record for %s: %s", host, e) + raise HttpSrvLookupException(f"Failed to resolve SRV record for {host}: {e}") from e + return _select_srv_target(answers) + @asynccontextmanager async def session(self, method: str | None = None) -> AsyncGenerator[AsyncHttpSession, None]: """ diff --git a/providers/http/tests/unit/http/hooks/test_http.py b/providers/http/tests/unit/http/hooks/test_http.py index f89532d7e2756..6a21feed95bdb 100644 --- a/providers/http/tests/unit/http/hooks/test_http.py +++ b/providers/http/tests/unit/http/hooks/test_http.py @@ -22,6 +22,7 @@ import json import logging import os +import sys from http import HTTPStatus from unittest import mock @@ -35,6 +36,7 @@ from airflow.models import Connection from airflow.providers.common.compat.sdk import AirflowException +from airflow.providers.http.exceptions import HttpSrvLookupException from airflow.providers.http.hooks.http import HttpAsyncHook, HttpHook, _process_extra_options_from_connection from tests_common.test_utils.aiohttp import MockAiohttpClientResponse @@ -615,6 +617,7 @@ def test_url_from_endpoint_lazy_initialization(self, mock_get_connection): mock_connection.host = "foo.bar.com" mock_connection.schema = "https" mock_connection.port = None + mock_connection.extra_dejson = {} mock_get_connection.return_value = mock_connection # Create hook without calling get_conn() and verify that base_url is not initialized @@ -654,6 +657,8 @@ def test_process_extra_options_from_connection(self): "allow_redirects": False, "max_redirects": 3, "trust_env": False, + "srv_lookup": True, + "srv_cache_ttl": 30, } )() @@ -673,6 +678,143 @@ def test_process_extra_options_from_connection(self): } assert actual_conn_extra == {"bearer": "test"} assert extra_options == {} + assert all(isinstance(value, str) for value in actual_conn_extra.values()) + + +class TestHttpHookSrvLookup: + """Test DNS SRV record resolution support in HttpHook.""" + + @staticmethod + def _make_srv_answer(priority: int, port: int, target: str): + answer = mock.Mock() + answer.priority = priority + answer.port = port + answer.target = target + return answer + + @mock.patch("airflow.providers.http.hooks.http.HttpHook.get_connection") + def test_set_base_url_enables_srv_lookup_from_extra(self, mock_get_connection): + conn = Connection( + conn_id="http_default", + conn_type="http", + host="_http._tcp.example.com", + schema="https", + extra=json.dumps({"srv_lookup": True, "srv_cache_ttl": 30}), + ) + mock_get_connection.return_value = conn + hook = HttpHook() + hook._set_base_url(conn) + assert hook._srv_lookup_enabled is True + assert hook._srv_name == "_http._tcp.example.com" + assert hook._srv_scheme == "https" + assert hook._srv_cache_ttl == 30 + + def test_set_base_url_srv_lookup_disabled_by_default(self): + conn = Connection(conn_id="http_default", conn_type="http", host="test.com") + hook = HttpHook() + hook._set_base_url(conn) + assert hook._srv_lookup_enabled is False + + @mock.patch("airflow.providers.http.hooks.http.HttpHook.get_connection") + def test_get_conn_only_puts_string_values_in_headers(self, mock_get_connection): + # Any consumed (non-header) extra option that isn't a string, e.g. srv_lookup/srv_cache_ttl, + # must never reach session.headers: requests rejects non-str/bytes header values outright. + mock_get_connection.return_value = Connection( + conn_id="http_default", + conn_type="http", + host="_http._tcp.example.com", + schema="https", + extra=json.dumps({"srv_lookup": True, "srv_cache_ttl": 30}), + ) + hook = HttpHook() + + session = hook.get_conn() + + assert all(isinstance(value, str) for value in session.headers.values()) + + @mock.patch("dns.resolver.resolve") + def test_resolve_srv_record_picks_lowest_priority(self, mock_resolve): + decoy = self._make_srv_answer(priority=20, port=9999, target="decoy.example.com.") + winner = self._make_srv_answer(priority=10, port=8080, target="svc-1.example.com.") + mock_resolve.return_value = [decoy, winner] + + hook = HttpHook() + host, port = hook._resolve_srv_record("_http._tcp.example.com") + + assert (host, port) == ("svc-1.example.com", 8080) + mock_resolve.assert_called_once_with("_http._tcp.example.com", "SRV") + + @mock.patch("random.choice") + @mock.patch("dns.resolver.resolve") + def test_resolve_srv_record_picks_randomly_among_ties(self, mock_resolve, mock_choice): + first = self._make_srv_answer(priority=10, port=8080, target="a.example.com.") + second = self._make_srv_answer(priority=10, port=8081, target="b.example.com.") + mock_resolve.return_value = [first, second] + mock_choice.return_value = second + + hook = HttpHook() + host, port = hook._resolve_srv_record("_http._tcp.example.com") + + mock_choice.assert_called_once_with([first, second]) + assert (host, port) == ("b.example.com", 8081) + + @mock.patch("dns.resolver.resolve") + def test_resolve_srv_record_dns_failure_raises(self, mock_resolve): + import dns.exception + + mock_resolve.side_effect = dns.exception.DNSException("boom") + hook = HttpHook() + + with pytest.raises(HttpSrvLookupException, match="Failed to resolve SRV record"): + hook._resolve_srv_record("_http._tcp.example.com") + + def test_resolve_srv_record_missing_dependency_raises(self): + hook = HttpHook() + with mock.patch.dict(sys.modules, {"dns": None, "dns.resolver": None, "dns.exception": None}): + with pytest.raises(HttpSrvLookupException, match="dnspython"): + hook._resolve_srv_record("_http._tcp.example.com") + + @mock.patch("airflow.providers.http.hooks.http.HttpHook._resolve_srv_record") + @mock.patch("airflow.providers.http.hooks.http.HttpHook.get_connection") + def test_url_from_endpoint_resolves_srv_record(self, mock_get_connection, mock_resolve): + conn = Connection( + conn_id="http_default", + conn_type="http", + host="_http._tcp.example.com", + schema="https", + extra=json.dumps({"srv_lookup": True}), + ) + mock_get_connection.return_value = conn + mock_resolve.return_value = ("svc-1.example.com", 8443) + + hook = HttpHook() + url = hook.url_from_endpoint("v1/test") + + assert url == "https://svc-1.example.com:8443/v1/test" + mock_resolve.assert_called_once_with("_http._tcp.example.com") + + @mock.patch("airflow.providers.http.hooks.http.time.monotonic") + @mock.patch("airflow.providers.http.hooks.http.HttpHook._resolve_srv_record") + @mock.patch("airflow.providers.http.hooks.http.HttpHook.get_connection") + def test_url_from_endpoint_caches_srv_resolution_within_ttl( + self, mock_get_connection, mock_resolve, mock_monotonic + ): + conn = Connection( + conn_id="http_default", + conn_type="http", + host="_http._tcp.example.com", + extra=json.dumps({"srv_lookup": True, "srv_cache_ttl": 60}), + ) + mock_get_connection.return_value = conn + mock_resolve.return_value = ("svc-1.example.com", 8080) + mock_monotonic.side_effect = [0.0, 10.0, 65.0] + + hook = HttpHook() + hook.url_from_endpoint("a") + hook.url_from_endpoint("b") # still within the 60s TTL: no re-resolution + hook.url_from_endpoint("c") # past the TTL: re-resolves + + assert mock_resolve.call_count == 2 class TestHttpAsyncHook: @@ -898,3 +1040,111 @@ async def test_build_request_url_from_endpoint_param(self): async with aiohttp.ClientSession() as session: await hook.run(session=session, endpoint="test.com:8080/v1/test") assert mocked_function.call_args.args[0] == "http://test.com:8080/v1/test" + + +class TestHttpAsyncHookSrvLookup: + """Test DNS SRV record resolution support in HttpAsyncHook.""" + + @pytest.fixture(autouse=True) + def setup_connections(self, create_connection_without_db): + create_connection_without_db( + Connection( + conn_id="http_async_srv_conn", + conn_type="http", + host="_http._tcp.example.com", + schema="https", + extra=json.dumps({"srv_lookup": True}), + ) + ) + + @staticmethod + def _make_srv_answer(priority: int, port: int, target: str): + answer = mock.Mock() + answer.priority = priority + answer.port = port + answer.target = target + return answer + + @pytest.mark.asyncio + @mock.patch("dns.asyncresolver.resolve", new_callable=mock.AsyncMock) + async def test_run_resolves_srv_record(self, mock_resolve): + mock_resolve.return_value = [self._make_srv_answer(10, 8443, "svc-1.example.com.")] + hook = HttpAsyncHook(http_conn_id="http_async_srv_conn", method="GET") + + with mock.patch("aiohttp.ClientSession.get", new_callable=mock.AsyncMock) as mocked_get: + mocked_get.return_value = MockAiohttpClientResponse( + status=200, + payload={"status": {"status": 200}}, + method="GET", + url="https://svc-1.example.com:8443/v1/test", + ) + async with aiohttp.ClientSession() as session: + await hook.run(session=session, endpoint="v1/test") + + assert mocked_get.call_args.args[0] == "https://svc-1.example.com:8443/v1/test" + mock_resolve.assert_called_once_with("_http._tcp.example.com", "SRV") + + @pytest.mark.asyncio + @mock.patch("dns.asyncresolver.resolve", new_callable=mock.AsyncMock) + async def test_run_caches_srv_resolution_within_ttl(self, mock_resolve): + # `time.monotonic` is also used by the asyncio event loop internals, so it can't be + # mocked wholesale here; instead, the cache timestamp is nudged directly to simulate + # TTL expiry. + mock_resolve.return_value = [self._make_srv_answer(10, 8080, "svc-1.example.com.")] + hook = HttpAsyncHook(http_conn_id="http_async_srv_conn", method="GET") + + with mock.patch("aiohttp.ClientSession.get", new_callable=mock.AsyncMock) as mocked_get: + mocked_get.return_value = MockAiohttpClientResponse( + status=200, payload={}, method="GET", url="https://svc-1.example.com:8080" + ) + async with aiohttp.ClientSession() as session: + await hook.run(session=session, endpoint="a") + await hook.run(session=session, endpoint="b") # within TTL: no re-resolution + + assert mock_resolve.call_count == 1 + + hook._srv_cache_time -= hook._srv_cache_ttl + 1 # simulate TTL expiry + + with mock.patch("aiohttp.ClientSession.get", new_callable=mock.AsyncMock) as mocked_get: + mocked_get.return_value = MockAiohttpClientResponse( + status=200, payload={}, method="GET", url="https://svc-1.example.com:8080" + ) + async with aiohttp.ClientSession() as session: + await hook.run(session=session, endpoint="c") # past TTL: re-resolves + + assert mock_resolve.call_count == 2 + + @pytest.mark.asyncio + async def test_config_only_puts_string_values_in_headers(self, create_connection_without_db): + create_connection_without_db( + Connection( + conn_id="http_async_srv_conn_with_ttl", + conn_type="http", + host="_http._tcp.example.com", + schema="https", + extra=json.dumps({"srv_lookup": True, "srv_cache_ttl": 30}), + ) + ) + hook = HttpAsyncHook(http_conn_id="http_async_srv_conn_with_ttl", method="GET") + + config = await hook.config() + + assert all(isinstance(value, str) for value in config.headers.values()) + + @pytest.mark.asyncio + async def test_resolve_srv_record_async_dns_failure_raises(self): + import dns.exception + + with mock.patch("dns.asyncresolver.resolve", new_callable=mock.AsyncMock) as mock_resolve: + mock_resolve.side_effect = dns.exception.DNSException("boom") + hook = HttpAsyncHook() + + with pytest.raises(HttpSrvLookupException, match="Failed to resolve SRV record"): + await hook._resolve_srv_record_async("_http._tcp.example.com") + + @pytest.mark.asyncio + async def test_resolve_srv_record_async_missing_dependency_raises(self): + hook = HttpAsyncHook() + with mock.patch.dict(sys.modules, {"dns": None, "dns.asyncresolver": None, "dns.exception": None}): + with pytest.raises(HttpSrvLookupException, match="dnspython"): + await hook._resolve_srv_record_async("_http._tcp.example.com") diff --git a/uv.lock b/uv.lock index b66ad967fd43c..4dd88047b9792 100644 --- a/uv.lock +++ b/uv.lock @@ -5903,6 +5903,11 @@ dependencies = [ { name = "requests-toolbelt" }, ] +[package.optional-dependencies] +srv = [ + { name = "dnspython" }, +] + [package.dev-dependencies] dev = [ { name = "apache-airflow" }, @@ -5921,10 +5926,12 @@ requires-dist = [ { name = "apache-airflow-providers-common-compat", editable = "providers/common/compat" }, { name = "asgiref", marker = "python_full_version < '3.14'", specifier = ">=2.3.0" }, { name = "asgiref", marker = "python_full_version >= '3.14'", specifier = ">=3.11.1" }, + { name = "dnspython", marker = "extra == 'srv'", specifier = ">=2.0.0" }, { name = "pydantic", specifier = ">=2.11.0" }, { name = "requests", specifier = ">=2.32.0,<3" }, { name = "requests-toolbelt", specifier = ">=1.0.0" }, ] +provides-extras = ["srv"] [package.metadata.requires-dev] dev = [