Merge branch 'feature/DBC22-5164-2' of https://github.com/bcgov/DriveBC.ca into feature/DBC22-5164-2

Author: bcgov-brwang

src/backend/apps/consumer/processor.py

@@ -64,10 +64,6 @@
 EXCHANGE_NAME = os.getenv("RABBITMQ_EXCHANGE_NAME")
 CAMERA_CACHE_REFRESH_SECONDS = int(os.getenv("CAMERA_CACHE_REFRESH_SECONDS", "60"))
 
-RABBITMQ_HEARTBEAT = int(os.getenv("RABBITMQ_HEARTBEAT", "60"))
-RABBITMQ_TIMEOUT = int(os.getenv("RABBITMQ_TIMEOUT", "30")) 
-RABBITMQ_RECONNECT_INTERVAL = int(os.getenv("RABBITMQ_RECONNECT_INTERVAL", "5"))
-
 
 #boto3.set_stream_logger('botocore', logging.DEBUG)
 
@@ -100,141 +96,28 @@
 last_camera_refresh = {}
 image_invalid = False
 
-last_activity = 0.0  # Track last message processing time
-health_check_interval = 60  # seconds
-max_idle_time = 300  # 5 minutes before warning
-
-tz_pst = 'America/Vancouver'
-
-async def consume_from(rb_url: str):
+async def run_consumer():
     """
-    Consumer instance bound to a single RabbitMQ URL.
+    Long-running RabbitMQ consumer that processes image messages and watermarks them.
+    Shuts down cleanly when stop_event is set.
     """
-    try:
-        db_data = await parse_rows(rb_url)
-    except Exception as exc:
-        logger.error("Failed to initialize consumer for %s: %s", rb_url, exc)
-        return
+    rb_url = os.getenv("RABBITMQ_URL")
+    if not rb_url:
+        raise RuntimeError("RABBITMQ_URL environment variable is not set.")
 
-async def on_reconnect(conn):
-    logger.info("RabbitMQ connection re-established")
-    global last_activity
-    last_activity = time.time()
-   
-async def on_close(conn, exc=None):
-    logger.warning(f"RabbitMQ connection closed: {exc}")
-
-async def on_channel_close(ch, exc=None):
-    logger.warning(f"RabbitMQ channel closed: {exc}")
-
-
-async def on_message(message: aio_pika.IncomingMessage):
-    """Process incoming RabbitMQ message."""
-    if stop_event.is_set():
-        logger.info("Stop requested. Dropping message.")
-    
-    async with message.process(ignore_processed=True):
-        filename = message.headers.get("filename", "unknown.jpg")
-        camera_id = filename.split("_")[0].split(".")[0]
-        timestamp_utc = message.headers.get("timestamp", "unknown")
-        camera_status = calculate_camera_status(camera_id, timestamp_utc)
-        
-        try:
-            timestamp_local = await sync_to_async(safe_db_call)(
-                generate_local_timestamp, db_data, camera_id, timestamp_utc
-            )
-            try:
-                await asyncio.wait_for(
-                    handle_image_message(camera_id, message.body, timestamp_local, camera_status),
-                    timeout=120
-                )
-            except asyncio.TimeoutError:
-                logger.error(f"Processing timeout for camera {camera_id}")
-            
-            logger.info("Processed message for camera %s.", camera_id)
-        except Exception as e:
-            logger.exception("Failed processing message (camera %s): %s", camera_id, e)
-            await asyncio.sleep(2)
-
-async def wait_for_timeout_or_signal(timeout_seconds: int):
-    """Wait for timeout or stop signal."""
-    try:
-        await asyncio.wait_for(stop_event.wait(), timeout=timeout_seconds)
-        return True  # Stop signal received
-    except asyncio.TimeoutError:
-        return False  # Timeout reached
+    rows = await sync_to_async(get_all_from_db)()
 
