How to use opensearch_client method in localstack

Best Python code snippet using localstack_python

provider.py

Source:provider.py

1from contextlib import contextmanager2from typing import Optional, cast3from botocore.exceptions import ClientError4from localstack import constants5from localstack.aws.api import RequestContext6from localstack.aws.api.es import (7    ARN,8    AccessDeniedException,9    AdvancedOptions,10    AdvancedSecurityOptionsInput,11    AutoTuneOptionsInput,12)13from localstack.aws.api.es import BaseException as EsBaseException14from localstack.aws.api.es import (15    CognitoOptions,16    CompatibleElasticsearchVersionsList,17    CompatibleVersionsMap,18    ConflictException,19    CreateElasticsearchDomainResponse,20    DeleteElasticsearchDomainResponse,21    DescribeElasticsearchDomainConfigResponse,22    DescribeElasticsearchDomainResponse,23    DescribeElasticsearchDomainsResponse,24    DisabledOperationException,25    DomainEndpointOptions,26    DomainInfoList,27    DomainName,28    DomainNameList,29    EBSOptions,30    ElasticsearchClusterConfig,31    ElasticsearchClusterConfigStatus,32    ElasticsearchDomainConfig,33    ElasticsearchDomainStatus,34    ElasticsearchVersionStatus,35    ElasticsearchVersionString,36    EncryptionAtRestOptions,37    EngineType,38    EsApi,39    GetCompatibleElasticsearchVersionsResponse,40    InternalException,41    InvalidPaginationTokenException,42    InvalidTypeException,43    LimitExceededException,44    ListDomainNamesResponse,45    ListElasticsearchVersionsResponse,46    ListTagsResponse,47    LogPublishingOptions,48    MaxResults,49    NextToken,50    NodeToNodeEncryptionOptions,51    OptionStatus,52    PolicyDocument,53    ResourceAlreadyExistsException,54    ResourceNotFoundException,55    SnapshotOptions,56    StringList,57    TagList,58    ValidationException,59    VPCOptions,60)61from localstack.aws.api.opensearch import (62    ClusterConfig,63    CompatibleVersionsList,64    DomainConfig,65    DomainStatus,66    VersionString,67)68from localstack.utils.analytics import event_publisher69from localstack.utils.aws import aws_stack70def _version_to_opensearch(71    version: Optional[ElasticsearchVersionString],72) -> Optional[VersionString]:73    if version is not None:74        if version.startswith("OpenSearch_"):75            return version76        else:77            return f"Elasticsearch_{version}"78def _version_from_opensearch(79    version: Optional[VersionString],80) -> Optional[ElasticsearchVersionString]:81    if version is not None:82        if version.startswith("Elasticsearch_"):83            return version.split("_")[1]84        else:85            return version86def _instancetype_to_opensearch(instance_type: Optional[str]) -> Optional[str]:87    if instance_type is not None:88        return instance_type.replace("elasticsearch", "search")89def _instancetype_from_opensearch(instance_type: Optional[str]) -> Optional[str]:90    if instance_type is not None:91        return instance_type.replace("search", "elasticsearch")92def _clusterconfig_from_opensearch(93    cluster_config: Optional[ClusterConfig],94) -> Optional[ElasticsearchClusterConfig]:95    if cluster_config is not None:96        # Just take the whole typed dict and typecast it to our target type97        result = cast(ElasticsearchClusterConfig, cluster_config)98        # Adjust the instance type names99        result["InstanceType"] = _instancetype_from_opensearch(cluster_config.get("InstanceType"))100        result["DedicatedMasterType"] = _instancetype_from_opensearch(101            cluster_config.get("DedicatedMasterType")102        )103        result["WarmType"] = _instancetype_from_opensearch(cluster_config.get("WarmType"))104        return result105def _domainstatus_from_opensearch(106    domain_status: Optional[DomainStatus],107) -> Optional[ElasticsearchDomainStatus]:108    if domain_status is not None:109        # Just take the whole typed dict and typecast it to our target type110        result = cast(ElasticsearchDomainStatus, domain_status)111        # Only specifically handle keys which are named differently or their values differ (version and clusterconfig)112        result["ElasticsearchVersion"] = _version_from_opensearch(113            domain_status.get("EngineVersion")114        )115        result["ElasticsearchClusterConfig"] = _clusterconfig_from_opensearch(116            domain_status.get("ClusterConfig")117        )118        result.pop("EngineVersion", None)119        result.pop("ClusterConfig", None)120        return result121def _clusterconfig_to_opensearch(122    elasticsearch_cluster_config: Optional[ElasticsearchClusterConfig],123) -> Optional[ClusterConfig]:124    if elasticsearch_cluster_config is not None:125        result = cast(ClusterConfig, elasticsearch_cluster_config)126        result["InstanceType"] = _instancetype_to_opensearch(result.get("InstanceType"))127        result["DedicatedMasterType"] = _instancetype_to_opensearch(128            result.get("DedicatedMasterType")129        )130        result["WarmType"] = _instancetype_to_opensearch(result.get("WarmType"))131        return result132def _domainconfig_from_opensearch(133    domain_config: Optional[DomainConfig],134) -> Optional[ElasticsearchDomainConfig]:135    if domain_config is not None:136        result = cast(ElasticsearchDomainConfig, domain_config)137        engine_version = domain_config.get("EngineVersion", {})138        result["ElasticsearchVersion"] = ElasticsearchVersionStatus(139            Options=_version_from_opensearch(engine_version.get("Options")),140            Status=cast(OptionStatus, engine_version.get("Status")),141        )142        cluster_config = domain_config.get("ClusterConfig", {})143        result["ElasticsearchClusterConfig"] = ElasticsearchClusterConfigStatus(144            Options=_clusterconfig_from_opensearch(cluster_config.get("Options")),145            Status=cluster_config.get("Status"),146        )147        result.pop("EngineVersion", None)148        result.pop("ClusterConfig", None)149        return result150def _compatible_version_list_from_opensearch(151    compatible_version_list: Optional[CompatibleVersionsList],152) -> Optional[CompatibleElasticsearchVersionsList]:153    if compatible_version_list is not None:154        return [155            CompatibleVersionsMap(156                SourceVersion=_version_from_opensearch(version_map["SourceVersion"]),157                TargetVersions=[158                    _version_from_opensearch(target_version)159                    for target_version in version_map["TargetVersions"]160                ],161            )162            for version_map in compatible_version_list163        ]164@contextmanager165def exception_mapper():166    """Maps an exception thrown by the OpenSearch client to an exception thrown by the ElasticSearch API."""167    try:168        yield169    except ClientError as err:170        exception_types = {171            "AccessDeniedException": AccessDeniedException,172            "BaseException": EsBaseException,173            "ConflictException": ConflictException,174            "DisabledOperationException": DisabledOperationException,175            "InternalException": InternalException,176            "InvalidPaginationTokenException": InvalidPaginationTokenException,177            "InvalidTypeException": InvalidTypeException,178            "LimitExceededException": LimitExceededException,179            "ResourceAlreadyExistsException": ResourceAlreadyExistsException,180            "ResourceNotFoundException": ResourceNotFoundException,181            "ValidationException": ValidationException,182        }183        mapped_exception_type = exception_types.get(err.response["Error"]["Code"], EsBaseException)184        raise mapped_exception_type(err.response["Error"]["Message"])185class EsProvider(EsApi):186    def create_elasticsearch_domain(187        self,188        context: RequestContext,189        domain_name: DomainName,190        elasticsearch_version: ElasticsearchVersionString = None,191        elasticsearch_cluster_config: ElasticsearchClusterConfig = None,192        ebs_options: EBSOptions = None,193        access_policies: PolicyDocument = None,194        snapshot_options: SnapshotOptions = None,195        vpc_options: VPCOptions = None,196        cognito_options: CognitoOptions = None,197        encryption_at_rest_options: EncryptionAtRestOptions = None,198        node_to_node_encryption_options: NodeToNodeEncryptionOptions = None,199        advanced_options: AdvancedOptions = None,200        log_publishing_options: LogPublishingOptions = None,201        domain_endpoint_options: DomainEndpointOptions = None,202        advanced_security_options: AdvancedSecurityOptionsInput = None,203        auto_tune_options: AutoTuneOptionsInput = None,204        tag_list: TagList = None,205    ) -> CreateElasticsearchDomainResponse:206        opensearch_client = aws_stack.connect_to_service("opensearch", region_name=context.region)207        # If no version is given, we set our default elasticsearch version208        engine_version = (209            _version_to_opensearch(elasticsearch_version)210            if elasticsearch_version211            else constants.ELASTICSEARCH_DEFAULT_VERSION212        )213        kwargs = {214            "DomainName": domain_name,215            "EngineVersion": engine_version,216            "ClusterConfig": _clusterconfig_to_opensearch(elasticsearch_cluster_config),217            "EBSOptions": ebs_options,218            "AccessPolicies": access_policies,219            "SnapshotOptions": snapshot_options,220            "VPCOptions": vpc_options,221            "CognitoOptions": cognito_options,222            "EncryptionAtRestOptions": encryption_at_rest_options,223            "NodeToNodeEncryptionOptions": node_to_node_encryption_options,224            "AdvancedOptions": advanced_options,225            "LogPublishingOptions": log_publishing_options,226            "DomainEndpointOptions": domain_endpoint_options,227            "AdvancedSecurityOptions": advanced_security_options,228            "AutoTuneOptions": auto_tune_options,229            "TagList": tag_list,230        }231        # Filter the kwargs to not set None values at all (boto doesn't like that)232        kwargs = {key: value for key, value in kwargs.items() if value is not None}233        with exception_mapper():234            domain_status = opensearch_client.create_domain(**kwargs)["DomainStatus"]235        # record event236        event_publisher.fire_event(237            event_publisher.EVENT_ES_CREATE_DOMAIN,238            payload={"n": event_publisher.get_hash(domain_name)},239        )240        status = _domainstatus_from_opensearch(domain_status)241        return CreateElasticsearchDomainResponse(DomainStatus=status)242    def delete_elasticsearch_domain(243        self, context: RequestContext, domain_name: DomainName244    ) -> DeleteElasticsearchDomainResponse:245        opensearch_client = aws_stack.connect_to_service("opensearch", region_name=context.region)246        with exception_mapper():247            domain_status = opensearch_client.delete_domain(248                DomainName=domain_name,249            )["DomainStatus"]250        # record event251        event_publisher.fire_event(252            event_publisher.EVENT_ES_DELETE_DOMAIN,253            payload={"n": event_publisher.get_hash(domain_name)},254        )255        status = _domainstatus_from_opensearch(domain_status)256        return DeleteElasticsearchDomainResponse(DomainStatus=status)257    def describe_elasticsearch_domain(258        self, context: RequestContext, domain_name: DomainName259    ) -> DescribeElasticsearchDomainResponse:260        opensearch_client = aws_stack.connect_to_service("opensearch", region_name=context.region)261        with exception_mapper():262            opensearch_status = opensearch_client.describe_domain(263                DomainName=domain_name,264            )["DomainStatus"]265        status = _domainstatus_from_opensearch(opensearch_status)266        return DescribeElasticsearchDomainResponse(DomainStatus=status)267    def describe_elasticsearch_domains(268        self, context: RequestContext, domain_names: DomainNameList269    ) -> DescribeElasticsearchDomainsResponse:270        opensearch_client = aws_stack.connect_to_service("opensearch", region_name=context.region)271        with exception_mapper():272            opensearch_status_list = opensearch_client.describe_domains(273                DomainNames=domain_names,274            )["DomainStatusList"]275        status_list = [_domainstatus_from_opensearch(s) for s in opensearch_status_list]276        return DescribeElasticsearchDomainsResponse(DomainStatusList=status_list)277    def list_domain_names(278        self, context: RequestContext, engine_type: EngineType = None279    ) -> ListDomainNamesResponse:280        opensearch_client = aws_stack.connect_to_service("opensearch", region_name=context.region)281        # Only hand the EngineType param to boto if it's set282        kwargs = {}283        if engine_type:284            kwargs["EngineType"] = engine_type285        with exception_mapper():286            domain_names = opensearch_client.list_domain_names(**kwargs)["DomainNames"]287        return ListDomainNamesResponse(DomainNames=cast(Optional[DomainInfoList], domain_names))288    def list_elasticsearch_versions(289        self,290        context: RequestContext,291        max_results: MaxResults = None,292        next_token: NextToken = None,293    ) -> ListElasticsearchVersionsResponse:294        opensearch_client = aws_stack.connect_to_service("opensearch", region_name=context.region)295        # Construct the arguments as kwargs to not set None values at all (boto doesn't like that)296        kwargs = {297            key: value298            for key, value in {"MaxResults": max_results, "NextToken": next_token}.items()299            if value is not None300        }301        with exception_mapper():302            versions = opensearch_client.list_versions(**kwargs)303        return ListElasticsearchVersionsResponse(304            ElasticsearchVersions=[305                _version_from_opensearch(version) for version in versions["Versions"]306            ],307            NextToken=versions.get(next_token),308        )309    def get_compatible_elasticsearch_versions(310        self, context: RequestContext, domain_name: DomainName = None311    ) -> GetCompatibleElasticsearchVersionsResponse:312        opensearch_client = aws_stack.connect_to_service("opensearch", region_name=context.region)313        # Only hand the DomainName param to boto if it's set314        kwargs = {}315        if domain_name:316            kwargs["DomainName"] = domain_name317        with exception_mapper():318            compatible_versions_response = opensearch_client.get_compatible_versions(**kwargs)319        compatible_versions = compatible_versions_response.get("CompatibleVersions")320        return GetCompatibleElasticsearchVersionsResponse(321            CompatibleElasticsearchVersions=_compatible_version_list_from_opensearch(322                compatible_versions323            )324        )325    def describe_elasticsearch_domain_config(326        self, context: RequestContext, domain_name: DomainName327    ) -> DescribeElasticsearchDomainConfigResponse:328        opensearch_client = aws_stack.connect_to_service("opensearch", region_name=context.region)329        with exception_mapper():330            domain_config = opensearch_client.describe_domain_config(DomainName=domain_name).get(331                "DomainConfig"332            )333        return DescribeElasticsearchDomainConfigResponse(334            DomainConfig=_domainconfig_from_opensearch(domain_config)335        )336    def add_tags(self, context: RequestContext, arn: ARN, tag_list: TagList) -> None:337        opensearch_client = aws_stack.connect_to_service("opensearch", region_name=context.region)338        with exception_mapper():339            opensearch_client.add_tags(ARN=arn, TagList=tag_list)340    def list_tags(self, context: RequestContext, arn: ARN) -> ListTagsResponse:341        opensearch_client = aws_stack.connect_to_service("opensearch", region_name=context.region)342        with exception_mapper():343            response = opensearch_client.list_tags(ARN=arn)344        return ListTagsResponse(TagList=response.get("TagList"))345    def remove_tags(self, context: RequestContext, arn: ARN, tag_keys: StringList) -> None:346        opensearch_client = aws_stack.connect_to_service("opensearch", region_name=context.region)347        with exception_mapper():...

