rabbitmq is now abstract channel

Author: tcnichol

backend/app/dependencies.py

@@ -3,6 +3,9 @@
 
 import boto3
 import pika
+import aio_pika
+from aio_pika.abc import AbstractChannel
+
 from app.config import settings
 from app.search.connect import connect_elasticsearch
 from minio import Minio
@@ -74,20 +77,27 @@ async def get_external_fs() -> AsyncGenerator[Minio, None]:
     yield file_system
 
 
-def get_rabbitmq() -> BlockingChannel:
+async def get_rabbitmq() -> AbstractChannel:
     """Client to connect to RabbitMQ for listeners/extractors interactions."""
-    credentials = pika.PlainCredentials(settings.RABBITMQ_USER, settings.RABBITMQ_PASS)
-    parameters = pika.ConnectionParameters(
-        settings.RABBITMQ_HOST, credentials=credentials
-    )
+    RABBITMQ_URL = f"amqp://{settings.RABBITMQ_USER}:{settings.RABBITMQ_PASS}@{settings.RABBITMQ_HOST}/"
+
     logger.debug("Connecting to rabbitmq at %s", settings.RABBITMQ_HOST)
-    connection = pika.BlockingConnection(parameters)
-    channel = connection.channel()
+    connection = await aio_pika.connect_robust(RABBITMQ_URL)
+    channel = await connection.channel()
+
     print(f"DEBUG: get_rabbitmq() called. Returning channel of type: {type(channel)}")
-    print(f"DEBUG: Channel object: {channel}")
     return channel
 
 
+# Keep the old function for compatibility if needed
+def get_blocking_rabbitmq() -> BlockingChannel:
+    """Legacy blocking RabbitMQ client (for extractors that need it)"""
+    credentials = pika.PlainCredentials(settings.RABBITMQ_USER, settings.RABBITMQ_PASS)
+    parameters = pika.ConnectionParameters(settings.RABBITMQ_HOST, credentials=credentials)
+    connection = pika.BlockingConnection(parameters)
+    return connection.channel()
+
+
 async def get_elasticsearchclient():
     es = await connect_elasticsearch()
     return es

backend/app/rabbitmq/listeners.py

@@ -41,33 +41,6 @@ async def create_reply_queue(channel: AbstractChannel):
 
     return queue.name
 
