"""Parse the SONY_STARS_MONTHLY_FINANCIAL_FEED_VIEW to files.""" import io # You know what this is. import os # basic os operations like paths import shutil # for archiving import sys # for sysexit() import time # for sleep() from datetime import datetime # for date formatting from typing import Optional, Sequence from pebble import concurrent # type: ignore # for async processing from src.logger import get_logger from src.utils.paths import create_path_with_todays_date from models.snowflake_models import SMEFeedFileHistoryPersister as feed_persist from constants import ( database as db_consts, files as file_consts, processing_modes as modes, ) import config from sftp_transfer import ( SFTPConfig, transfer_non_zipped_files_if_enabled, resolve_password_interactive, ) from s3_transfer import ( S3Config, transfer_to_s3_if_enabled, ) from src.utils.compression import gzip_files_in_directory logger = get_logger() # Scheduler workspace base URL for generating output links SCHEDULER_WS_BASE_URL = ( 'https://scheduler.theorchard.io/job/sme-feed-file-exporter/ws/' ) def get_scheduler_output_link(output_path: str) -> str: """Generate a scheduler workspace link for the given output path. Extracts the relative path from 'output/' onwards and constructs a full URL to the scheduler workspace. Args: output_path: The absolute or relative path to the output folder. Returns: A full URL to the scheduler workspace for the output folder. """ # Normalize and find the 'output/' portion of the path normalized = output_path.replace('\\', '/') # Find 'output/' in the path and extract from there output_marker = 'output/' if output_marker in normalized: relative_path = normalized[normalized.find(output_marker):] else: # Fallback: use the last few path components relative_path = normalized.lstrip('./') # Ensure trailing slash for directory links if not relative_path.endswith('/'): relative_path += '/' return f'{SCHEDULER_WS_BASE_URL}{relative_path}' # Dev-only local SFTP helper: enable ephemeral filesystem-backed SFTP when # explicitly targeting localhost in a developer context. This avoids network # usage while allowing end-to-end verification. def _enable_local_sftp_if_requested() -> None: try: import os import paramiko # type: ignore from pathlib import Path import shutil except Exception: return if os.environ.get('SFTP_HOST') != 'localhost': return # Guard to avoid unintended activation outside dev/local contexts env = (os.environ.get('ENVIRONMENT') or '').upper() if env not in ('', 'DEV', 'LOCAL'): return user = os.environ.get('SFTP_USER', 'devuser') pw = os.environ.get('SFTP_PASSWORD', 'devpass') root = Path('tmp/local_sftp_root') root.mkdir(parents=True, exist_ok=True) class _FSBackedSFTPClient: def __init__(self, base: Path): self.base = base def put(self, local: str, remote: str): dest = self.base / remote dest.parent.mkdir(parents=True, exist_ok=True) shutil.copy2(local, dest) def close(self): pass class _MockSSHClient: def __init__(self, base: Path, expect_user: str, expect_pass: str): self.base = base self.expect_user = expect_user self.expect_pass = expect_pass self._policy = None def load_system_host_keys(self): return None def set_missing_host_key_policy(self, policy): self._policy = policy def connect( self, hostname: str, username: str, password: str, port: int, timeout: int, **kwargs, ): if username != self.expect_user or password != self.expect_pass: raise PermissionError('Invalid local SFTP credentials') return None def open_sftp(self): return _FSBackedSFTPClient(self.base) def close(self): return None # Monkeypatch paramiko for the current process paramiko.SSHClient = lambda: _MockSSHClient(root, user, pw) # type: ignore # TODO: Refactor out the printed logs to a returned object def row_to_dict(row): """Convert a SQLAlchemy Row (or RowProxy) to a dict. If it's already a dict, return it as is. """ if isinstance(row, dict): return row # SQLAlchemy 2.x: Row objects have _mapping attribute for dict-like access if hasattr(row, '_mapping'): return dict(row._mapping) # SQLAlchemy 1.x: RowProxy has keys() method if hasattr(row, 'keys') and callable(row.keys): return dict(zip(row.keys(), row)) # Fallback: try to convert directly return dict(row) def get_period_list(table_name): """Retrieve a list of all periods present in table. Returns: list """ logger.info("Finding periods in view.") period_list = feed_persist.get_all_periods(table_name).message logger.info("Periods found: {}".format(len(period_list))) return period_list # store_list, period_list def get_affiliate_list(table_name): """Retrieve a list of all booking affiliates present in table. Returns: list """ logger.info("Finding booking affiliates in view.") affiliate_list = feed_persist.get_all_booking_affiliates(table_name).message logger.info("Booking affiliates found: {}".format(len(affiliate_list))) return affiliate_list # store_list, period_list def get_group_name_list(table_name): """Retrieve a list of all booking affiliates present in table. Returns: list """ logger.info("Finding group names in view.") group_list = feed_persist.get_all_group_names(table_name).message logger.info("Number of group names found: {}".format(len(group_list))) return group_list # store_list, period_list def get_territory_list(table_name): """Retrieve a list of all territories present in table. Returns: list """ logger.info("Finding territories in view.") territory_list = feed_persist.get_all_territories(table_name).message logger.info("Territories found: {}".format(len(territory_list))) return territory_list # store_list, period_list def get_group_indicator_list(table_name): """Retrieve a list of all territories present in table. Returns: list """ logger.info("Finding territories in view.") territory_list = feed_persist.get_all_group_indicators(table_name).message logger.info("Group indicators found: {}".format(len(territory_list))) return territory_list # store_list, period_list def get_label_list(table_name): """Retrieve a list of all labels present in table. Returns: list """ logger.info("Finding labels in view.") label_list = feed_persist.get_all_labels(table_name).message logger.info("Labels found: {}".format(len(label_list))) return label_list # store_list, period_list def get_store_list(table_name): """Retrieve a list of all labels present in table. Returns: list """ logger.info("Finding stores in view.") store_list = feed_persist.get_all_stores(table_name).message logger.info("Stores found: {}".format(len(store_list))) return store_list # store_list, period_list def create_temp_table_from_period_id(table_name, period_id): """Creates a table by selecting a set filtered by period_id, by using the same table name, and appending a period_id on to it. Args: table_name: The table to duplicate period_id: The period ID by which to filter Returns: """ temp_table_name = "{}_{}".format(table_name, period_id) logger.info( "Creating temp table '{}' from {} using period id: {}.".format( temp_table_name, table_name, period_id ) ) feed_persist.create_table_as_select_by_period( source_table=table_name, target_table=temp_table_name, period_id=period_id ) return temp_table_name def process_table_to_files(table_name, temp_path, process=None): """Iterate through stores and periods, and write local files. Args: table_name (str): Table to process temp_path (str): temporary file storage (file write target) process (str): Whether the input set reflects us or ex-us sales """ if not process: process = modes.US logger.info("Counting rows in view.") # Count all rows count_retrieved_rows = feed_persist.get_count_of_rows(table_name=table_name).message logger.info( "Begin processing {} rows in `{}`".format(count_retrieved_rows, table_name) ) # Process Rows ------------------ count_processed_rows = 0 # Default file_name_list = [] # Default if process == modes.US: logger.info("Processing file as 'US'.") count_processed_rows, file_name_list = process_regular_feed( table_name, temp_path ) elif process == modes.GB: logger.info("Processing file as 'GB'.") count_processed_rows, file_name_list = process_regular_feed( table_name, temp_path ) elif process == modes.EX_US: logger.info("Processing file as 'Ex-US'.") count_processed_rows, file_name_list = process_ex_us_feed(table_name, temp_path) elif process == modes.AGGREGATE: logger.info('Processing file as \'AGGREGATE\'.') count_processed_rows, file_name_list = process_aggregate_feed( table_name, temp_path) elif process == modes.SAP_SETTLEMENT: # SAP Settlement uses same processing as AGGREGATE but routes to S3 logger.info('Processing file as \'SAP_SETTLEMENT\'.') count_processed_rows, file_name_list = process_aggregate_feed( table_name, temp_path ) logger.info("{} rows retrieved from `{}`".format(count_retrieved_rows, table_name)) logger.info("{} rows processed".format(count_processed_rows)) # Compare count_processed_rows to count_retrieved_rows if not count_retrieved_rows == count_processed_rows: logger.error( "The count of rows in the table does not match " "the count of rows processed to files. " "Retrieved {} =/= {} processed. Diff = {}".format( count_retrieved_rows, count_processed_rows, abs(count_retrieved_rows - count_processed_rows), ) ) return file_name_list def slice_rows(rows, size): """Yield successive n-sized chunks from rows. Args: rows (list): The list of rows to slice size (int): The size of the slices Returns: generator - A generator of slices of the rows """ row_count = len(rows) for i in range(0, row_count, size): yield rows[i : i + size] def write_rows_to_file(output_path, file_name, rows): """Write a list of dictionaries, representing a query result, into a file. Uses batched join writes with explicit file buffering to reduce syscall overhead for large datasets (INT-2638 optimization). Args: output_path: (str) The output path name file_name: (str) The output file name rows: (dict) The rows to write to the file, with keys: - header: Header content - body: List of body content strings - tail: Tail content Returns: (str) The full path to the written file """ local_file_full_path = os.path.join(output_path, os.path.basename(file_name)) # Local refs for speed in tight loop line_term = file_consts.LINE_TERMINATOR buffer_size = config.WRITE_BUFFER_SIZE # Write file with explicit large buffering for performance with io.open( local_file_full_path, "a", encoding="utf8", buffering=config.FILE_BUFFER_SIZE ) as outfile: # Write header outfile.write(str(rows["header"]) + line_term) # Write body in batches using join instead of list comprehension # This reduces memory allocations and syscalls for chunk in slice_rows(rows["body"], buffer_size): # Use generator expression with join for efficient batch write outfile.write(''.join( str(line) + line_term for line in chunk )) # Write tail (no trailing terminator per original behavior) outfile.write(str(rows["tail"])) return local_file_full_path def split_rows_to_files(file_rows, output_path): """Split a list of dictionaries, representing a query result, into files, by reading the content of the body_content field of the rows. Note: When multiple H/T pairs are present in file_rows, all body content is aggregated and written to ONE file using the LAST header/tail. This is the original aggregation behavior preserved for backward compatibility. The underlying write_rows_to_file() uses buffered batch writes for performance (INT-2638 optimization). Args: file_rows: (list) A list of dicts representing rows of data output_path: (str) The output path name Returns: (tuple) count_file_processed_rows, upload_file_list, generated_file_list - The count of processed rows, the list of files to upload, and the generated file list """ # Init vars count_file_processed_rows = 0 # file row count file_name = None # output file name header = None # header unit from db tail = None # tail unit from db body = [] # Accumulate body content across all H/T pairs upload_file_list = [] generated_file_list = [] # Send rows to helper funcs for parsing for row in file_rows: period_id = row.get(db_consts.PERIOD_ID) store_id = row.get(db_consts.STORE_ID) group_name = row.get(db_consts.GROUP_NAME) label_id = row.get(db_consts.LABEL_ID) affiliate_id = row.get(db_consts.BOOKING_AFFILIATE) if row[db_consts.RECORD_TYPE] == "H": header = row[db_consts.BODY_CONTENT] # Preserve raw Snowflake filename for H/T validation (INT-2649) raw_file_name = row[db_consts.FILE_NAME] file_name = raw_file_name # Add timestamp to filename if available timestamp = os.environ.get('FEED_FILE_TIMESTAMP') if timestamp and file_name: # Insert timestamp before file extension name_parts = os.path.splitext(file_name) if name_parts[1]: # Has extension file_name = f"{name_parts[0]}_{timestamp}{name_parts[1]}" else: # No extension file_name = f"{file_name}_{timestamp}" # Add file name to list for S3 upload upload_file_list.append(file_name) if row[db_consts.RECORD_TYPE] in file_consts.BODY_RECORD_TYPES: body.append(row[db_consts.BODY_CONTENT]) if row[db_consts.RECORD_TYPE] == "T": tail = row[db_consts.BODY_CONTENT] # Compare raw Snowflake filenames (pre-timestamp-mutation) so that # FEED_FILE_TIMESTAMP injection does not cause a spurious mismatch if row[db_consts.FILE_NAME] != raw_file_name: error_msg = ( "There was a mismatch between Header and Tail:" " Header: {} - Tail: {}" ) error_msg = error_msg.format( raw_file_name, row[db_consts.FILE_NAME], ) if affiliate_id: error_msg += " - Booking Affiliate: {}".format(affiliate_id) if period_id: error_msg += " - Period: {}".format(period_id) if store_id: error_msg += " - Store: {}".format(store_id) if group_name: error_msg += " - Group Name: {}".format(group_name) if label_id: error_msg += " - Label ID: {}".format(label_id) raise ValueError(error_msg) # Increment file row count count_file_processed_rows += 1 if count_file_processed_rows > 0: # Report writing file logger.info( "Writing File ({} rows): '{}'".format(count_file_processed_rows, file_name) ) # Combine rows to new object output_rows = {"header": header, "body": body, "tail": tail} # Create path to local file using optimized write_rows_to_file # which uses buffered batch writes (INT-2638 optimization) local_file_full_path = write_rows_to_file(output_path, file_name, output_rows) generated_file_list.append(local_file_full_path) return count_file_processed_rows, upload_file_list, generated_file_list def split_res_to_files(result_proxy, output_path): """Split a SQLAlchemy resultset into files, by reading the content of the body_content field of the rows. Uses buffered batch writes to reduce syscall overhead for large datasets. Rows are accumulated in a buffer and flushed when threshold is reached or when the tail record is encountered. Args: result_proxy: (ResultProxy) A result object representing rows of data output_path: (str) The output path name Returns: (tuple) count_file_processed_rows, generated_file_list - The count of processed rows and the list of generated local file paths """ # Init vars count_file_processed_rows = 0 # file row count generated_file_list = [] local_file_full_path = None # Buffer configuration for batched writes (INT-2638 optimization) write_buffer = [] # Accumulate rows before flushing buffer_threshold = config.WRITE_BUFFER_SIZE line_term = file_consts.LINE_TERMINATOR # Local ref for speed row = result_proxy.fetchone() row = row_to_dict(row) if not row[db_consts.RECORD_TYPE] == "H": raise ValueError("First row must be a header row.") header = row[db_consts.BODY_CONTENT] # Preserve raw Snowflake filename for H/T validation raw_file_name = row[db_consts.FILE_NAME] file_name = raw_file_name # Add timestamp to filename if available (consistent with split_rows_to_files) timestamp = os.environ.get('FEED_FILE_TIMESTAMP') if timestamp and file_name: name_parts = os.path.splitext(file_name) if name_parts[1]: # Has extension file_name = f"{name_parts[0]}_{timestamp}{name_parts[1]}" else: # No extension file_name = f"{file_name}_{timestamp}" local_file_full_path = os.path.join(output_path, os.path.basename(file_name)) logger.info("Writing File: {}".format(file_name)) # Write header to file with explicit large buffering for performance with io.open( local_file_full_path, "a", encoding="utf8", buffering=config.FILE_BUFFER_SIZE ) as outfile: outfile.write(str(header)) count_file_processed_rows += 1 # Send rows to helper funcs for parsing for row in result_proxy: # zip result tuple into dict row = row_to_dict(row) period_id = row.get(db_consts.PERIOD_ID) store_id = row.get(db_consts.STORE_ID) label_id = row.get(db_consts.LABEL_ID) affiliate_id = row.get(db_consts.BOOKING_AFFILIATE) if row[db_consts.RECORD_TYPE] in file_consts.BODY_RECORD_TYPES: body = row[db_consts.BODY_CONTENT] # Accumulate body rows in buffer instead of per-row writes write_buffer.append(line_term + str(body)) # Flush buffer when threshold reached if len(write_buffer) >= buffer_threshold: outfile.write(''.join(write_buffer)) write_buffer.clear() if row[db_consts.RECORD_TYPE] == "T": tail = row[db_consts.BODY_CONTENT] # If row header file_name doesn't match tail file_name # then throw exception if row[db_consts.FILE_NAME] != raw_file_name: error_msg = ( "There was a mismatch between Header and Tail:" " Header: {} - Tail: {}" ) error_msg = error_msg.format( raw_file_name, row[db_consts.FILE_NAME], ) if affiliate_id: error_msg += " - Booking Affiliate: {}".format(affiliate_id) if period_id: error_msg += " - Period: {}".format(period_id) if store_id: error_msg += " - Store: {}".format(store_id) if label_id: error_msg += " - Label ID: {}".format(label_id) raise ValueError(error_msg) # Flush any remaining buffer before writing tail if write_buffer: outfile.write(''.join(write_buffer)) write_buffer.clear() outfile.write(line_term + str(tail)) # Increment file row count count_file_processed_rows += 1 # Final buffer flush (should be empty if tail was processed) if write_buffer: outfile.write(''.join(write_buffer)) write_buffer.clear() generated_file_list.append(local_file_full_path) return count_file_processed_rows, generated_file_list def process_regular_feed(table_name, output_path): """Process US feed table by period_id and store_id. Args: table_name: (str) The name of the source table output_path: (str) The file output path. Returns: (tuple) count_processed_rows, file_name_list The count of processed rows, and the list of generated filenames """ file_name_list = [] count_total_processed_rows = 0 promise = None if os.environ.get('ENVIRONMENT') != 'TEST': logger.info('Spawning async keepalive process.') # Use pebble to create a new async process promise = write_blank_file() # Retrieve a list of all periods present in table period_list = get_period_list(table_name) # For each store/label/period in store_dict # for label in label_list: for period_id in period_list: group_name_list = feed_persist.get_group_names_for_period( period_id=period_id, table_name=table_name ).message logger.info( "{} group names found for period {}".format(len(group_name_list), period_id) ) for num, group_name in enumerate(group_name_list, start=1): params = { "period_id": period_id, "group_name": group_name, "table_name": table_name, } logger.info( "Finding rows for group name {} in period {}.".format( group_name, period_id ) ) # Get file rows by period_store file_rows = feed_persist.get_by_period_group_name(**params).message if not file_rows: logger.info( "No rows for group name {} in period {}".format( group_name, period_id ) ) continue count_file_processed_rows, _upload_file_list, new_file_list = ( split_rows_to_files(file_rows, output_path) ) count_total_processed_rows += count_file_processed_rows file_name_list = file_name_list + new_file_list if promise: logger.info('Terminating keepalive process.') promise.cancel() return count_total_processed_rows, file_name_list # TODO - Test this before resuming use def process_ex_us_feed(table_name, output_path): """Process EX-US feed table by territory ID and Period ID, and label ID. Args: table_name: (str) The name of the source table output_path: (str) The file output path. Returns: (tuple) count_processed_rows, file_name_list The count of processed rows, and the list of generated filenames """ file_name_list = [] count_processed_rows = 0 promise = None if os.environ.get('ENVIRONMENT') != 'TEST': logger.info('Spawning async keepalive process.') promise = write_blank_file() # Retrieve a list of all periods present in table period_list = get_period_list(table_name) # Retrieve a list of all territories present in table territory_list = get_territory_list(table_name) # Retrieve a list of all labels present in table label_list = get_label_list(table_name) # For each store/label/period in store_dict # TODO: Is this still correct?! - Need to test before resuming use. for territory in territory_list: for period_id in period_list: for label_id in label_list: params = { "period_id": period_id, "table_name": table_name, "label_id": label_id, } logger.info( "Finding rows for label {} in period {}.".format( label_id, period_id ) ) file_rows = feed_persist.get_by_period_id(**params).message if not file_rows: logger.info( "No rows for label {} in period {}".format(label_id, period_id) ) continue count_file_processed_rows, _upload_file_list, new_file_list = ( split_rows_to_files(file_rows, output_path) ) count_processed_rows += count_file_processed_rows file_name_list = file_name_list + new_file_list if promise: logger.info('Terminating keepalive process.') promise.cancel() return count_processed_rows, file_name_list def process_aggregate_feed_streaming(table_name, output_path, group_name): """Process a single group using streaming mode.""" params = { "group_name": group_name, "table_name": table_name, } count_processed_rows = 0 file_name_list = [] with feed_persist.get_by_group_name_stream(**params) as res: # split_res_to_files consumes the cursor count, new_files = split_res_to_files(res, output_path) count_processed_rows += count file_name_list.extend(new_files) return count_processed_rows, file_name_list def process_aggregate_feed_batch(table_name, output_path, group_name, mode="CHUNKED"): """Process a single group using batch mode (CHUNKED or CURSOR). Uses streaming row conversion and buffered batch writes to reduce memory usage and syscall overhead for large datasets (INT-2638 optimization). """ params = { "group_name": group_name, "table_name": table_name, "chunk_size": config.CHUNK_SIZE, } count_processed_rows = 0 file_name_list = [] if mode == "CURSOR": batch_gen = feed_persist.get_by_group_name_cursor(**params) else: batch_gen = feed_persist.get_by_group_name_chunked(**params) # State tracking for file writing across chunks current_file_path = None file_handle = None current_file_name = None # Buffer configuration for batched writes (INT-2638 optimization) write_buffer = [] # Accumulate body rows before flushing buffer_threshold = config.WRITE_BUFFER_SIZE line_term = file_consts.LINE_TERMINATOR # Local ref for speed try: for chunk in batch_gen: # Stream row conversion on-demand instead of list comprehension # This reduces memory allocation from O(chunk_size) to O(1) for raw_row in chunk: row = row_to_dict(raw_row) row_type = row.get(db_consts.RECORD_TYPE) if row_type == "H": # Flush any remaining buffer before closing file if write_buffer and file_handle: file_handle.write(''.join(write_buffer)) write_buffer.clear() # Close previous file if open if file_handle: file_handle.close() # Preserve raw Snowflake filename for H/T validation raw_current_file_name = row[db_consts.FILE_NAME] current_file_name = raw_current_file_name # Add timestamp to filename if available timestamp = os.environ.get('FEED_FILE_TIMESTAMP') if timestamp and current_file_name: name_parts = os.path.splitext(current_file_name) if name_parts[1]: # Has extension current_file_name = f"{name_parts[0]}_{timestamp}{name_parts[1]}" else: # No extension current_file_name = f"{current_file_name}_{timestamp}" current_file_path = os.path.join( output_path, os.path.basename(current_file_name) ) file_name_list.append(current_file_path) logger.info("Writing File: {}".format(current_file_name)) # Open with explicit large buffering for performance file_handle = io.open( current_file_path, "a", encoding="utf8", buffering=config.FILE_BUFFER_SIZE ) file_handle.write(str(row[db_consts.BODY_CONTENT])) count_processed_rows += 1 elif row_type in file_consts.BODY_RECORD_TYPES: if file_handle: # Accumulate body rows in buffer instead of per-row writes write_buffer.append( line_term + str(row[db_consts.BODY_CONTENT]) ) count_processed_rows += 1 # Flush buffer when threshold reached if len(write_buffer) >= buffer_threshold: file_handle.write(''.join(write_buffer)) write_buffer.clear() elif row_type == "T": if file_handle: # Verify footer matches header file if row[db_consts.FILE_NAME] != raw_current_file_name: raise ValueError( f"Header/Tail Mismatch: {raw_current_file_name} != {row[db_consts.FILE_NAME]}" ) # Flush remaining buffer before writing tail if write_buffer: file_handle.write(''.join(write_buffer)) write_buffer.clear() file_handle.write( line_term + str(row[db_consts.BODY_CONTENT]) ) count_processed_rows += 1 finally: # Final buffer flush and file close if write_buffer and file_handle: file_handle.write(''.join(write_buffer)) write_buffer.clear() if file_handle: file_handle.close() return count_processed_rows, file_name_list def process_aggregate_feed(table_name, output_path): """Process Aggregate feed table by store_id and affiliate_id. Args: table_name: (str) The name of the source table output_path: (str) The file output path. Returns: (tuple) count_processed_rows, file_name_list The count of processed rows, and the list of generated filenames """ file_name_list = [] count_processed_rows = 0 fetch_mode = config.FETCH_MODE logger.info(f"Processing Aggregate Feed with mode: {fetch_mode}") # Retrieve a list of all labels present in table group_name_list = get_group_name_list(table_name) # For each store/affiliate_id in store_dict # for territory in territory_list: for group_name in group_name_list: params = { "group_name": group_name, "table_name": table_name, } row_count = feed_persist.get_count_by_group_name(**params).message if not row_count: logger.info("No rows for group_name ({})".format(group_name)) continue logger.info( "Processing {} rows for group_name ({}).".format( # territory, row_count, group_name, ) ) if fetch_mode == "STREAMING": c, f = process_aggregate_feed_streaming( table_name, output_path, group_name ) elif fetch_mode in ["CHUNKED", "CURSOR"]: c, f = process_aggregate_feed_batch( table_name, output_path, group_name, mode=fetch_mode ) else: # Fallback to legacy behavior file_rows = feed_persist.get_by_group_name(**params).message logger.info("Writing file for group_name: {}".format(group_name)) count_file_processed_rows, _upload_file_list, new_file_list = ( split_rows_to_files(file_rows, output_path) ) c = count_file_processed_rows f = new_file_list count_processed_rows += c file_name_list.extend(f) # Loop each Store logger.info("------") # separator return count_processed_rows, file_name_list # Create a new async process to keep env alive while the main process runs @concurrent.process def write_blank_file(timeout=None): if not timeout: timeout = 30 while True: with open("keepalive.remove", "w") as f: f.write(" ") time.sleep(timeout) # Begin main control def parse_dev_cli_args(argv: Optional[Sequence[str]] = None): """Parse development-oriented CLI arguments and apply env overrides. This helper is extracted to enable unit testing of CLI behavior without invoking the full module main sequence (which touches databases and file system). Jenkins usage remains unaffected. Returns the argparse namespace for further introspection if needed. """ import argparse import os parser = argparse.ArgumentParser( description='SME Stars feed file exporter (with optional SFTP).', add_help=True, ) parser.add_argument( '--sftp-password', dest='sftp_password', help='Inject SFTP password (DEV only; never commit).' ) parser.add_argument( '--prompt-sftp-password', dest='prompt_pw', action='store_true', help='Prompt interactively for password if missing.' ) parser.add_argument( '--mock-sftp', dest='mock_sftp', action='store_true', help='Enable mock SFTP mode (no real network).' ) parser.add_argument( '--enable-sftp', dest='enable_sftp', action='store_true', help='Force enable SFTP transfer regardless of env TRANSFER_FILES.' ) parser.add_argument( '--clear-local-sftp', dest='clear_local_sftp', action='store_true', help='Clear tmp/local_sftp_root before run (DEV only).' ) parser.add_argument( '--sftp-auto-confirm', dest='sftp_auto_confirm', action='store_true', help='Skip interactive confirmation prompt before SFTP upload.' ) ns, _unknown = parser.parse_known_args(argv) if ns.sftp_password: os.environ['SFTP_PASSWORD'] = ns.sftp_password if ns.prompt_pw: os.environ['SFTP_PROMPT_PASSWORD'] = 'True' if ns.mock_sftp: os.environ['SFTP_MOCK'] = 'True' if ns.enable_sftp: os.environ['TRANSFER_FILES'] = 'True' if getattr(ns, 'clear_local_sftp', False): os.environ['SFTP_CLEAR_LOCAL'] = 'True' if getattr(ns, 'sftp_auto_confirm', False): os.environ['SFTP_AUTO_CONFIRM'] = 'True' return ns def _is_interactive_tty() -> bool: """Return True if stdin/stdout are TTY (interactive shell).""" return sys.stdin.isatty() and sys.stdout.isatty() REQUIRED_ENV_FIELDS = [ 'SNOWFLAKE_SOURCE_TABLE', 'PERIOD_ID', ] # Secret fields (use getpass) SECRET_ENV_FIELDS = [ 'SNOWFLAKE_KEY_PASSPHRASE', # 'SNOWFLAKE_PASSWORD', # optional; included for completeness ] def prompt_missing_env(fields: Sequence[str]) -> None: """Prompt user for any missing required environment variable values. Only prompts when: - Session is interactive (TTY) - ENVIRONMENT indicates local/dev (ENVIRONMENT in {DEV, LOCAL, ''}) Skips prompting in CI/non-interactive contexts to avoid blocking. Empty answers are left unset. """ env = os.environ.get('ENVIRONMENT', '').upper() if not _is_interactive_tty() or env not in {'', 'DEV', 'LOCAL'}: return for key in fields: if os.environ.get(key): continue try: val = input(f"Enter value for {key} (leave blank to skip): ") except EOFError: # Non-interactive after all; abort prompting. return if val: os.environ[key] = val logger.info("Captured value for %s interactively.", key) else: logger.warning( "%s left unset; downstream operations may fail.", key ) def prompt_missing_secrets(fields: Sequence[str]) -> None: """Prompt for secret fields using getpass when missing. Uses getpass to avoid echoing sensitive values. Only performed in interactive DEV/LOCAL sessions. Blank responses leave variable unset. """ env = os.environ.get('ENVIRONMENT', '').upper() if not _is_interactive_tty() or env not in {'', 'DEV', 'LOCAL'}: return try: import getpass except ImportError: # pragma: no cover return for key in fields: if os.environ.get(key): continue try: val = getpass.getpass( f"Enter secret for {key} (leave blank to skip): " ) except Exception: # pragma: no cover - defensive return if val: os.environ[key] = val logger.info("Secret value for %s captured interactively.", key) else: logger.warning("Secret %s left unset.", key) def prepare_s3_transfer_files( generated_files: Sequence[str], temp_file_path: str, gzip_enabled: bool, continue_on_error: bool, ) -> list[str]: """Return the files that should be uploaded to S3.""" if not generated_files: return [] if not gzip_enabled: logger.info( 'S3 gzip disabled. Uploading %d original files.', len(generated_files), ) return list(generated_files) logger.info('Gzipping files before S3 transfer.') try: success_results, failed_results = gzip_files_in_directory( directory=temp_file_path, pattern='*.[Tt][Xx][Tt]', delete_originals=False, compression_level=9, ) gzipped_files = [str(r.output_path) for r in success_results] logger.info( 'Gzip complete: %d files compressed, %d failed.', len(success_results), len(failed_results), ) if failed_results and not continue_on_error: logger.error('Gzip failures encountered. Aborting.') sys.exit('Gzip compression failed.') return gzipped_files except Exception as exc: logger.error('Failed to gzip files: %s', exc) if not continue_on_error: sys.exit('Gzip compression raised fatal exception.') return [] def resolve_s3_transfer_mode(transfer_mode: str) -> tuple[bool, bool]: """Return whether S3 should upload individual files and/or zip.""" return ( transfer_mode in ['Individual Files', 'Both'], transfer_mode in ['Zip File', 'Both'], ) def build_zipfile_basename( temp_file_path: str, proc_mode: str, current_time: Optional[datetime] = None, ) -> str: """Return the basename to use for the generated workspace zip file.""" normalized_path = temp_file_path.rstrip('/') default_basename = os.path.basename(normalized_path) if proc_mode != modes.SAP_SETTLEMENT: return default_basename current_time = current_time or datetime.now() zip_date = current_time.strftime('%Y_%m_%d') return f'sony_settlement_feed-{zip_date}' if __name__ == '__main__': # Invoke CLI parsing first to allow overrides before reading config # driven flags later. # Dev-only: enable local SFTP mock if targeting localhost _enable_local_sftp_if_requested() # Optional: clear local SFTP root if requested via CLI try: from pathlib import Path is_local = os.environ.get('SFTP_HOST') == 'localhost' wants_clear = os.environ.get('SFTP_CLEAR_LOCAL') == 'True' if is_local and wants_clear: root = Path('tmp/local_sftp_root') if root.exists(): shutil.rmtree(root) except Exception: pass try: # pragma: no cover - exercised via dedicated unit test parse_dev_cli_args(sys.argv[1:]) except Exception as _arg_exc: # pragma: no cover - defensive logger.debug('Argparse initialization failed: %s', _arg_exc) # Prompt for missing required fields (dev/interactive only) before # resolving config values. prompt_missing_env(REQUIRED_ENV_FIELDS) prompt_missing_secrets(SECRET_ENV_FIELDS) # Lazy Snowflake key load and session init AFTER interactive prompts. try: from src.config import snowflake as _cfg # type: ignore _cfg.load_snowflake_private_key() except Exception as _key_err: # pragma: no cover logger.debug('Deferred key load skipped: %s', _key_err) try: from models import snowflake_models as _sfmod # type: ignore _sfmod.initialize_snowflake_session() except Exception as _init_err: # pragma: no cover logger.debug('Deferred Snowflake session init skipped: %s', _init_err) # Grab env vars source_table = os.environ.get( 'SNOWFLAKE_SOURCE_TABLE', config.SNOWFLAKE_SOURCE_TABLE ) period_id = os.environ.get('PERIOD_ID', config.PERIOD_ID) copy_table = config.COPY_TABLE proc_mode = config.PROCESS_MODE # Change output path based on processing mode date_str = datetime.now().strftime('%Y%m%d') base_path = config.FILE_OUTPUT_PATH.replace('{date}', date_str) temp_file_path = \ os.path.join( base_path, proc_mode.lower().replace('-', '_') ) # Make Output Folder logger.info('Using FILE_OUTPUT_PATH: {}'.format(temp_file_path)) temp_file_path = create_path_with_todays_date(temp_file_path) logger.info('Created file path: {}'.format(temp_file_path)) # Preserve source filenames for output/transfer and clear stale state. os.environ.pop('FEED_FILE_TIMESTAMP', None) logger.info('Output filenames will preserve source names.') # Optional create temp table. if copy_table: temp_table_name = create_temp_table_from_period_id(source_table, period_id) else: # Use passed source table temp_table_name = source_table # Process entire table try: list_file_names = process_table_to_files( temp_table_name, temp_file_path, proc_mode ) except RuntimeError as e: sys.exit("Unhandled exception: {}".format(str(e))) except ValueError as e: sys.exit(str(e)) logger.info('%s files created.', len(list_file_names)) logger.debug("Files created:") for fn in list_file_names: logger.debug(fn) # Display scheduler workspace link for output (before any network ops) scheduler_link = get_scheduler_output_link(temp_file_path) logger.info('Output available at: %s', scheduler_link) print(f'\n*** Output Location (Individual Files) ***\n{scheduler_link}\n') # Optional SFTP transfer BEFORE compressing & deleting source files. # Re-resolve config booleans after potential CLI overrides. # (Importing config earlier captured original env; update dynamic # fields manually.) # Only update the flags affected by CLI overrides to avoid side effects. # Check both env var names: SFTP_TRANSFER_ENABLED (Jenkins) and # TRANSFER_FILES (legacy CLI name, set by --enable-sftp flag). _sftp_enabled_env = ( os.environ.get('SFTP_TRANSFER_ENABLED') or os.environ.get('TRANSFER_FILES') ) if _sftp_enabled_env is not None: config.SFTP_TRANSFER_ENABLED = int( str(_sftp_enabled_env).lower() in {'true', '1', 'yes'} ) # type: ignore if 'SFTP_PASSWORD' in os.environ: config.SFTP_PASSWORD = os.environ.get('SFTP_PASSWORD') # type: ignore if 'SFTP_PROMPT_PASSWORD' in os.environ: # env var already set; presence ensures prompt logic can trigger. os.environ['SFTP_PROMPT_PASSWORD'] = os.environ['SFTP_PROMPT_PASSWORD'] if 'SFTP_MOCK' in os.environ: os.environ['SFTP_MOCK'] = os.environ['SFTP_MOCK'] # ensure presence # Get transfer mode (Individual Files, Zip File, or Both) transfer_mode = os.environ.get('TRANSFER_MODE', 'Individual Files') logger.info('Transfer mode: %s', transfer_mode) # Determine what to transfer transfer_individual = transfer_mode in ['Individual Files', 'Both'] transfer_zip = transfer_mode in ['Zip File', 'Both'] # Transfer individual files if enabled and requested if config.SFTP_TRANSFER_ENABLED and transfer_individual: logger.info( 'SFTP transfer enabled. Preparing to upload individual files.' ) # Prompt for confirmation before upload unless auto-confirm is set auto_confirm = str( os.environ.get('SFTP_AUTO_CONFIRM', 'False') ).lower() in {'1', 'true', 'yes'} proceed_with_upload = True # Track whether user confirmed if not auto_confirm and _is_interactive_tty(): logger.info( 'Pausing before SFTP upload. %d files ready to transfer.', len(list_file_names), ) try: dest = f'{config.SFTP_HOST}:{config.SFTP_REMOTE_DIR}' response = input( '\n*** SFTP Upload Confirmation ***\n' f'Files to upload: {len(list_file_names)}\n' f'Destination: {dest}\n' 'Proceed with SFTP upload? [Y/n]: ' ).strip().lower() if response and response not in {'y', 'yes', ''}: logger.info('SFTP upload cancelled by user.') print('SFTP upload cancelled.') proceed_with_upload = False except (EOFError, KeyboardInterrupt): logger.info('SFTP upload cancelled (no input/interrupt).') print('\nSFTP upload cancelled.') proceed_with_upload = False if proceed_with_upload: try: # TODO: DEPRECATE - Remove password resolution when migrating # to key-only authentication. # Resolve password interactively if allowed and missing. password = resolve_password_interactive(config.SFTP_PASSWORD) sftp_cfg = SFTPConfig( host=config.SFTP_HOST, port=config.SFTP_PORT, username=config.SFTP_USERNAME, password=password, remote_dir=config.SFTP_REMOTE_DIR, retries=config.SFTP_RETRIES, backoff_sec=config.SFTP_RETRY_BACKOFF_SEC, timeout=config.SFTP_TIMEOUT_SEC, strict_host_key=bool(config.SFTP_STRICT_HOST_KEY), continue_on_error=bool(config.SFTP_CONTINUE_ON_ERROR), auth_mode=config.SFTP_AUTH_MODE, private_key_path=config.SFTP_PRIVATE_KEY_PATH, ) upload_result = transfer_non_zipped_files_if_enabled( list_file_names, sftp_cfg, ) if upload_result: if upload_result.success(): logger.info( 'All %d files transferred via SFTP.', upload_result.uploaded, ) else: logger.warning( 'SFTP incomplete. Up=%d Fail=%d Total=%d', upload_result.uploaded, len(upload_result.failed), upload_result.total, ) for f in upload_result.failed: logger.error('Failed SFTP upload: %s', f.name) if not config.SFTP_CONTINUE_ON_ERROR: logger.error( 'Aborting subsequent steps due to SFTP failure.' ) sys.exit('SFTP transfer failed.') except Exception as exc: logger.error('Unhandled SFTP transfer exception: %s', exc) if not config.SFTP_CONTINUE_ON_ERROR: sys.exit('SFTP transfer raised fatal exception.') elif config.SFTP_TRANSFER_ENABLED: logger.info('Individual file transfer disabled by transfer mode.') else: logger.info('SFTP transfer disabled. Skipping file upload stage.') # -------------------- S3 Transfer Whitelist Guard -------------------- # S3 transfer is gated to whitelisted processing modes. The allowed set # is driven by S3_ALLOWED_MODES (comma-separated env var, see config.py). # If an operator enables S3_TRANSFER_ENABLED for a non-whitelisted mode # the job logs a warning and continues without transferring. _S3_ALLOWED_MODES = config.S3_ALLOWED_MODES if config.S3_TRANSFER_ENABLED and proc_mode not in _S3_ALLOWED_MODES: logger.warning( 'S3_TRANSFER_ENABLED is set but proc_mode "%s" is not ' 'in the S3 transfer whitelist (%s). ' 'Skipping S3 transfer.', proc_mode, ', '.join(sorted(_S3_ALLOWED_MODES)), ) # -------------------- S3 Transfer ------------------------------------ if config.S3_TRANSFER_ENABLED and proc_mode in _S3_ALLOWED_MODES: logger.info('Preparing S3 transfer for mode %s.', proc_mode) logger.info( 'S3 target environment: %s (role: %s)', config.S3_TARGET_ENV, config.S3_ROLE_ARN, ) s3_transfer_mode = config.S3_TRANSFER_MODE logger.info('S3 transfer mode: %s', s3_transfer_mode) s3_transfer_individual, s3_transfer_zip = ( resolve_s3_transfer_mode(s3_transfer_mode) ) s3_transfer_files = [] if s3_transfer_individual: s3_transfer_files = prepare_s3_transfer_files( generated_files=list_file_names, temp_file_path=temp_file_path, gzip_enabled=bool(config.S3_GZIP_ENABLED), continue_on_error=bool(config.S3_CONTINUE_ON_ERROR), ) if not s3_transfer_files: logger.info( 'No files available for S3 individual transfer. ' 'generated_files=%d transfer_mode=%s', len(list_file_names), s3_transfer_mode, ) if s3_transfer_individual and s3_transfer_files: logger.info( 'Transferring %d files to S3.', len(s3_transfer_files), ) try: s3_cfg = S3Config( bucket=config.S3_BUCKET, folder=config.S3_FOLDER, region=config.S3_REGION, retries=config.S3_RETRIES, backoff_sec=config.S3_RETRY_BACKOFF_SEC, timeout=config.S3_TIMEOUT_SEC, continue_on_error=bool(config.S3_CONTINUE_ON_ERROR), ) s3_result = transfer_to_s3_if_enabled( s3_transfer_files, s3_cfg, ) if s3_result: if s3_result.success(): logger.info( 'All %d files transferred to S3.', s3_result.uploaded, ) for s3_key in s3_result.s3_keys: logger.info( 'S3: s3://%s/%s', config.S3_BUCKET, s3_key, ) else: logger.warning( 'S3 transfer incomplete. Up=%d Fail=%d Total=%d', s3_result.uploaded, len(s3_result.failed), s3_result.total, ) for f in s3_result.failed: logger.error('Failed S3 upload: %s', f.name) if not config.S3_CONTINUE_ON_ERROR: logger.error( 'Aborting due to S3 transfer failure.' ) sys.exit('S3 transfer failed.') except Exception as exc: logger.error('Unhandled S3 transfer exception: %s', exc) if not config.S3_CONTINUE_ON_ERROR: sys.exit('S3 transfer raised fatal exception.') # Note: Zip transfer to S3 is handled after zip creation below elif proc_mode in _S3_ALLOWED_MODES: logger.info('S3 transfer disabled. Skipping S3 upload stage.') # Get zip generation mode zip_mode = os.environ.get('ZIP_GENERATION_MODE', 'Both') logger.info('Zip generation mode: %s', zip_mode) # Determine what to generate generate_zip = zip_mode in ['Zip File Only', 'Both'] keep_individual_files = zip_mode in ['Individual Files Only', 'Both'] # Only create zip if requested if generate_zip: # Compress the output folder (the timestamped run folder) # Normalize path to remove trailing slashes for consistent behavior normalized_path = temp_file_path.rstrip('/') zipfile_basename = build_zipfile_basename( temp_file_path=temp_file_path, proc_mode=proc_mode, ) parent_dir = os.path.dirname(normalized_path) if proc_mode == modes.SAP_SETTLEMENT: logger.info( 'SAP Settlement zip naming: %s.zip', zipfile_basename ) target_path = os.path.join(parent_dir, zipfile_basename) # Make the zip file zipfile_name = shutil.make_archive( target_path, 'zip', temp_file_path ) logger.info('Zipped archive created at %s.', zipfile_name) # Transfer zip file if enabled and requested if config.SFTP_TRANSFER_ENABLED and transfer_zip: logger.info('SFTP transfer enabled. Uploading zip file.') # Prompt for confirmation before upload unless auto-confirm is set auto_confirm = str( os.environ.get('SFTP_AUTO_CONFIRM', 'False') ).lower() in {'1', 'true', 'yes'} if not auto_confirm and _is_interactive_tty(): logger.info('Pausing before SFTP upload of zip file.') try: dest = f'{config.SFTP_HOST}:{config.SFTP_REMOTE_DIR}' response = input( '\n*** SFTP Upload Confirmation (Zip) ***\n' f'Zip file: {zipfile_name}\n' f'Destination: {dest}\n' 'Proceed with SFTP upload? [Y/n]: ' ).strip().lower() if response and response not in {'y', 'yes', ''}: logger.info('Zip SFTP upload cancelled by user.') print('SFTP upload cancelled.') transfer_zip = False # Skip this upload except (EOFError, KeyboardInterrupt): logger.info('Zip SFTP upload cancelled (no input).') print('\nSFTP upload cancelled.') transfer_zip = False # Skip this upload if config.SFTP_TRANSFER_ENABLED and transfer_zip: try: # TODO: DEPRECATE - Remove password resolution when migrating # to key-only authentication. # Resolve password interactively if allowed and missing. password = resolve_password_interactive(config.SFTP_PASSWORD) sftp_cfg = SFTPConfig( host=config.SFTP_HOST, port=config.SFTP_PORT, username=config.SFTP_USERNAME, password=password, # TODO: DEPRECATE - Remove param remote_dir=config.SFTP_REMOTE_DIR, retries=config.SFTP_RETRIES, backoff_sec=config.SFTP_RETRY_BACKOFF_SEC, timeout=config.SFTP_TIMEOUT_SEC, strict_host_key=bool(config.SFTP_STRICT_HOST_KEY), continue_on_error=bool(config.SFTP_CONTINUE_ON_ERROR), auth_mode=config.SFTP_AUTH_MODE, private_key_path=config.SFTP_PRIVATE_KEY_PATH, ) upload_result = transfer_non_zipped_files_if_enabled( [zipfile_name], sftp_cfg, ) if upload_result: if upload_result.success(): logger.info('Zip file transferred via SFTP.') else: logger.warning('Zip file transfer failed.') if not config.SFTP_CONTINUE_ON_ERROR: logger.error( 'Aborting due to zip SFTP failure.' ) sys.exit('Zip SFTP transfer failed.') except Exception as exc: logger.error('Unhandled zip SFTP transfer exception: %s', exc) if not config.SFTP_CONTINUE_ON_ERROR: sys.exit('Zip SFTP transfer raised fatal exception.') # Transfer zip file to S3 if S3 is enabled for this mode _s3_zip_eligible = ( config.S3_TRANSFER_ENABLED and proc_mode in _S3_ALLOWED_MODES ) if _s3_zip_eligible: _, s3_zip_flag = resolve_s3_transfer_mode( config.S3_TRANSFER_MODE ) if s3_zip_flag: logger.info('Transferring zip file to S3.') try: s3_cfg = S3Config( bucket=config.S3_BUCKET, folder=config.S3_FOLDER, region=config.S3_REGION, retries=config.S3_RETRIES, backoff_sec=config.S3_RETRY_BACKOFF_SEC, timeout=config.S3_TIMEOUT_SEC, continue_on_error=bool( config.S3_CONTINUE_ON_ERROR ), ) s3_result = transfer_to_s3_if_enabled( [zipfile_name], s3_cfg ) if s3_result: if s3_result.success(): logger.info( 'Zip file transferred to S3: s3://%s/%s', config.S3_BUCKET, s3_result.s3_keys[0] if s3_result.s3_keys else 'unknown', ) else: logger.warning('Zip file S3 transfer failed.') if not config.S3_CONTINUE_ON_ERROR: logger.error( 'Aborting due to zip S3 failure.' ) sys.exit('Zip S3 transfer failed.') except Exception as exc: logger.error( 'Unhandled zip S3 transfer exception: %s', exc ) if not config.S3_CONTINUE_ON_ERROR: sys.exit('Zip S3 transfer raised fatal exception.') else: logger.info( 'S3 zip transfer not selected by transfer mode: %s', config.S3_TRANSFER_MODE, ) else: logger.info('Zip file creation skipped (mode: %s)', zip_mode) # Delete individual files if only zip was requested if not keep_individual_files: logger.info('Deleting individual files (keeping zip only)') shutil.rmtree(os.path.dirname(temp_file_path)) else: logger.info('Keeping individual files (mode: %s)', zip_mode) # Optional drop temp table if config.COPY_TABLE and config.DELETE_COPY: feed_persist.drop_table(table_name=temp_table_name) logger.info("Temp table {} dropped.".format(temp_table_name)) # Final summary: display links to all generated artifacts print('\n*** Run Complete - Output Artifacts ***') if keep_individual_files: ind_link = get_scheduler_output_link(temp_file_path) print(f'Individual Files: {ind_link}') logger.info('Individual files at: %s', ind_link) if generate_zip: # Zip is in the parent directory of the timestamped folder zip_parent = os.path.dirname(temp_file_path) zip_link = get_scheduler_output_link(zip_parent) print(f'Zip File Location: {zip_link}') logger.info('Zip file at: %s', zip_link) print() logger.info("Operation completed.")