opensearch.py

Source:opensearch.py

...8from opensearchpy.exceptions import RequestError9from typing import Any, Optional10logger = logging.getLogger(__name__)11opensearch_client = None12async def get_opensearch_client() -> Optional[OpenSearch]:13    """14    Create or return an OpenSearch client connected to the local15    OpenSearch server.16    :return: a connected OpenSearch client instance17    """18    global opensearch_client19    if not opensearch_client:20        opensearch_client = await create_opensearch_client()21    return opensearch_client22async def create_opensearch_client() -> Optional[OpenSearch]:23    """24    Create an OpenSearch client, connected to the configured OpenSearch server.25    :return: a connected OpenSearch client instance26    """27    settings = get_settings()28    # Create the client with SSL/TLS enabled, but hostname verification disabled29    client = OpenSearch(30        hosts=[{"host": settings.opensearch_server, "port": settings.opensearch_port}],31        http_compress=True,  # enables gzip compression for request bodies32        http_auth=(settings.opensearch_user, settings.opensearch_password),33        use_ssl=True,34        verify_certs=settings.certificate_verify,35        ssl_assert_hostname=False,36        ssl_show_warn=False,37        ca_certs=settings.certificate_authority_path,38    )39    logger.info("Created OpenSearch client")40    return client41async def add_index(index: str):42    """43    Create a new index.44    """45    client = await get_opensearch_client()46    try:47        response = client.indices.create(index=index, body={})48        logger.trace(f"Added index={index}, response={response}")49    except RequestError as re:50        if re.error == "resource_already_exists_exception":51            logger.trace(f"index={index} already exists, no need to create")52        else:53            logger.error(f"Exception creating index={index}, exception={re}")54            raise55async def add_document(index: str, document: dict):56    """57    Add a document to an index.58    :param index: Name of the index to add the document to59    :param document: The dict specifying the document to add.60        For example, to add a patient document for a patient with id = "001"61        and kafka and ipfs storage locations, use:62        {63            "patient_id": "001"",64            "data_record_location": message["data_record_location"],65            "ipfs_uri": message["ipfs_uri"],66        }67        where the LinuxForHealth message[] contains the kafka and ipfs data storage locations.68    :return:69    """70    client = await get_opensearch_client()71    response = client.index(index=index, body=document, refresh=True)72    logger.trace(f"Added document to index={index}, response={response}")73async def search_by_query(index: str, query: dict) -> Optional[dict]:74    """75    Search via a provided query.76    :param index: Name of the index to search77    :param query: The dict specifying the query search terms.78        For example, to search an index for a patient_id of "001", use:79        {"query": {"term": {"patient_id": "001"}}}80    :return: dict containing the search results81    """82    client = await get_opensearch_client()83    response = client.search(body=query, index=index)84    logger.trace(f"Search results for query={query} in index={index}: {response}")85    return response86async def delete_document(index: str, doc_id: str):87    """88    Delete a document from an index.89    :param index: Name of the index to add the document to90    :param doc_id: Id of the document to delete91    One way to use this function is to first query for records, then use the92    query results to get the ids of the documents to delete.93    """94    client = await get_opensearch_client()95    response = client.delete(index=index, id=doc_id)96    logger.trace(97        f"Deleted document with id={doc_id} from index={index}, response={response}"98    )99async def delete_index(index: str):100    """101    Delete an index.102    :param index: Name of the index to delete103    """104    client = await get_opensearch_client()105    response = client.indices.delete(index=index)...

