Pass plugin module directly to FunctionRegistry.initialize()

Instead of scanning all sys.modules, pass the imported plugin module
directly so discovery is targeted and avoids false positives from
infrastructure modules. Also hardens the sys.modules fallback path
against RuntimeError (dict changed during iteration) and TypeError
from problematic module attributes.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: kesmit13

singlestoredb/functions/ext/plugin/registry.py

@@ -144,15 +144,29 @@ class FunctionRegistry:
     def __init__(self) -> None:
         self.functions: Dict[str, Dict[str, Any]] = {}
 
-    def initialize(self) -> None:
+    def initialize(self, plugin_module: Any = None) -> None:
         """Initialize and discover UDF functions from loaded modules.
 
-        Scans sys.modules for any module containing @udf-decorated
+        If plugin_module is provided, only that module is scanned.
+        Otherwise scans sys.modules for any module containing @udf-decorated
         functions. No _exports.py is needed -- modules just need to be
         imported before initialize() is called (componentize-py captures
         them at build time).
         """
-        self._discover_udf_functions()
+        if plugin_module is not None:
+            self._extract_functions(plugin_module)
+            if self.functions:
+                logger.info(
+                    f'Discovered UDF functions from module: '
+                    f'{plugin_module.__name__}',
+                )
+            else:
+                logger.warning(
+                    f'No @udf functions found in module: '
+                    f'{plugin_module.__name__}',
+                )
+        else:
+            self._discover_udf_functions()
 
     @staticmethod
     def _is_stdlib_or_infra(mod_name: str, mod_file: str) -> bool:
@@ -164,8 +178,9 @@ def _is_stdlib_or_infra(mod_name: str, mod_file: str) -> bool:
         """
         _infra = frozenset({
             'udf_handler',
+            'singlestoredb',
         })
-        if mod_name in _infra:
+        if mod_name in _infra or mod_name.startswith('singlestoredb.'):
             return True
 
         real_file = os.path.realpath(mod_file)
@@ -197,10 +212,18 @@ def _discover_udf_functions(self) -> None:
             if self._is_stdlib_or_infra(mod_name, mod_file):
                 continue
 
-            if not any(
-                hasattr(obj, '_singlestoredb_attrs')
-                for obj in vars(mod).values()
-            ):
+            def _has_udf_marker(obj: object) -> bool:
+                try:
+                    return hasattr(obj, '_singlestoredb_attrs')
+                except (TypeError, Exception):
+                    return False
+
+            try:
+                mod_vars = list(vars(mod).values())
+            except RuntimeError:
+                continue
+
+            if not any(_has_udf_marker(obj) for obj in mod_vars):
                 continue
 
             self._extract_functions(mod)

singlestoredb/functions/ext/plugin/server.py

@@ -487,11 +487,11 @@ def _initialize_registry(self) -> FunctionRegistry:
 
         # Import the plugin module
         logger.info(f'Importing plugin module: {plugin_name}')
-        importlib.import_module(plugin_name)
+        mod = importlib.import_module(plugin_name)
 
-        # Initialize registry (discovers @udf functions from sys.modules)
+        # Initialize registry (discovers @udf functions from the plugin module)
         registry = FunctionRegistry()
-        registry.initialize()
+        registry.initialize(plugin_module=mod)
 
         func_count = len(registry.functions)
         if func_count == 0: