Skip to content

Commit

Permalink
Handle empty login and password with opensearch client (apache#39982)
Browse files Browse the repository at this point in the history
* Don't pass auth to opensearch client with empty login and password

This updates the opensearch hook to only pass the http_auth argument
to the opensearch client if a login and password are part of the
connection.

* Add typed dict for opensearch client arguments to satisfy mypy
  • Loading branch information
pdebelak authored Oct 4, 2024
1 parent 0711581 commit b0a18d9
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 4 deletions.
19 changes: 15 additions & 4 deletions airflow/providers/opensearch/hooks/opensearch.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@

import json
from functools import cached_property
from typing import TYPE_CHECKING, Any
from typing import TYPE_CHECKING, Any, TypedDict

from opensearchpy import OpenSearch, RequestsHttpConnection

Expand All @@ -31,6 +31,16 @@
from airflow.utils.strings import to_boolean


class OpenSearchClientArguments(TypedDict, total=False):
"""Typed arguments to the open search client."""

hosts: str | list[dict] | None
use_ssl: bool
verify_certs: bool
connection_class: type[OpenSearchConnectionClass] | None
http_auth: tuple[str, str]


class OpenSearchHook(BaseHook):
"""
Provide a thin wrapper around the OpenSearch client.
Expand Down Expand Up @@ -67,14 +77,15 @@ def conn(self):
@cached_property
def client(self) -> OpenSearch:
"""This function is intended for Operators that forward high level client objects."""
auth = (self.conn.login, self.conn.password)
client = OpenSearch(
client_args: OpenSearchClientArguments = dict(
hosts=[{"host": self.conn.host, "port": self.conn.port}],
http_auth=auth,
use_ssl=self.use_ssl,
verify_certs=self.verify_certs,
connection_class=self.connection_class,
)
if self.conn.login and self.conn.password:
client_args["http_auth"] = (self.conn.login, self.conn.password)
client = OpenSearch(**client_args)
return client

def search(self, query: dict, index_name: str, **kwargs: Any) -> Any:
Expand Down
29 changes: 29 additions & 0 deletions tests/providers/opensearch/hooks/test_opensearch.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
# under the License.
from __future__ import annotations

import json
from unittest import mock

import pytest
Expand Down Expand Up @@ -75,3 +76,31 @@ def test_load_conn_param(self, mock_hook):
open_search_conn_class=Urllib3HttpConnection,
)
assert hook_Urllib3.connection_class == Urllib3HttpConnection

def test_hook_with_auth(self, monkeypatch):
monkeypatch.setenv(
"AIRFLOW_CONN_OPENSEARCH_DEFAULT",
json.dumps(
{
"conn_type": "opensearch hook",
"host": "testhost",
"login": "testuser",
"password": "testpass",
}
),
)
hook = OpenSearchHook(open_search_conn_id="opensearch_default", log_query=True)
assert hook.client.transport.kwargs["http_auth"] == ("testuser", "testpass")

def test_hook_no_auth(self, monkeypatch):
monkeypatch.setenv(
"AIRFLOW_CONN_OPENSEARCH_DEFAULT",
json.dumps(
{
"conn_type": "opensearch hook",
"host": "testhost",
}
),
)
hook = OpenSearchHook(open_search_conn_id="opensearch_default", log_query=True)
assert "http_auth" not in hook.client.transport.kwargs

0 comments on commit b0a18d9

Please sign in to comment.