test_open_search.py

Source:test_open_search.py

1from django.test import TestCase, override_settings2from core.opensearch import get_open_search, OSWrapper, OSWrapperError3from opensearchpy import OpenSearch4class OpenSearchTest(TestCase):5    """Test the OpenSearch object instantiation process"""6    def test_get_open_search(self):7        """Tests that the get_open_search() function returns a working OpenSearch object"""8        opensearch_client = get_open_search()9        self.assertIsInstance(opensearch_client, OpenSearch)10    def test_get_client_cache(self):11        """Tests that the OSWrapper.get_client() method caches previous OpenSearch objects and returns them"""12        opensearch_client = get_open_search()13        self.assertIs(opensearch_client, OSWrapper._os_client)14        OSWrapper._os_client = None15        new_opensearch_client = get_open_search()16        self.assertFalse(opensearch_client is new_opensearch_client)17    @override_settings(OPENSEARCH_HOST=None, OPENSEARCH_PORT=None, OPENSEARCH_URI=None)18    def test_get_open_search_error(self):19        """Tests that without the correct environment variables, an OSWrapperError() is raised"""20        with self.assertRaises(OSWrapperError):...

Automation Testing Tutorials

Learn to execute automation testing from scratch with LambdaTest Learning Hub. Right from setting up the prerequisites to run your first automation test, to following best practices and diving deeper into advanced test scenarios. LambdaTest Learning Hubs compile a list of step-by-step guides to help you be proficient with different test automation frameworks i.e. Selenium, Cypress, TestNG etc.