Merge pull request #10 from Aignosi/SIENTIAPDE-1166-alterar-orquestrador-para-criar-collections-com-ttl-no-mongo]

feat: add TTL index functionality to MongoDB activities and update requirements

orchestrator/activities/activities.py

@@ -55,6 +55,7 @@ class Activities(  # Couchbase,
         MongoDB.__init__(self,
                          connection_string=mongodb_config['connection_string'],
                          database_name=mongodb_config['database_name'],
+                         ttl_index_seconds=mongodb_config['ttl_index_seconds'],
                          logger=logger,
                          notification_handler=notification_handler)
 

orchestrator/activities/mongo_db.py

@@ -41,7 +41,7 @@ def clear_mongo_id(docs: list) -> list:
 
 
 class MongoDB(BaseActivity):
-    def __init__(self, connection_string: str, database_name: str,
+    def __init__(self, connection_string: str, database_name: str, ttl_index_seconds: int,
                  logger: Logger,
                  notification_handler: NotificationHandler):
         self.connection_string = connection_string
@@ -53,6 +53,8 @@ class MongoDB(BaseActivity):
 
         self.database = self.client[self.database_name]
 
+        self.ttl_index_seconds = ttl_index_seconds
+
         # Initialize MongoDB client here (omitted for brevity)
         logger.info("MongoDB connection initialized")
 
@@ -295,3 +297,58 @@ class MongoDB(BaseActivity):
             )
             self.error(trace, metadata=metadata)
             raise e
+
+    @activity.defn(name="create_collection_with_ttl_index")
+    async def create_collection_with_ttl_index(self, input_data: dict[str, Any]) -> None:
+        """
+        Create a collection with a TTL index.
+        input_data:
+            - collection_name (str): The name of the collection to create.
+            - ttl_index (str): The name of the TTL index to create.
+        """
+        pipelines = input_data.get("pipelines", {})
+        metadata = input_data.get("metadata", {})
+
+        self.info(
+            f"Creating collection with TTL index for pipelines: {list(pipelines.keys())}",
+            metadata=metadata)
+
+        collection_names = self.database.list_collection_names()
+
+        for pipeline_name, pipeline_config in pipelines.items():
+            collection = pipeline_config["topic"]
+
+            try:
+                # Check if collection exists
+                if collection not in collection_names:
+                    self.database.create_collection(collection)
+
+                collection = self.database[collection]
+                # Check if TTL index exists
+                existing_indexes = collection.list_indexes()
+                ttl_index_exists = False
+                for index in existing_indexes:
+                    if "inserted_at" in index["key"] and index.get("expireAfterSeconds") is not None:
+                        ttl_index_exists = True
+                        break
+
+                # Create TTL index if it doesn't exist
+                if not ttl_index_exists:
+                    collection.create_index(
+                        "inserted_at",
+                        expireAfterSeconds=self.ttl_index_seconds,
+                        background=True
+                    )
+
+            except Exception as e:
+                trace = traceback.format_exc()
+                self.send_notification(
+                    metadata=metadata,
+                    notification_id="MONGODB_CREATE_COLLECTION_ERROR",
+                    message=f"Failed to create collection {collection} with TTL index: {e}",
+                    block="create_collection_with_ttl_index",
+                    level=NotificationLevel.ERROR,
+                    attachment_content=trace
+                )
+                self.error(trace, metadata=metadata)
+                raise e

orchestrator/utils/connectors_config.py

@@ -18,7 +18,8 @@ def build_mongodb_config():
     connection_string = f'mongodb://{username}:{password}@{uri}'
     return {
         'connection_string': connection_string,
-        'database_name': getenv('MONGODB_DATABASE_NAME', 'sientia')
+        'database_name': getenv('MONGODB_DATABASE_NAME', 'sientia'),
+        'ttl_index_seconds': int(getenv('MONGODB_TTL_INDEX_HOURS', '1')) * 3600
     }
 
 

orchestrator/worker/worker.py

@@ -75,6 +75,7 @@ async def main():
                 activities.update_pipelines_timestamps,
                 activities.create_pipelines_timestamps,
                 activities.delete_pipelines_timestamps,
+                activities.create_collection_with_ttl_index,
                 # Temporal
                 activities.create_schedules,
                 activities.update_schedules,

orchestrator/workflows/orchestrator.py

@@ -147,9 +147,20 @@ class Orchestrator:
             start_to_close_timeout=timedelta(seconds=60)
         )
 
