gitea_admin/sientia-dataops-orchestrator_temporal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
feat: add TTL index functionality to MongoDB activities and update requirements
This commit is contained in:
Matheus Demoner
2025-07-24 10:06:25 -03:00
committed by GitHub
parent f72735857a ff683eea14
commit 6ba85a755a
11 changed files with 219 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
orchestrator/activities/activities.py
View File

@@ -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)

59
orchestrator/activities/mongo_db.py
View File

@@ -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

3
orchestrator/utils/connectors_config.py
View File

@@ -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
}

1
orchestrator/worker/worker.py
View File

@@ -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,

11
orchestrator/workflows/orchestrator.py
View File

@@ -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,

2
requirements.txt
View File

@@ -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

4
tests/orchestrator/activities/test_activities.py
View File

@@ -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
)

120
tests/orchestrator/activities/test_mongo_db.py
View File

@@ -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"

10
tests/orchestrator/utils/test_connectors_config.py
View File

@@ -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
}

12
tests/orchestrator/workflows/test_orchestrator.py
View File

@@ -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,

6
values.yaml
View File

@@ -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"