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):...
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.
You could also refer to video tutorials over LambdaTest YouTube channel to get step by step demonstration from industry experts.
Get 100 minutes of automation test minutes FREE!!