"""Opensearch Connector.""" from os import environ import random import boto3 from aws_requests_auth.aws_auth import AWSRequestsAuth from opensearchpy import exceptions as opensearch_exceptions from opensearchpy import OpenSearch from opensearchpy import RequestsHttpConnection from content_utils.exceptions import IndexingFailedError from content_utils.exceptions import IneligibleEventError def _get_current_region(): """Get the current AWS region.""" if 'AWS_REGION' in environ: return environ['AWS_REGION'] if 'AWS_DEFAULT_REGION' in environ: return environ['AWS_DEFAULT_REGION'] return boto3.session.Session().region_name def _get_assumed_role_credentials(assume_role_arn): """Get credentials for the assumed role.""" assumed_role = boto3.client('sts').assume_role( RoleArn=assume_role_arn, RoleSessionName=f'LambdaContentESAssumedRole{random.randint(1, 1000)}' ) return assumed_role['Credentials'] class LambdaOpensearchConnector: """Opensearch client.""" def __init__( self, opensearch_endpoint, logger, username=None, password=None, assume_role_arn=None, port=443, index_name='review.v01', alias_name='review_write', ): """Opensearch connector.""" self.os = None self.opensearch_endpoint = opensearch_endpoint self.assume_role_arn = assume_role_arn self.logger = logger self.port = port self.index_name = index_name self.alias_name = alias_name self._initialize_client(self._get_auth(username, password)) def _get_auth(self, username=None, password=None): """Get client credentials using an assumed role.""" if username and password: return (username, password) credentials = {} if self.assume_role_arn: session_creds = _get_assumed_role_credentials(self.assume_role_arn) credentials = { 'aws_access_key': session_creds['AccessKeyId'], 'aws_secret_access_key': session_creds['SecretAccessKey'], 'aws_token': session_creds['SessionToken'] } else: session_creds = boto3.Session().get_credentials() credentials = { 'aws_access_key': session_creds.access_key, 'aws_secret_access_key': session_creds.secret_key, 'aws_token': session_creds.token } return AWSRequestsAuth( **credentials, aws_host=self.opensearch_endpoint, aws_region=_get_current_region(), aws_service='es' ) def _initialize_client(self, http_auth_credentials): """Initialize Opensearch Client.""" self.os = OpenSearch( hosts=[{'host': self.opensearch_endpoint, 'port': self.port}], http_auth=http_auth_credentials, use_ssl=self.port == 443, verify_certs=self.port == 443, connection_class=RequestsHttpConnection ) def drop_index(self): """Drop the alias and index.""" if self.os.indices.exists_alias(name=self.alias_name): self.os.aliases.delete(name=self.alias_name, ignore=[400, 404]) return self.os.indices.delete( index=self.index_name, ignore=[400, 404] ) def clear_index_via_alias(self): """Delete all documents in the index.""" return self.os.delete_by_query( index=self.alias_name, body={ 'query': { 'match_all': {} } }, refresh=True, wait_for_completion=True, ignore=[400, 404] ) def index_product(self, index_record): """Create an indexed document.""" response = self._call( 'index', index=self.alias_name, id=index_record.get('product_id'), body=index_record ) created = response.get('result') == 'created' or response.get('created') is not None updated = response.get('result') == 'updated' or response.get('updated') is not None if not created and not updated: self.logger.error(response) raise IndexingFailedError('Failed to index product.') self.logger.debug(f'ADDED: product {index_record.get("product_id")}') def patch_product(self, patch_data): """Patch indexed product data.""" response = self._call( 'update', index=self.alias_name, id=patch_data.get('product_id'), ignore=[404], body={'doc': patch_data} ) patched = response.get('result') == 'updated' or response.get('updated') is not None if patched: self.logger.debug(f'PATCHED: product {patch_data.get("product_id")}') return if response.get('status') == 404: self.logger.debug( f'NOOP: nothing to patch, product {patch_data.get("product_id")} not found') return if not response.get('_shards', {}).get('failed'): self.logger.debug( f'NOOP: nothing to patch, product {patch_data.get("product_id")} has no updates') return self.logger.error(response) raise IndexingFailedError(f'Failed to patch product. ID {patch_data.get("product_id")}') def remove_product(self, product_id): """Remove a product from the index.""" try: self._call( 'delete', index=self.alias_name, id=product_id, ignore=[404] ) self.logger.debug(f'DELETED: product {product_id}') except opensearch_exceptions.NotFoundError: self.logger.warning( f'failed to remove product {product_id} ' 'from index. product was not in index.' ) def check_product_existing_status(self, row_data, expected_status=None): """Check if a product exists. Raise exception if response doesn't match expected status.""" try: search_hit = self._call( 'get', index=self.alias_name, id=row_data.get('product_id'), _source=True ) self.logger.debug(search_hit) if search_hit.get('found'): if expected_status == 'not found': raise IneligibleEventError( 'Product already in queue: ' f'queue id: {row_data.get("review_queue_id")} ' f'product id: {row_data.get("product_id")}' ) self.logger.debug(f'FOUND: product {row_data.get("product_id")}') return search_hit.get('_source') except opensearch_exceptions.NotFoundError: if expected_status == 'found': raise IneligibleEventError( 'Product not found in queue: ' f'product id: {row_data.get("product_id")}' ) self.logger.debug(f'NOT FOUND: product {row_data.get("product_id")}') def _call(self, method, retry_on_expired=True, **params): """Call es client, retrying once on expired session exceptions.""" if not hasattr(self.os, method) or not callable(getattr(self.os, method)): raise Exception(f'Unknown es method: {method}') try: return getattr(self.os, method)(**params) except opensearch_exceptions.AuthorizationException as ex: if not retry_on_expired or \ not self.assume_role_arn or \ 'message' not in ex.info or \ 'expired' not in ex.info['message']: raise ex # Call errored out with expired session token # recreate the client and try again self._initialize_client(self._get_auth()) return self._call(method, retry_on_expired=False, **params)