init commit to add support for service-level max queue memory size

Author: mattnowzari

connectors/config.py

@@ -9,12 +9,12 @@
 from envyaml import EnvYAML
 
 from connectors.logger import logger
-
-DEFAULT_ELASTICSEARCH_MAX_RETRIES = 5
-DEFAULT_ELASTICSEARCH_RETRY_INTERVAL = 10
-
-DEFAULT_MAX_FILE_SIZE = 10485760  # 10MB
-
+from connectors.utils import (
+    DEFAULT_MAX_FILE_SIZE,
+    DEFAULT_QUEUE_MEM_SIZE,
+    DEFAULT_ELASTICSEARCH_MAX_RETRIES,
+    DEFAULT_ELASTICSEARCH_RETRY_INTERVAL
+)
 
 def load_config(config_file):
     logger.info(f"Loading config from {config_file}")
@@ -67,7 +67,7 @@ def _default_config():
             "verify_certs": True,
             "bulk": {
                 "queue_max_size": 1024,
-                "queue_max_mem_size": 25,
+                "queue_max_mem_size": 25 * 1024 * 1024, # 25 MB
                 "queue_refresh_interval": 1,
                 "queue_refresh_timeout": 600,
                 "display_every": 100,
@@ -108,6 +108,7 @@ def _default_config():
             "max_concurrent_access_control_syncs": 1,
             "max_concurrent_scheduling_tasks": 4,
             "max_file_download_size": DEFAULT_MAX_FILE_SIZE,
+            "queue_max_mem_size": DEFAULT_QUEUE_MEM_SIZE,
             "job_cleanup_interval": 300,
             "log_level": "INFO",
         },
@@ -220,19 +221,28 @@ class DataSourceFrameworkConfig:
     preventing them from requiring substantial changes to access new configs that may be added.
     """
 
-    def __init__(self, max_file_size):
+    def __init__(self, max_file_size, queue_max_mem_size):
         """
         Should not be called directly. Use the Builder.
         """
         self.max_file_size = max_file_size
+        self.max_queue_mem_size = queue_max_mem_size
 
     class Builder:
         def __init__(self):
             self.max_file_size = DEFAULT_MAX_FILE_SIZE
+            self.max_queue_mem_size = DEFAULT_QUEUE_MEM_SIZE
 
         def with_max_file_size(self, max_file_size):
             self.max_file_size = max_file_size
             return self
 
+        def with_max_queue_memory_size(self, queue_max_mem_size):
+            self.max_queue_mem_size = queue_max_mem_size
+            return self
+
         def build(self):
-            return DataSourceFrameworkConfig(self.max_file_size)
+            return DataSourceFrameworkConfig(
+                self.max_file_size,
+                self.max_queue_mem_size
+            )

connectors/es/sink.py

@@ -1023,7 +1023,7 @@ async def async_bulk(
 
         stream = MemQueue(
             maxsize=queue_size,
-            maxmemsize=queue_mem_size * 1024 * 1024,
+            maxmemsize=queue_mem_size,
             refresh_timeout=mem_queue_refresh_timeout,
             refresh_interval=mem_queue_refresh_interval,
         )

connectors/sources/box.py

@@ -475,7 +475,21 @@ async def _consumer(self):
             else:
                 yield item
 
-    async def get_docs(self, filtering=None):
+    def _update_queue_max_mem_size(self):
+        """
+        Helper function to update the queue's max memory size
+        if the configured value exceeds the constant defined in the data source
+        """
+        max_mem_size_from_config = self.framework_config.max_queue_mem_size
+
+        if max_mem_size_from_config > QUEUE_MEM_SIZE:
+            self.queue.update_maxmemsize(max_mem_size_from_config)
+            logger.debug(f"MemQueue max memory size updated from {QUEUE_MEM_SIZE} to {max_mem_size_from_config}")
+
+    async def get_docs(self, filtering=None): # type: ignore
+        # update queue max memory size if needed
+        self._update_queue_max_mem_size()
+
         seen_ids = set()
         root_folder = "0"
         if self.is_enterprise == BOX_ENTERPRISE:

connectors/sources/confluence.py

@@ -767,7 +767,7 @@ async def get_user(self):
                     user=user
                 )
 
-    async def get_access_control(self):
+    async def get_access_control(self): # type: ignore
         """Get access control documents for active Atlassian users.
 
         This method fetches access control documents for active Atlassian users when document level security (DLS)
@@ -1277,14 +1277,28 @@ async def _consumer(self):
             else:
                 yield item
 
