import json import logging import uuid from contextlib import contextmanager from dataclasses import dataclass from typing import Any import mysql.connector from mysql.connector import Error logger = logging.getLogger(__name__) @dataclass class DatabaseConfig: """MySQL database configuration""" host: str database: str username: str password: str port: int = 3306 charset: str = "utf8mb4" pool_size: int = 10 class MySQLRequestStore: """MySQL-based request storage""" def __init__(self, config: DatabaseConfig): self.config = config self.connection_pool = None self.init_connection_pool() self.init_database() def init_connection_pool(self): """Initialize MySQL connection pool""" try: self.connection_pool = mysql.connector.pooling.MySQLConnectionPool( pool_name="sftp_pipeline_pool", pool_size=self.config.pool_size, host=self.config.host, port=self.config.port, database=self.config.database, user=self.config.username, password=self.config.password, charset=self.config.charset, ) logger.info( f"MySQL connection pool initialized: {self.config.host}:{self.config.port}" ) except Error as e: logger.error(f"Error creating MySQL connection pool: {e}") raise @contextmanager def get_connection(self): """Get connection from pool""" connection = None try: connection = self.connection_pool.get_connection() yield connection except Error as e: logger.error(f"Database connection error: {e}") if connection: connection.rollback() raise finally: if connection: connection.close() def init_database(self): """Initialize database tables""" with self.get_connection() as conn: cursor = conn.cursor() try: cursor.execute( """ CREATE TABLE IF NOT EXISTS connectivity_requests ( consultation_id VARCHAR(100) PRIMARY KEY, reference_id VARCHAR(100) NOT NULL, message_id VARCHAR(100) NOT NULL, original_filename VARCHAR(255) NOT NULL, original_request JSON NOT NULL, rx_request JSON NOT NULL, status ENUM( 'pending', 'sent_to_1rx', 'sent_to_service_s1', 'received_1rx_response', 'processing_complete', 'failed_1rx', 'failed_processing', 'completed' ) DEFAULT 'pending', created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, sftp_path VARCHAR(500) NULL, retry_count INT DEFAULT 0, last_error TEXT NULL, rx_response JSON NULL, egress_api_status VARCHAR(32) NULL, egress_tracking_id VARCHAR(64) NULL, egress_api_error TEXT NULL, egress_api_at TIMESTAMP NULL, tracking_id VARCHAR(64) NULL, is_test_data TINYINT(1) NULL, INDEX idx_reference_id (reference_id), INDEX idx_message_id (message_id), INDEX idx_status (status), INDEX idx_created_at (created_at) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_general_ci """ ) # Existing installs: add the Amazon-egress columns in place. # MySQL has no ADD COLUMN IF NOT EXISTS, so a duplicate # column (errno 1060) means the migration already ran. for clause in ( "ADD COLUMN rx_response JSON NULL", "ADD COLUMN egress_api_status VARCHAR(32) NULL", "ADD COLUMN egress_tracking_id VARCHAR(64) NULL", "ADD COLUMN egress_api_error TEXT NULL", "ADD COLUMN egress_api_at TIMESTAMP NULL", "ADD COLUMN tracking_id VARCHAR(64) NULL", "ADD COLUMN is_test_data TINYINT(1) NULL", ): try: cursor.execute(f"ALTER TABLE connectivity_requests {clause}") except Error as e: if getattr(e, "errno", None) != 1060: raise # The real-order API ramp counts acknowledged real orders on # every webhook; without this index that is a full scan of # the JSON-heavy table. Duplicate key name is errno 1061. try: cursor.execute( "ALTER TABLE connectivity_requests " "ADD INDEX idx_egress_ramp (is_test_data, egress_api_status)" ) except Error as e: if getattr(e, "errno", None) != 1061: raise conn.commit() logger.info("Database tables initialized successfully") except Error as e: logger.error(f"Error initializing database: {e}") conn.rollback() raise finally: cursor.close() def store_request(self, request_data: dict[str, Any]) -> str: """Store consultation request context""" consultation_id = request_data["consultation_id"] with self.get_connection() as conn: cursor = conn.cursor() try: cursor.execute( """ INSERT IGNORE INTO connectivity_requests ( consultation_id, reference_id, message_id, original_filename, original_request, rx_request, sftp_path, tracking_id, is_test_data ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s) """, ( consultation_id, request_data["reference_id"], request_data["message_id"], request_data["original_filename"], json.dumps(request_data["original_request"]), json.dumps(request_data["rx_request"]), request_data.get("sftp_path"), request_data.get("tracking_id"), request_data.get("is_test_data"), ), ) is_new = cursor.rowcount > 0 if is_new: logger.info( f"Stored new request with consultation_id: {consultation_id}" ) else: logger.info( f"Request with consultation_id {consultation_id} already exists" ) conn.commit() return is_new except Error as e: logger.error(f"Error storing request: {e}") conn.rollback() raise finally: cursor.close() def reclaim_failed_request(self, request_data: dict[str, Any]) -> bool: """Atomically reclaim a request whose 1Rx call failed so a retry from Amazon (same order, new trackingId) can be reprocessed. Only 'failed_1rx' rows may be reclaimed — any later status means the request already reached 1Rx and the retry is a true duplicate. The stored payload is refreshed so responses echo the retry's identifiers. Returns True if the row was claimed by this call. """ consultation_id = request_data["consultation_id"] with self.get_connection() as conn: cursor = conn.cursor() try: cursor.execute( """ UPDATE connectivity_requests SET status = 'pending', reference_id = %s, message_id = %s, original_filename = %s, original_request = %s, rx_request = %s, tracking_id = %s, is_test_data = %s, last_error = NULL WHERE consultation_id = %s AND status = 'failed_1rx' """, ( request_data["reference_id"], request_data["message_id"], request_data["original_filename"], json.dumps(request_data["original_request"]), json.dumps(request_data["rx_request"]), request_data.get("tracking_id"), request_data.get("is_test_data"), consultation_id, ), ) reclaimed = cursor.rowcount > 0 conn.commit() if reclaimed: logger.info( f"Reclaimed failed request for retry, consultation_id: {consultation_id}" ) return reclaimed except Error as e: logger.error(f"Error reclaiming failed request: {e}") conn.rollback() return False finally: cursor.close() def get_request_by_consultation_id( self, consultation_id: str ) -> dict[str, Any] | None: """Get request by consultation ID (primary key)""" with self.get_connection() as conn: cursor = conn.cursor(dictionary=True) try: cursor.execute( """ SELECT * FROM connectivity_requests WHERE consultation_id = %s """, (consultation_id,), ) result = cursor.fetchone() if result: result["original_request"] = json.loads(result["original_request"]) result["rx_request"] = json.loads(result["rx_request"]) if result.get("rx_response"): result["rx_response"] = json.loads(result["rx_response"]) return result except Error as e: logger.error(f"Error getting request: {e}") return None finally: cursor.close() def save_rx_response(self, consultation_id: str, rx_response: dict[str, Any]): """Persist the 1Rx webhook result so a failed Amazon API delivery can be rebuilt and resent later without replaying the webhook.""" with self.get_connection() as conn: cursor = conn.cursor() try: cursor.execute( """ UPDATE connectivity_requests SET rx_response = %s WHERE consultation_id = %s """, (json.dumps(rx_response), consultation_id), ) conn.commit() except Error as e: logger.error(f"Error saving 1Rx response: {e}") conn.rollback() finally: cursor.close() def count_real_orders_delivered_via_api(self) -> int: """Real (non-test) orders whose response Amazon acknowledged over the Shuttle API. Drives the EGRESS_API_REAL_ORDER_LIMIT ramp; legacy rows with no is_test_data flag are real orders.""" with self.get_connection() as conn: cursor = conn.cursor() try: cursor.execute( """ SELECT COUNT(*) FROM connectivity_requests WHERE (is_test_data IS NULL OR is_test_data = 0) AND egress_api_status = 'acknowledged' """ ) row = cursor.fetchone() return int(row[0]) if row else 0 finally: cursor.close() def record_egress_api_result( self, consultation_id: str, tracking_id: str, status: str, error_message: str = None, ): """Record the outcome of one Amazon Consume API delivery attempt.""" with self.get_connection() as conn: cursor = conn.cursor() try: cursor.execute( """ UPDATE connectivity_requests SET egress_api_status = %s, egress_tracking_id = %s, egress_api_error = %s, egress_api_at = CURRENT_TIMESTAMP WHERE consultation_id = %s """, (status, tracking_id, error_message, consultation_id), ) conn.commit() except Error as e: logger.error(f"Error recording egress API result: {e}") conn.rollback() finally: cursor.close() def update_request_status( self, consultation_id: str, status: str, error_message: str = None ): """Update request status using consultation_id""" with self.get_connection() as conn: cursor = conn.cursor() try: if error_message: cursor.execute( """ UPDATE connectivity_requests SET status = %s, last_error = %s, retry_count = retry_count + 1 WHERE consultation_id = %s """, (status, error_message, consultation_id), ) else: cursor.execute( """ UPDATE connectivity_requests SET status = %s WHERE consultation_id = %s """, (status, consultation_id), ) conn.commit() except Error as e: logger.error(f"Error updating request status: {e}") conn.rollback() raise finally: cursor.close()