Hugo van Kemenade
6e3f1a1194
* Upgrade syntax with pyupgrade --py38-plus Signed-off-by: Hugo van Kemenade <1324225+hugovk@users.noreply.github.com> * Convert to f-strings with flynt Signed-off-by: Hugo van Kemenade <1324225+hugovk@users.noreply.github.com> * Format with Black Signed-off-by: Hugo van Kemenade <1324225+hugovk@users.noreply.github.com> * Remove redundant mock backport dependency Signed-off-by: Hugo van Kemenade <1324225+hugovk@users.noreply.github.com> * isort imports Signed-off-by: Hugo van Kemenade <1324225+hugovk@users.noreply.github.com> * Add changelog entry Signed-off-by: Hugo van Kemenade <1324225+hugovk@users.noreply.github.com> --------- Signed-off-by: Hugo van Kemenade <1324225+hugovk@users.noreply.github.com>
116 lines
3.8 KiB
Python
116 lines
3.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# The OpenSearch Contributors require contributions made to
|
|
# this file be licensed under the Apache-2.0 license or a
|
|
# compatible open source license.
|
|
#
|
|
# Modifications Copyright OpenSearch Contributors. See
|
|
# GitHub history for details.
|
|
|
|
from typing import Any
|
|
|
|
import opensearchpy
|
|
from opensearchpy._async.helpers.actions import aiter
|
|
from opensearchpy.serializer import serializer
|
|
|
|
|
|
class AsyncConnections:
|
|
_conns: Any
|
|
|
|
"""
|
|
Class responsible for holding connections to different clusters. Used as a
|
|
singleton in this module.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self._kwargs: Any = {}
|
|
self._conns: Any = {}
|
|
|
|
async def configure(self, **kwargs: Any) -> None:
|
|
"""
|
|
Configure multiple connections at once, useful for passing in config
|
|
dictionaries obtained from other sources, like Django's settings or a
|
|
configuration management tool.
|
|
|
|
Example::
|
|
|
|
async_connections.configure(
|
|
default={'hosts': 'localhost'},
|
|
dev={'hosts': ['opensearchdev1.example.com:9200'], 'sniff_on_start': True},
|
|
)
|
|
|
|
Connections will only be constructed lazily when requested through
|
|
``get_connection``.
|
|
"""
|
|
async for k in aiter(list(self._conns)):
|
|
# try and preserve existing client to keep the persistent connections alive
|
|
if k in self._kwargs and kwargs.get(k, None) == self._kwargs[k]:
|
|
continue
|
|
del self._conns[k]
|
|
self._kwargs = kwargs
|
|
|
|
async def add_connection(self, alias: str, conn: Any) -> None:
|
|
"""
|
|
Add a connection object, it will be passed through as-is.
|
|
"""
|
|
self._conns[alias] = conn
|
|
|
|
async def remove_connection(self, alias: str) -> None:
|
|
"""
|
|
Remove connection from the registry. Raises ``KeyError`` if connection
|
|
wasn't found.
|
|
"""
|
|
errors = 0
|
|
async for d in aiter((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}.")
|
|
|
|
async def create_connection(self, alias: str = "default", **kwargs: Any) -> Any:
|
|
"""
|
|
Construct an instance of ``opensearchpy.AsyncOpenSearch`` and register
|
|
it under given alias.
|
|
"""
|
|
kwargs.setdefault("serializer", serializer)
|
|
conn = self._conns[alias] = opensearchpy.AsyncOpenSearch(**kwargs)
|
|
return conn
|
|
|
|
async def get_connection(self, alias: str = "default") -> Any:
|
|
"""
|
|
Retrieve a connection, construct it if necessary (only configuration
|
|
was passed to us). If a non-string alias has been passed through we
|
|
assume it's already a client instance and will just return it as-is.
|
|
|
|
Raises ``KeyError`` if no client (or its definition) is registered
|
|
under the alias.
|
|
"""
|
|
# do not check isinstance(AsyncOpenSearch) so that people can wrap their
|
|
# clients
|
|
if not isinstance(alias, str):
|
|
return alias
|
|
|
|
# connection already established
|
|
try:
|
|
return self._conns[alias]
|
|
except KeyError:
|
|
pass
|
|
|
|
# if not, try to create it
|
|
try:
|
|
return await self.create_connection(alias, **self._kwargs[alias])
|
|
except KeyError:
|
|
# no connection and no kwargs to set one up
|
|
raise KeyError(f"There is no connection with alias {alias!r}.")
|
|
|
|
|
|
async_connections = AsyncConnections()
|
|
configure = async_connections.configure
|
|
add_connection = async_connections.add_connection
|
|
remove_connection = async_connections.remove_connection
|
|
create_connection = async_connections.create_connection
|
|
get_connection = async_connections.get_connection
|