-# async def create_reply_queue(channel: BlockingChannel):
-#     # channel: BlockingChannel = dependencies.get_rabbitmq()
-#
-#     if (
-#         config_entry := await ConfigEntryDB.find_one({"key": "instance_id"})
-#     ) is not None:
-#         instance_id = config_entry.value
-#     else:
-#         # If no ID has been generated for this instance, generate a 10-digit alphanumeric identifier
-#         instance_id = "".join(
-#             random.choice(
-#                 string.ascii_uppercase + string.ascii_lowercase + string.digits
-#             )
-#             for _ in range(10)
-#         )
-#         config_entry = ConfigEntryDB(key="instance_id", value=instance_id)
-#         await config_entry.insert()
-#
-#     queue_name = "clowder.%s" % instance_id
-#     channel.exchange_declare(exchange="clowder", durable=True)
-#     result = channel.queue_declare(
-#         queue=queue_name, durable=True, exclusive=False, auto_delete=False
-#     )
-#     queue_name = result.method.queue
-#     channel.queue_bind(exchange="clowder", queue=queue_name)
-#     return queue_name
-
 
 async def submit_file_job(
     file_out: FileOut,
@@ -102,7 +75,6 @@ async def submit_file_job(
 
     # Use aio_pika publishing
     # Get the existing clowder exchange
-    exchange = await rabbitmq_client.get_exchange("clowder")
     reply_to = await create_reply_queue(rabbitmq_client)
     print("RABBITMQ_CLIENT: " + str(rabbitmq_client))
     await rabbitmq_client.default_exchange.publish(
@@ -122,7 +94,7 @@ async def submit_dataset_job(
     routing_key: str,
     parameters: dict,
     user: UserOut,
-    rabbitmq_client: BlockingChannel = Depends(dependencies.get_rabbitmq),
+    rabbitmq_client: AbstractChannel,
 ):
     # Create an entry in job history with unique ID
     job = EventListenerJobDB(
@@ -142,13 +114,14 @@ async def submit_dataset_job(
         job_id=str(job.id),
         parameters=parameters,
     )
-    reply_to = await create_reply_queue()
-    rabbitmq_client.basic_publish(
-        exchange="",
-        routing_key=routing_key,
-        body=json.dumps(msg_body.dict(), ensure_ascii=False),
-        properties=pika.BasicProperties(
-            content_type="application/json", delivery_mode=1, reply_to=reply_to
+    reply_to = await create_reply_queue(rabbitmq_client)
+    await rabbitmq_client.default_exchange.publish(
+        aio_pika.Message(
+            body=json.dumps(msg_body.dict(), ensure_ascii=False).encode('utf-8'),
+            content_type="application/json",
+            delivery_mode=aio_pika.DeliveryMode.PERSISTENT,
+            reply_to=reply_to,
         ),
+        routing_key=routing_key,
     )
     return str(job.id)

backend/app/routers/files.py

@@ -35,6 +35,8 @@
 from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
 from minio import Minio
 from pika.adapters.blocking_connection import BlockingChannel
+import aio_pika
+from aio_pika.abc import AbstractChannel
 
 router = APIRouter()
 security = HTTPBearer()
@@ -49,7 +51,7 @@ def default(self, obj):
 
 async def _resubmit_file_extractors(
     file: FileOut,
-    rabbitmq_client: BlockingChannel,
+    rabbitmq_client: AbstractChannel,
     user: UserOut,
     credentials: HTTPAuthorizationCredentials = Security(security),
 ):
@@ -94,7 +96,7 @@ async def add_file_entry(
     user: UserOut,
     fs: Minio,
     es: Elasticsearch,
-    rabbitmq_client: BlockingChannel,
+    rabbitmq_client: AbstractChannel,
     file: Optional[io.BytesIO] = None,
     content_type: Optional[str] = None,
     public: bool = False,
@@ -146,22 +148,28 @@ async def add_file_entry(
 
     # Publish a message when indexing is complete
 
+
+    # FIXED: Use aio_pika publishing
     message_body = {
         "event_type": "file_indexed",
         "file_data": json.loads(new_file.json()),
-        "user": json.loads(user.json()),# This handles ObjectID serialization
+        "user": json.loads(user.json()),
         "timestamp": datetime.now().isoformat()
     }
 
-    rabbitmq_client.basic_publish(
-        exchange='clowder',
-        routing_key='file_indexed_events',
-        body=json.dumps(message_body).encode('utf-8')
+    # Get the exchange first
+    exchange = await rabbitmq_client.get_exchange("clowder")
+
+    # Use aio_pika publish method
+    await exchange.publish(
+        aio_pika.Message(
+            body=json.dumps(message_body).encode('utf-8'),
+            content_type="application/json",
+            delivery_mode=aio_pika.DeliveryMode.PERSISTENT,
+        ),
+        routing_key="file_indexed_events",
     )
 
-    # TODO - timing issue here, check_feed_listeners needs to happen asynchronously.
-    time.sleep(1)
-
     # Submit file job to any qualifying feeds
     # await check_feed_listeners(
     #     es,
@@ -175,7 +183,7 @@ async def add_local_file_entry(
     new_file: FileDB,
     user: UserOut,
     es: Elasticsearch,
-    rabbitmq_client: BlockingChannel,
+    rabbitmq_client: AbstractChannel,
     content_type: Optional[str] = None,
 ):
     """Insert FileDB object into MongoDB (makes Clowder ID). Bytes are not stored in DB and versioning not supported
@@ -188,22 +196,27 @@ async def add_local_file_entry(
     # Add entry to the file index
     await index_file(es, FileOut(**new_file.dict()))
     # Publish a message when indexing is complete
+
     message_body = {
         "event_type": "file_indexed",
         "file_data": json.loads(new_file.json()),
-        "user": json.loads(user.json()),# This handles ObjectID serialization
+        "user": json.loads(user.json()),
         "timestamp": datetime.now().isoformat()
     }
 
-    rabbitmq_client.basic_publish(
-        exchange='clowder',
-        routing_key='file_indexed_events',
-        body=json.dumps(message_body).encode('utf-8')
+    # Get the exchange first
+    exchange = await rabbitmq_client.get_exchange("clowder")
+
+    # Use aio_pika publish method
+    await exchange.publish(
+        aio_pika.Message(
+            body=json.dumps(message_body).encode('utf-8'),
+            content_type="application/json",
+            delivery_mode=aio_pika.DeliveryMode.PERSISTENT,
+        ),
+        routing_key="file_indexed_events",
     )
 
-    # TODO - timing issue here, check_feed_listeners needs to happen asynchronously.
-    time.sleep(1)
-
     # Submit file job to any qualifying feeds
     # await check_feed_listeners(
     #     es,
@@ -255,7 +268,7 @@ async def update_file(
     file: UploadFile = File(...),
     es: Elasticsearch = Depends(dependencies.get_elasticsearchclient),
     credentials: HTTPAuthorizationCredentials = Security(security),
-    rabbitmq_client: BlockingChannel = Depends(dependencies.get_rabbitmq),
+    rabbitmq_client: AbstractChannel = Depends(dependencies.get_rabbitmq),
     allow: bool = Depends(FileAuthorization("uploader")),
 ):
     # Check all connection and abort if any one of them is not available
@@ -593,7 +606,7 @@ async def post_file_extract(
     parameters: dict = None,
     user=Depends(get_current_user),
     credentials: HTTPAuthorizationCredentials = Security(security),
-    rabbitmq_client: BlockingChannel = Depends(dependencies.get_rabbitmq),
+    rabbitmq_client: AbstractChannel = Depends(dependencies.get_rabbitmq),
     allow: bool = Depends(FileAuthorization("uploader")),
 ):
     if extractorName is None:
@@ -620,7 +633,7 @@ async def resubmit_file_extractions(
     file_id: str,
     user=Depends(get_current_user),
     credentials: HTTPAuthorizationCredentials = Security(security),
-    rabbitmq_client: BlockingChannel = Depends(dependencies.get_rabbitmq),
+    rabbitmq_client: AbstractChannel = Depends(dependencies.get_rabbitmq),
     allow: bool = Depends(FileAuthorization("editor")),
 ):
     """This route will check metadata. We get the extractors run from metadata from extractors.