"""Application logic layer.""" import json import time from typing import Any from typing import Dict from typing import List from kafka import KafkaProducer from kafka.consumer.fetcher import ConsumerRecord from kafka.errors import NoBrokersAvailable from kafka.errors import NodeNotReadyError from kafka.structs import TopicPartition from lxml import etree from dbdeploy.base import config from dbdeploy.base import constants from dbdeploy.dtos import ChangeSet from dbdeploy.dtos import Column from dbdeploy.dtos import ExecutionJob from dbdeploy.dtos import JsonSchema from dbdeploy.dtos import TableSchema from dbdeploy.util.aws.helpers import get_msk_bootstrap_brokers from dbdeploy.util.aws.helpers import change_fargate_task_count from dbdeploy.util.aws.helpers import check_running_task_count from dbdeploy.util.aws.helpers import delete_table_lock_item from dbdeploy.util.aws.helpers import get_dynamodb_client from dbdeploy.util.aws.helpers import get_ecs_client from dbdeploy.util.aws.helpers import get_ecs_cluster_arn from dbdeploy.util.aws.helpers import get_table_item from dbdeploy.util.aws.helpers import put_table_item from dbdeploy.util.aws.helpers import put_table_lock_item from dbdeploy.util.common import threadify from dbdeploy.util.common import get_string_hash from dbdeploy.util.exceptions import KafkaDLQError from dbdeploy.util.kafka.connect import create_connector from dbdeploy.util.kafka.connect import delete_connector from dbdeploy.util.kafka.connect import get_connector_list from dbdeploy.util.kafka.connect import get_neo4j_sink_connector_config from dbdeploy.util.kafka.connect import get_jdbc_sink_connector_config from dbdeploy.util.kafka.executor import KafkaExecutor from dbdeploy.util.kafka.helpers import serialize_to_json from dbdeploy.util.kafka.helpers import serialize_to_string from dbdeploy.util.kafka.helpers import check_dlq_topic_records from dbdeploy.util.kafka.helpers import compare_offsets from dbdeploy.util.kafka.helpers import get_offsets_per_partiton from dbdeploy.util.kafka.helpers import on_send_error_callback from dbdeploy.util.kafka.helpers import on_send_success_callback from dbdeploy.util.kafka.structs import AutoOffsetReset from dbdeploy.util.snowflake.executor import sfcursor from dbdeploy.util.xml_helpers import get_schema_file_path from dbdeploy.util.xml_helpers import validate log = config.LOGGER def acquire_deploy_lock( timeout: int = 60 * 15, force_lock_release: bool = False) -> bool: """Get db-deploy execution lock.""" log.info('Acquiring an execution lock...') ddb_client = get_dynamodb_client(aws_region=config.AWS_REGION) while True: if timeout < 0: return False is_lock_acquired = put_table_lock_item( dynamodb_client=ddb_client, table_name=config.DYNAMODB_TABLE) if is_lock_acquired or force_lock_release: log.info('Execution lock acquired.') return True log.info('Another PR is getting executed right now. ' 'Waiting 60s before next lock attempt...') timeout -= config.LOCK_ACQUISITION_BACKOFF time.sleep(config.LOCK_ACQUISITION_BACKOFF) def clean_up_previous_runs() -> None: """Clean up previuos runs. In case any of the previous deploy jobs were not able to delete connectors because of too short jenkins timeout it might be needed to delete other connector instances. This is temporary until we can implement this: https://theorchard.atlassian.net/browse/SYS-19443 """ connector_list, msg = get_connector_list( uri=config.CONNECT_CLUSTER_URL) for connector in connector_list: connector_deleted, msg = delete_connector( uri=config.CONNECT_CLUSTER_URL, connector_name=connector) if not connector_deleted: log.error(msg) raise Exception(msg) def parse_db_pr_file(xml_contents: str) -> List[ChangeSet]: """Validate and parse XML file. Args: xml_contents: contents of XML file with changesets. Returns: changeset_list: A list of ChangeSet DTOs """ log.info('Parsing XML file.') log.info(f'Validating XML schema for flow: {config.FLOW_NAME}.') xml_schema_file_path = get_schema_file_path(config.FLOW_NAME) if xml_schema_file_path is None: raise FileNotFoundError( f'XML schema not found for flow: {config.FLOW_NAME}.') xml_is_valid, err = validate(xml_schema_file_path, xml_contents) if not xml_is_valid: raise ValueError(f'Could not validate XML file. Error: {err}') # create XML document tree root = etree.fromstring(xml_contents.encode()) changesets = list() for changeset in root.xpath('//changeset'): changeset_id = str(changeset.xpath('@id')[0]) run_on_change = ( True if changeset.xpath('@run-on-change') else False) run_always = ( True if changeset.xpath('@run-always') else False) topic_name = ( None if not changeset.xpath('@topic') else changeset.xpath('@topic')[0]) kafka_cluster_name = ( None if not changeset.xpath('@kafka-cluster-name') else changeset.xpath('@kafka-cluster-name')[0]) kafka_message_root = ( None if not changeset.xpath('./kafkamessage') else changeset.xpath('./kafkamessage')[0]) snowflake_account = ( config.SnowflakeAccount.ORCHARD if not changeset.xpath('@snowflake-account') else changeset.xpath('@snowflake-account')[0]) neo4j_server = ( config.Neo4jServerNames.MUSIC_GRAPH if not changeset.xpath('@neo4j-server') else changeset.xpath('@neo4j-server')[0]) sql_query_string, precondition = None, False kafka_message, kafka_message_key = None, None key_serializer, value_serializer = None, None kafka_tombstone_message = False if sql_query := changeset.xpath('./precondition/sqlquery/text()'): # get contents of the first found XML tag # and strip [:space:] characters sql_query_string = ' '.join([ word.strip() for word in sql_query[0].split()]) precondition = True elif kafka_message_root is not None: # get contents of the first found XML tag # and strip [:space:] characters precondition = True kafka_message_str = kafka_message_root.text kafka_message_key = ( None if not kafka_message_root.xpath('@key') else str(kafka_message_root.xpath('@key')[0]).strip()) key_serializer = ( None if not kafka_message_root.xpath('@key_serializer') else str(kafka_message_root.xpath('@key_serializer')[0]).strip()) value_serializer = ( None if not kafka_message_root.xpath('@value_serializer') else str(kafka_message_root.xpath('@value_serializer')[0]).strip()) kafka_tombstone_message = ( False if not kafka_message_root.xpath('@tombstone') else True) if kafka_tombstone_message and not kafka_message_key: raise ValueError('A `key` attribute is required in tag for tombstone messages.') # noqa: E501 # validate values if key_serializer and key_serializer not in constants.KEY_SERIALIZER_VALUES: # noqa: E501 raise ValueError( 'Invalid key_serializer. Value can be `string` or `json`.') if value_serializer and value_serializer not in constants.VALUE_SERIALIZER_VALUES: # noqa: E501 raise ValueError( 'Invalid value_serializer. Value can be `string` or `json`') if not kafka_tombstone_message: kafka_message_str = ' '.join([ word.strip() for word in kafka_message_str.split()]) if value_serializer and value_serializer == 'string': kafka_message = kafka_message_str else: try: kafka_message = json.loads(kafka_message_str) except Exception: if value_serializer and value_serializer == 'json': raise ValueError('Invalid kafkamessage. Not a valid json for json serialization.') # noqa: E501 else: log.warning('Message is not a json. Will write it as string.') # noqa: E501 kafka_message = kafka_message_str if not kafka_message: raise ValueError('XML tag is empty.') if kafka_message_key and key_serializer and key_serializer == 'json': try: # serializer needs python object to write json. kafka_message_key = json.loads(kafka_message_key) except Exception: log.warning('Key is not a json. Will write it as string.') # noqa: E501 table_schema = None cypher_query_string = None skip_sink_connector = False # noqa: E501 insert_mode = None if cypher_query := changeset.xpath('./cypherquery/text()'): # get contents of the first found XML tag # and strip [:space:] characters cypher_query_string = ' '.join([ word.strip() for word in cypher_query[0].split()]) topic_name, kafka_cluster_name = None, None elif not precondition: raise ValueError('XML tag / is empty.') # noqa: E501 elif table_schema := changeset.xpath('./tableschema/@*'): table_name = changeset.xpath('./tableschema/@table-name')[0] primary_key = changeset.xpath('./tableschema/@primary-key')[0] insert_mode = ( 'update' if not changeset.xpath('./tableschema/@insert-mode') else changeset.xpath('./tableschema/@insert-mode')[0]) if insert_mode not in constants.JDBC_INSERT_MODES: raise ValueError( f"Invalid insert_mode '{insert_mode}' in tableschema. " f"Must be one of: {constants.JDBC_INSERT_MODES}") fields = list() for column in changeset.xpath('./tableschema/column'): name, type = column.xpath('./*/text()') fields.append(Column(field=name.strip(), type=type.strip())) table_schema = TableSchema( fields=fields, primary_key=primary_key, table_name=table_name) elif not topic_name: raise ValueError("To skip sink step 'topic_name' property must be set.") # noqa: E501 else: skip_sink_connector = True log.info('Done parsing XML file.') log.debug('SQL: %s' % sql_query_string) log.debug('Cypher: %s' % cypher_query_string) changeset = ChangeSet( changeset_id=changeset_id, precondition=precondition, sql_query=sql_query_string, kafka_message=kafka_message, kafka_message_key=kafka_message_key, kafka_message_key_serializer=key_serializer, kafka_message_value_serializer=value_serializer, cypher_query=cypher_query_string, run_on_change=run_on_change, run_always=run_always, skip_sink_connector=skip_sink_connector, topic_name=topic_name, snowflake_account=snowflake_account, neo4j_server=neo4j_server, insert_mode=insert_mode,) if skip_sink_connector and kafka_cluster_name is not None: changeset.kafka_cluster_name = kafka_cluster_name elif (kafka_message or kafka_message_key) and kafka_cluster_name is not None: # noqa: E501 changeset.kafka_cluster_name = kafka_cluster_name elif not skip_sink_connector: changeset.kafka_cluster_name = None if table_schema: changeset.table_schema = table_schema changesets.append(changeset) return changesets def get_jobs_statuses(changeset_list: List[ChangeSet]) -> List[ChangeSet]: """Check job execution status. Make sure there are no duplicate runs. Args: changeset_list: The list of execution job DTOs Returns: changeset_list: DTO with a list of execution job DTOs """ log.info('Getting changeset statuses.') ddb_client = get_dynamodb_client(aws_region=config.AWS_REGION) changesets_to_deploy = list() for changeset in changeset_list: changeset_status = get_table_item( changeset_id=changeset.changeset_id, dynamodb_client=ddb_client, table_name=config.DYNAMODB_TABLE) if not changeset_status or changeset.run_always: changesets_to_deploy.append(changeset) continue if changeset.cypher_query: cypher_has_changed = ( changeset_status.get('cypher_hash') != get_string_hash(changeset.cypher_query)) if cypher_has_changed and changeset.run_on_change: changesets_to_deploy.append(changeset) continue if changeset.sql_query: sql_has_changed = ( changeset_status.get('sql_hash') != get_string_hash(changeset.sql_query)) if sql_has_changed and changeset.run_on_change: changesets_to_deploy.append(changeset) if changeset.kafka_message: msg = json.dumps(changeset.kafka_message) \ if isinstance(changeset.kafka_message, (object,)) \ else str(changeset.kafka_message) json_has_changed = ( changeset_status.get('json_hash') != get_string_hash(msg)) if json_has_changed and changeset.run_on_change: changesets_to_deploy.append(changeset) continue if changeset.kafka_message_key: # incase this is tombstone with empty message key = json.dumps(changeset.kafka_message_key) \ if isinstance(changeset.kafka_message_key, (object,)) \ else str(changeset.kafka_message_key) json_has_changed = ( changeset_status.get('json_hash') != get_string_hash(key)) if json_has_changed and changeset.run_on_change: changesets_to_deploy.append(changeset) continue if not changesets_to_deploy: raise ValueError('XML file hasn\'t changed. ' 'Consider adding runonchange or runalways flags.') log.info('Got changeset statuses. ' f'{len(changesets_to_deploy)} / {len(changeset_list)} ' 'changeset(s) to execute.') return changesets_to_deploy def prepare_job_list(changeset_list: List[ChangeSet]) -> List[ExecutionJob]: """Add generated connector and topic names to each execution Job. Args: changeset_list: The list of execution job DTOs Returns: job_list: DTO with a list of execution job DTOs """ jobs = [] for changeset in changeset_list: job = ExecutionJob(**changeset.dict()) # if we are only to get data from Snowflake into a topic # let's skip name generation and use topic name # from XML property if job.skip_sink_connector and job.topic_name: job.topic.name = job.topic_name jobs.append(job) continue # generate topic names for each job job.topic.name = f'{config.KAFKA_TOPIC_PREFIX}{job.id}' job.dlq_topic.name = f'{config.KAFKA_DLQ_TOPIC_PREFIX}{job.id}' log.debug('Generated changeset topic names: ' f'"{job.topic.name}", ' f'"{job.dlq_topic.name}"') # generate connector name for each job job.connector_name = (constants.CONNECTOR_NAME_TEMPLATE .format(job_id=job.id) .replace('-', '_')) log.debug('Generated changeset connector name: ' f'"{job.connector_name}"') jobs.append(job) return jobs def start_fargate_task(desired_count: int) -> bool: """Bump the number of fargate task instances. Args: desired_count: The desired number of running task instances Returns: boolean: The result of tasks start attempt. """ log.info(f'Starting {desired_count} Fargate ' 'task(s) for neo-sink-db service...') start = time.perf_counter() try: ecs_client = get_ecs_client(config.AWS_REGION) cluster_arn = get_ecs_cluster_arn( ecs_client=ecs_client, cluster_name=config.CONNECT_CLUSTER_NAME) if not cluster_arn: raise RuntimeError('Was not able to get ECS Cluster ARN.') change_fargate_task_count( ecs_client=ecs_client, desired_count=desired_count, cluster_name=config.CONNECT_CLUSTER_NAME, cluster_arn=cluster_arn) tasks_running = check_running_task_count( ecs_client=ecs_client, desired_count=desired_count, grace_period=config.FARGATE_TASK_HEALTH_CHECK_GRACE_PERIOD, backoff_timeout=config.FARGATE_TASK_HEALTH_CHECK_BACKOFF_TIMEOUT, cluster_arn=cluster_arn, cluster_url=config.CONNECT_CLUSTER_URL) if not tasks_running: raise RuntimeError('Something went wrong. ' 'Was not able to start ECS Fargate tasks.') except Exception as err: log.error('Caught an exception while starting fargate tasks. ' 'Stopping...') stop_fargate_task() raise err end = time.perf_counter() log.info('Succesfully started fargate task(s) ' f'in {end - start} seconds.') return True def create_kafka_topics(job: ExecutionJob) -> ExecutionJob: """Create kafka topics. Args: job: An execution job DTO Returns: job: Updated execution job DTO """ log.info('Creating topics.') retries = config.KAFKA_CONNECTION_RETRIES client = None while not client and retries > 0: client, err = ( KafkaExecutor(config.KAFKA_CLUSTER) .admin_client(client_id=job.id)) if not client: log.debug(f'Kafka is not ready, retry in 60 sec. Error: {err}.') # Wait till kafka is ready retries -= 1 time.sleep(config.RETRY_BACKOFF) continue break if not client: raise TimeoutError('Kafka connection error. ' 'Was not able to connect to the Kafka cluster.') client.create_topics([job.topic, job.dlq_topic]) job.topics_created = True log.info('Done creating topics.') log.debug(f'Topics created: "{job.topic.name}", "{job.dlq_topic.name}".') return job def update_kafka_cluster_config(job: ExecutionJob) -> None: """Update Kafka cluster configuration. Args: job: An execution job DTO """ full_cluster_name = f'{config.ENVIRONMENT}-{job.kafka_cluster_name}' log.info(f'Updating kafka cluster configuration for {full_cluster_name}') config.KAFKA_CLUSTER.name = full_cluster_name config.KAFKA_CLUSTER.bootstrap_brokers = get_msk_bootstrap_brokers(full_cluster_name) # noqa: E501 log.info('Kafka cluster config successfully updated. ' f'Cluster-name: {full_cluster_name}') def check_kafka_topic_exists(job: ExecutionJob) -> ExecutionJob: """If sink connector is skipped make sure topics exsist. Args: job: An execution job DTO Returns: job: Updated execution job DTO """ log.info('Getting topic info.') retries = config.KAFKA_CONNECTION_RETRIES client = None while not client and retries > 0: client, err = ( KafkaExecutor(config.KAFKA_CLUSTER) .admin_client(client_id=job.id)) if not client: log.debug(f'Kafka is not ready, retry in 60 sec. Error: {err}.') # Wait till kafka is ready retries -= 1 time.sleep(config.RETRY_BACKOFF) continue break if not client: raise TimeoutError('Kafka connection error. ' 'Was not able to connect to the Kafka cluster.') if job.topic_name not in client.list_topics(): raise ValueError( f"Please make sure topic '{job.topic_name}' " f'was created in {config.KAFKA_CLUSTER} ' 'before running db-deploy job.') log.info(f"Topic exists in kafka: '{job.topic_name}'") return job def start_kafka_connector(job: ExecutionJob) -> ExecutionJob: """Generate config and create connector in Kafka-Connect cluster. Args: job: An execution job DTO Returns: job: Update execution job DTO """ log.info('Getting connector configuration.') if job.cypher_query: neo4j_server_config = config.NEO4J_SERVERS[job.neo4j_server or ''] connector_config = get_neo4j_sink_connector_config( # ignore type when skipping sink cypher_query=job.cypher_query, topic=job.topic.name, dlq_topic_name=job.dlq_topic.name, environment=config.ENVIRONMENT, aws_region=config.AWS_REGION, tasks_max=config.CONNECTOR_TASKS_MAX, kafka_bootstrap_brokers=config.KAFKA_CLUSTER.bootstrap_brokers, kafka_security_protocol=config.KAFKA_CLUSTER.security_protocol, neo4j_server_uri=neo4j_server_config.uri, batch_size=config.CONNECTOR_BATCH_SIZE, service_name=config.SERVICE_NAME, max_poll_interval=( config.CONNECTOR_MAX_POLL_INTERVAL if not job.precondition else None), secret_key=neo4j_server_config.secret_name) elif job.table_schema: connector_config = get_jdbc_sink_connector_config( topic=job.topic.name, dlq_topic_name=job.dlq_topic.name, environment=config.ENVIRONMENT, aws_region=config.AWS_REGION, table_name=job.table_schema.table_name, primary_keys=job.table_schema.primary_key, insert_mode=job.insert_mode, tasks_max=config.CONNECTOR_TASKS_MAX, batch_size=config.CONNECTOR_BATCH_SIZE, service_name=config.SERVICE_NAME, secret_key=config.MYSQL_SECRET_KEY) log.debug('Configuration: %s' % connector_config) connector_created, msg = create_connector( uri=config.CONNECT_CLUSTER_URL, connector_name=job.connector_name, connector_config=connector_config) if not connector_created: raise Exception(msg) job.connector_started = True log.info('Connector was successfully created.') return job def get_producer(job_id: str, **kwargs: Any) -> KafkaProducer: """Get a Producer instance.""" producer = None while not producer: producer, err = ( KafkaExecutor(config.KAFKA_CLUSTER) .producer(client_id=job_id, **kwargs)) if not producer: log.debug(f'Kafka is not ready, retry in 60 sec. Error: {err}.') # Wait till kafka is ready time.sleep(config.RETRY_BACKOFF) continue break return producer def push_to_kafka(job: ExecutionJob) -> ExecutionJob: """Push data into Kafka topic. This produces messages to kafka topic from snowflake, neo4j or json. Args: job: An execution job DTO Returns: job: Updated execution job DTO """ latest_offsets: Dict[int, int] = {} log.info(f'changeset id: {job.changeset_id} Starting to push data ' f'to {job.topic.name}') message_key = {'DB_DEPLOY_ID': job.id} if not job.sql_query and job.precondition is False: # Type neo4j: # there are no sql query and no messages because of that # only single cypher should be executed # this message only required to trigger the cypher producer = get_producer(job.id) (producer .send(job.topic.name, key=message_key, value={'MESSAGE': 'execute single cypher'}) .add_callback(on_send_success_callback, latest_offsets) .add_errback(on_send_error_callback)) producer.flush() job.offsets = latest_offsets return job if ( (job.kafka_message or job.kafka_message_key) and job.precondition is True ): # Type kafka: # there are no sql query or cypher. It contains kafka_message # and kafka_message_key that is written to kafka topic. # kafka_message_key_serializer and value_serializer are only used for # this format. Default key_serializer=string and value_serializer=json params = { 'key_serializer': serialize_to_string, 'value_serializer': serialize_to_json } # overwrite defaults. if job.kafka_message_key_serializer and job.kafka_message_key_serializer == 'json': params['key_serializer'] = serialize_to_json if job.kafka_message_value_serializer and job.kafka_message_value_serializer == 'string': params['value_serializer'] = serialize_to_string producer = get_producer(job.id, **params) (producer .send( job.topic.name, key=job.kafka_message_key, value=job.kafka_message ) .add_callback(on_send_success_callback, latest_offsets) .add_errback(on_send_error_callback)) producer.flush() job.offsets = latest_offsets return job # Type snowflake-*: producer = get_producer(job.id) total_row_count = 0 batch_size = config.SNOWFLAKE_BATCH_SIZE with sfcursor(job.snowflake_account) as cursor: cursor.execute(str(job.sql_query)) while rows := cursor.fetchmany(batch_size): if job.dlq_error is not None: # raise on first found error raise job.dlq_error row_count = len(rows) total_row_count += row_count log.info( f'Pushing data to "{job.topic.name}" ' f'with {row_count} rows to execute...') for row in rows: if job.skip_sink_connector: # Type snowflake-kafka: message_key = row.get('KEY', message_key) row = json.loads(row.get('VALUE')) if row is None: raise ValueError( 'Please use Snowflake object_construct() ' 'to serialize the value in your SQL ' 'and use "KEY" and "VALUE" as column names ' 'where "KEY" is optional.') if job.table_schema: # Type snowflake-mysql: message_key = None row = dict( schema=JsonSchema(**job.table_schema.dict()).dict(), payload=row) (producer .send(job.topic.name, key=message_key, value=row) .add_callback(on_send_success_callback, latest_offsets) .add_errback(on_send_error_callback)) producer.flush() log.info( f'Completed pushing {total_row_count} messages to "{job.topic.name}".') job.offsets = latest_offsets return job def track_kafka_consumer_offsets(job: ExecutionJob) -> ExecutionJob: """Track Sink Connector's consumer offsets in Kafka. Args: job: An execution job DTO Returns: job: An execution job DTO """ start = time.perf_counter() latest_offsets = job.offsets log.info( 'Waiting for Sink connector to process ' f'data in "{job.topic.name}".') client = None while not client: client, err = ( KafkaExecutor(config.KAFKA_CLUSTER) .admin_client(client_id=job.id)) if not client: log.debug(f'Kafka is not ready, retry in 60 sec. Error: {err}.') # Wait till kafka is ready time.sleep(config.RETRY_BACKOFF) continue break while latest_offsets: if job.dlq_error is not None: # raise on first found error raise job.dlq_error log.debug('Comparing consumer offsets with producer offsets.') try: group_offsets = client.list_consumer_group_offsets( group_id=(constants.CONSUMER_GROUP_NAME_TEMPLATE .format(connector_name=job.connector_name)) ) except (NoBrokersAvailable, NodeNotReadyError) as err: log.debug(f'Kafka is not ready, retry in 60 sec. Error: {err}.') # Wait till kafka is ready time.sleep(config.RETRY_BACKOFF) continue # Get sorted list of (partition, offset) namedtuples for the topic current_offsets = get_offsets_per_partiton( job.topic.name, group_offsets) # Compare current consumed latest offsets with the latest pushed offest compare_offsets(latest_offsets, current_offsets) # All partitions are done, break the loop if len(latest_offsets) < 1: break log.debug('Consumer is not done yet. Wait for another 60 sec.') # Wait for 1 minute and continue while loop # if there is unconsumed data left time.sleep(config.RETRY_BACKOFF) # Offests are beyond those sent, Sink connector done consuming end = time.perf_counter() # Do the last check before marking as done to avoid race conditions if job.dlq_error is not None: # raise on first found error raise job.dlq_error job.is_done = True log.info(f'Sink is completed for "{job.topic.name}" ' f'in {end - start} seconds.') return job @threadify def monitor_dlq(job: ExecutionJob) -> ExecutionJob: """Track Sink exceptions in Kafka. Args: job: An execution job DTO Returns: job: Updated execution job DTO """ dlq_topic = job.dlq_topic.name log.info( f'Started listening DLQ "{dlq_topic}" topic in kafka with ' f'CONTINUE_ON_DLQ={"true" if config.CONTINUE_ON_DLQ else "false"}.') consumer = None while not consumer: consumer, err = KafkaExecutor(config.KAFKA_CLUSTER).consumer( client_id=job.id, group_id=config.SERVICE_NAME, enable_auto_commit=False, auto_offset_reset=AutoOffsetReset.EARLIEST) if not consumer: log.debug(f'Kafka is not ready, retry in 60 sec. Error: {err}.') # Wait till kafka is ready time.sleep(config.RETRY_BACKOFF) continue break try: partitions_id_set = consumer.partitions_for_topic(dlq_topic) if not partitions_id_set: msg = (f'Not able to find any partitions for: {dlq_topic}.' 'Please check if the topic exists.') log.error(msg) job.dlq_error = KafkaDLQError(msg) return job partitions = [] for partition_id in partitions_id_set: partitions.append( TopicPartition(dlq_topic, partition_id)) consumer.assign(partitions=partitions) consumer.poll(timeout_ms=0) # Ensure group is rebalanced consumer.seek_to_beginning() found_some_dlq_messages = False while not job.is_done and not job.main_error: polled_records: Dict[ TopicPartition, List[ConsumerRecord] ] = consumer.poll(timeout_ms=200) result = check_dlq_topic_records( polled_records=polled_records, partitions=partitions) if result: found_some_dlq_messages = True msg = ('Sink Connector has written DLQ messages ' f'for ExecutionJob with id: "{job.id}". ') log.error(msg) if result.get('stop'): log.info('Job is configured to immediately stop on DLQ, so throw error and stop.') msg += f'DLQ Exception: {result.get("record")}' job.dlq_error = KafkaDLQError(msg) break else: log.error(f'DLQ Exception: {result.get("record")}') log.info('Job is configured to continue on DLQ, so continue processing.') if not found_some_dlq_messages and not job.dlq_error: log.info('No exceptions found in DLQ topic.') return job finally: consumer.unsubscribe() # close() can produce 'Fetch to node %i failed: Cancelled' error # which is transient in current state until this PR is merged # https://github.com/dpkp/kafka-python/pull/2172/files consumer.close() def update_job_status(job: ExecutionJob) -> ExecutionJob: """Check job execution status. Make sure there are no duplicate runs. Args: job_list: The list of execution job DTOs Returns: job_list: DTO with a list of execution job DTOs """ ddb_client = get_dynamodb_client(aws_region=config.AWS_REGION) put_table_item( dynamodb_client=ddb_client, changeset_id=job.changeset_id, table_name=config.DYNAMODB_TABLE, cypher_hash=(get_string_hash(job.cypher_query) if job.cypher_query else None), # noqa: E501 sql_hash=(get_string_hash(job.sql_query) if job.sql_query else None)) return job def stop_kafka_connector(job: ExecutionJob) -> bool: """Delete connector from Kafka-Connect cluster. Args: job: An execution job DTO Returns: boolean: True for successfully deleted connector """ log.info(f'Deleting {job.connector_name} connector...') connector_deleted, msg = delete_connector( uri=config.CONNECT_CLUSTER_URL, connector_name=job.connector_name) if not connector_deleted: log.error(msg) raise Exception(msg) log.info('Connector was successfully deleted.') return True def delete_kafka_topics(job: ExecutionJob) -> bool: """Delete data and dlq topics when done consuming. Args: job: An execution job DTO Returns: boolean: True for successfully deleted topics """ log.info('Deleting topics.') client = None while not client: client, err = ( KafkaExecutor(config.KAFKA_CLUSTER) .admin_client(client_id=job.id)) if not client: log.debug(f'Kafka is not ready, retry in 60 sec. Error: {err}.') # Wait till kafka is ready time.sleep(config.RETRY_BACKOFF) continue break # DLQ consumer is already closed (or closing) by this point # and the counter is not really needed but as a safe measure. retries = config.TOPIC_DELETE_RETRIES topics_to_delete = [job.topic.name, job.dlq_topic.name] while topics_to_delete: response = client.delete_topics(topics_to_delete, timeout_ms=3600) # Example response schema: # DeleteTopicsResponse_v3( # throttle_time_ms=0, # topic_error_codes=[ # (topic='topic_name', error_code=0), # (topic='topic2_name', error_code=0)]) topic_name = 0 error_code = 1 topics_to_delete = [ topic[topic_name] for topic in response.topic_error_codes if topic[error_code] != 0] if topics_to_delete: log.debug('Failed to delete these topics from the' f' first attempt: {", ".join(topics_to_delete)}.') retries -= 1 if not retries: break time.sleep(config.RETRY_BACKOFF) if not topics_to_delete: log.info('Topics were successfully deleted.') else: log.info('Topics deletedion failed.') return True def stop_fargate_task() -> bool: """Stop all fargate task instances. Returns: boolean: The result of tasks start attempt. """ log.info('Stopping Fargate tasks for neo-sink-db service...') STOP_ALL_TAKSKS = 0 ecs_client = get_ecs_client(config.AWS_REGION) cluster_arn = get_ecs_cluster_arn( ecs_client=ecs_client, cluster_name=config.CONNECT_CLUSTER_NAME) if not cluster_arn: log.error('Was not able to get ECS Cluster ARN.') return False change_fargate_task_count( ecs_client=ecs_client, desired_count=STOP_ALL_TAKSKS, cluster_name=config.CONNECT_CLUSTER_NAME, cluster_arn=cluster_arn) tasks_stopped = check_running_task_count( ecs_client=ecs_client, desired_count=STOP_ALL_TAKSKS, grace_period=config.FARGATE_TASK_HEALTH_CHECK_GRACE_PERIOD, backoff_timeout=config.FARGATE_TASK_HEALTH_CHECK_BACKOFF_TIMEOUT, cluster_arn=cluster_arn, cluster_url=config.CONNECT_CLUSTER_URL) if not tasks_stopped: log.error('Was not able to stop running tasks.') return False log.info('Succesfully stopped fargate task(s)...') return True def release_deploy_lock() -> bool: """Release db-deploy job lock.""" log.debug('==== release_deploy_lock ====') ddb_client = get_dynamodb_client(aws_region=config.AWS_REGION) is_lock_released = delete_table_lock_item( dynamodb_client=ddb_client, table_name=config.DYNAMODB_TABLE) return is_lock_released