"""Salesforce consumer.""" from copy import deepcopy import json import time from datetime import datetime from datetime import date from datetime import timedelta from typing import List from confluent_kafka import DeserializingConsumer from confluent_kafka import SerializingProducer from confluent_kafka import TopicPartition from confluent_kafka.admin import AdminClient from confluent_kafka.cimpl import Message from confluent_kafka.cimpl import NewTopic import requests from simple_salesforce import format_soql from simple_salesforce import Salesforce from simple_salesforce import SalesforceError from tenacity import retry from tenacity import stop_after_attempt from tenacity import wait_exponential import config # noqa from src import constants # noqa logger = config.logger class SalesforceConsumer: """SalesForce API client that reads messages from SalesForce and saves them into Kafka.""" def __init__( self, sf_domain: str = config.SALESFORCE_DOMAIN, bootstrap_servers: str = config.BOOTSTRAP_SERVERS, kafka_topic: str = config.KAFKA_TOPIC, group_id: str = config.GROUP_ID, end_date: str = None, with_delay: bool = True, external_test: bool = False, ): """Create a new instance of SalesforceConsumer.""" self._kafka_config = { 'bootstrap.servers': bootstrap_servers, 'security.protocol': 'SSL', 'group.id': group_id } self._kafka_topic = kafka_topic self._with_delay = with_delay self._external_test = external_test if self._with_delay: self._offset_topic = config.KAFKA_OFFSET_TOPIC_WITH_DELAY else: self._offset_topic = config.KAFKA_OFFSET_TOPIC_WITHOUT_DELAY if self._external_test: self._offset_topic = config.KAFKA_OFFSET_TOPIC_EXTERNAL_TEST sf_consumer_key = config.secrets_manager_client.get_cred( 'SALESFORCE_CONSUMER_KEY') sf_consumer_secret = config.secrets_manager_client.get_cred( 'SALESFORCE_CONSUMER_SECRET') sf_refresh_token = config.secrets_manager_client.get_cred( 'SALESFORCE_REFRESH_TOKEN') oauth_url = f'https://{sf_domain}.salesforce.com/services/oauth2/token' # noqa: E231 payload = '&'.join([ f'client_id={sf_consumer_key}', f'client_secret={sf_consumer_secret}', 'grant_type=refresh_token', f'refresh_token={sf_refresh_token}' ]) headers = {'content-type': 'application/x-www-form-urlencoded'} response = requests.request( 'POST', oauth_url, data=payload, headers=headers) credentials = response.json() self._sf = Salesforce( instance_url=credentials['instance_url'], session_id=credentials['access_token'], version=config.SALESFORCE_VERSION) self._producer = SerializingProducer({ **self._kafka_config, 'value.serializer': lambda v, _: json.dumps(v).encode(), }) self._offset_producer = SerializingProducer({ **self._kafka_config, 'value.serializer': lambda v, _: json.dumps(v).encode(), }) self._consumer = DeserializingConsumer({ **self._kafka_config, 'value.deserializer': lambda v, _: json.loads(v), }) self._admin_client = AdminClient(self._kafka_config) self._received_records = 0 self._sent_records_counter = 0 end_date = datetime.combine( date.today(), datetime.max.time()) - timedelta(hours=config.IMPORT_INTERVAL_HOURS) self._end_date = end_date.strftime('%Y-%m-%dT%H:%M:%SZ') if self._with_delay: query_file = constants.QUERY_FILE_WITH_DELAY else: query_file = constants.QUERY_FILE_WITHOUT_DELAY if self._external_test: query_file = constants.QUERY_FILE_EXTERNAL_TEST with open(query_file, 'r') as query: self._soql_query = query.read().replace('\n', '') def __enter__(self): """Context manager enter.""" return self def __exit__(self, exc_type, exc_value, traceback): """Context manager exit.""" self._producer.flush() logger.info( f'Finished. {self._received_records} records received from SalesForce. ' f'{self._sent_records_counter} records sent to Kafka.' ) @staticmethod def _massage_contact_for_downstream_requirements(contact): """Prepare a contact payload for consuming by gda-account-creation.""" massaged_contact = deepcopy(contact) massaged_contact.update(massaged_contact.get('Account')) massaged_contact['ArtistName__c'] = massaged_contact.get('Name') massaged_contact.pop('Account') massaged_contact.pop('Name') return massaged_contact @retry( stop=stop_after_attempt(config.RETRY_COUNT), wait=wait_exponential(max=config.RETRY_MAX_WAIT), reraise=True, ) def _load_contacts( self, end_date: str, last_auto_id: str ) -> List[dict]: """Load a list of contacts from SalesForce.""" if self._with_delay: logger.info(f'Fetching records created before {end_date}') query = format_soql( self._soql_query, end_date, int(last_auto_id)) else: query = format_soql(self._soql_query, int(last_auto_id)) # this is redundant, but let's leave it here in sake of possible future changes if self._external_test: query = format_soql(self._soql_query, int(last_auto_id)) records = [] logger.info('Loading records.') try: res = self._sf.query(query) if res['totalSize']: records = res['records'] logger.info(f'Received: {res["totalSize"]} records.') # noqa: Q000 self._received_records = len(records) else: logger.info('No records found.') except SalesforceError as ex: logger.exception('Failed to query SalesForce.. retrying.') raise ex return records @retry( stop=stop_after_attempt(config.RETRY_COUNT), wait=wait_exponential(max=config.RETRY_MAX_WAIT), reraise=True, ) def _create_offset_topic(self): logger.info(f'Check for existing offset topic "{self._offset_topic}"..') existing_offset_topic = self._admin_client.list_topics(topic=self._offset_topic) if not existing_offset_topic.topics: logger.info('Offset topic does not exist. Creating it...') self._admin_client.create_topics( [ NewTopic( self._offset_topic, num_partitions=config.OFFSET_TOPIC_PARTITIONS, replication_factor=config.OFFSET_TOPIC_REPLICAS, ) ] ) logger.info(f'Offset topic {self._offset_topic} created.') else: logger.info(f'Offset topic {self._offset_topic} already exists.') @retry( stop=stop_after_attempt(config.RETRY_COUNT), wait=wait_exponential(max=config.RETRY_MAX_WAIT), reraise=True, ) def _read_last_contact_id(self) -> str: """Read last processed contact id from Kafka offset topic.""" last_id = '0' logger.info('Loading last event id from Kafka..') tp = TopicPartition(self._offset_topic, 0) first_offset, last_offset = self._consumer.get_watermark_offsets(tp) self._consumer.assign([TopicPartition(self._offset_topic, 0, last_offset - 1)]) offsets_data = self._consumer.poll(config.OFFSETS_POLL_WAIT_TIME) if offsets_data: if offsets_data.error(): logger.error(offsets_data.error().str()) else: data = offsets_data.value() last_id = data.get('last_id') or '0' logger.info(f'Last contact id found: {last_id}') else: logger.info('Last contact id not found. Use 0 as last contact id.') return last_id @retry( stop=stop_after_attempt(config.RETRY_COUNT), wait=wait_exponential(max=config.RETRY_MAX_WAIT), reraise=True, ) def _save_last_contact_id(self, last_id: str): """Save the last processed contat id into kafka offset topic.""" offset_data = {'ts_sec': int(time.time()), 'last_id': last_id} logger.info(f'Storing last contact id: {last_id}...') self._offset_producer.produce( self._offset_topic, value=offset_data, key=constants.OFFSET_KEY, ) self._offset_producer.flush() logger.info('Last contact id saved.') def _delivery_report(self, err: Exception, msg: Message): """Kafka message delivery report handler.""" if err is not None: logger.info(f'Message delivery failed: {err}') else: self._sent_records_counter += 1 if self._sent_records_counter == self._received_records: logger.info( f'Kafka delivery reports for all {self._received_records} records received!' ) def process_contacts(self): """Sync SalesForce contacts to Kafka.""" self._create_offset_topic() last_auto_id = self._read_last_contact_id() contacts = self._load_contacts(self._end_date, last_auto_id) if contacts: logger.info(f'Sending {len(contacts)} records to Kafka') for contact in contacts: massaged_contact = self._massage_contact_for_downstream_requirements(contact) self._producer.produce( self._kafka_topic, value=massaged_contact, key=contact[constants.ID_FIELD], on_delivery=self._delivery_report, ) logger.info(f'{len(contacts)} records sent.') self._producer.flush() last_id = contacts[-1][constants.ID_FIELD] self._save_last_contact_id(last_id)