refactoring code

Author: Avi-Robusta

src/robusta/core/sinks/slack/slack_sink_params.py

@@ -1,23 +1,46 @@
 from robusta.core.sinks.sink_base_params import SinkBaseParams
 from robusta.core.sinks.sink_config import SinkConfigBase
 from robusta.core.sinks.common import ChannelTransformer
-
+from enum import Enum
 from typing import Optional, Dict, Literal
 from pydantic import validator
 
 
+class SlackTemplateStyle(str, Enum):
+    DEFAULT = "default"
+    LEGACY = "legacy"
+
+
 class SlackSinkParams(SinkBaseParams):
     slack_channel: str
     api_key: str
     channel_override: Optional[str] = None
     max_log_file_limit_kb: int = 1000
     investigate_link: bool = True
-    send_svg: bool = True
-    prefer_redirect_to_platform: bool = True
-    
-    # Template selection and customization options
-    template_style: Literal["default", "legacy"] = "default"  # Use "legacy" for old-style formatting
-    custom_templates: Optional[Dict[str, str]] = None  # Template name -> custom template content
+
+    template_style: SlackTemplateStyle = SlackTemplateStyle.DEFAULT  # Use "legacy" for old-style formatting
+    slack_custom_templates: Optional[Dict[str, str]] = None  # Template name -> custom template content
+    template_name: Optional[str] = None
+
+    def get_effective_template_name(self) -> str:
+        """
+        Returns the template name to use for this sink. If template_name is set, use it.
+        Otherwise, use 'legacy.j2' if template_style is legacy, else 'header.j2'.
+        """
+        if self.template_name:
+            return self.template_name
+        if self.template_style == SlackTemplateStyle.LEGACY:
+            return "legacy.j2"
+        return "header.j2"
+
+    def get_custom_template(self) -> Optional[str]:
+        """
+        Returns the custom template string for the effective template name, if it exists.
+        """
+        template_name = self.get_effective_template_name()
+        if self.slack_custom_templates and template_name in self.slack_custom_templates:
+            return self.slack_custom_templates[template_name]
+        return None
 
     @classmethod
     def _supports_grouping(cls):

src/robusta/core/sinks/slack/templates/template_loader.py

@@ -92,32 +92,33 @@ def render_to_blocks(self, template_name: str, context: Dict[str, Any]) -> List[
             logging.error(f"Error rendering template {template_name}: {e}")
             return []
 
-    def render_custom_or_file_template_to_blocks(self, template_name: str, context: Dict[str, Any], custom_template: str = None) -> List[Dict[str, Any]]:
+    def render_custom_template_to_blocks(self, custom_template: str, context: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
-        Render a custom Jinja template string (if provided) or a file-based template to Slack blocks.
+        Render a custom Jinja template string to Slack blocks.
         Args:
-            template_name: The name of the file-based template (e.g., "header.j2")
+            custom_template: Jinja template string
             context: Dictionary of variables to pass to the template
-            custom_template: Optional Jinja template string to use instead of file-based template
         Returns:
             List of Slack block objects (dictionaries)
         """
-        if custom_template:
-            try:
-                template = Template(custom_template)
-                rendered_blocks = []
-                for block_str in template.render(**context).strip().split("\n\n"):
-                    if block_str.strip():
-                        block_str_fixed = escape_raw_newlines_in_json_strings(block_str)
-                        block = json.loads(block_str_fixed)
-                        rendered_blocks.append(block)
-                return rendered_blocks
-            except Exception as e:
-                logging.error(f"Error rendering custom template: {e}")
-                # Fall back to file-based template
+        try:
+            template = Template(custom_template)
+            return self.render_to_blocks(template, context)
+        except Exception as e:
+            logging.error(f"Error rendering custom template: {e}")
+            return []
 
-        # Use file-based template
-        return self.render_to_blocks(template_name, context)
+    def render_file_template_to_blocks(self, template_name: str, context: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+        Render a file-based Jinja template to Slack blocks.
+        Args:
+            template_name: The name of the file-based template (e.g., "header.j2")
+            context: Dictionary of variables to pass to the template
+        Returns:
+            List of Slack block objects (dictionaries)
+        """
+        template = self.get_template(template_name)
+        return self.render_to_blocks(template, context)
 
 
 # Singleton instance

src/robusta/integrations/slack/sender.py

@@ -437,17 +437,8 @@ def __create_finding_header(
         title, mention = self.extract_mentions(title)
 
         sev = finding.severity
+
 
-        # Select appropriate template based on user preference
-        template_name = "header.j2"  # default template
-        if sink_params and hasattr(sink_params, "template_style") and sink_params.template_style == "legacy":
-            template_name = "legacy.j2"
-        
-        # Check if the user has provided a custom template
-        custom_template = None
-        if sink_params and sink_params.custom_templates and template_name in sink_params.custom_templates:
-            custom_template = sink_params.custom_templates[template_name]
-            
         # Prepare data for template
         status_text = "Firing" if status == FindingStatus.FIRING else "Resolved"
         status_emoji = "⚠️" if status == FindingStatus.FIRING else "✅"
@@ -508,10 +499,17 @@ def __create_finding_header(
             "finding": finding
         }
 
-        # Use the new template loader method for both custom and file-based templates
-        return template_loader.render_custom_or_file_template_to_blocks(
-            template_name, template_context, custom_template
-        )
+        # Determine the template name to use
+        template_name = sink_params.get_effective_template_name() if sink_params else "header.j2"
+
+        # Get the custom template for this template name, if any
+        custom_template = sink_params.get_custom_template() if sink_params else None
+
+        # Use the new template loader methods for custom or file-based templates
+        if custom_template:
+            return template_loader.render_custom_template_to_blocks(custom_template, template_context)
+        else:
+            return template_loader.render_file_template_to_blocks(template_name, template_context)
 
     def __create_links(
         self,