added sample collection to volatile tables

Author: remi-td

docs/developer_guide/PROGRESSIVE_DISCLOSURE.md

@@ -2,13 +2,23 @@
 
 ## Overview
 
-Progressive disclosure is an optimization technique for MCP servers with a large number of tools (100+). Instead of listing all tools in the `tools/list` response (which consumes significant context window space), tools are dynamically discovered and executed through a catalog system.
-
-## Context Window Savings
-
-- **Static Mode** (traditional): All 100+ tools listed → ~50,000 tokens
-- **Progressive Disclosure Mode**: 3 proxy tools + 1 core tool → ~500 tokens
-- **Savings**: 99% reduction in initial context window usage
+We enable progressive disclosure of MCP tools and other assets (to be implemented) to optimize the context window usage. 
+
+This is particularly relevant considering the large number of tools that this server offers by default, and the pace at which new custom tools are created after deployment.
+
+Instead of listing all tools in the `tools/list` response (which consumes significant context window space), tools are dynamically discovered and executed through a catalog system.
+
+```mermaid
+%%{init: {'theme':'neutral', 'themeVariables': { 'fontSize':'12px'}, 'flowchart':{'htmlLabels':true, 'curve':'basis', 'nodeSpacing': 30, 'rankSpacing': 30}}}%%
+flowchart TD
+    A[User Input] --> B[Current Context]
+    B --> C[Start Reasoning]
+    C --> D{Do I need more<br/>information?}
+    D -->|No| G[Answer]
+    D -->|Yes| E[Search for<br/>relevant tools]
+    E --> F[Execute tool and<br/>add to context]
+    F --> B
+```    
 
 ## Usage
 

docs/developer_guide/REGISTRY_IMPLEMENTATION.md

@@ -102,23 +102,37 @@ class RegistryLoader:
 - `DATE` → `str`
 - Default → `str`
 
-#### 2. SQL Builder (`registry_tools.py`)
+#### 2. SQL Builder & Type Casting (`registry_tools.py`)
 
-Generates SQL statements for tool execution:
+Generates SQL statements with named parameter placeholders and handles type casting:
 
 ```python
-def build_registry_sql(tool_def: Dict[str, Any], params: Dict[str, Any]) -> str:
-    """Build SQL statement to execute a database-registered tool (UDF or Macro)."""
-    # For UDFs: SELECT * FROM db_object(param1, param2, ...)
-    # For Macros: EXEC db_object(param1, param2, ...)
-    # Handles parameter ordering, type formatting, NULL handling
+def build_registry_sql(tool_def: Dict[str, Any]) -> str:
+    """Build SQL with named parameter placeholders for safe binding."""
+    # For UDFs: SELECT database.function_name(:param1, :param2)
+    # For Macros: EXEC database.macro_name(:param1, :param2)
+    # Parameters sorted by position, using named placeholders
+
+def cast_parameters(params: Dict[str, Any], tool_def: Dict[str, Any]) -> Dict[str, Any]:
+    """Cast parameter values to correct Python types before SQLAlchemy binding."""
+    # Ensures integer parameters stay as int, not string
+    # Critical for Teradata operations like TOP N
 ```
 
 **SQL Generation**:
-- **UDF**: `SELECT * FROM database.function_name(param1, param2)`
-- **Macro**: `EXEC database.macro_name(param1, param2)`
+- **UDF**: `SELECT database.function_name(:param1, :param2)`
+- **Macro**: `EXEC database.macro_name(:param1, :param2)`
 - Parameters sorted by position
-- Values formatted by type (strings quoted, numbers unquoted, NULL handling)
+- Uses named placeholders (`:param_name`) for safe SQLAlchemy parameter binding
+- No value formatting needed - database handles escaping
+- Eliminates SQL injection risks
+
+**Type Casting**:
+- Parameters received from clients (JSON/HTTP) are often strings
+- `cast_parameters()` converts them to correct Python types based on tool definition
+- Example: String `'10'` → integer `10` for TOP N parameters
+- Essential for Teradata type-sensitive operations
+- Handles int, float, str, bool, and None/NULL values
 
 #### 3. Main Integration (`app.py`)
 
@@ -167,28 +181,51 @@ def execute_db_tool_with_registry(*args, **kwargs):
 
 ### Tool Registration
 
