diff --git a/smart_kit/message/smartapp_to_message.py b/smart_kit/message/smartapp_to_message.py index b777de94..3ab1a35d 100644 --- a/smart_kit/message/smartapp_to_message.py +++ b/smart_kit/message/smartapp_to_message.py @@ -1,8 +1,9 @@ from functools import cached_property -from typing import Iterable +from typing import Iterable, Optional, List import json from copy import copy +from core.basic_models.actions.command import Command from core.utils.masking_message import masking from core.message.msg_validator import MessageValidator from smart_kit.request.kafka_request import SmartKitKafkaRequest @@ -13,8 +14,9 @@ class SmartAppToMessage: ROOT_NODES_KEY = "root_nodes" PAYLOAD = "payload" - def __init__(self, command, message, request: SmartKitKafkaRequest, forward_fields=None, masking_fields=None, - validators: Iterable[MessageValidator] = (), **kwargs): + def __init__(self, command: Command, message, request: SmartKitKafkaRequest, + forward_fields=None, masking_fields=None, validators: Iterable[MessageValidator] = (), + masking_white_list: Optional[List[str]] = None, **kwargs): root_nodes = command.payload.pop(self.ROOT_NODES_KEY, None) self.command = command self.root_nodes = root_nodes or {} @@ -22,6 +24,7 @@ def __init__(self, command, message, request: SmartKitKafkaRequest, forward_fiel self.request = request self.forward_fields = forward_fields or () self.masking_fields = masking_fields + self.masking_white_list = masking_white_list self.validators = validators self._kwargs = kwargs