This is an automated email from the ASF dual-hosted git repository.
potiuk 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 dc4742bea24 opensearch: generalize extra parameters (#65613)
dc4742bea24 is described below
commit dc4742bea24bb0b01d3adf6f08b3bc66a3339a50
Author: João Ramiro <[email protected]>
AuthorDate: Fri Jul 3 14:55:23 2026 +0200
opensearch: generalize extra parameters (#65613)
* opensearch: generalize extra parameters
* opensearch: specify all arguments
---
.../providers/opensearch/hooks/opensearch.py | 32 ++++++++++++++++++--
.../tests/unit/opensearch/hooks/test_opensearch.py | 34 ++++++++++++++++++----
2 files changed, 59 insertions(+), 7 deletions(-)
diff --git
a/providers/opensearch/src/airflow/providers/opensearch/hooks/opensearch.py
b/providers/opensearch/src/airflow/providers/opensearch/hooks/opensearch.py
index 2728ad0bfa5..0c3d8155ee9 100644
--- a/providers/opensearch/src/airflow/providers/opensearch/hooks/opensearch.py
+++ b/providers/opensearch/src/airflow/providers/opensearch/hooks/opensearch.py
@@ -36,6 +36,12 @@ class OpenSearchClientArguments(TypedDict, total=False):
hosts: str | list[dict] | None
use_ssl: bool
verify_certs: bool
+ url_prefix: str
+ timeout: int
+ headers: dict[str, str] | None
+ http_compress: bool | None
+ opaque_id: str | None
+ scheme: str
connection_class: type[OpenSearchConnectionClass] | None
http_auth: tuple[str, str]
@@ -64,8 +70,18 @@ class OpenSearchHook(BaseHook):
self.conn_id = open_search_conn_id
self.log_query = log_query
- self.use_ssl = to_boolean(str(self.conn.extra_dejson.get("use_ssl",
False)))
- self.verify_certs =
to_boolean(str(self.conn.extra_dejson.get("verify_certs", False)))
+ extra = self.conn.extra_dejson
+ self.use_ssl = to_boolean(str(extra.get("use_ssl", False)))
+ self.verify_certs = to_boolean(str(extra.get("verify_certs", False)))
+ self.url_prefix = extra.get("url_prefix")
+ self.timeout = int(extra["timeout"]) if extra.get("timeout") is not
None else None
+ self.headers = extra.get("headers")
+ self.http_compress = (
+ to_boolean(str(extra["http_compress"])) if
extra.get("http_compress") is not None else None
+ )
+ self.opaque_id = extra.get("opaque_id")
+ self.scheme = extra.get("scheme")
+
self.connection_class = open_search_conn_class
self.__SERVICE = "es"
@@ -82,6 +98,18 @@ class OpenSearchHook(BaseHook):
verify_certs=self.verify_certs,
connection_class=self.connection_class,
)
+ if self.scheme:
+ client_args["scheme"] = self.scheme
+ if self.url_prefix:
+ client_args["url_prefix"] = self.url_prefix
+ if self.timeout is not None:
+ client_args["timeout"] = self.timeout
+ if self.headers is not None:
+ client_args["headers"] = self.headers
+ if self.http_compress is not None:
+ client_args["http_compress"] = self.http_compress
+ if self.opaque_id is not None:
+ client_args["opaque_id"] = self.opaque_id
if self.conn.login and self.conn.password:
client_args["http_auth"] = (self.conn.login, self.conn.password)
client = OpenSearch(**client_args)
diff --git
a/providers/opensearch/tests/unit/opensearch/hooks/test_opensearch.py
b/providers/opensearch/tests/unit/opensearch/hooks/test_opensearch.py
index 3ab71d8e62f..bfddf075091 100644
--- a/providers/opensearch/tests/unit/opensearch/hooks/test_opensearch.py
+++ b/providers/opensearch/tests/unit/opensearch/hooks/test_opensearch.py
@@ -62,15 +62,39 @@ class TestOpenSearchHook:
hook.delete(index_name="test_index")
@mock.patch(f"{BASEHOOK_PATCH_PATH}.get_connection")
- def test_hook_param_bool(self, mock_get_connection):
+ @mock.patch("airflow.providers.opensearch.hooks.opensearch.OpenSearch")
+ def test_hook_extra_params(self, mock_opensearch, mock_get_connection):
mock_conn = Connection(
- conn_id="opensearch_default", extra={"use_ssl": "True",
"verify_certs": "True"}
+ conn_id="opensearch_default",
+ host="opensearch.local",
+ port=9200,
+ extra={
+ "use_ssl": True,
+ "verify_certs": False,
+ "scheme": "https",
+ "timeout": 30,
+ "http_compress": True,
+ "url_prefix": "os",
+ "headers": {"x-trace-id": "abc123"},
+ "opaque_id": "request-1",
+ },
)
mock_get_connection.return_value = mock_conn
hook = OpenSearchHook(open_search_conn_id="opensearch_default",
log_query=True)
-
- assert isinstance(hook.use_ssl, bool)
- assert isinstance(hook.verify_certs, bool)
+ hook.client
+
+ mock_opensearch.assert_called_once_with(
+ hosts=[{"host": "opensearch.local", "port": 9200}],
+ use_ssl=True,
+ verify_certs=False,
+ scheme="https",
+ timeout=30,
+ http_compress=True,
+ url_prefix="os",
+ headers={"x-trace-id": "abc123"},
+ opaque_id="request-1",
+ connection_class=DEFAULT_CONN,
+ )
def test_load_conn_param(self, mock_hook):
hook_default =
OpenSearchHook(open_search_conn_id="opensearch_default", log_query=True)
