Fix user agent handling in DSL

Author: pquentin

elasticsearch/dsl/connections.py

@@ -117,15 +117,15 @@ def get_connection(self, alias: Union[str, _T] = "default") -> _T:
 
     def _with_user_agent(self, conn: _T) -> _T:
         # try to inject our user agent
-        if hasattr(conn, "_headers"):
-            is_frozen = conn._headers.frozen
+        if hasattr(conn, "_base_client") and hasattr(conn._base_client, "_headers"):
+            is_frozen = conn._base_client._headers.frozen
             if is_frozen:
-                conn._headers = conn._headers.copy()
-            conn._headers.update(
+                conn._base_client._headers = conn._base_client._headers.copy()
+            conn._base_client._headers.update(
                 {"user-agent": f"elasticsearch-dsl-py/{__versionstr__}"}
             )
             if is_frozen:
-                conn._headers.freeze()
+                conn._base_client._headers.freeze()
         return conn
 
 

test_elasticsearch/test_dsl/test_connections.py

@@ -130,12 +130,16 @@ def test_connection_has_correct_user_agent() -> None:
     my_client = Elasticsearch(hosts=["http://localhost:9200"])
     my_client = my_client.options(headers={"user-agent": "my-user-agent/1.0"})
     c.add_connection("default", my_client)
-    assert c.get_connection()._headers["user-agent"].startswith("elasticsearch-dsl-py/")
+    assert (
+        c.get_connection()
+        ._base_client._headers["user-agent"]
+        .startswith("elasticsearch-dsl-py/")
+    )
 
     my_client = Elasticsearch(hosts=["http://localhost:9200"])
     assert (
         c.get_connection(my_client)
-        ._headers["user-agent"]
+        ._base_client._headers["user-agent"]
         .startswith("elasticsearch-dsl-py/")
     )