-Registry tools are registered using the same pattern as YAML tools:
+Registry tools are registered using the same pattern as YAML tools, with safe parameter binding:
 
 ```python
-# Create executor that generates SQL and executes
-def make_executor(tool_def_captured=tool_def, tool_name_captured=tool_name):
-    def executor(**kwargs):
-        sql = build_registry_sql(tool_def_captured, kwargs)
-        return execute_db_tool(td.handle_base_readQuery, sql, tool_name=tool_name_captured, **kwargs)
-    return executor
-
-# Create MCP tool
-tool_func = create_mcp_tool(
-    executor_func=make_executor(),
-    signature=sig,
-    validate_required=True,
-    tool_name=tool_name,
-)
+# Create handler function (like handle_* functions)
+def handler(conn, tool_name=None, **kwargs):
+    """Registry-defined database tool handler."""
+    # Build SQL with named parameter placeholders
+    sql = build_registry_sql(tool_def)
+
+    # Extract tool parameters (excluding special params)
+    special_params = {'persist', 'tool_name'}
+    tool_params = {k: v for k, v in kwargs.items() if k not in special_params}
+
+    # Extract special params for handle_base_readQuery
+    persist = kwargs.get('persist', False)
+
+    # Pass SQL with named placeholders and parameter values for safe binding
+    return execute_db_tool(
+        td.handle_base_readQuery,
+        sql,
+        tool_name=tool_name,
+        persist=persist,
+        **tool_params  # These will be safely bound by SQLAlchemy
+    )
+
+# Wrap handler as MCP tool
+wrapped = make_tool_wrapper(handler)
 
 # Register with FastMCP
-mcp.tool(name=tool_name, description=description)(tool_func)
+mcp.tool(name=tool_name, description=description)(wrapped)
 ```
 
+**Key Points**:
+- SQL is generated once with placeholders (`:param1`, `:param2`)
+- Parameter values are passed separately as kwargs
+- SQLAlchemy safely binds parameters, preventing SQL injection
+- Same secure pattern used by YAML tools
+
+**Important Limitations**:
+- **Registry tools do not support the `persist` parameter**:
+  - The `persist` parameter is not available for registry tools (both UDFs and Macros)
+  - Macros: Cannot be wrapped in `CREATE VOLATILE TABLE` (invalid: `CREATE VOLATILE TABLE vt AS (EXEC mydb.my_macro(param)) WITH DATA`) ❌
+  - UDFs: While technically possible, persist is disabled for consistency and simplicity
+  - If you need to persist results from registry tools, use YAML tools or Python tools instead
+
 ## Configuration
 
 ### Profile Setup

src/teradata_mcp_server/app.py

@@ -1013,14 +1013,13 @@ def create_registry_handler(tool_name, tool_def):
 
         # Build docstring with parameters
         docstring_parts = [description]
-        if param_defs or True:  # Always show Arguments section to include persist
+        if param_defs:
             docstring_parts.append("\nArguments:")
             for param_name, p in sorted(param_defs.items(), key=lambda x: x[1].get('position', 0)):
                 param_desc = p.get("description", "")
                 docstring_parts.append(f"  {param_name} - {param_desc}")
-            # Add persist parameter documentation
-            docstring_parts.append(f"  persist - If True, materializes result as a volatile table and returns table name")
         docstring_parts.append(f"\nRegistry tool: {tool_def['object_type']} {tool_def['db_object']}")
+        docstring_parts.append(f"Note: Registry tools do not support the 'persist' parameter")
 
         # Add required 'conn' parameter at the beginning (for catalog compatibility)
         parameters = [
@@ -1032,13 +1031,6 @@ def create_registry_handler(tool_name, tool_def):
             inspect.Parameter("tool_name", kind=inspect.Parameter.POSITIONAL_OR_KEYWORD, default=None)
         )
 
-        # Add persist parameter (for materializing results as volatile table)
-        persist_description = "If True, materializes result as a volatile table and returns table name"
-        parameters.append(
-            inspect.Parameter("persist", kind=inspect.Parameter.POSITIONAL_OR_KEYWORD,
-                            default=False, annotation=Annotated[bool, persist_description])
-        )
-
         # Add registry parameters - separate required and optional
         required_params = []
         optional_params = []
@@ -1057,14 +1049,43 @@ def create_registry_handler(tool_name, tool_def):
             else:
                 optional_params.append(param)
 
-        # Build signature with correct order: conn, required_params, tool_name, persist (both with defaults), optional_params
-        sig = inspect.Signature([parameters[0]] + required_params + [parameters[1], parameters[2]] + optional_params)
+        # Build signature: conn, required_params, tool_name (with default), optional_params
+        sig = inspect.Signature([parameters[0]] + required_params + [parameters[1]] + optional_params)
 
         # Create the handler function (like handle_* functions)
         def handler(conn, tool_name=None, **kwargs):
             """Registry-defined database tool handler."""
