sientia-dataops-laborious_temporal/laborious/activities/opc.py

from temporalio import activity, workflow

with workflow.unsafe.imports_passed_through():
    import traceback
    from collections.abc import Hashable
    from typing import Any

    from pandas import DataFrame
    from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
    from sientia_do.notifications.models import NotificationLevel
    from sientia_do.observability.logger import Logger
    from sientia_do.observability.metrics_controller import MetricsController
    from sientia_do.observability.sientia_monitoring import SientiaMonitoring

    from laborious.utils.repository.opc_repository import OpcRepository

OPC_WRITTING_ERROR_CONFIDENCE = 12
OPC_SESSION_BAD_CONFIDENCE = 14
OPC_SESSION_BAD_COMMENT_PREFIX = 'OPC UA session/channel error:'
OPC_WRITTING_ERROR_MESSAGE = 'Some data could not be written to OPC servers'
OPC_RECONNECT_IN_PROGRESS_COMMENT = 'OPC UA reconnect in progress'
OPC_COMMENT_SEPARATOR = ' | '


def _opc_session_bad_comment(opc_status: str | None) -> str:
    status = opc_status or 'Unknown'
    return f'{OPC_SESSION_BAD_COMMENT_PREFIX} {status}'


def _apply_opc_write_error(
    error_info: dict[str, Any] | None,
    session_bad_seen: bool,
    session_bad_status: str | None,
    reconnect_in_progress_seen: bool,
) -> tuple[bool, str | None, bool]:
    """
    Update session/reconnect flags from an OPC write error payload.

    Args:
        error_info: Repository error details, or None when the write succeeded.
        session_bad_seen: Whether a session_bad error was seen so far.
        session_bad_status: Last known OPC status for session errors.
        reconnect_in_progress_seen: Whether reconnect_in_progress was seen so far.

    Return:
        Updated (session_bad_seen, session_bad_status, reconnect_in_progress_seen).
    """
    if not error_info:
        return session_bad_seen, session_bad_status, reconnect_in_progress_seen

    kind = error_info.get('opc_error_kind')
    if kind == 'session_bad':
        return True, error_info.get('opc_status', session_bad_status), reconnect_in_progress_seen
    if kind == 'reconnect_in_progress':
        return session_bad_seen, session_bad_status, True
    return session_bad_seen, session_bad_status, reconnect_in_progress_seen


