1st iter

backend/src/plugins/search.py

@@ -0,0 +1,161 @@
+# [DEF:SearchPluginModule:Module]
+# @SEMANTICS: plugin, search, datasets, regex, superset
+# @PURPOSE: Implements a plugin for searching text patterns across all datasets in a specific Superset environment.
+# @LAYER: Plugins
+# @RELATION: Inherits from PluginBase. Uses SupersetClient from core.
+# @CONSTRAINT: Must use belief_scope for logging.
+
+# [SECTION: IMPORTS]
+import re
+from typing import Dict, Any, List, Optional
+from ..core.plugin_base import PluginBase
+from ..core.superset_client import SupersetClient
+from ..core.logger import logger, belief_scope
+# [/SECTION]
+
+# [DEF:SearchPlugin:Class]
+# @PURPOSE: Plugin for searching text patterns in Superset datasets.
+class SearchPlugin(PluginBase):
+    """
+    Plugin for searching text patterns in Superset datasets.
+    """
+
+    @property
+    def id(self) -> str:
+        return "search-datasets"
+
+    @property
+    def name(self) -> str:
+        return "Search Datasets"
+
+    @property
+    def description(self) -> str:
+        return "Search for text patterns across all datasets in a specific environment."
+
+    @property
+    def version(self) -> str:
+        return "1.0.0"
+
+    # [DEF:SearchPlugin.get_schema:Function]
+    # @PURPOSE: Returns the JSON schema for the search plugin parameters.
+    def get_schema(self) -> Dict[str, Any]:
+        return {
+            "type": "object",
+            "properties": {
+                "env": {
+                    "type": "string",
+                    "title": "Environment",
+                    "description": "The Superset environment to search in (e.g., 'dev', 'prod')."
+                },
+                "query": {
+                    "type": "string",
+                    "title": "Search Query (Regex)",
+                    "description": "The regex pattern to search for."
+                }
+            },
+            "required": ["env", "query"]
+        }
+    # [/DEF:SearchPlugin.get_schema:Function]
+
+    # [DEF:SearchPlugin.execute:Function]
+    # @PURPOSE: Executes the dataset search logic.
+    # @PRE: Params contain valid 'env' and 'query'.
+    # @POST: Returns a dictionary with count and results list.
+    async def execute(self, params: Dict[str, Any]) -> Dict[str, Any]:
+        with belief_scope("SearchPlugin.execute", f"params={params}"):
+            env_name = params.get("env")
+            search_query = params.get("query")
+            
+            if not env_name or not search_query:
+                logger.error("[SearchPlugin.execute][State] Missing required parameters.")
+                raise ValueError("Missing required parameters: env, query")
+
+            # Get config and initialize client
+            from ..dependencies import get_config_manager
+            config_manager = get_config_manager()
+            env_config = config_manager.get_environment(env_name)
+            if not env_config:
+                logger.error(f"[SearchPlugin.execute][State] Environment '{env_name}' not found.")
+                raise ValueError(f"Environment '{env_name}' not found in configuration.")
+
+            client = SupersetClient(env_config)
+            client.authenticate()
+
+            logger.info(f"[SearchPlugin.execute][Action] Searching for pattern: '{search_query}' in environment: {env_name}")
+            
+            try:
+                # Ported logic from search_script.py
+                _, datasets = client.get_datasets(query={"columns": ["id", "table_name", "sql", "database", "columns"]})
+
+                if not datasets:
+                    logger.warning("[SearchPlugin.execute][State] No datasets found.")
+                    return {"count": 0, "results": []}
+
+                pattern = re.compile(search_query, re.IGNORECASE)
+                results = []
+                
+                for dataset in datasets:
+                    dataset_id = dataset.get('id')
+                    dataset_name = dataset.get('table_name', 'Unknown')
+                    if not dataset_id:
+                        continue
+
+                    for field, value in dataset.items():
+                        value_str = str(value)
+                        if pattern.search(value_str):
+                            match_obj = pattern.search(value_str)
+                            results.append({
+                                "dataset_id": dataset_id,
+                                "dataset_name": dataset_name,
+                                "field": field,
+                                "match_context": self._get_context(value_str, match_obj.group() if match_obj else ""),
+                                "full_value": value_str
+                            })
+
+                logger.info(f"[SearchPlugin.execute][Success] Found matches in {len(results)} locations.")
+                return {
+                    "count": len(results),
+                    "results": results
+                }
+
+            except re.error as e:
+                logger.error(f"[SearchPlugin.execute][Failure] Invalid regex pattern: {e}")
+                raise ValueError(f"Invalid regex pattern: {e}")
+            except Exception as e:
+                logger.error(f"[SearchPlugin.execute][Failure] Error during search: {e}")
+                raise
+    # [/DEF:SearchPlugin.execute:Function]
+
+    # [DEF:SearchPlugin._get_context:Function]
+    # @PURPOSE: Extracts a small context around the match for display.
+    def _get_context(self, text: str, match_text: str, context_lines: int = 1) -> str:
+        """
+        Extracts a small context around the match for display.
+        """
+        if not match_text:
+            return text[:100] + "..." if len(text) > 100 else text
+            
+        lines = text.splitlines()
+        match_line_index = -1
+        for i, line in enumerate(lines):
+            if match_text in line:
+                match_line_index = i
+                break
+        
+        if match_line_index != -1:
+            start = max(0, match_line_index - context_lines)
+            end = min(len(lines), match_line_index + context_lines + 1)
+            context = []
+            for i in range(start, end):
+                line_content = lines[i]
+                if i == match_line_index:
+                    context.append(f"==> {line_content}")
+                else:
+                    context.append(f"    {line_content}")
+            return "\n".join(context)
+        
+        return text[:100] + "..." if len(text) > 100 else text
+    # [/DEF:SearchPlugin._get_context:Function]
+
+# [/DEF:SearchPlugin:Class]
+# [/DEF:SearchPluginModule:Module]