fix: close the search client a handle owns, not the default alias

_ElasticsearchHandle.close() and _OpenSearchHandle.close() re-resolved
the client library's "default" connection alias at close time. That is
wrong on the SIGHUP reload path, which is not a shutdown path:
_main's reload block calls _init_output_clients() -- and therefore
set_hosts() -> connections.create_connection() -- before
_close_output_clients() on the old clients, deliberately, so a broken
new config leaves the old clients running. By then the alias names the
*new* client, so the old handle closed the new client and then dropped
the alias entirely. Every later save raised
KeyError("There is no connection with alias 'default'.") until
parsedmarc was restarted, and the original client was leaked.

Verified against the installed SDK sources rather than the docs
(AGENTS.md, "Verify bug claims against authoritative sources", point 3):

  elasticsearch/dsl/connections.py (elasticsearch 9.x) and
  opensearchpy/connection/connections.py (opensearch-py 3.x) both define

      def create_connection(self, alias="default", **kwargs):
          conn = self._conns[alias] = <Client>(**kwargs)

  -- so create_connection() overwrites whatever the alias held, and
  returns the very object it stored (elasticsearch.dsl passes it through
  _with_user_agent(), which mutates headers and returns the same object,
  so identity is preserved) --

      def remove_connection(self, alias):
          errors = 0
          for d in (self._conns, self._kwargs):
              try:
                  del d[alias]
              except KeyError:
                  errors += 1
          if errors == 2:
              raise KeyError(f"There is no connection with alias {alias!r}.")

  -- i.e. it deletes the alias outright, with no regard for which client
  currently holds it, and

      def get_connection(self, alias="default"):
          if not isinstance(alias, str):
              return alias            # (opensearch-py; elasticsearch.dsl
                                      #  returns _with_user_agent(alias))
          try:
              return self._conns[alias]
          ...
          raise KeyError(f"There is no connection with alias {alias!r}.")

  -- a str only comes *back* out of get_connection() when a str was
  passed *in* as the alias, which these call sites never do; the client
  stored under "default" is always the object create_connection() built.
  The old `isinstance(conn, str)` guard was therefore dead code and is
  removed rather than carried forward (AGENTS.md: delete unreachable
  branches, don't hide them). add_connection(alias, conn) stores an
  arbitrary object as-is, which is what the new tests use.

set_hosts() now returns the client it registered, _init_output_clients()
passes that client to the handle, and close() closes that object and
gives up the alias only when the registry's current "default" *is* that
same object (identity check). The init/close order in the reload block
is unchanged.

tests/test_cli.py gains TestSearchBackendHandles, which drives the real
connection registries (the SDK boundary) with fake client objects, so no
network is touched: per backend, a reload-shaped test asserts the old
client is closed, the new one is not, and the alias still resolves to the
new client; a shutdown-shaped test asserts the alias is released when it
is still ours; and a best-effort test asserts a raising close() neither
propagates nor blocks the alias release. Each test restores whatever the
process-wide registry held beforehand. All six fail against the unfixed
handles.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Sean Whalen
2026-09-09 19:06:37 -04:00
co-authored by Claude Opus 5
parent 77045c2df0
commit 3d387c29e8
5 changed files with 247 additions and 19 deletions
+180
View File
@@ -24,6 +24,7 @@ from unittest.mock import MagicMock, patch
import httpx
from azure.core.exceptions import ClientAuthenticationError
from elasticsearch import Elasticsearch
from kiota_abstractions.api_error import APIError
from msgraph.generated.models.o_data_errors.inner_error import InnerError
from msgraph.generated.models.o_data_errors.main_error import MainError
@@ -4173,6 +4174,185 @@ watch = true
es_client.close.assert_called_once()
class _FakeSearchConnection:
"""Stand-in for an Elasticsearch/OpenSearch client held in the real
connection registry. Records close() calls and nothing else, so no
network is touched while the registry itself stays real -- it is the
SDK boundary these handles are written against."""
def __init__(self, name: str):
self.name = name
self.close_count = 0
def __repr__(self):
return f"<fake {self.name} connection, closed {self.close_count}x>"
def close(self):
self.close_count += 1
class _RaisingSearchConnection(_FakeSearchConnection):
"""A client whose close() fails, e.g. because the transport is already
broken. The handles treat closing as best-effort."""
def close(self):
super().close()
raise RuntimeError("transport already closed")
class TestSearchBackendHandles(unittest.TestCase):
"""_ElasticsearchHandle / _OpenSearchHandle must close the client they
were built for, not whatever holds the ``default`` alias at close time.
A SIGHUP reload calls _init_output_clients -- which re-registers the
``default`` alias with the replacement client -- before it calls
_close_output_clients on the old clients (the order is deliberate, so a
broken new config leaves the old clients running). A handle that
re-resolved the alias when closing therefore closed the *new* client and
then dropped the alias entirely: both SDKs' Connections.remove_connection()
deletes the alias from ``_conns`` and ``_kwargs`` outright, so every later
save raised KeyError("There is no connection with alias 'default'.") and
the original client was never closed.
"""
def _isolate_default_alias(self, connections):
"""Restore whatever the process-wide registry held under the
``default`` alias, so these tests cannot leak into other tests
regardless of ordering."""
try:
previous = connections.get_connection("default")
except KeyError:
previous = None
def restore():
try:
connections.remove_connection("default")
except KeyError:
# Already gone: the test under this cleanup removed the
# alias itself, which is the state we are restoring to.
pass
if previous is not None:
connections.add_connection("default", previous)
self.addCleanup(restore)
def testElasticsearchHandleClosesItsOwnConnectionAfterReload(self):
connections = parsedmarc.elastic.connections
self._isolate_default_alias(connections)
old_conn = _FakeSearchConnection("old")
new_conn = _FakeSearchConnection("new")
connections.add_connection("default", cast(Elasticsearch, old_conn))
handle = parsedmarc.cli._ElasticsearchHandle(
connections.get_connection("default")
)
# What a reload does before the old handle is closed: the new
# client takes over the alias.
connections.add_connection("default", cast(Elasticsearch, new_conn))
handle.close()
self.assertEqual(old_conn.close_count, 1)
self.assertEqual(new_conn.close_count, 0)
self.assertIs(connections.get_connection("default"), new_conn)
def testElasticsearchHandleReleasesAliasOnShutdown(self):
"""The normal teardown path still frees the alias it owns."""
connections = parsedmarc.elastic.connections
self._isolate_default_alias(connections)
conn = _FakeSearchConnection("only")
connections.add_connection("default", cast(Elasticsearch, conn))
handle = parsedmarc.cli._ElasticsearchHandle(
connections.get_connection("default")
)
handle.close()
self.assertEqual(conn.close_count, 1)
with self.assertRaises(KeyError):
connections.get_connection("default")
def testElasticsearchHandleCloseIsBestEffort(self):
"""A client whose close() raises must not keep the handle from
releasing its alias. A second close() still calls the client's
close() again -- which raises and is swallowed -- even though the
alias-release step is now a no-op, since get_connection() raises
KeyError with the alias already gone."""
connections = parsedmarc.elastic.connections
self._isolate_default_alias(connections)
conn = _RaisingSearchConnection("raiser")
connections.add_connection("default", cast(Elasticsearch, conn))
handle = parsedmarc.cli._ElasticsearchHandle(
connections.get_connection("default")
)
handle.close()
self.assertEqual(conn.close_count, 1)
with self.assertRaises(KeyError):
connections.get_connection("default")
handle.close()
self.assertEqual(conn.close_count, 2)
def testOpenSearchHandleClosesItsOwnConnectionAfterReload(self):
connections = opensearch_module.connections
self._isolate_default_alias(connections)
old_conn = _FakeSearchConnection("old")
new_conn = _FakeSearchConnection("new")
connections.add_connection("default", old_conn)
handle = parsedmarc.cli._OpenSearchHandle(connections.get_connection("default"))
connections.add_connection("default", new_conn)
handle.close()
self.assertEqual(old_conn.close_count, 1)
self.assertEqual(new_conn.close_count, 0)
self.assertIs(connections.get_connection("default"), new_conn)
def testOpenSearchHandleReleasesAliasOnShutdown(self):
"""The normal teardown path still frees the alias it owns."""
connections = opensearch_module.connections
self._isolate_default_alias(connections)
conn = _FakeSearchConnection("only")
connections.add_connection("default", conn)
handle = parsedmarc.cli._OpenSearchHandle(connections.get_connection("default"))
handle.close()
self.assertEqual(conn.close_count, 1)
with self.assertRaises(KeyError):
connections.get_connection("default")
def testOpenSearchHandleCloseIsBestEffort(self):
"""A client whose close() raises must not keep the handle from
releasing its alias. A second close() still calls the client's
close() again -- which raises and is swallowed -- even though the
alias-release step is now a no-op, since get_connection() raises
KeyError with the alias already gone."""
connections = opensearch_module.connections
self._isolate_default_alias(connections)
conn = _RaisingSearchConnection("raiser")
connections.add_connection("default", conn)
handle = parsedmarc.cli._OpenSearchHandle(connections.get_connection("default"))
handle.close()
self.assertEqual(conn.close_count, 1)
with self.assertRaises(KeyError):
connections.get_connection("default")
handle.close()
self.assertEqual(conn.close_count, 2)
def _domain_map_tls_reports():
"""Four SMTP TLS reports for the index_prefix_domain_map tests: two
whose policy domains fold to the mapped base domain example.com (one of