class OPC(SientiaMonitoring):
    """
    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,
        metrics_controller: MetricsController,
    ):
        self.logger = logger
        self.notification_handler = notification_handler
        self.opc_servers = opc_servers

        SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)

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

    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 opc_id, server in self.opc_servers.items():
            self.opc_repository[opc_id] = OpcRepository(
                opc_id=server['id'],
                server_name=server['server_name'],
                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'],
                metrics_controller=self.metrics_controller,
            )
            is_connected, error_data = await self.opc_repository[opc_id].connect()
            if not is_connected:
                await self.send_notification_async(
                    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 {opc_id}:{server["server_name"]} connected successfully.'
                )

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

        Return:
            tuple[float | None, dict[str, Any] | None]: Response time on success, or
                (None, error info_data) on repository failure.
        """

        try:
            is_success, info_data = await self.opc_repository[server_id].write_data(
                tag, data, data_type, metadata
            )
            if not is_success:
                await self.send_notification_async(
                    metadata=metadata,
                    notification_id=info_data['notification_id'],
                    message=info_data['message'],
                    block=info_data['block'],
                    level=info_data.get('level', NotificationLevel.ERROR),
                    attachment_content=info_data.get('attachment_content', None),
                )
                return None, info_data
            return info_data['response_time'], None
        except Exception as e:
            trace = traceback.format_exc()
            await self.send_notification_async(
                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

    async 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.'
            await self.send_notification_async(
                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 _write_tags_from_config(
        self,
        server_id: str,
        tags_config: dict[str, dict[str, Any]],
        data: DataFrame,
        data_column: str,
        tag_type: str,
        log_label: str,
        metadata: dict[str, Any],
    ) -> tuple[dict[str, float | None], bool, str | None, bool]:
        """
        Write a group of OPC tags and collect response times and error flags.

        Args:
            server_id: Target OPC server identifier.
            tags_config: Tag name to configuration mapping.
            data: DataFrame with prediction/confidence columns.
            data_column: Column name whose first row value is written.
            tag_type: Tag category passed to write_data ('prediction' or 'confidence').
            log_label: Human-readable label for success logs.
            metadata: Context metadata for logging and notifications.

        Return:
            (response_times, session_bad_seen, session_bad_status, reconnect_in_progress_seen)
        """
        response_times: dict[str, float | None] = {}
        session_bad_seen = False
        session_bad_status: str | None = None
        reconnect_in_progress_seen = False

        for tag, tag_config in tags_config.items():
            response_time, error_info = await self.write_data(
                server_id=server_id,
                tag=tag,
                data=data.head(1)[data_column].values[0],
                data_type=tag_config['data_type'],
                tag_type=tag_type,
                metadata=metadata,
            )
            session_bad_seen, session_bad_status, reconnect_in_progress_seen = (
                _apply_opc_write_error(
                    error_info,
                    session_bad_seen,
                    session_bad_status,
                    reconnect_in_progress_seen,
                )
            )
            if response_time is not None:
                self.info(
                    f'{log_label} written to OPC server {server_id} for tag {tag}.',
                    metadata,
                )
            response_times[tag] = response_time

        return response_times, session_bad_seen, session_bad_status, reconnect_in_progress_seen

    async def manage_output_tags(
        self,
        server_id: str,
        config: dict[str, Any],
        data: DataFrame,
        metadata: dict[str, Any],
    ) -> tuple[bool, dict[str, float | None], bool, str | None, bool]:
        """
        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
        """
        response_times: dict[str, float | None] = {}
        session_bad_seen = False
        session_bad_status: str | None = None
        reconnect_in_progress_seen = False

        tag_groups = (
            ('prediction_tags', 'prediction', 'prediction', 'Prediction data'),
            ('confidence_tags', 'prediction_confidence', 'confidence', 'Confidence data'),
        )
        for config_key, data_column, tag_type, log_label in tag_groups:
            if config_key not in config:
                continue
            (
                group_times,
                group_session_bad,
                group_status,
                group_reconnect,
            ) = await self._write_tags_from_config(
                server_id=server_id,
                tags_config=config[config_key],
                data=data,
                data_column=data_column,
                tag_type=tag_type,
                log_label=log_label,
                metadata=metadata,
            )
            response_times.update(group_times)
            if group_session_bad:
                session_bad_seen = True
                session_bad_status = group_status or session_bad_status
            if group_reconnect:
                reconnect_in_progress_seen = True

        success = None not in response_times.values()
        return (
            success,
            response_times,
            session_bad_seen,
            session_bad_status,
            reconnect_in_progress_seen,
        )

    @activity.defn(name='write_opc_data')
    async def write_opc_data(
        self, input_data: dict[str, Any]
    ) -> tuple[dict[Hashable, Any], dict[str, dict[str, float | None]]]:
        """
        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
        session_bad_seen = False
        session_bad_status: str | None = None
        reconnect_in_progress_seen = False

        metrics: dict[str, dict[str, float | None]] = {}

        for server_id, config in opc_output_config.items():
            if not await self.validate_server(server_id, metadata):
                success = False
                continue

            (
                local_success,
                local_response_times,
                local_session_bad,
                local_status,
                local_reconnect_in_progress,
            ) = await self.manage_output_tags(server_id, config, data, metadata)
            metrics[server_id] = local_response_times
            local_count = len(local_response_times)
            success = success and local_success
            if local_session_bad:
                session_bad_seen = True
                session_bad_status = local_status or session_bad_status
            if local_reconnect_in_progress:
                reconnect_in_progress_seen = True

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

        return (
            self.process_confidence(
                data,
                success,
                metadata,
                session_bad=session_bad_seen,
                opc_status=session_bad_status,
                reconnect_in_progress=reconnect_in_progress_seen,
            ),
            metrics,
        )

    def process_confidence(
        self,
        data: DataFrame,
        success: bool,
        metadata: dict[str, Any],
        *,
        session_bad: bool = False,
        opc_status: str | None = None,
        reconnect_in_progress: bool = False,
    ) -> dict[Hashable, 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:
            comment_parts: list[str] = []
            confidence = OPC_WRITTING_ERROR_CONFIDENCE

            if session_bad:
                comment_parts.append(_opc_session_bad_comment(opc_status))
                confidence = OPC_SESSION_BAD_CONFIDENCE
            if reconnect_in_progress:
                comment_parts.append(OPC_RECONNECT_IN_PROGRESS_COMMENT)
                confidence = OPC_SESSION_BAD_CONFIDENCE
            if not comment_parts:
                comment_parts.append(OPC_WRITTING_ERROR_MESSAGE)

            comments = OPC_COMMENT_SEPARATOR.join(comment_parts)
            data['prediction_confidence'] = confidence
            data['comments'] = comments
            self.debug(
                f'OPC write issues, confidence={confidence}, comments={comments}',
                metadata,
            )
        else:
            self.debug('Data written to OPC servers successfully.', metadata)

        return data.to_dict()

    async def aclose(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()