+        create_collection_with_ttl_index_handler = workflow.start_activity_method(
+            Activities.create_collection_with_ttl_index,
+            {
+                **metadata,
+                'pipelines': schedules_config['scouter']
+            },
+            retry_policy=retry_policy,
+            start_to_close_timeout=timedelta(seconds=60)
+        )
+
         schedule_actions = await schedule_actions_handler
         slot_actions = await slot_actions_handler
         await normalize_schedules_handler
+        await create_collection_with_ttl_index_handler
 
         slot_deletion_report_handler = workflow.start_activity_method(
             Activities.delete_slots,

requirements.txt

@@ -4,4 +4,4 @@ sqlalchemy
 redis
 couchbase
 pymongo
-git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.3.3
+git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.3.4

tests/orchestrator/activities/test_activities.py

@@ -18,7 +18,8 @@ def test___init__(mock_formatters_init, mock_slot_manager_init,
 
     mongo_db_config = {
         'connection_string': 'mongodb://localhost:27017',
-        'database_name': 'test_db'
+        'database_name': 'test_db',
+        'ttl_index_seconds': 3600
     }
 
     redis_config = {
@@ -65,6 +66,7 @@ def test___init__(mock_formatters_init, mock_slot_manager_init,
         ANY,
         connection_string='mongodb://localhost:27017',
         database_name='test_db',
+        ttl_index_seconds=3600,
         logger=logger,
         notification_handler=notification_handler
     )

tests/orchestrator/activities/test_mongo_db.py

@@ -55,6 +55,7 @@ def mongo_db(mongo_mock):
         MongoDB(
             connection_string="mongodb://localhost:27017",
             database_name="test_db",
+            ttl_index_seconds=3600,
             logger=MagicMock(),
             notification_handler=MagicMock()
         )
@@ -68,6 +69,7 @@ def test___init__(mongo_mock):
     mongo_db = MongoDB(
         connection_string="mongodb://localhost:27017",
         database_name="test_db",
+        ttl_index_seconds=3600,
         logger=MagicMock(),
         notification_handler=MagicMock()
     )
@@ -398,3 +400,121 @@ async def test_delete_pipelines_timestamps_failure(datetime_mock, mongo_db):
 
     else:
         assert False, "Expected an exception to be raised"
+
+
+@mark.asyncio
+async def test_create_collection_with_ttl_index_success(mongo_db):
+    input_data = {
+        **metadata,
+        "pipelines": {
+            "scouter-pipeline": {
+                "topic": "raw_scouter_pipeline"
+            },
+            "scouter-pipeline-2": {
+                "topic": "raw_scouter_pipeline_2"
+            },
+            "scouter-pipeline-3": {
+                "topic": "raw_scouter_pipeline_3"
+            }
+        }
+    }
+
+    mongo_db.database.list_collection_names.return_value = [
+        "raw_scouter_pipeline_2",
+        "raw_scouter_pipeline_3"
+    ]
+
+    collection_1 = MagicMock(
+        list_indexes=MagicMock(
+            return_value=[
+                {
+                    "key": "asdad",
+                }
+            ]
+        )
+    )
+
+    collection_2 = MagicMock(
+        list_indexes=MagicMock(
+            return_value=[
+                {
+                    "key": "inserted_at",
+                    "expireAfterSeconds": None
+                }
+            ]
+        )
+    )
+    collection_3 = MagicMock(
+        list_indexes=MagicMock(
+            return_value=[
+                {
+                    "key": "inserted_at",
+                    "expireAfterSeconds": 3600
+                }
+            ]
+        )
+    )
+
+    mongo_db.database.__getitem__ = MagicMock(
+        side_effect=[
+            collection_1,
+            collection_2,
+            collection_3
+        ]
+    )
+
+    await mongo_db.create_collection_with_ttl_index(input_data)
+
+    mongo_db.database.list_collection_names.assert_called_once_with()
+
+    mongo_db.database.create_collection.assert_called_once_with(
+        "raw_scouter_pipeline"
+    )
+
+    collection_1.list_indexes.assert_called_once()
+    collection_1.create_index.assert_called_once_with(
+        "inserted_at",
+        expireAfterSeconds=3600,
+        background=True
+    )
+
+    collection_2.list_indexes.assert_called_once()
+    collection_2.create_index.assert_called_once_with(
+        "inserted_at",
+        expireAfterSeconds=3600,
+        background=True
+    )
+
+    collection_3.list_indexes.assert_called_once()
+    collection_3.create_index.assert_not_called()
+
+
+@mark.asyncio
+async def test_create_collection_with_ttl_index_failure(mongo_db):
+    input_data = {
+        **metadata,
+        "pipelines": {
+            "scouter-pipeline": {
+                "topic": "raw_scouter_pipeline"
+            }
+        }
+    }
+
+    mongo_db.database.list_collection_names.return_value = []
+
+    mongo_db.database.create_collection.side_effect = Exception("Error")
+
+    try:
+        await mongo_db.create_collection_with_ttl_index(input_data)
+    except Exception as e:
+        assert str(e) == "Error"
+        mongo_db.send_notification.assert_called_once_with(
+            metadata=metadata['metadata'],
+            notification_id="MONGODB_CREATE_COLLECTION_ERROR",
+            message="Failed to create collection raw_scouter_pipeline with TTL index: Error",
+            block="create_collection_with_ttl_index",
+            level=NotificationLevel.ERROR,
+            attachment_content=ANY
+        )
+    else:
+        assert False, "Expected an exception to be raised"

tests/orchestrator/utils/test_connectors_config.py

@@ -57,10 +57,11 @@ def test_build_mongo_db_config_with_env_vars():
     environ['MONGODB_PASSWORD'] = 'sientia1'
     environ['MONGODB_URL'] = 'localhost:27018'
     environ['MONGODB_DATABASE_NAME'] = 'test_db'
-
+    environ['MONGODB_TTL_INDEX_HOURS'] = '2'
     assert build_mongodb_config() == {
         'connection_string': 'mongodb://sientia1:sientia1@localhost:27018',
-        'database_name': 'test_db'
+        'database_name': 'test_db',
+        'ttl_index_seconds': 7200
     }
 
 
@@ -69,10 +70,11 @@ def test_build_mongo_db_config_with_defaults():
     environ.pop('MONGODB_PASSWORD', None)
     environ.pop('MONGODB_DATABASE_NAME', None)
     environ.pop('MONGODB_URL', None)
-
+    environ.pop('MONGODB_TTL_INDEX_HOURS', None)
     assert build_mongodb_config() == {
         'connection_string': 'mongodb://root:wKZDbMNU1c@localhost:27018',
-        'database_name': 'sientia'
+        'database_name': 'sientia',
+        'ttl_index_seconds': 3600
     }
 
 

tests/orchestrator/workflows/test_orchestrator.py

@@ -166,6 +166,18 @@ async def test_run(workflow_mock, orchestrator):
         )
     ])
 
