This is an automated email from the ASF dual-hosted git repository.
ashb 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 b6acedbab23 Add SRV DNS record support to HTTP Operator (#70960)
b6acedbab23 is described below
commit b6acedbab23f9a3bbfb307b5381ac2ffbcfc4255
Author: Gengtao Xu <[email protected]>
AuthorDate: Wed Sep 23 16:24:52 2026 -0400
Add SRV DNS record support to HTTP Operator (#70960)
* Add DNS SRV record support to HTTP Operator
* Fix fail-over mechanism never uses backup(lower-priority) SRV targets
* Document the srv extra required for HTTP SRV lookup
The srv extra is only listed in the provider's dependency table, so
nothing on the connection page tells you SRV lookup needs it. The
srv_lookup field default is left unset to match mongo's srv field.
---
providers/http/docs/connections/http.rst | 26 +-
providers/http/docs/index.rst | 17 +
providers/http/provider.yaml | 11 +-
providers/http/pyproject.toml | 7 +
.../http/src/airflow/providers/http/exceptions.py | 4 +
.../airflow/providers/http/get_provider_info.py | 8 +-
.../http/src/airflow/providers/http/hooks/http.py | 201 +++++++++-
providers/http/tests/unit/http/hooks/test_http.py | 409 ++++++++++++++++++++-
uv.lock | 7 +
9 files changed, 671 insertions(+), 19 deletions(-)
diff --git a/providers/http/docs/connections/http.rst
b/providers/http/docs/connections/http.rst
index ccaf6b55069..28414696985 100644
--- a/providers/http/docs/connections/http.rst
+++ b/providers/http/docs/connections/http.rst
@@ -47,12 +47,25 @@ 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.
+
+ Targets are tried in `RFC 2782 <https://www.rfc-editor.org/rfc/rfc2782>`__
order: lowest priority
+ first, weighted random within the same priority. If a connection to a
target cannot be established,
+ the request fails over to the next target.
+
+ Resolving SRV records requires the ``srv`` extra: ``pip install
apache-airflow-providers-http[srv]``.
+
Extra (optional)
Specify headers and default requests parameters in json format.
Following default requests parameters are taken into account:
@@ -64,6 +77,10 @@ Extra (optional)
* ``allow_redirects``
* ``max_redirects``
+ "Use DNS SRV Lookup" above is stored as the ``srv_lookup`` key in this
same Extra field, so it
+ 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 +92,10 @@ For example:
.. code-block:: bash
export
AIRFLOW_CONN_HTTP_DEFAULT='http://username:[email protected]: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='http://_http._tcp.example.com/https?srv_lookup=true'
diff --git a/providers/http/docs/index.rst b/providers/http/docs/index.rst
index 51df82172ef..44f274ff1f9 100644
--- a/providers/http/docs/index.rst
+++ b/providers/http/docs/index.rst
@@ -112,6 +112,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 1a53b392909..2f96ecda669 100644
--- a/providers/http/provider.yaml
+++ b/providers/http/provider.yaml
@@ -127,4 +127,13 @@ 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"
diff --git a/providers/http/pyproject.toml b/providers/http/pyproject.toml
index 197b91b0c6b..bafeb26952b 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 3c5f52cf655..43ea5c5a23a 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 93d137842de..ee468a44fc6 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,13 @@ 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"]},
+ }
+ },
}
],
}
diff --git a/providers/http/src/airflow/providers/http/hooks/http.py
b/providers/http/src/airflow/providers/http/hooks/http.py
index af27f886091..e4f5de03ed8 100644
--- a/providers/http/src/airflow/providers/http/hooks/http.py
+++ b/providers/http/src/airflow/providers/http/hooks/http.py
@@ -18,7 +18,8 @@
from __future__ import annotations
import copy
-from collections.abc import AsyncGenerator, Awaitable, Callable
+import random
+from collections.abc import AsyncGenerator, Awaitable, Callable, Iterable
from contextlib import asynccontextmanager
from typing import TYPE_CHECKING, Any, cast
from urllib.parse import urlparse
@@ -33,13 +34,16 @@ from requests.exceptions import ConnectionError, HTTPError
from requests.models import DEFAULT_REDIRECT_LIMIT
from requests_toolbelt.adapters.socket_options import TCPKeepAliveAdapter
from tenacity import retry_if_exception
+from urllib3.exceptions import ConnectTimeoutError
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
+ from dns.rdtypes.IN.SRV import SRV
from requests.adapters import HTTPAdapter
from airflow.models import Connection
@@ -52,6 +56,45 @@ def _url_from_endpoint(base_url: str | None, endpoint: str |
None) -> str:
return (base_url or "") + (endpoint or "")
+def _order_srv_targets(answers: Iterable[SRV]) -> list[tuple[str, int]]:
+ """
+ Order resolved DNS SRV records into the sequence of ``(host, port)``
targets to try.
+
+ Follows RFC 2782: lower priorities first, and within a priority a weighted
random order.
+ """
+ records_by_priority: dict[int, list[SRV]] = {}
+ for record in answers:
+ records_by_priority.setdefault(record.priority, []).append(record)
+
+ targets: list[tuple[str, int]] = []
+ for priority in sorted(records_by_priority):
+ remaining = records_by_priority[priority]
+ random.shuffle(remaining)
+ # RFC 2782 places zero-weight records first so they keep a small
chance of being picked first.
+ remaining.sort(key=lambda record: record.weight > 0)
+ while remaining:
+ threshold = random.randint(0, sum(record.weight for record in
remaining))
+ running_weight = 0
+ for index, record in enumerate(remaining):
+ running_weight += record.weight
+ if running_weight >= threshold:
+ chosen = remaining.pop(index)
+ break
+ targets.append((str(chosen.target).rstrip("."), chosen.port))
+ return targets
+
+
+def _is_connect_failure(exception: ConnectionError) -> bool:
+ """
+ Whether ``requests`` failed before a connection was established.
+
+ Only then is it safe to fail over to another SRV target: the request was
never sent, so
+ non-idempotent methods cannot be replayed.
+ """
+ reason = getattr(exception.args[0], "reason", None) if exception.args else
None
+ return isinstance(reason, ConnectTimeoutError)
+
+
def _process_extra_options_from_connection(
conn, extra_options: dict[str, Any]
) -> tuple[dict[str, Any], dict[str, Any]]:
@@ -76,6 +119,8 @@ 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)
+
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 +180,10 @@ 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``.
"""
conn_name_attr = "http_conn_id"
@@ -162,6 +211,9 @@ class HttpHook(BaseHook):
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"
# If no adapter is provided, use TCPKeepAliveAdapter (default behavior)
self.adapter = adapter
@@ -218,6 +270,8 @@ class HttpHook(BaseHook):
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)))
# RFC 3986 (https://www.rfc-editor.org/rfc/rfc3986.html#page-16)
if "://" in host:
self.base_url = host
@@ -228,8 +282,42 @@ class HttpHook(BaseHook):
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._base_url_initialized = True
+ def _get_srv_base_urls(self) -> list[str]:
+ """Resolve the SRV record into base URLs, in the order they should be
tried."""
+ return [
+ f"{self._srv_scheme}://{target_host}:{target_port}"
+ for target_host, target_port in
self._resolve_srv_targets(cast("str", self._srv_name))
+ ]
+
+ def _resolve_srv_targets(self, host: str) -> list[tuple[str, int]]:
+ """
+ Resolve a DNS SRV record to ``(host, port)`` targets, in the order
they should be tried.
+
+ 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 _order_srv_targets(answers)
+
def _configure_session_from_auth(self, session: Session, connection:
Connection) -> Session:
session.auth = self._extract_auth(connection)
return session
@@ -306,23 +394,33 @@ class HttpHook(BaseHook):
"""
extra_options = extra_options or {}
session = self.get_conn(headers, extra_options) # This sets
self.merged_extra, which is used later
- url = self.url_from_endpoint(endpoint)
+ if self._srv_lookup_enabled:
+ urls = [_url_from_endpoint(base_url, endpoint) for base_url in
self._get_srv_base_urls()]
+ else:
+ urls = [self.url_from_endpoint(endpoint)]
if self.method == "GET":
# GET uses params
- req = Request(self.method, url, params=data, headers=headers,
**request_kwargs)
+ req = Request(self.method, urls[0], params=data, headers=headers,
**request_kwargs)
elif self.method == "HEAD":
# HEAD doesn't use params
- req = Request(self.method, url, headers=headers, **request_kwargs)
+ req = Request(self.method, urls[0], headers=headers,
**request_kwargs)
else:
# Others use data
- req = Request(self.method, url, data=data, headers=headers,
**request_kwargs)
+ req = Request(self.method, urls[0], data=data, headers=headers,
**request_kwargs)
- prepped_request = session.prepare_request(req)
- self.log.debug("Sending '%s' to url: %s", self.method, url)
+ for target_number, url in enumerate(urls, start=1):
+ req.url = url
+ prepped_request = session.prepare_request(req)
+ self.log.debug("Sending '%s' to url: %s", self.method, url)
- # This is referencing self.merged_extra, which is update by _process
...
- return self.run_and_check(session, prepped_request, self.merged_extra)
+ try:
+ # This is referencing self.merged_extra, which is update by
_process ...
+ return self.run_and_check(session, prepped_request,
self.merged_extra)
+ except ConnectionError as ex:
+ if target_number == len(urls) or not _is_connect_failure(ex):
+ raise
+ self.log.warning("Could not connect to %s, failing over to the
next SRV target", url)
def check_response(self, response: Response) -> None:
"""
@@ -407,12 +505,18 @@ class HttpHook(BaseHook):
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 SRV record is
re-resolved and the
+ first target is combined 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)
+ base_url = self._get_srv_base_urls()[0] if self._srv_lookup_enabled
else self.base_url
+ return _url_from_endpoint(base_url=base_url, endpoint=endpoint)
def test_connection(self):
"""Test HTTP Connection."""
@@ -507,11 +611,33 @@ class AsyncHttpSession(LoggingMixin):
For example, ``run(json=obj)`` is passed as
``aiohttp.ClientSession().get(json=obj)``.
"""
- from tenacity import AsyncRetrying, stop_after_attempt, wait_fixed
-
- url = _url_from_endpoint(self.base_url, endpoint)
merged_headers = {**(self.headers or {}), **(headers or {})}
extra_options = {**(self.extra_options or {}), **(extra_options or {})}
+ if self._hook._srv_lookup_enabled:
+ base_urls = await self._hook._get_srv_base_urls_async()
+ else:
+ base_urls = [self.base_url]
+
+ for target_number, base_url in enumerate(base_urls, start=1):
+ url = _url_from_endpoint(base_url, endpoint)
+ try:
+ return await self._run_with_retries(url, data, json,
merged_headers, extra_options)
+ except (aiohttp.ClientConnectorError,
aiohttp.ConnectionTimeoutError):
+ if target_number == len(base_urls):
+ raise
+ self.log.warning("Could not connect to %s, failing over to the
next SRV target", url)
+
+ raise NotImplementedError # should not reach this, but makes mypy
happy
+
+ async def _run_with_retries(
+ self,
+ url: str,
+ data: dict[str, Any] | str | None,
+ json: dict[str, Any] | str | None,
+ merged_headers: dict[str, Any],
+ extra_options: dict[str, Any],
+ ) -> ClientResponse:
+ from tenacity import AsyncRetrying, stop_after_attempt, wait_fixed
async def request_func() -> ClientResponse:
response = await self._request(
@@ -558,6 +684,10 @@ 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``.
"""
conn_name_attr = "http_conn_id"
@@ -583,6 +713,9 @@ class HttpAsyncHook(BaseHook):
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"
def _get_request_func(
self, session: aiohttp.ClientSession, method: str | None = None
@@ -634,6 +767,15 @@ class HttpAsyncHook(BaseHook):
)
headers.update(conn_extra_options)
+ extra = conn.extra_dejson
+ self._srv_lookup_enabled =
to_boolean(str(extra.get("srv_lookup", False)))
+ 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 +784,35 @@ class HttpAsyncHook(BaseHook):
)
return self._config
+ async def _get_srv_base_urls_async(self) -> list[str]:
+ """Resolve the SRV record into base URLs, in the order they should be
tried."""
+ return [
+ f"{self._srv_scheme}://{target_host}:{target_port}"
+ for target_host, target_port in await
self._resolve_srv_targets_async(cast("str", self._srv_name))
+ ]
+
+ async def _resolve_srv_targets_async(self, host: str) -> list[tuple[str,
int]]:
+ """
+ Resolve a DNS SRV record to ordered ``(host, port)`` targets 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 _order_srv_targets(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 31cdb649a39..f3fc043ca78 100644
--- a/providers/http/tests/unit/http/hooks/test_http.py
+++ b/providers/http/tests/unit/http/hooks/test_http.py
@@ -19,9 +19,11 @@ from __future__ import annotations
import contextlib
import functools
+import importlib
import json
import logging
import os
+import sys
from http import HTTPStatus
from unittest import mock
@@ -32,13 +34,26 @@ import tenacity
from requests.adapters import HTTPAdapter, Response
from requests.auth import AuthBase, HTTPBasicAuth
from requests.models import DEFAULT_REDIRECT_LIMIT
+from urllib3.connection import HTTPConnection
+from urllib3.connectionpool import HTTPConnectionPool
+from urllib3.exceptions import ConnectTimeoutError, MaxRetryError,
NewConnectionError, ProtocolError
from airflow.models import Connection
from airflow.providers.common.compat.sdk import AirflowException
-from airflow.providers.http.hooks.http import HttpAsyncHook, HttpHook,
_process_extra_options_from_connection
+from airflow.providers.http.exceptions import HttpSrvLookupException
+from airflow.providers.http.hooks.http import (
+ HttpAsyncHook,
+ HttpHook,
+ _order_srv_targets,
+ _process_extra_options_from_connection,
+)
from tests_common.test_utils.aiohttp import MockAiohttpClientResponse
+# urllib3 requires a pool/connection on these exceptions; the fail-over path
never reads them.
+_DUMMY_POOL = HTTPConnectionPool("example.com")
+_DUMMY_CONN = HTTPConnection("example.com")
+
def get_airflow_dummy_connection(conn_id: str = "http_default"):
return Connection(conn_id=conn_id, conn_type="http", host="test:8080/")
@@ -615,6 +630,7 @@ class TestHttpHook:
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 +670,7 @@ class TestHttpHook:
"allow_redirects": False,
"max_redirects": 3,
"trust_env": False,
+ "srv_lookup": True,
}
)()
@@ -673,6 +690,238 @@ class TestHttpHook:
}
assert actual_conn_extra == {"bearer": "test"}
assert extra_options == {}
+ assert all(isinstance(value, str) for value in
actual_conn_extra.values())
+
+
[email protected]
+def stable_dns_import():
+ """
+ Re-import the ``dns`` submodules so ``sys.modules`` and the package
attributes agree.
+
+ In CI, Python 3.10's ``mock.patch`` resolves dotted targets
attribute-first, so a stale attribute left
+ by another test's ``patch.dict(sys.modules, ...)`` gets patched while the
hook re-imports a
+ fresh module — bypassing the mock and hitting real DNS.
+ """
+ importlib.import_module("dns.resolver")
+ importlib.import_module("dns.asyncresolver")
+
+
[email protected]("stable_dns_import")
+class TestHttpHookSrvLookup:
+ """Test DNS SRV record resolution support in HttpHook."""
+
+ @staticmethod
+ def _make_srv_answer(priority: int, port: int, target: str, weight: int =
0):
+ answer = mock.Mock()
+ answer.priority = priority
+ answer.port = port
+ answer.target = target
+ answer.weight = weight
+ 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}),
+ )
+ 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"
+
+ 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, 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}),
+ )
+ 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_targets_orders_by_priority(self, mock_resolve):
+ backup = self._make_srv_answer(priority=20, port=9999,
target="backup.example.com.")
+ primary = self._make_srv_answer(priority=10, port=8080,
target="primary.example.com.")
+ mock_resolve.return_value = [backup, primary]
+
+ hook = HttpHook()
+ targets = hook._resolve_srv_targets("_http._tcp.example.com")
+
+ assert targets == [("primary.example.com", 8080),
("backup.example.com", 9999)]
+ mock_resolve.assert_called_once_with("_http._tcp.example.com", "SRV")
+
+ @pytest.mark.parametrize(
+ ("threshold", "expected_targets"),
+ [
+ pytest.param(
+ 0, [("zero.example.com", 8080), ("heavy.example.com", 8081)],
id="zero-weight-first"
+ ),
+ pytest.param(5, [("heavy.example.com", 8081), ("zero.example.com",
8080)], id="weighted-first"),
+ ],
+ )
+ @mock.patch("random.randint")
+ @mock.patch("random.shuffle")
+ def test_order_srv_targets_uses_weighted_order_within_priority(
+ self, mock_shuffle, mock_randint, threshold, expected_targets
+ ):
+ heavy = self._make_srv_answer(priority=10, port=8081,
target="heavy.example.com.", weight=10)
+ zero = self._make_srv_answer(priority=10, port=8080,
target="zero.example.com.", weight=0)
+ mock_randint.side_effect = lambda low, high: min(threshold, high)
+
+ assert _order_srv_targets([heavy, zero]) == expected_targets
+
+ @mock.patch("dns.resolver.resolve")
+ def test_resolve_srv_targets_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_targets("_http._tcp.example.com")
+
+ def test_resolve_srv_targets_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_targets("_http._tcp.example.com")
+
+
@mock.patch("airflow.providers.http.hooks.http.HttpHook._resolve_srv_targets")
+ @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),
("svc-2.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.HttpHook._resolve_srv_targets")
+ @mock.patch("airflow.providers.http.hooks.http.HttpHook.get_connection")
+ def test_url_from_endpoint_resolves_srv_record_on_every_call(self,
mock_get_connection, mock_resolve):
+ conn = Connection(
+ conn_id="http_default",
+ conn_type="http",
+ host="_http._tcp.example.com",
+ extra=json.dumps({"srv_lookup": True}),
+ )
+ mock_get_connection.return_value = conn
+ mock_resolve.return_value = [("svc-1.example.com", 8080)]
+
+ hook = HttpHook()
+ hook.url_from_endpoint("a")
+ hook.url_from_endpoint("b")
+
+ assert mock_resolve.call_count == 2
+
+ @pytest.mark.parametrize(
+ ("primary_error", "expected_urls"),
+ [
+ pytest.param(
+ requests.exceptions.ConnectionError(
+ MaxRetryError(
+ _DUMMY_POOL, "/",
reason=NewConnectionError(_DUMMY_CONN, "Connection refused")
+ )
+ ),
+ ["https://primary.example.com:8443/v1/test";,
"https://backup.example.com:8443/v1/test";],
+ id="connect-failure-fails-over",
+ ),
+ pytest.param(
+ requests.exceptions.ConnectTimeout(
+ MaxRetryError(_DUMMY_POOL, "/",
reason=ConnectTimeoutError("Connect timed out"))
+ ),
+ ["https://primary.example.com:8443/v1/test";,
"https://backup.example.com:8443/v1/test";],
+ id="connect-timeout-fails-over",
+ ),
+ pytest.param(
+ requests.exceptions.ConnectionError(
+ ProtocolError("Connection aborted", ConnectionResetError())
+ ),
+ ["https://primary.example.com:8443/v1/test";],
+ id="failure-after-connect-does-not-fail-over",
+ ),
+ ],
+ )
+
@mock.patch("airflow.providers.http.hooks.http.HttpHook._resolve_srv_targets")
+ @mock.patch("airflow.providers.http.hooks.http.HttpHook.get_connection")
+ def test_run_fails_over_to_next_srv_target(
+ self, mock_get_connection, mock_resolve, requests_mock, primary_error,
expected_urls
+ ):
+ 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}),
+ )
+ mock_resolve.return_value = [("primary.example.com", 8443),
("backup.example.com", 8443)]
+ requests_mock.post("https://primary.example.com:8443/v1/test";,
exc=primary_error)
+ requests_mock.post("https://backup.example.com:8443/v1/test";,
status_code=200)
+
+ hook = HttpHook()
+ if len(expected_urls) == 1:
+ with pytest.raises(type(primary_error)):
+ hook.run("v1/test")
+ else:
+ assert hook.run("v1/test").status_code == 200
+
+ assert [request.url for request in requests_mock.request_history] ==
expected_urls
+
+
@mock.patch("airflow.providers.http.hooks.http.HttpHook._resolve_srv_targets")
+ @mock.patch("airflow.providers.http.hooks.http.HttpHook.get_connection")
+ def test_run_raises_when_all_srv_targets_unreachable(
+ self, mock_get_connection, mock_resolve, requests_mock
+ ):
+ 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}),
+ )
+ mock_resolve.return_value = [("primary.example.com", 8443),
("backup.example.com", 8443)]
+ for host in ("primary", "backup"):
+ requests_mock.post(
+ f"https://{host}.example.com:8443/v1/test";,
+ exc=requests.exceptions.ConnectionError(
+ MaxRetryError(_DUMMY_POOL, "/",
reason=NewConnectionError(_DUMMY_CONN, f"{host} refused"))
+ ),
+ )
+
+ hook = HttpHook()
+ with pytest.raises(requests.exceptions.ConnectionError, match="backup
refused"):
+ hook.run("v1/test")
+
+ assert requests_mock.call_count == 2
class TestHttpAsyncHook:
@@ -922,3 +1171,161 @@ class TestHttpAsyncHook:
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";
+
+
[email protected]("stable_dns_import")
+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, weight: int =
0):
+ answer = mock.Mock()
+ answer.priority = priority
+ answer.port = port
+ answer.target = target
+ answer.weight = weight
+ 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_resolves_srv_record_on_every_call(self, mock_resolve):
+ 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")
+
+ assert mock_resolve.call_count == 2
+
+ @pytest.mark.parametrize(
+ "primary_error",
+ [
+ pytest.param(aiohttp.ClientConnectorError(mock.Mock(),
OSError("refused")), id="connector-error"),
+ pytest.param(aiohttp.ConnectionTimeoutError("timed out"),
id="connection-timeout"),
+ ],
+ )
+ @pytest.mark.asyncio
+ @mock.patch("dns.asyncresolver.resolve", new_callable=mock.AsyncMock)
+ async def test_run_fails_over_to_next_srv_target(self, mock_resolve,
primary_error):
+ mock_resolve.return_value = [
+ self._make_srv_answer(20, 8443, "backup.example.com."),
+ self._make_srv_answer(10, 8443, "primary.example.com."),
+ ]
+ hook = HttpAsyncHook(http_conn_id="http_async_srv_conn", method="GET")
+ backup_response = MockAiohttpClientResponse(
+ status=200, payload={}, method="GET",
url="https://backup.example.com:8443/v1/test";
+ )
+
+ with mock.patch("aiohttp.ClientSession.get",
new_callable=mock.AsyncMock) as mocked_get:
+ mocked_get.side_effect = [primary_error, backup_response]
+ async with aiohttp.ClientSession() as session:
+ response = await hook.run(session=session, endpoint="v1/test")
+
+ assert response is backup_response
+ assert [call.args[0] for call in mocked_get.call_args_list] == [
+ "https://primary.example.com:8443/v1/test";,
+ "https://backup.example.com:8443/v1/test";,
+ ]
+
+ @pytest.mark.parametrize(
+ ("errors", "expected_error"),
+ [
+ pytest.param(
+ [aiohttp.ClientConnectorError(mock.Mock(),
OSError("refused"))] * 2,
+ aiohttp.ClientConnectorError,
+ id="all-targets-unreachable",
+ ),
+ pytest.param(
+ [aiohttp.ServerDisconnectedError()],
+ aiohttp.ServerDisconnectedError,
+ id="failure-after-connect-does-not-fail-over",
+ ),
+ ],
+ )
+ @pytest.mark.asyncio
+ @mock.patch("dns.asyncresolver.resolve", new_callable=mock.AsyncMock)
+ async def test_run_raises_without_further_failover(self, mock_resolve,
errors, expected_error):
+ mock_resolve.return_value = [
+ self._make_srv_answer(10, 8443, "primary.example.com."),
+ self._make_srv_answer(20, 8443, "backup.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.side_effect = errors
+ async with aiohttp.ClientSession() as session:
+ with pytest.raises(expected_error):
+ await hook.run(session=session, endpoint="v1/test")
+
+ assert mocked_get.call_count == len(errors)
+
+ @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_extra",
+ conn_type="http",
+ host="_http._tcp.example.com",
+ schema="https",
+ extra=json.dumps({"srv_lookup": True}),
+ )
+ )
+ hook = HttpAsyncHook(http_conn_id="http_async_srv_conn_extra",
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_targets_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_targets_async("_http._tcp.example.com")
+
+ @pytest.mark.asyncio
+ async def test_resolve_srv_targets_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_targets_async("_http._tcp.example.com")
diff --git a/uv.lock b/uv.lock
index 373adbbe4c0..b32aa190e52 100644
--- a/uv.lock
+++ b/uv.lock
@@ -6197,6 +6197,11 @@ dependencies = [
{ name = "requests-toolbelt" },
]
+[package.optional-dependencies]
+srv = [
+ { name = "dnspython" },
+]
+
[package.dev-dependencies]
dev = [
{ name = "apache-airflow" },
@@ -6215,10 +6220,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 = [