-    async def get_docs(self, filtering=None):
+    def _update_queue_max_mem_size(self):
+        """
+        Helper function to update the queue's max memory size
+        if the configured value exceeds the constant defined in the data source
+        """
+        max_mem_size_from_config = self.framework_config.max_queue_mem_size
+
+        if max_mem_size_from_config > QUEUE_MEM_SIZE:
+            self.queue.update_maxmemsize(max_mem_size_from_config)
+            logger.debug(f"MemQueue max memory size updated from {QUEUE_MEM_SIZE} to {max_mem_size_from_config}")
+
+    async def get_docs(self, filtering=None): # type: ignore
         """Executes the logic to fetch Confluence content in async manner.
 
         Args:
             filtering (Filtering): Object of Class Filtering
         Yields:
             dictionary: dictionary containing meta-data of the content.
         """
+        # update queue max memory size if needed
+        self._update_queue_max_mem_size()
+
         self._logger.info("Successfully connected to Confluence")
         if filtering and filtering.has_advanced_rules():
             advanced_rules = filtering.get_advanced_rules()

connectors/sources/jira.py

@@ -729,7 +729,7 @@ async def _issue_access_control(self, issue_key, project):
                 return project_access_controls
         return list(access_control)
 
-    async def get_access_control(self):
+    async def get_access_control(self): # type: ignore
         """Get access control documents for active Atlassian users.
 
         This method fetches access control documents for active Atlassian users when document level security (DLS)
@@ -1017,7 +1017,18 @@ async def _consumer(self):
             else:
                 yield item
 
-    async def get_docs(self, filtering=None):
+    def _update_queue_max_mem_size(self):
+        """
+        Helper function to update the queue's max memory size
+        if the configured value exceeds the constant defined in the data source
+        """
+        max_mem_size_from_config = self.framework_config.max_queue_mem_size
+
+        if max_mem_size_from_config > QUEUE_MEM_SIZE:
+            self.queue.update_maxmemsize(max_mem_size_from_config)
+            logger.debug(f"MemQueue max memory size updated from {QUEUE_MEM_SIZE} to {max_mem_size_from_config}")
+
+    async def get_docs(self, filtering=None): # type: ignore
         """Executes the logic to fetch jira objects in async manner
 
         Args:
@@ -1026,6 +1037,9 @@ async def get_docs(self, filtering=None):
         Yields:
             dictionary: dictionary containing meta-data of the files.
         """
+        # update queue max memory size if needed
+        self._update_queue_max_mem_size()
+
         self.custom_fields = await anext(self.jira_client.get_jira_fields())
 
         if filtering and filtering.has_advanced_rules():

connectors/sources/microsoft_teams.py

@@ -452,7 +452,7 @@ async def _get(self, absolute_url, use_token=True):
                     url=absolute_url,
                 ) as resp:
                     yield resp
-        except aiohttp.client_exceptions.ClientOSError:
+        except aiohttp.client_exceptions.ClientOSError: # type: ignore
             self._logger.error(
                 "Graph API dropped the connection. It might indicate, that connector makes too many requests - decrease concurrency settings, otherwise Graph API can block this app."
             )
@@ -1267,7 +1267,18 @@ async def calendars_producer(self, user):
                 )
         await self.queue.put(EndSignal.CALENDAR_TASK_FINISHED)
 
-    async def get_docs(self, filtering=None):
+    def _update_queue_max_mem_size(self):
+        """
+        Helper function to update the queue's max memory size
+        if the configured value exceeds the constant defined in the data source
+        """
+        max_mem_size_from_config = self.framework_config.max_queue_mem_size
+
+        if max_mem_size_from_config > QUEUE_MEM_SIZE:
+            self.queue.update_maxmemsize(max_mem_size_from_config)
+            logger.debug(f"MemQueue max memory size updated from {QUEUE_MEM_SIZE} to {max_mem_size_from_config}")
+
+    async def get_docs(self, filtering=None): # type: ignore
         """Executes the logic to fetch Microsoft Teams objects in async manner
 
         Args:
@@ -1276,6 +1287,9 @@ async def get_docs(self, filtering=None):
         Yields:
             dictionary: dictionary containing meta-data of the files.
         """
+        # update queue max memory size if needed
+        self._update_queue_max_mem_size()
+
         async for chats in self.client.get_user_chats():
             for chat in chats:
                 await self.fetchers.put(partial(self.user_chat_producer, chat))

connectors/sources/outlook.py

@@ -50,8 +50,6 @@
 RETRIES = 3
 RETRY_INTERVAL = 2
 
-QUEUE_MEM_SIZE = 5 * 1024 * 1024  # Size in Megabytes
-
 OUTLOOK_SERVER = "outlook_server"
 OUTLOOK_CLOUD = "outlook_cloud"
 API_SCOPE = "https://graph.microsoft.com/.default"
@@ -821,7 +819,7 @@ def _dls_enabled(self):
 
         return self.configuration["use_document_level_security"]
 
-    async def get_access_control(self):
+    async def get_access_control(self): # type: ignore
         if not self._dls_enabled():
             self._logger.warning("DLS is not enabled. Skipping")
             return
@@ -1064,7 +1062,7 @@ async def ping(self):
         await self.client.ping()
         self._logger.info("Successfully connected to Outlook")
 
-    async def get_docs(self, filtering=None):
+    async def get_docs(self, filtering=None): # type: ignore
         """Executes the logic to fetch outlook objects in async manner
 
         Args:

