Update MongoDB connection handling in DataManager and enhance IngestorManager configuration
- Updated DataManager to simplify MongoDB connection logic by removing retry loop and adding logging for connection attempts. - Enhanced IngestorManager to accept MongoDB connection parameters and export_to_kafka flag as part of its initialization. - Updated requirements.txt to use the latest version of the sientia-dataops-library.
This commit is contained in:
@@ -75,26 +75,17 @@ class DataManager:
|
||||
logger.info(
|
||||
f"DataManager initialized with Kafka servers: {kafka_servers}")
|
||||
|
||||
logger.info(
|
||||
f"Trying to initializing DataManager with MongoDB servers: {mongo_connection_string}"
|
||||
)
|
||||
|
||||
self.connection_string = mongo_connection_string
|
||||
self.database = mongo_database
|
||||
|
||||
for i in range(0, 3):
|
||||
logger.info(
|
||||
f"Trying ({i}) to initializing DataManager with MongoDB servers: {self.connection_string}"
|
||||
)
|
||||
try:
|
||||
self.mongo_client = MongoClient(self.connection_string)
|
||||
self.mongo_client.server_info()
|
||||
self.mongo_client = MongoClient(self.connection_string)
|
||||
self.mongo_client.server_info()
|
||||
|
||||
self.mongo_db = self.mongo_client[self.database]
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error(f"Error connecting to MongoDB: {e}")
|
||||
sleep(5)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Failed to connect to MongoDB servers {self.connection_string} after 3 attempts."
|
||||
)
|
||||
self.mongo_db = self.mongo_client[self.database]
|
||||
|
||||
logger.info(
|
||||
f"DataManager initialized with MongoDB servers: {self.connection_string}"
|
||||
@@ -123,6 +114,9 @@ class DataManager:
|
||||
self.mongo_client.close()
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error closing MongoDB client: {e}")
|
||||
else:
|
||||
self.logger.warning(
|
||||
"MongoDB client is already closed or not initialized.")
|
||||
|
||||
def __del__(self):
|
||||
self.shutdown()
|
||||
|
||||
@@ -162,7 +162,8 @@ class IngestorManager():
|
||||
_a = [self.managed_tags[slot].pop(server, None)
|
||||
for slot, _value in current_managed_tags.items()]
|
||||
|
||||
for server in self.opc_managers.keys():
|
||||
servers = list(self.opc_managers.keys())
|
||||
for server in servers:
|
||||
if server not in registered_servers:
|
||||
self.logger.warning(
|
||||
f"Server {server} not found in managed tags. "
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
asyncua==1.1.5
|
||||
redis
|
||||
git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.1.14
|
||||
git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.2.0
|
||||
prometheus_client
|
||||
pymongo
|
||||
@@ -7,20 +7,28 @@ from ingestor.managers.data_manager import DataManager
|
||||
|
||||
@fixture
|
||||
@patch("ingestor.managers.data_manager.KafkaProducer")
|
||||
def data_manager(kafka):
|
||||
@patch("ingestor.managers.data_manager.MongoClient")
|
||||
def data_manager(mongo, kafka):
|
||||
return DataManager(
|
||||
kafka_servers="localhost:9092",
|
||||
mongo_connection_string="mongodb://localhost:27017",
|
||||
mongo_database="sientia",
|
||||
export_to_kafka=True,
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock()
|
||||
)
|
||||
|
||||
|
||||
@patch("ingestor.managers.data_manager.KafkaProducer")
|
||||
def test___init___success(kafka):
|
||||
@patch("ingestor.managers.data_manager.MongoClient")
|
||||
def test___init___success(mongo, kafka):
|
||||
logger_mock = MagicMock()
|
||||
|
||||
data_manager = DataManager(
|
||||
kafka_servers="localhost:9092",
|
||||
mongo_connection_string="mongodb://localhost:27017",
|
||||
mongo_database="sientia",
|
||||
export_to_kafka=True,
|
||||
logger=logger_mock,
|
||||
notification_handler=MagicMock()
|
||||
)
|
||||
@@ -38,16 +46,20 @@ def test___init___success(kafka):
|
||||
"DataManager initialized with Kafka servers: localhost:9092"
|
||||
)
|
||||
logger_mock.error.assert_not_called()
|
||||
assert logger_mock.info.call_count == 2
|
||||
assert logger_mock.info.call_count == 4
|
||||
|
||||
|
||||
@patch("ingestor.managers.data_manager.KafkaProducer")
|
||||
def test___init___second_attempt(kafka):
|
||||
@patch("ingestor.managers.data_manager.MongoClient")
|
||||
def test___init___second_attempt(mongo, kafka):
|
||||
kafka.side_effect = [NoBrokersAvailable, MagicMock()]
|
||||
logger_mock = MagicMock()
|
||||
|
||||
data_manager = DataManager(
|
||||
kafka_servers="localhost:9092",
|
||||
mongo_connection_string="mongodb://localhost:27017",
|
||||
mongo_database="sientia",
|
||||
export_to_kafka=True,
|
||||
logger=logger_mock,
|
||||
notification_handler=MagicMock()
|
||||
)
|
||||
@@ -71,17 +83,21 @@ def test___init___second_attempt(kafka):
|
||||
logger_mock.error.assert_called_once_with(
|
||||
"Kafka servers localhost:9092 are not available. Retrying..."
|
||||
)
|
||||
assert logger_mock.info.call_count == 3
|
||||
assert logger_mock.info.call_count == 5
|
||||
|
||||
|
||||
@patch("ingestor.managers.data_manager.KafkaProducer")
|
||||
def test___init___failure_max_attempts(kafka):
|
||||
@patch("ingestor.managers.data_manager.MongoClient")
|
||||
def test___init___failure_max_attempts(mongo, kafka):
|
||||
kafka.side_effect = NoBrokersAvailable
|
||||
logger_mock = MagicMock()
|
||||
|
||||
try:
|
||||
DataManager(
|
||||
kafka_servers="localhost:9092",
|
||||
mongo_connection_string="mongodb://localhost:27017",
|
||||
mongo_database="sientia",
|
||||
export_to_kafka=True,
|
||||
logger=logger_mock,
|
||||
notification_handler=MagicMock()
|
||||
)
|
||||
@@ -130,6 +146,16 @@ def test_shutdown_no_producer(data_manager):
|
||||
)
|
||||
|
||||
|
||||
def test_shutdown_no_mongo_client(data_manager):
|
||||
data_manager.mongo_client = None
|
||||
|
||||
data_manager.shutdown()
|
||||
|
||||
data_manager.logger.warning.assert_any_call(
|
||||
"MongoDB client is already closed or not initialized."
|
||||
)
|
||||
|
||||
|
||||
def test_shutdown_exception(data_manager):
|
||||
data_manager.kafka_producer.flush = MagicMock(
|
||||
side_effect=Exception("Test error"))
|
||||
@@ -141,6 +167,17 @@ def test_shutdown_exception(data_manager):
|
||||
)
|
||||
|
||||
|
||||
def test_shutdown_exception_mongo(data_manager):
|
||||
data_manager.mongo_client.close = MagicMock(
|
||||
side_effect=Exception("Test error"))
|
||||
|
||||
data_manager.shutdown()
|
||||
|
||||
data_manager.logger.error.assert_called_once_with(
|
||||
"Error closing MongoDB client: Test error"
|
||||
)
|
||||
|
||||
|
||||
def test___del__(data_manager):
|
||||
data_manager.shutdown = MagicMock()
|
||||
data_manager.__del__()
|
||||
@@ -190,6 +227,18 @@ def test_publish(data_manager):
|
||||
data_manager.kafka_producer.flush.assert_called_once()
|
||||
|
||||
|
||||
def test_publish_no_kafka(data_manager):
|
||||
data_manager.export_to_kafka = False
|
||||
topic = "test_topic"
|
||||
data = {"key": "value"}
|
||||
|
||||
data_manager.publish(topic, data)
|
||||
|
||||
data_manager.logger.debug.assert_any_call(
|
||||
f"Skipping message to topic {topic}: {data}"
|
||||
)
|
||||
|
||||
|
||||
@patch("ingestor.managers.data_manager.traceback")
|
||||
def test_publish_error(traceback, data_manager):
|
||||
topic = "test_topic"
|
||||
@@ -215,3 +264,19 @@ def test_publish_error(traceback, data_manager):
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=traceback.format_exc.return_value
|
||||
)
|
||||
|
||||
|
||||
def test_publish_error_mongo(data_manager):
|
||||
data_manager.export_to_kafka = False
|
||||
data_manager.mongo_db.__getitem__.return_value.insert_one = MagicMock(
|
||||
side_effect=Exception("Test error"))
|
||||
|
||||
data_manager.publish("test_topic", {"key": "value"})
|
||||
|
||||
data_manager.notification_handler.build_and_send_notification.assert_called_once_with(
|
||||
notification_id="MONGO_PRODUCER_ERROR_test_topic",
|
||||
message="Error inserting message to MongoDB: Test error",
|
||||
block="mongo_producer",
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=ANY
|
||||
)
|
||||
|
||||
@@ -10,12 +10,17 @@ from ingestor.managers.ingestor_manager import IngestorManager
|
||||
def ingestor_manager(data_manager_mock, resource_manager_mock):
|
||||
return IngestorManager(
|
||||
kafka_servers="localhost:9092",
|
||||
redis_host="localhost",
|
||||
redis_port=6379,
|
||||
redis_data={
|
||||
"host": "localhost",
|
||||
"port": 6379
|
||||
},
|
||||
lease_ttl=60,
|
||||
heartbeat_ttl=60,
|
||||
pod_id="test_pod",
|
||||
poll_interval=5,
|
||||
mongo_connection_string="mongodb://localhost:27017",
|
||||
mongo_database="sientia",
|
||||
export_to_kafka=False,
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock()
|
||||
)
|
||||
@@ -29,19 +34,25 @@ def test___init__(notification_handler_mock, resource_manager_mock, data_manager
|
||||
|
||||
ingestor = IngestorManager(
|
||||
kafka_servers="localhost:9092",
|
||||
redis_host="localhost",
|
||||
redis_port=6379,
|
||||
redis_data={
|
||||
"host": "localhost",
|
||||
"port": 6379
|
||||
},
|
||||
lease_ttl=60,
|
||||
heartbeat_ttl=60,
|
||||
pod_id="test_pod",
|
||||
poll_interval=5,
|
||||
mongo_connection_string="mongodb://localhost:27017",
|
||||
mongo_database="sientia",
|
||||
export_to_kafka=False,
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock()
|
||||
)
|
||||
|
||||
opc_manager_mock.assert_not_called()
|
||||
data_manager_mock.assert_called_once_with(
|
||||
"localhost:9092", ingestor.logger, ingestor.notification_handler)
|
||||
"localhost:9092", "mongodb://localhost:27017", "sientia", False,
|
||||
ingestor.logger, ingestor.notification_handler)
|
||||
resource_manager_mock.assert_called_once_with(
|
||||
"localhost", 6379, 60, 60, "test_pod", None, None)
|
||||
assert ingestor.poll_interval == 5
|
||||
@@ -176,8 +187,10 @@ def test_update_opc_servers(metrics, opc_manager, ingestor_manager):
|
||||
|
||||
assert ingestor_manager.opc_managers['server2'] != mock
|
||||
|
||||
metrics.OPC_MANAGERS_ACTIVE.labels.assert_called_once_with(pod_id=ingestor_manager.pod_id)
|
||||
metrics.OPC_MANAGERS_ACTIVE.labels.return_value.set.assert_called_once_with(len(ingestor_manager.opc_managers))
|
||||
metrics.OPC_MANAGERS_ACTIVE.labels.assert_called_once_with(
|
||||
pod_id=ingestor_manager.pod_id)
|
||||
metrics.OPC_MANAGERS_ACTIVE.labels.return_value.set.assert_called_once_with(
|
||||
len(ingestor_manager.opc_managers))
|
||||
|
||||
|
||||
def test_declare_active(ingestor_manager):
|
||||
@@ -359,7 +372,8 @@ def test_drop_slot_leases(metrics, ingestor_manager):
|
||||
ingestor_manager.resource_manager.drop_tag_lease.assert_any_call("1")
|
||||
ingestor_manager.resource_manager.drop_tag_lease.assert_any_call("2")
|
||||
|
||||
metrics.SLOTS_RELEASED.labels.assert_any_call(pod_id=ingestor_manager.pod_id)
|
||||
metrics.SLOTS_RELEASED.labels.assert_any_call(
|
||||
pod_id=ingestor_manager.pod_id)
|
||||
metrics.SLOTS_RELEASED.labels.return_value.inc.assert_any_call()
|
||||
|
||||
|
||||
@@ -540,8 +554,10 @@ def test_check_opc_servers_integrity_all_healthy(metrics, ingestor_manager):
|
||||
# Verify that no reinitialization was needed
|
||||
ingestor_manager.initialize_opc_from_config.assert_not_called()
|
||||
|
||||
metrics.OPC_MANAGERS_ACTIVE.labels.assert_called_once_with(pod_id=ingestor_manager.pod_id)
|
||||
metrics.OPC_MANAGERS_ACTIVE.labels.return_value.set.assert_called_once_with(len(ingestor_manager.opc_managers))
|
||||
metrics.OPC_MANAGERS_ACTIVE.labels.assert_called_once_with(
|
||||
pod_id=ingestor_manager.pod_id)
|
||||
metrics.OPC_MANAGERS_ACTIVE.labels.return_value.set.assert_called_once_with(
|
||||
len(ingestor_manager.opc_managers))
|
||||
|
||||
|
||||
def test_check_opc_servers_integrity_server_lost(ingestor_manager):
|
||||
|
||||
@@ -11,6 +11,7 @@ from ingestor.ingestor import Ingestor
|
||||
def test___init__(notification_handler, init_logger, getenv):
|
||||
getenv.side_effect = [
|
||||
"localhost:9092,localhost:35", # KAFKA_SERVERS
|
||||
"true", # EXPORT_TO_KAFKA
|
||||
"localhost1", # REDIS_HOST
|
||||
'63790', # REDIS_PORT
|
||||
"user", # REDIS_USERNAME
|
||||
@@ -18,7 +19,11 @@ def test___init__(notification_handler, init_logger, getenv):
|
||||
'100', # LEASE_TTL
|
||||
'200', # HEARTBEAT_TTL
|
||||
"localhost1", # HOSTNAME
|
||||
'50' # POLL_INTERVAL
|
||||
'50', # POLL_INTERVAL
|
||||
"mongodb://localhost:27017", # MONGODB_URL
|
||||
"sientia", # MONGODB_USERNAME
|
||||
"sientia", # MONGODB_PASSWORD
|
||||
"sientia" # MONGODB_DATABASE
|
||||
]
|
||||
|
||||
ingestor = Ingestor()
|
||||
@@ -114,16 +119,21 @@ def test_prepare_ingestor(ingestor_manager_mock, ingestor):
|
||||
|
||||
ingestor_manager_mock.assert_called_once_with(
|
||||
ingestor.kafka_servers,
|
||||
ingestor.redis_host,
|
||||
ingestor.redis_port,
|
||||
{
|
||||
"host": ingestor.redis_host,
|
||||
"port": ingestor.redis_port,
|
||||
"username": ingestor.redis_username,
|
||||
"password": ingestor.redis_password
|
||||
},
|
||||
ingestor.lease_ttl,
|
||||
ingestor.heartbeat_ttl,
|
||||
ingestor.pod_id,
|
||||
ingestor.poll_interval,
|
||||
ingestor.mongo_connection_string,
|
||||
ingestor.mongo_database,
|
||||
ingestor.logger,
|
||||
ingestor.notification_handler,
|
||||
ingestor.redis_username,
|
||||
ingestor.redis_password
|
||||
ingestor.export_to_kafka
|
||||
)
|
||||
ingestor_manager.declare_active.assert_called_once()
|
||||
ingestor_manager.get_slot_leases.assert_called_once()
|
||||
|
||||
Reference in New Issue
Block a user