sientia-dataops-model-manager/model-manager/activities/opc.py

from temporalio import activity, workflow


with workflow.unsafe.imports_passed_through():
    from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
    from sientia_do.notifications.models import NotificationLevel
    from sientia_do.temporal.activities.base import BaseActivity
    from sientia_do.observability.logger import Logger
    from laborious.utils.repository.opc_repository import OpcRepository
    from typing import Any
    import traceback
    from pandas import DataFrame

OPC_WRITTING_ERROR_CONFIDENCE = 12


class OPC(BaseActivity):
    """
    OPC server integration activities for real-time data export.

    This class provides comprehensive OPC UA client functionality for connecting
    to multiple OPC servers and writing prediction data in real-time. It implements
    secure communication with certificate-based authentication and automatic
    reconnection capabilities.

    The class supports multiple OPC servers with individual configurations and
    provides robust error handling and monitoring for production environments.

    Attributes:
        opc_servers (dict): Configuration for multiple OPC servers
        opc_repository (dict): Active OPC repository connections
        logger (Logger): Logging and observability instance
        notification_handler (NotificationHandler): Notification management instance
    """

    def __init__(self, opc_servers: dict[str, dict[str, Any]],
                 logger: Logger, notification_handler: NotificationHandler):

        self.logger = logger
        self.notification_handler = notification_handler
        self.opc_servers = opc_servers

        BaseActivity.__init__(
            self, logger, notification_handler, set_error_counter=True)

        self.opc_repository: dict[str, OpcRepository] = {}
        self.opc_servers = opc_servers

    async def init_opc(self):
        """
        Initialize OPC server connections and establish communication channels.

        This method iterates through all configured OPC servers and attempts to
        establish secure connections using certificate-based authentication.
        Each server connection is managed independently, and connection failures
        are reported through the notification system.

        The method performs the following operations:
        1. Creates OpcRepository instances for each configured server
        2. Establishes secure connections with certificate validation
        3. Reports connection success/failure through notifications
        4. Logs connection status for operational visibility

        Raises:
            Exception: If OPC repository initialization fails or connection
                establishment encounters critical errors

        Note:
            Connection failures are logged and reported but do not prevent
            the initialization of other OPC servers. Each server is handled
            independently to ensure maximum availability.
        """
        self.logger.info("Initializing OPC servers...")
        for id, server in self.opc_servers.items():
            self.opc_repository[id] = OpcRepository(
                id=server['id'],
                url=server['url'],
                logger=self.logger,
                server_uri=server['server_uri'],
                cert_path=server['cert_path'],
                private_key_path=server['private_key_path'],
                server_cert_path=server['server_cert_path'],
                notification_handler=self.notification_handler,
                reconnection_interval=server['reconnection_interval'],
                pod_id=self.pod_id
            )
            is_connected, error_data = await self.opc_repository[id].connect()
            if not is_connected:
                self.send_notification(
                    metadata={
                        'model_id': '-',
                        'model_name': '-',
                        'workflow_name': '-',
                        'schedule_name': 'INITIALIZATION'
                    },
                    notification_id=error_data['notification_id'],
                    message=error_data['message'],
                    block=error_data['block'],
                    level=error_data.get('level', NotificationLevel.ERROR),
                    attachment_content=error_data.get(
                        'attachment_content', None)
                )
            else:
                self.logger.info(
                    f"OPC server {id} connected successfully.")

    async def write_data(self, server_id: str, tag: str, data: Any,
                         data_type: str, tag_type: str, metadata: dict[str, Any]) -> bool:
        """
        Write data to a specific OPC server tag with comprehensive error handling.

        This method provides a secure and reliable way to write data to OPC servers
        with automatic error handling, notification integration, and detailed logging.
        It validates server availability before attempting write operations and
        provides comprehensive error reporting for operational monitoring.

        Args:
            - server_id (str): The id of the OPC server.
            - tag (str): The tag to write to.
            - data (Any): The data to write.
            - data_type (str): The data type.
            - tag_type (str): The tag type.

        Returns:
            - bool: True if the data was written successfully, False otherwise.
        """

        try:
            is_success, error_data = await self.opc_repository[server_id].write_data(
                tag, data, data_type, self.logger, metadata)
            if not is_success:
                self.send_notification(
                    metadata=metadata,
                    notification_id=error_data['notification_id'],
                    message=error_data['message'],
                    block=error_data['block'],
                    level=error_data.get('level', NotificationLevel.ERROR),
                    attachment_content=error_data.get(
                        'attachment_content', None)
                )
                return False
            return True
        except Exception as e:
            trace = traceback.format_exc()
            self.send_notification(
                metadata=metadata,
                notification_id=f"WRITE_OPC_{tag_type.upper()}_ERROR",
                message=f"Error writing data to OPC server: {e}",
                block="write_opc_data",
                level=NotificationLevel.ERROR,
                attachment_content=trace
            )
            raise e

    def validate_server(self, server_id: str, metadata: dict[str, Any]) -> bool:
        """
        Validate that an OPC server is available and configured for write operations.

        This method checks if the specified OPC server exists in the active
        repository and is available for data writing operations. It provides
        immediate feedback for server availability and logs validation failures
        for operational monitoring.

        Args:
            server_id (str): Unique identifier for the OPC server to validate
            metadata (dict[str, Any]): Context metadata for logging and notifications

        Returns:
            bool: True if server is available, False otherwise

        Note:
            Server validation failures are automatically reported through the
            notification system with detailed information about available servers.
            This helps operators quickly identify configuration issues.
        """
        if self.opc_repository.get(server_id) is None:
            message = f"OPC server {server_id} not found to perform write operation."
            self.send_notification(
                metadata=metadata,
                notification_id="OPC_SERVER_NOT_FOUND",
                message=message,
                block="write_opc_data",
                level=NotificationLevel.ERROR,
                attachment_content=f"OPC servers: {list(self.opc_repository.keys())}"
            )
            return False
        return True

    async def manage_output_tags(
            self, server_id: str, config: dict[str, Any], data: DataFrame,
            metadata: dict[str, Any], success: bool) -> tuple[bool, int]:
        """
        Manage the writing of prediction and confidence data to OPC server tags.

        This method orchestrates the writing of multiple data types to OPC servers
        based on configuration. It handles both prediction data and confidence
        values independently, allowing for flexible tag configuration and
        comprehensive error handling.

        The method supports two main tag types:
        1. Prediction tags: Write actual prediction values to configured OPC tags
        2. Confidence tags: Write confidence scores to separate OPC tags

        Args:
            server_id (str): Unique identifier for the target OPC server
            config (dict[str, Any]): OPC tag configuration containing:
                - prediction_tags (dict, optional): Prediction tag configurations
                - confidence_tags (dict, optional): Confidence tag configurations
            data (DataFrame): DataFrame containing prediction and confidence data
            metadata (dict[str, Any]): Context metadata for logging and notifications
            success (bool): Current success status to maintain across operations

        Returns:
            tuple[bool, int]: (overall_success, total_tags_written)
                - overall_success: True if all configured tags were written successfully
                - total_tags_written: Count of successfully written tags
        """

        count = 0
        if 'prediction_tags' in config:
            for tag, tag_config in config['prediction_tags'].items():
                local_success = await self.write_data(
                    server_id=server_id,
                    tag=tag,
                    data=data.head(1)['prediction'].values[0],
                    data_type=tag_config['data_type'],
                    tag_type='prediction',
                    metadata=metadata
                )
                if local_success:
                    self.info(
                        f"Prediction data written to OPC server {server_id} for tag {tag}.", metadata)
                    count += 1
                success = success and local_success

        if 'confidence_tags' in config:
            for tag, tag_config in config['confidence_tags'].items():
                local_success = await self.write_data(
                    server_id=server_id,
                    tag=tag,
                    data=data.head(1)['prediction_confidence'].values[0],
                    data_type=tag_config['data_type'],
                    tag_type='confidence',
                    metadata=metadata
                )
                if local_success:
                    self.info(
                        f"Confidence data written to OPC server {server_id} for tag {tag}.", metadata)
                    count += 1
                success = success and local_success

        return success, count

    @activity.defn(name='write_opc_data')
    async def write_opc_data(self, input_data: dict[str, Any]) -> dict[Any, Any]:
        """
        Write prediction and confidence data to OPC servers. The two writing
        operations are optional and independent of each other.

        Args:
            - input_data(dict[str, Any]): The input data. Contains the following keys:
                - data(dict[str, Any]): The dataframe that contains the data to write
                to the OPC servers.
                - opc_output_config(dict[str, Any]): The OPC writing configuration.
                The keys are the OPC server names and the values contain:
                    - prediction_tags(dict[str, Any]): The tags to write to the OPC servers.
                    - confidence_tags(dict[str, Any]): The tags to write to the OPC servers.

        Returns:
            - dict[Any, Any]: The data that was written to the OPC servers.

        """
        metadata = input_data['metadata']
        self.info("Writing data to OPC servers...", metadata)
        data = DataFrame(input_data['data'])
        opc_output_config = input_data['opc_output_config']
        self.info(f"Data to write: {data.size} rows", metadata)

        success = True

        for server_id, config in opc_output_config.items():

            if not self.validate_server(server_id, metadata):
                success = False
                continue

            local_success, local_count = await self.manage_output_tags(
                server_id, config, data, metadata, success)
            success = success and local_success

            self.info(
                f"Process completed for OPC server {server_id}: {local_count} of {len(config['prediction_tags'])} prediction tags and {len(config['confidence_tags'])} confidence tags", metadata)

        return self.process_confidence(data, success, metadata)

    def process_confidence(self, data: DataFrame, success: bool, metadata: dict[str, Any]) -> dict[Any, Any]:
        """
        Process prediction confidence based on OPC write operation success.

        This method updates the prediction confidence values in the DataFrame
        based on the success status of OPC server write operations. If any
        write operations failed, it sets the confidence to a predefined error
        value to indicate data quality issues.

        The method implements a confidence degradation strategy:
        - Success: Maintains original confidence values
        - Failure: Sets confidence to error value for operational awareness

        Args:
            data (DataFrame): DataFrame containing prediction and confidence data
            success (bool): Overall success status of OPC write operations
            metadata (dict[str, Any]): Context metadata for logging and notifications

        Returns:
            dict[Any, Any]: Processed data as a dictionary with updated confidence values

        Note:
            The error confidence value (OPC_WRITTING_ERROR_CONFIDENCE = 12) is
            used to indicate that data was not successfully exported to OPC servers.
            This allows downstream systems to handle data quality appropriately.
        """

        if not success:
            data['prediction_confidence'] = OPC_WRITTING_ERROR_CONFIDENCE
            self.debug(
                f"Some data could not be written to OPC servers, setting confidence to  {OPC_WRITTING_ERROR_CONFIDENCE}.",
                metadata
            )

        else:
            self.debug("Data written to OPC servers successfully.", metadata)

        return data.to_dict()

    async def shutdown(self):
        """
        Gracefully shutdown all OPC server connections and cleanup resources.

        This method ensures proper cleanup of all active OPC server connections
        by calling the disconnect method on each repository instance. It's
        designed to be called during application shutdown to prevent resource
        leaks and ensure clean termination.

        The method performs the following cleanup operations:
        1. Iterates through all active OPC repository connections
        2. Calls disconnect() on each repository instance
        3. Allows for graceful connection termination
        4. Prevents resource leaks and connection hanging

        Note:
            This method should be called during application shutdown to ensure
            proper cleanup. It handles all active connections regardless of
            their current state and provides a clean shutdown experience.
        """
        for opc in self.opc_repository.values():
            await opc.disconnect()