Airflow task failing intermittently when using SnowflakeHook - amazon-web-services

Question

I'm using AWS managed Airflow (MWAA) to orchestrate some dags, but some of them are failing in an intermittent way in tasks that use SnowflakeHook provider.
Here is the dag code:
# Standard Library Imports
import logging
from datetime import timedelta, datetime
# 3rd Party Imports
from airflow.models import DAG
from airflow.operators.python_operator import PythonOperator
# Internal Imports
from DagTasksImpl import raw_ingest, preprocess, stage, get_parameter, send_failed_email
from operators.CustomS3ToSnowflakeOperator import CustomS3ToSnowflakeOperator
from operators.PostprocessingOperator import PostprocessingOperator
from airflow.contrib.operators.sns_publish_operator import SnsPublishOperator
LOGGER = logging.getLogger(__name__)
EXECUTION_DATE = '{{ macros.ds_format(macros.ds_add(yesterday_ds, 7), "%Y-%m-%d", "%Y_%m_%d") }}'
SOURCE_BUCKET = "{{params.get_parameter('/durable/buckets/ce_signal_channel')}}"
RAW_INGEST_BUCKET = "{{params.get_parameter('/durable/buckets/raw')}}"
PROCESSED_BUCKET = "{{params.get_parameter('/durable/buckets/processed')}}"
STAGED_BUCKET = "{{params.get_parameter('/durable/buckets/staged')}}"
DAY_SYM_BRAND_CHANNEL_EMAX_KEY = f"{EXECUTION_DATE}/day_sym_brand_channel_emax_{EXECUTION_DATE}.csv"
DAY_SYM_CHANNEL_EMAX_KEY = f"{EXECUTION_DATE}/day_sym_channel_emax_{EXECUTION_DATE}.csv"
PERIOD_SYM_BRAND_CHANNEL_EMAX_KEY = f"{EXECUTION_DATE}/period_sym_brand_channel_emax_{EXECUTION_DATE}.csv"
PERIOD_SYM_CHANNEL_EMAX_KEY = f"{EXECUTION_DATE}/period_sym_channel_emax_{EXECUTION_DATE}.csv"
SIGNAL_CHANNEL_FILES = [DAY_SYM_BRAND_CHANNEL_EMAX_KEY, DAY_SYM_CHANNEL_EMAX_KEY,
PERIOD_SYM_BRAND_CHANNEL_EMAX_KEY, PERIOD_SYM_CHANNEL_EMAX_KEY]
IS_GZIP = False
DEFAULT_ARGS = {
'owner': 'airflow',
'depends_on_past': False,
'start_date': datetime(2021, 5, 16),
'on_failure_callback': send_failed_email,
'retries': 1,
'retry_delay': timedelta(minutes=5),
}
def raw_ingest_task_callable(source_bucket, target_bucket, keys, **kwargs):
for key in keys:
LOGGER.info(f"Preparing to ingest data from source: {source_bucket}/{key} to target: {target_bucket}/{key}")
raw_ingest(source_bucket, target_bucket, key, IS_GZIP)
TASK_RAW_INGEST = PythonOperator(
task_id="raw_ingest",
provide_context=True,
python_callable=raw_ingest_task_callable,
op_kwargs={"source_bucket": SOURCE_BUCKET, "target_bucket": RAW_INGEST_BUCKET, "keys": SIGNAL_CHANNEL_FILES},
params={"get_parameter": get_parameter},
dag=CE_TRANSACT_WEEKLY_SIGNAL_CHANNEL_DAG,
)
def preprocess_task_callable(source_bucket, target_bucket, keys, **kwargs):
for key in keys:
LOGGER.info(f"Preparing to ingest data from source: {source_bucket}/{key} to target: {target_bucket}/{key}")
preprocess(source_bucket, target_bucket, key, IS_GZIP)
TASK_PREPROCESS = PythonOperator(
task_id="preprocess",
provide_context=True,
python_callable=preprocess_task_callable,
op_kwargs={"source_bucket": RAW_INGEST_BUCKET, "target_bucket": PROCESSED_BUCKET, "keys": SIGNAL_CHANNEL_FILES},
params={"get_parameter": get_parameter},
dag=CE_TRANSACT_WEEKLY_SIGNAL_CHANNEL_DAG,
)
def stage_task_callable(source_bucket, target_bucket, keys, **kwargs):
for key in keys:
LOGGER.info(f"Preparing to ingest data from source: {source_bucket}/{key} to target: {target_bucket}/{key}")
stage(source_bucket, target_bucket, key, IS_GZIP)
TASK_STAGE = PythonOperator(
task_id="stage",
provide_context=True,
python_callable=stage_task_callable,
op_kwargs={"source_bucket": PROCESSED_BUCKET, "target_bucket": STAGED_BUCKET, "keys": SIGNAL_CHANNEL_FILES},
params={"get_parameter": get_parameter},
dag=CE_TRANSACT_WEEKLY_SIGNAL_CHANNEL_DAG,
)
TASK_COPY_INTO_TABLE_DAY_SYM_BRAND_CHANNEL_EMAX = CustomS3ToSnowflakeOperator(
task_id='copy_into_table_day_sym_brand_channel_emax',
s3_keys=[DAY_SYM_BRAND_CHANNEL_EMAX_KEY],
table="day_sym_brand_channel_emax",
dag=CE_TRANSACT_WEEKLY_SIGNAL_CHANNEL_DAG,
pool="snowflake_pool"
)
TASK_COPY_INTO_TABLE_PERIOD_SYM_CHANNEL_EMAX = CustomS3ToSnowflakeOperator(
task_id='copy_into_table_period_sym_channel_emax',
s3_keys=[PERIOD_SYM_CHANNEL_EMAX_KEY],
table="period_sym_channel_emax",
dag=CE_TRANSACT_WEEKLY_SIGNAL_CHANNEL_DAG,
pool="snowflake_pool"
)
SEND_EMAIL = SnsPublishOperator(
task_id='send_email',
target_arn=get_parameter("/sns/mwaa"),
subject='Transact Weekly Signal Channel Notification',
message=f"Transact Weekly Signal Channel workflow has been successfully completed for the CE Data uploaded on {EXECUTION_DATE}",
dag=CE_TRANSACT_WEEKLY_SIGNAL_CHANNEL_DAG
)
TASK_RAW_INGEST >> TASK_PREPROCESS >> TASK_STAGE >>
TASK_COPY_INTO_TABLE_DAY_SYM_BRAND_CHANNEL_EMAX >> TASK_POSTPROCESS_SYM_BRAND_CHANNEL_EMAX >> SEND_EMAIL
CustomS3ToSnowflakeOperator
from typing import Any, Optional
from airflow.models import BaseOperator
from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
from airflow.utils.decorators import apply_defaults
import boto3
class CustomS3ToSnowflakeOperator(BaseOperator):
template_fields = ('s3_keys',)
sum_column = "spend_amount_90D"
#apply_defaults
def __init__(
self,
*,
s3_keys,
table: str,
stage: Any = 'consumer_edge_stage',
file_format: str = "(type = csv field_delimiter = ',' skip_header = 1 field_optionally_enclosed_by = '\"')",
schema: str = 'PUBLIC',
gzip: Optional[bool] = False,
autocommit: bool = True,
snowflake_conn_id: str = 'snowflake_connection_id',
**kwargs,
) -> None:
super().__init__(**kwargs)
self.s3_keys = s3_keys
self.table = table
self.stage = stage
self.file_format = file_format
self.schema = schema
self.gzip = gzip
self.autocommit = autocommit
self.snowflake_conn_id = snowflake_conn_id
self.temp_table = f"{table}_tmp"
def execute(self, context: Any) -> None:
snowflake_hook = SnowflakeHook(snowflake_conn_id=self.snowflake_conn_id)
self.wipe_table(snowflake_hook)
if "day" in self.table:
self.create_temp_table(snowflake_hook)
self.copy_into_table(self.get_files(), snowflake_hook, self.temp_table)
else:
self.copy_into_table(self.get_files(), snowflake_hook, self.table)
def wipe_table(self, snowflake_hook):
wipe_query = f"truncate table if exists {self.table}"
self.log.info('Wiping table...')
snowflake_hook.run(wipe_query, self.autocommit)
self.log.info('Truncate command complete...')
def create_temp_table(self, snowflake_hook):
create_temp_table_query = f"create or replace table {self.temp_table} clone {self.table}"
self.log.info('Creating table...')
snowflake_hook.run(create_temp_table_query, self.autocommit)
self.log.info('Create table command complete...')
drop_column_query = f"alter table {self.temp_table} drop column {self.sum_column}"
snowflake_hook.run(drop_column_query, self.autocommit)
def get_files(self):
if self.gzip:
bucket_and_key = str(self.s3_keys)
bucket_and_key = bucket_and_key.replace('[', '')
bucket_and_key = bucket_and_key.replace(']', '')
bucket_and_key = bucket_and_key.replace("'", '')
s3 = boto3.resource('s3')
path_parts = bucket_and_key.replace("s3://", "").split("/")
bucket = path_parts.pop(0)
key = "/".join(path_parts)
source_bucket = s3.Bucket(bucket)
files_list = []
for obj in source_bucket.objects.filter(Prefix=key):
files_list.append(obj.key)
files = str(files_list)
else:
files = str(self.s3_keys)
files = files.replace('[', '(')
files = files.replace(']', ')')
return files
def copy_into_table(self, files, snowflake_hook, table):
base_sql = """
FROM #{stage}/
files={files}
file_format={file_format}
""".format(
stage=self.stage, files=files, file_format=self.file_format
)
copy_query = """
COPY INTO {schema}.{table} {base_sql}
""".format(
schema=self.schema, table=table, base_sql=base_sql
)
self.log.info('Executing COPY command...')
snowflake_hook.run(copy_query, self.autocommit)
self.log.info("COPY command completed")
Here are the logs from the failing tasks:
*** Reading remote log from Cloudwatch log_group: airflow-ConsumerEdgeAirflowEnvironment-Task log_stream: ce_transact_weekly_signal_channel_2/copy_into_table_day_sym_brand_channel_emax/2022-12-05T14_00_00+00_00/1.log.
[2022-12-12 14:02:13,653] {{taskinstance.py:670}} INFO - Dependencies all met for <TaskInstance: ce_transact_weekly_signal_channel_2.copy_into_table_day_sym_brand_channel_emax 2022-12-05T14:00:00+00:00 [queued]>
[2022-12-12 14:02:13,887] {{taskinstance.py:670}} INFO - Dependencies all met for <TaskInstance: ce_transact_weekly_signal_channel_2.copy_into_table_day_sym_brand_channel_emax 2022-12-05T14:00:00+00:00 [queued]>
[2022-12-12 14:02:13,929] {{taskinstance.py:880}} INFO -
--------------------------------------------------------------------------------
[2022-12-12 14:02:13,946] {{taskinstance.py:881}} INFO - Starting attempt 1 of 2
[2022-12-12 14:02:13,967] {{taskinstance.py:882}} INFO -
--------------------------------------------------------------------------------
[2022-12-12 14:02:14,052] {{taskinstance.py:901}} INFO - Executing <Task(CustomS3ToSnowflakeOperator): copy_into_table_day_sym_brand_channel_emax> on 2022-12-05T14:00:00+00:00
[2022-12-12 14:02:14,088] {{standard_task_runner.py:54}} INFO - Started process 12690 to run task
[2022-12-12 14:02:14,165] {{standard_task_runner.py:77}} INFO - Running: ['airflow', 'run', 'ce_transact_weekly_signal_channel_2', 'copy_into_table_day_sym_brand_channel_emax', '2022-12-05T14:00:00+00:00', '--job_id', '11386', '--pool', 'default_pool', '--raw', '-sd', 'DAGS_FOLDER/CeTransactWeeklySignalChannel.py', '--cfg_path', '/tmp/tmptw5qhtdm']
[2022-12-12 14:02:14,975] {{standard_task_runner.py:78}} INFO - Job 11386: Subtask copy_into_table_day_sym_brand_channel_emax
[2022-12-12 14:02:15,640] {{logging_mixin.py:112}} INFO - Running %s on host %s <TaskInstance: ce_transact_weekly_signal_channel_2.copy_into_table_day_sym_brand_channel_emax 2022-12-05T14:00:00+00:00 [running]> ip-10-192-21-170.us-east-2.compute.internal
[2022-12-12 14:02:15,943] {{CustomS3ToSnowflakeOperator.py:49}} INFO - Wiping table...
[2022-12-12 14:02:16,037] {{connection.py:280}} INFO - Snowflake Connector for Python Version: 2.8.2, Python Version: 3.7.10, Platform: Linux-4.14.294-220.533.amzn2.x86_64-x86_64-with-glibc2.2.5
[2022-12-12 14:02:16,055] {{connection.py:935}} INFO - This connection is in OCSP Fail Open Mode. TLS Certificates would be checked for validity and revocation status. Any other Certificate Revocation related exceptions or OCSP Responder failures would be disregarded in favor of connectivity.
[2022-12-12 14:02:16,073] {{connection.py:952}} INFO - Setting use_openssl_only mode to False
[2022-12-12 14:02:28,275] {{logging_mixin.py:112}} WARNING - Exception ignored in: <function BaseFileLock.__del__ at 0x7fba82bbf5f0>
[2022-12-12 14:02:28,467] {{logging_mixin.py:112}} WARNING - Traceback (most recent call last):
[2022-12-12 14:02:28,527] {{logging_mixin.py:112}} WARNING - File "/usr/local/airflow/.local/lib/python3.7/site-packages/filelock/_api.py", line 240, in __del__
[2022-12-12 14:02:28,548] {{logging_mixin.py:112}} WARNING - self.release(force=True)
[2022-12-12 14:02:28,571] {{logging_mixin.py:112}} WARNING - File "/usr/local/airflow/.local/lib/python3.7/site-packages/filelock/_api.py", line 201, in release
[2022-12-12 14:02:28,628] {{logging_mixin.py:112}} WARNING - with self._thread_lock:
[2022-12-12 14:02:28,646] {{logging_mixin.py:112}} WARNING - AttributeError: 'UnixFileLock' object has no attribute '_thread_lock'
Does anyone faced some similar error?
Usually after one or two retries of the failed tasks make them succeed. Because of this intermittent intermittent character and the BaseLock mention in the logs I thogth it was related to some concurrency problem. But even eliminating all the concurrency between tasks on the airflow instance (by limiting the pool slots to 1) didn't solved the problem.