connectors/sources/servicenow.py

@@ -384,7 +384,7 @@ async def validate(self, advanced_rules):
 
     async def _remote_validation(self, advanced_rules):
         try:
-            ServiceNowAdvancedRulesValidator.SCHEMA(advanced_rules)
+            ServiceNowAdvancedRulesValidator.SCHEMA(advanced_rules) # type: ignore
         except fastjsonschema.JsonSchemaValueException as e:
             return SyncRuleValidationResult(
                 rule_id=SyncRuleValidationResult.ADVANCED_RULES,
@@ -567,7 +567,7 @@ async def _fetch_users_by_roles(self, role):
         ):
             yield user
 
-    async def get_access_control(self):
+    async def get_access_control(self): # type: ignore
         if not self._dls_enabled():
             self._logger.warning("DLS is not enabled. Skipping")
             return
@@ -858,7 +858,18 @@ async def _consumer(self):
             else:
                 yield item
 
-    async def get_docs(self, filtering=None):
+    def _update_queue_max_mem_size(self):
+        """
+        Helper function to update the queue's max memory size
+        if the configured value exceeds the constant defined in the data source
+        """
+        max_mem_size_from_config = self.framework_config.max_queue_mem_size
+
+        if max_mem_size_from_config > QUEUE_MEM_SIZE:
+            self.queue.update_maxmemsize(max_mem_size_from_config)
+            logger.debug(f"MemQueue max memory size updated from {QUEUE_MEM_SIZE} to {max_mem_size_from_config}")
+
+    async def get_docs(self, filtering=None): # type: ignore
         """Get documents from ServiceNow.
 
         Args:
@@ -867,6 +878,8 @@ async def get_docs(self, filtering=None):
         Yields:
             dict: Documents from ServiceNow.
         """
+        # update queue max memory size if needed
+        self._update_queue_max_mem_size()
 
         self._logger.info("Fetching ServiceNow data")
         if filtering and filtering.has_advanced_rules():

connectors/sync_job_runner.py

@@ -218,9 +218,14 @@ async def execute(self):
                 await self.data_provider.close()
 
     def _data_source_framework_config(self):
-        builder = DataSourceFrameworkConfig.Builder().with_max_file_size(
+        builder = DataSourceFrameworkConfig.Builder()
+
+        builder.with_max_file_size(
             self.service_config.get("max_file_download_size")
         )
+        builder.with_max_queue_memory_size(
+            self.service_config.get("queue_max_mem_size")
+        )
         return builder.build()
 
     async def _execute_access_control_sync_job(self, job_type, bulk_options):

connectors/utils.py

@@ -36,10 +36,13 @@
 DEFAULT_CHUNK_SIZE = 500
 DEFAULT_QUEUE_SIZE = 1024
 DEFAULT_DISPLAY_EVERY = 100
-DEFAULT_QUEUE_MEM_SIZE = 5
+DEFAULT_QUEUE_MEM_SIZE = 5242880 # 5 MB
 DEFAULT_CHUNK_MEM_SIZE = 25
 DEFAULT_MAX_CONCURRENCY = 5
 DEFAULT_CONCURRENT_DOWNLOADS = 10
+DEFAULT_MAX_FILE_SIZE = 10485760 # 10 MB
+DEFAULT_ELASTICSEARCH_MAX_RETRIES = 5
+DEFAULT_ELASTICSEARCH_RETRY_INTERVAL = 10
 
 # Regular expression pattern to match a basic email format (no whitespace, valid domain)
 EMAIL_REGEX_PATTERN = r"^\S+@\S+\.\S+$"
@@ -297,9 +300,16 @@ def __init__(
         self._current_memsize = 0
         self.refresh_timeout = refresh_timeout
 
+        logger.debug(f"MemQueue created with maxsize={maxsize}, maxmemsize={maxmemsize}, "
+                     f"refresh_interval={refresh_interval}, refresh_timeout={refresh_timeout}")
+
     def qmemsize(self):
         return self._current_memsize
 
+    def update_maxmemsize(self, new_maxmemsize):
+        current_maxmemsize = self.maxmemsize
+        self.maxmemsize = new_maxmemsize
+
     def _get(self):
         item_size, item = self._queue.popleft()  # pyright: ignore
         self._current_memsize -= item_size