limit batchsize (#138)

* limit batchsize

* change error message

---------

Co-authored-by: ido <ido@ip-192-168-1-78.eu-central-1.compute.internal>

Author: idonaaman123

memphis/consumer.py

@@ -9,6 +9,7 @@
 
 
 class Consumer:
+    MAX_BATCH_SIZE = 5000
     def __init__(
         self,
         connection,
@@ -108,14 +109,17 @@ async def __consume_dls(self):
                         self.context,
                     )
         except Exception as e:
-            await self.dls_callback_func([], MemphisError(str(e)), self.context)
-            return
+            if self.dls_callback_func != None:
+                await self.dls_callback_func([], MemphisError(str(e)), self.context)
+                return
 
     async def fetch(self, batch_size: int = 10):
         """Fetch a batch of messages."""
         messages = []
         if self.connection.is_connection_active:
             try:
+                if batch_size > self.MAX_BATCH_SIZE:
+                    raise MemphisError(f"Batch size can not be greater than {self.MAX_BATCH_SIZE}")
                 self.batch_size = batch_size
                 if len(self.dls_messages) > 0:
                     if len(self.dls_messages) <= batch_size:

memphis/memphis.py

@@ -40,6 +40,7 @@
 
 
 class Memphis:
+    MAX_BATCH_SIZE = 5000
     def __init__(self):
         self.is_connection_active = False
         self.schema_updates_data = {}
@@ -499,6 +500,8 @@ async def consumer(
         try:
             if not self.is_connection_active:
                 raise MemphisError("Connection is dead")
+            if batch_size > self.MAX_BATCH_SIZE:
+                raise MemphisError(f"Batch size can not be greater than {self.MAX_BATCH_SIZE}")
             real_name = consumer_name.lower()
             if generate_random_suffix:
                 consumer_name = self.__generateRandomSuffix(consumer_name)
@@ -639,6 +642,8 @@ async def fetch_messages(
             consumer = None
             if not self.is_connection_active:
                 raise MemphisError("Cant fetch messages without being connected!")
+            if batch_size > self.MAX_BATCH_SIZE:
+                raise MemphisError(f"Batch size can not be greater than {self.MAX_BATCH_SIZE}")
             internal_station_name = get_internal_name(station_name)
             consumer_map_key = internal_station_name + "_" + consumer_name.lower()
             if consumer_map_key in self.consumers_map: