From 4905e612144300591050740b10223f433f327bc2 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Tue, 22 Jul 2025 15:09:15 -0300 Subject: [PATCH 1/3] SIENTIAPDE-1166 feat: add TTL index functionality to MongoDB activities and update requirements - Updated MongoDB class to accept TTL index duration in seconds. - Implemented create_collection_with_ttl_index method to create collections with TTL indexes. - Added environment variable for TTL index duration in connectors_config. - Updated tests to cover new TTL index functionality and ensure proper behavior. - Bumped sientia-dataops-library version to 1.3.4 in requirements.txt. --- orchestrator/activities/activities.py | 1 + orchestrator/activities/mongo_db.py | 58 ++++++++- orchestrator/utils/connectors_config.py | 3 +- orchestrator/worker/worker.py | 1 + orchestrator/workflows/orchestrator.py | 11 ++ requirements.txt | 2 +- .../activities/test_activities.py | 4 +- .../orchestrator/activities/test_mongo_db.py | 120 ++++++++++++++++++ .../utils/test_connectors_config.py | 10 +- .../workflows/test_orchestrator.py | 12 ++ values.yaml | 2 + 11 files changed, 216 insertions(+), 8 deletions(-) diff --git a/orchestrator/activities/activities.py b/orchestrator/activities/activities.py index 23c0ace..6c769cc 100644 --- a/orchestrator/activities/activities.py +++ b/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) diff --git a/orchestrator/activities/mongo_db.py b/orchestrator/activities/mongo_db.py index 1c1510f..3d8f246 100644 --- a/orchestrator/activities/mongo_db.py +++ b/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,57 @@ 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 diff --git a/orchestrator/utils/connectors_config.py b/orchestrator/utils/connectors_config.py index 103d20b..519731a 100644 --- a/orchestrator/utils/connectors_config.py +++ b/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 } diff --git a/orchestrator/worker/worker.py b/orchestrator/worker/worker.py index f7c37d4..525a376 100644 --- a/orchestrator/worker/worker.py +++ b/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, diff --git a/orchestrator/workflows/orchestrator.py b/orchestrator/workflows/orchestrator.py index 816919b..f61587d 100644 --- a/orchestrator/workflows/orchestrator.py +++ b/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, diff --git a/requirements.txt b/requirements.txt index 8c82e94..23d31f9 100644 --- a/requirements.txt +++ b/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 diff --git a/tests/orchestrator/activities/test_activities.py b/tests/orchestrator/activities/test_activities.py index 353382a..1d76bae 100644 --- a/tests/orchestrator/activities/test_activities.py +++ b/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 ) diff --git a/tests/orchestrator/activities/test_mongo_db.py b/tests/orchestrator/activities/test_mongo_db.py index 355070b..707b866 100644 --- a/tests/orchestrator/activities/test_mongo_db.py +++ b/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" diff --git a/tests/orchestrator/utils/test_connectors_config.py b/tests/orchestrator/utils/test_connectors_config.py index 86f5154..a1c219e 100644 --- a/tests/orchestrator/utils/test_connectors_config.py +++ b/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 } diff --git a/tests/orchestrator/workflows/test_orchestrator.py b/tests/orchestrator/workflows/test_orchestrator.py index a55362e..c3e0257 100644 --- a/tests/orchestrator/workflows/test_orchestrator.py +++ b/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, diff --git a/values.yaml b/values.yaml index c505922..03b5b62 100644 --- a/values.yaml +++ b/values.yaml @@ -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" From 2317bd404be2e36fd4a2f937600347b6e3f238b1 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Wed, 23 Jul 2025 07:45:51 -0300 Subject: [PATCH 2/3] SIENTIAPDE-1166 fix: add missing newline in MongoDB class for improved readability --- orchestrator/activities/mongo_db.py | 1 + 1 file changed, 1 insertion(+) diff --git a/orchestrator/activities/mongo_db.py b/orchestrator/activities/mongo_db.py index 3d8f246..515a9b4 100644 --- a/orchestrator/activities/mongo_db.py +++ b/orchestrator/activities/mongo_db.py @@ -339,6 +339,7 @@ class MongoDB(BaseActivity): expireAfterSeconds=self.ttl_index_seconds, background=True ) + except Exception as e: trace = traceback.format_exc() self.send_notification( From ff683eea14a73d242fe381c19d5e8521272514dd Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Wed, 23 Jul 2025 08:04:21 -0300 Subject: [PATCH 3/3] SIENTIAPDE-1166 Update image tag to 0.2.5 and change GITHUB_BRANCH to SIENTIAPDE-1166-alterar-orquestrador-para-criar-collections-com-ttl-no-mongo in values.yaml --- values.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/values.yaml b/values.yaml index 03b5b62..bb7ac0d 100644 --- a/values.yaml +++ b/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"