-            sql = build_registry_sql(tool_def, kwargs)
-            return execute_db_tool(td.handle_base_readQuery, sql, tool_name=tool_name or tool_name, **kwargs)
+            from teradata_mcp_server.tools.registry.registry_tools import cast_parameters, build_registry_sql_with_values
+
+            logger.info(f"[REGISTRY_HANDLER] Starting handler for tool '{tool_name}'")
+            logger.info(f"[REGISTRY_HANDLER] Received kwargs: {kwargs}")
+
+            # Extract tool parameters (excluding special params)
+            # Note: persist is not supported for registry tools
+            special_params = {'tool_name'}
+            tool_params = {k: v for k, v in kwargs.items() if k not in special_params}
+            logger.info(f"[REGISTRY_HANDLER] Tool params before casting: {tool_params}")
+
+            # Cast parameters to their correct types based on tool definition
+            # This ensures values are properly typed before formatting into SQL
+            cast_params = cast_parameters(tool_params, tool_def)
+            logger.info(f"[REGISTRY_HANDLER] Tool params after casting: {cast_params}")
+
+            # Build SQL with values formatted as literals
+            # This approach is necessary because SQLAlchemy parameter binding doesn't
+            # preserve type information correctly with the Teradata driver
+            sql = build_registry_sql_with_values(tool_def, cast_params)
+            logger.info(f"[REGISTRY_HANDLER] Generated SQL: {sql}")
+
+            # Execute the SQL without parameters (values already in SQL)
+            # Note: persist=False for registry tools (not supported)
+            return execute_db_tool(
+                td.handle_base_readQuery,
+                sql,  # SQL string with values already formatted
+                tool_name=tool_name or tool_name,
+                persist=False  # Registry tools do not support persist
+                # No **kwargs here - values are already in the SQL string
+            )
 
         # Set metadata on the handler
         handler.__name__ = f"handle_{tool_name}"

src/teradata_mcp_server/tools/base/base_tools.py

@@ -24,7 +24,7 @@ def handle_base_readQuery(
 
     Arguments:
       sql     - SQL text, with optional bind-parameter placeholders
-      persist - Set to True to materializes results as a table and reuse it later
+      persist - Set to True to persist the results as a table and reuse it later. Recommended for large result sets.
 
     Returns:
       ResponseType: formatted response with query results + metadata
@@ -41,6 +41,9 @@ def handle_base_readQuery(
 
         # Strip trailing semicolons from the SQL
         sql_clean = sql.rstrip().rstrip(';')
+        
+        #Remove the final ORDER BY clause if present
+        sql_clean = re.sub(r'ORDER BY [\w\W\s\S]*$', '', sql_clean, flags=re.IGNORECASE).strip()
 
         # Wrap in CREATE VOLATILE TABLE statement
         sql = f"CREATE VOLATILE TABLE {volatile_table_name} AS ({sql_clean}) WITH DATA ON COMMIT PRESERVE ROWS"
@@ -53,6 +56,11 @@ def handle_base_readQuery(
     result = conn.execute(stmt, kwargs) if kwargs else conn.execute(stmt)
 
     # 3. Fetch rows & column metadata
+
+    # If we persisted in a volatile table, we won't get any rows back, we sample the resulting voltile table instead
+    if volatile_table_name:
+        result = conn.execute(text(f'select top 10 * from {volatile_table_name}'))
+
     cursor = result.cursor  # underlying DB-API cursor
     raw_rows = cursor.fetchall() or []
     data = rows_to_json(cursor.description, raw_rows)
@@ -63,7 +71,7 @@ def handle_base_readQuery(
         }
         for col in (cursor.description or [])
     ]
-
+        
     # 4. Compile the statement with literal binds for “final SQL”
     #    Fallback to DefaultDialect if conn has no `.dialect`
     dialect = getattr(conn, "dialect", default.DefaultDialect())
@@ -83,12 +91,11 @@ def handle_base_readQuery(
 
     # Add volatile table name if persisted
     if volatile_table_name:
-        metadata["columns"] = None
         metadata["row_count"] = None
+        metadata["sample_size"] = 10
         metadata["volatile_table"] = volatile_table_name
         metadata["persist"] = True
         logger.info(f"Query results persisted to volatile table: {volatile_table_name}")
-        data = [{"results stored in volatile_table": volatile_table_name}]
 
     logger.debug(f"Tool: handle_base_readQuery: metadata: {metadata}")
     return create_response(data, metadata)