-async def run_consumer():
-    """Long-running RabbitMQ consumer with automatic reconnection."""
-    
-    while not stop_event.is_set():
-        connection = None
-        channel = None
-        queue = None
-        
-        try:
-            # 1. Connect to RabbitMQ (auto-reconnects on RabbitMQ restart)
-            connection = await aio_pika.connect_robust(
-                RABBITMQ_URL,
-                heartbeat=RABBITMQ_HEARTBEAT,
-                timeout=RABBITMQ_TIMEOUT,
-                reconnect_interval=RABBITMQ_RECONNECT_INTERVAL,
-                fail_fast=False,
-            )
-            logger.info("RabbitMQ connection established.")
-            connection.reconnect_callbacks.add(on_reconnect)
-            connection.close_callbacks.add(on_close)
-            
-            # 2. Setup channel and queue
-            channel = await connection.channel()
-            logger.info("RabbitMQ channel created.")
-            channel.close_callbacks.add(on_channel_close)
-            
-            exchange = await channel.declare_exchange(
-                EXCHANGE_NAME,
-                type=aio_pika.ExchangeType.FANOUT,
-                durable=True,
-            )
-            logger.info(f"RabbitMQ exchange '{EXCHANGE_NAME}' declared.")
-            
-            # # Delete existing queue if wrong type, then recreate as quorum
-            # try:
-            #     await channel.queue_delete(QUEUE_NAME)
-            #     logger.info(f"Deleted existing queue '{QUEUE_NAME}' for recreation")
-            # except Exception:
-            #     pass
-            
-            queue = await channel.declare_queue(
-                QUEUE_NAME,
-                durable=True,
-                exclusive=False,
-                auto_delete=False,
-                arguments={
-                    "x-max-length-bytes": QUEUE_MAX_BYTES,
-                    "x-queue-type": "quorum"
-                },
-            )
-            logger.info(f"RabbitMQ queue '{QUEUE_NAME}' declared.")
-            
-            await queue.bind(exchange)
-            logger.info(f"RabbitMQ queue '{QUEUE_NAME}' bound to exchange '{EXCHANGE_NAME}'.")
-            
-            logger.info("Starting message consumption...")
-            
-            # 3. Consume messages using iterator (blocks waiting for messages)
-            async with queue.iterator() as queue_iter:
-                async for message in queue_iter:
-                    if stop_event.is_set():
-                        logger.info("Stop requested. Breaking consume loop.")
-                        break
-                    
-                    try:
-                        await process_message(message)
+    global db_data, last_camera_refresh
+    db_data = process_camera_rows(rows)
+    last_camera_refresh = time.time()
+    if not db_data:
+        logger.error("No camera data available for watermarking. Consumer exiting.")
+        return
 
-    logger.info("Starting RabbitMQ consumer for %s", rb_url)
+    logger.info("Starting RabbitMQ consumer.")
 
-    connection = None
-    channel = None
+    connection: Optional[aio_pika.RobustConnection] = None
+    channel: Optional[aio_pika.RobustChannel] = None
     queue = None
 
     try:
@@ -331,42 +214,7 @@ async def run_consumer():
         if connection:
             try: 
                 await connection.close()
-        
-        if not stop_event.is_set():
-            logger.info(f"Waiting {RABBITMQ_RECONNECT_INTERVAL}s before reconnecting...")
-            await asyncio.sleep(RABBITMQ_RECONNECT_INTERVAL)
-    
-    logger.info("Consumer stopped.")
-    
-async def process_message(message: aio_pika.IncomingMessage):
-    """Process a single message from the queue."""
-    async with message.process(ignore_processed=True):
-        filename = message.headers.get("filename", "unknown.jpg")
-        camera_id = filename.split("_")[0].split(".")[0]
-        timestamp_utc = message.headers.get("timestamp", "unknown")
-        camera_status = await sync_to_async(calculate_camera_status)(camera_id, timestamp_utc)
-        
-        timestamp_local = await sync_to_async(safe_db_call)(
-            generate_local_timestamp, db_data, camera_id, timestamp_utc
-        )
-        
-        try:
-            await asyncio.wait_for(
-                handle_image_message(camera_id, message.body, timestamp_local, camera_status),
-                timeout=120
-            )
-        except asyncio.TimeoutError:
-            logger.error(f"Processing timeout for camera {camera_id}")
-
-        logger.info("Processed message for camera %s.", camera_id)
-
-def safe_db_call(func, *args, **kwargs):
-    """Call a sync ORM function safely after dropping dead connections"""
-    close_old_connections()
-    # This forces Django to open a fresh connection if needed
-    connection.ensure_connection()
-    return func(*args, **kwargs)
-            except Exception as e: 
+            except Exception as e:
                 logger.warning("Error closing connection: %s", e)
 
         logger.info("Consumer stopped for %s.", rb_url)