+    workflow_mock.start_activity_method.assert_has_calls([
+        call(
+            Activities.create_collection_with_ttl_index,
+            {
+                **metadata,
+                'pipelines': workflow_mock.start_local_activity_method.return_value['scouter']
+            },
+            retry_policy=ANY,
+            start_to_close_timeout=ANY
+        )
+    ])
+
     workflow_mock.start_activity_method.assert_has_calls([
         call(
             Activities.delete_slots,

values.yaml

@@ -11,7 +11,7 @@ image:
   # This sets the pull policy for images.
   pullPolicy: Always
   # Overrides the image tag whose default is the chart appVersion.
-  tag: "0.2.4"
+  tag: "0.2.5"
 
 # This is for the secrets for pulling an image from a private repository more information can be found here: https://kubernetes.io/docs/tasks/configure-pod-container/pull-image-private-registry/
 imagePullSecrets:
@@ -132,7 +132,7 @@ env:
   - name: GITHUB_REPO_URL
     value: "git@github.com:Aignosi/sientia-dataops-orchestrator_temporal.git"
   - name: GITHUB_BRANCH
-    value: "SIENTIAPDE-1151-criar-testes-de-stress"
+    value: "SIENTIAPDE-1166-alterar-orquestrador-para-criar-collections-com-ttl-no-mongo]"
   - name: PYTHON_APP
     value: "orchestrator.worker.worker"
 
@@ -167,6 +167,8 @@ env:
     value: "my-release-mongodb.mongodb.svc.cluster.local:27017"
   - name: MONGODB_DATABASE
     value: "sientia"
+  - name: MONGODB_TTL_INDEX_HOURS
+    value: "1"
 
   - name: KAFKA_BOOTSTRAP_SERVERS
     value: "kafka.kafka.svc.cluster.local:9092"