"""Tasks of the GfK Streaming Ingestion Workflow.""" import datetime import ftplib import os import re import shutil import socket import ssl import tempfile import time import zipfile import boto3 from garcon import task from garcon_contrib.dynamo_feed_status import garcon_feed_status from feed_ingestion.flows.gfk_streaming import config from feed_ingestion.tasks import check_status _FTPS_CONNECT_TIMEOUT = 30 _DOWNLOAD_MAX_RETRIES = 30 _DOWNLOAD_RETRY_WAIT = 5 # GfK CSV files use backslash to escape embedded quotes (e.g. \"word\"), # but also contain literal trailing backslashes before closing quotes # (e.g. "value\"). These two patterns are incompatible with a single # Snowflake ESCAPE setting, so we normalise in-place: convert \" that is # NOT a field terminator to "" (standard CSV double-quote escaping), # leaving \" before ; or a line ending untouched. _EMBEDDED_BACKSLASH_QUOTE = re.compile(rb'\\"(?![;\r\n])') def _normalize_csv_quotes(path): tmp_path = path + '.tmp' with open(path, 'rb') as f_in, open(tmp_path, 'wb') as f_out: for line in f_in: f_out.write(_EMBEDDED_BACKSLASH_QUOTE.sub(b'""', line)) os.replace(tmp_path, path) class _RobustFTP_TLS(ftplib.FTP_TLS): """FTP_TLS that suppresses the malformed TLS shutdown this server sends.""" def retrbinary(self, cmd, callback, blocksize=8192, rest=None): self.voidcmd('TYPE I') with self.transfercmd(cmd, rest) as conn: while True: data = conn.read(blocksize) if not data: break callback(data) try: if isinstance(conn, ssl.SSLSocket): conn.unwrap() except ssl.SSLError: pass # server closes TLS without proper shutdown handshake return self.voidresp() def _get_ftps_connection(ftps_creds): """Open an authenticated FTPS connection with data channel TLS. Args: ftps_creds (dict): Keys: host, port, username, password. Returns: _RobustFTP_TLS: Connected and logged-in FTPS client. """ ftp = _RobustFTP_TLS(timeout=_FTPS_CONNECT_TIMEOUT) ftp.connect(ftps_creds['host'], ftps_creds['port']) ftp.login(ftps_creds['username'], ftps_creds['password']) ftp.prot_p() return ftp @task.decorate(timeout=700) def bootstrap(activity, date, reload): """Bootstrap workflow by getting the correct configurations. Args: activity (ActivityWorker): The activity worker. date (str): Reporting date (YYYY-MM-DD). reload (str or None): If 'True' reload status. Returns: dict: Context. """ date_obj = ( datetime.datetime.strptime(date, '%Y-%m-%d').date() if date else datetime.date.today() ) date = date_obj.strftime('%Y-%m-%d') if reload == 'True': activity.logger.info( f'Delete status for feed: {config.feed_name} {date}' ) garcon_feed_status.delete_status(config.feed_name, date) elif ( garcon_feed_status.get_overall_status(config.feed_name, date) == garcon_feed_status.STATUS_INGESTED ): activity.logger.info(f'Feed already ingested for {date}') return { 'stop': True, 'message': f'{config.feed_name} is already ingested for {date}', } activity.logger.info(f'Bootstrap flow for {date}') yyyy = date_obj.strftime('%Y') mm = date_obj.strftime('%m') dd = date_obj.strftime('%d') filename = config.file_template.format(yyyy=yyyy, mm=mm, dd=dd) archive_dir = config.s3['archive_path'].format(date=date) processed_dir = config.s3['processed_path'].format(date=date) processed_filename = filename.replace('.zip', '.csv') s3_dir_path = f's3://{config.data_bucket}/{processed_dir}' return dict( feed_name=config.feed_name, date=date, filename=filename, archive_path=f'{archive_dir}{filename}', processed_path=f'{processed_dir}{processed_filename}', staging_raw_table=config.staging_raw_table, s3_dir_path=s3_dir_path, ) @task.decorate(timeout=10800) @check_status() def fetch_from_ftps( activity, feed_name, date, filename, archive_path, processed_path ): """Download ZIP from FTPS, extract CSV, and upload both to S3. Args: activity (ActivityWorker): The activity worker. feed_name (str): The name of the feed. date (str): Reporting date (YYYY-MM-DD). filename (str): Expected ZIP file name on the FTPS server. archive_path (str): S3 key for the raw ZIP archive. processed_path (str): S3 key for the extracted CSV file. """ remote_dir = config.ftps.get('path', '/FEED').rstrip('/') remote_file_path = f'{remote_dir}/{filename}' activity.logger.info(f'Checking FTPS for {filename}') ftp = _get_ftps_connection(config.ftps) try: ftp.size(remote_file_path) except ftplib.error_perm: activity.logger.info(f'File not available on FTPS: {filename}') garcon_feed_status.set_overall_status( feed_name, date, garcon_feed_status.STATUS_NOT_AVAILABLE ) ftp.quit() return {'stop': True, 'message': f'Missing file {filename}'} ftp.quit() local_zip_path = os.path.join(tempfile.gettempdir(), filename) resume_from = ( os.path.getsize(local_zip_path) if os.path.exists(local_zip_path) else 0 ) for attempt in range(1, _DOWNLOAD_MAX_RETRIES + 1): try: activity.logger.info( f'Downloading {filename} via FTPS ' f'(attempt {attempt}/{_DOWNLOAD_MAX_RETRIES}, ' f'resume from {resume_from / 1_048_576:.1f} MB)' ) ftp = _get_ftps_connection(config.ftps) with open(local_zip_path, 'ab' if resume_from else 'wb') as f: ftp.retrbinary( f'RETR {remote_file_path}', f.write, rest=resume_from or None, ) ftp.quit() break except ( socket.error, ftplib.Error, ssl.SSLError, EOFError, OSError ) as e: resume_from = ( os.path.getsize(local_zip_path) if os.path.exists(local_zip_path) else 0 ) if attempt < _DOWNLOAD_MAX_RETRIES: activity.logger.warning( f'Transfer interrupted ({type(e).__name__}: {e}). ' f'Retrying in {_DOWNLOAD_RETRY_WAIT}s ' f'({resume_from / 1_048_576:.1f} MB saved so far).' ) time.sleep(_DOWNLOAD_RETRY_WAIT) else: raise RuntimeError( f'FTPS download of {filename} failed after ' f'{_DOWNLOAD_MAX_RETRIES} attempts' ) from e activity.logger.info(f'Downloaded {filename} to {local_zip_path}') s3_client = boto3.client('s3') activity.logger.info(f'Uploading ZIP to S3 archive: {archive_path}') s3_client.upload_file(local_zip_path, config.data_bucket, archive_path) local_csv_path = os.path.join( tempfile.gettempdir(), filename.replace('.zip', '.csv') ) activity.logger.info(f'Extracting CSV from ZIP: {filename}') with zipfile.ZipFile(local_zip_path) as zf: csv_names = [n for n in zf.namelist() if n.upper().endswith('.CSV')] if not csv_names: raise ValueError(f'No CSV file found inside ZIP: {filename}') with zf.open(csv_names[0]) as src, open(local_csv_path, 'wb') as dst: shutil.copyfileobj(src, dst) os.remove(local_zip_path) _normalize_csv_quotes(local_csv_path) activity.logger.info(f'Uploading extracted CSV to S3: {processed_path}') s3_client.upload_file(local_csv_path, config.data_bucket, processed_path) os.remove(local_csv_path) garcon_feed_status.set_overall_status( feed_name, date, garcon_feed_status.STATUS_DOWNLOADED ) activity.logger.info(f'Successfully processed {filename}') @task.decorate(timeout=300) def delete_processed_file(activity, processed_path): """Delete the extracted CSV from S3 after it has been loaded. Args: activity (ActivityWorker): The activity worker. processed_path (str): S3 key of the extracted CSV to delete. """ activity.logger.info(f'Deleting processed file from S3: {processed_path}') boto3.client('s3').delete_object( Bucket=config.data_bucket, Key=processed_path, ) activity.logger.info(f'Deleted {processed_path}')