updated dynamic "handle_" function wrapper to support most EFS tools and moved tool logic from server to to the tool/fs module

Author: remi-td

src/teradata_mcp_server/server.py

@@ -56,6 +56,19 @@
     else:
         config = all_profiles.get(profile_name)
 
+# Check if the EFS or EVS tools are enabled in the profiles
+if any(re.match(pattern, 'fs_*') for pattern in config.get('tool',[])):
+    _enableEFS = True
+else:
+    _enableEFS = False
+
+if any(re.match(pattern, 'evs_*') for pattern in config.get('tool',[])):
+    _enableEVS = True
+else:
+    _enableEVS = False
+
+_requireTdmlContext = _enableEFS or _enableEVS
+
 # Set up logging
 os.makedirs("logs", exist_ok=True)
 logging.basicConfig(
@@ -73,13 +86,16 @@
 #global shutdown flag
 shutdown_in_progress = False
 
-# Initiate connection to Teradata
-_tdconn = td.TDConn()
+# Connect to the Teradata server
+# Initiate base connection
+_tdconn = td.TDConn(tdml_context=_requireTdmlContext)
+
+# If the feature store is enabled, set it up
+if any(re.match(pattern, 'fs_*') for pattern in config.get('tool',[])):
+    fs_config = td.FeatureStoreConfig() 
 
-#afm-defect:
-_enableEVS = False
-# Only attempt to connect to EVS is the system has an EVS installed/configured
-if (len(os.getenv("VS_NAME", "").strip()) > 0):
+# If the enterprise vector store is enabled, set it up
+if _enableEVS and (len(os.getenv("VS_NAME", "").strip()) > 0):
     try:
         _evs    = td.get_evs()
         _enableEVS = True
@@ -126,7 +142,7 @@ def execute_db_tool(tool, *args, **kwargs):
     # (Re)initialize if needed
     if not getattr(_tdconn, "engine", None):
         logger.info("Reinitializing TDConn")
-        _tdconn = td.TDConn()
+        _tdconn = td.TDConn(tdml_context=_requireTdmlContext)
 
     # Check is the first argument of the tool is a SQLAlchemy Connection
     sig = inspect.signature(tool)
@@ -180,17 +196,22 @@ def execute_vs_tool(tool, *args, **kwargs) -> ResponseType:
 def make_tool_wrapper(func):
     """
     Given a handle_* function, return an async wrapper with:
-    - the same signature minus the first 'connection' param
-    - a call into execute_db_tool
+    - the same signature minus any 'conn' or 'fs_config' params
+    - injection of fs_config when declared, while conn injection is handled by execute_db_tool
     """
-
     sig = inspect.signature(func)
-    # Drop first param (connection) and skip 'tool_name' if present
-    first_param = next(iter(sig.parameters))
+
+    # Determine which parameters to inject and remove from the exposed signature
+    inject_kwargs = {}
+    removable = {"conn", "tool_name"}
+    if "fs_config" in sig.parameters:
+        inject_kwargs["fs_config"] = fs_config
+        removable.add("fs_config")
+
+    # Build the new signature without injected params
     params = [
         p for name, p in sig.parameters.items()
-        if name != first_param
-        and name != "tool_name" # skip tool_name if present as this allows to override the tool name for query tools declared in # the *_objects.yml files
+        if name not in removable
         and p.kind in (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY)
     ]
     new_sig = sig.replace(parameters=params)
@@ -199,8 +220,8 @@ def make_tool_wrapper(func):
     async def wrapper(*args, **kwargs):
         ba = new_sig.bind_partial(*args, **kwargs)
         ba.apply_defaults()
-
-        return execute_db_tool(func, **ba.arguments)
+        # execute_db_tool handles injecting `conn`; we add fs_config if needed
+        return execute_db_tool(func, **inject_kwargs, **ba.arguments)
 
     wrapper.__signature__ = new_sig
     return wrapper
@@ -209,7 +230,7 @@ async def wrapper(*args, **kwargs):
 def register_td_tools(config, td, mcp):
     patterns = config.get('tool', [])
     for name, func in inspect.getmembers(td, inspect.isfunction):
-        if not name.startswith("handle_") or name.startswith("handle_fs_"):
+        if not name.startswith("handle_"): #or name.startswith("handle_fs_"):
             continue
 
         tool_name = name[len("handle_"):]
@@ -456,23 +477,21 @@ def get_glossary_term(term_name: str)  -> dict:
 
 
 
+@mcp.tool(description="Reconnect to the Teradata database if the connection is lost.")
+async def reconnect_to_database() -> ResponseType:
+    """Reconnect to Teradata database if connection is lost."""
+    global _tdconn
+    try:
+        _tdconn = td.TDConn(tdml_context=_requireTdmlContext)  # Recreate TDConn with tdml_context enabled
+        return format_text_response("Reconnected to Teradata database successfully.")
+    except Exception as e:
+        logger.error(f"Error reconnecting to database: {e}")
+        return format_error_response(str(e))
+        
 #--------------- Feature Store Tools ---------------#
 # Feature tools leveraging the tdfs4ds package.
 # Run only if the EFS tools are defined in the config
-if any(re.match(pattern, 'fs_*') for pattern in config.get('tool',[])):
-    fs_config = td.FeatureStoreConfig() 
-
-    @mcp.tool(description="Reconnect to the Teradata database if the connection is lost.")
-    async def reconnect_to_database() -> ResponseType:
-        """Reconnect to Teradata database if connection is lost."""
-        global _tdconn
-        try:
-            _tdconn = td.TDConn()
-            td.teradataml_connection()
-            return format_text_response("Reconnected to Teradata database successfully.")
-        except Exception as e:
-            logger.error(f"Error reconnecting to database: {e}")
-            return format_error_response(str(e))
+if _enableEFS:
 
     @mcp.tool(description="Set or update the feature store configuration (database and data domain).")
     async def fs_setFeatureStoreConfig(
@@ -519,53 +538,6 @@ async def fs_setFeatureStoreConfig(
     async def fs_getFeatureStoreConfig() -> ResponseType:
         return format_text_response(f"Current feature store config: {fs_config.dict(exclude_none=True)}")
 
-    @mcp.tool(description=("Check if a feature store is present in the specified database." ))
-    async def fs_isFeatureStorePresent(
-        db_name: str = Field(..., description="Name of the database to check for a feature store.")
-    ) -> ResponseType:
-        return execute_db_tool(td.handle_fs_isFeatureStorePresent, db_name=db_name)
-
-    @mcp.tool(description=("Returns a summary of the feature store content. Use this to understand what data is available in the feature store"))
-    async def fs_featureStoreContent(
-    ) -> ResponseType:
-        return execute_db_tool(td.handle_fs_featureStoreContent, fs_config=fs_config)
-
-    @mcp.tool(description=( "List the available data domains. Requires a configured `db_name`  in the feature store config. Use this to explore which entities can be used when building a dataset."))
-    async def fs_getDataDomains(
-        entity: str = Field(..., description="The .")
-    ) -> ResponseType:
-        return execute_db_tool(td.handle_fs_getDataDomains, fs_config=fs_config)
-
-    @mcp.tool(description=("List the list of features. Requires a configured `db_name` and  `data_domain` in the feature store config. Use this to explore the features available ."))
-    async def fs_getFeatures(
-    ) -> ResponseType:
-        return execute_db_tool(td.handle_fs_getFeatures, fs_config=fs_config)
-
-    @mcp.tool(description=("List the list of available datasets.Requires a configured `db_name` in the feature store config.Use this to explore the datasets that are available ."))
-    async def fs_getAvailableDatasets(
-    ) -> ResponseType:
-        return execute_db_tool(td.handle_fs_getAvailableDatasets, fs_config=fs_config)
-
-    @mcp.tool(description=("Return the schema of the feature store.Requires a feature store in the configured database (`db_name`)."))
-    async def fs_getFeatureDataModel(
-    ) -> ResponseType:
-        return execute_db_tool(td.handle_fs_getFeatureDataModel, fs_config=fs_config)
-
-
-    @mcp.tool(description=("List the available entities for a given data domain. Requires a configured `db_name` and `data_domain` and  `entity` in the feature store config. Use this to explore which entities can be used when building a dataset."))
-    async def fs_getAvailableEntities(
-    ) -> ResponseType:
-        return execute_db_tool(td.handle_fs_getAvailableEntities, fs_config=fs_config)
-
-    @mcp.tool( description=("Create a dataset using selected features and an entity from the feature store. The dataset is created in the specified target database under the given name. Requires a configured feature store and data domain. Registers the dataset in the catalog automatically. Use this when you want to build and register a new dataset for analysis or modeling." ) )
-    async def fs_createDataset(
-        entity_name: str = Field(..., description="Entity for which the dataset will be created. Available entities are reported in the feature catalog."),
-        feature_selection: list[str] = Field(..., description="List of features to include in the dataset. Available features are reported in the feature catalog."),
-        dataset_name: str = Field(..., description="Name of the dataset to create."),
-        target_database: str = Field(..., description="Target database where the dataset will be created.")
-    ) -> ResponseType:
-        return execute_db_tool(td.handle_fs_createDataset, fs_config=fs_config, entity_name=entity_name, feature_selection=feature_selection, dataset_name=dataset_name, target_database=target_database)
-
 #------------------ Main ------------------#
 # Main function to start the MCP server
 #     Description: Initializes the MCP server and sets up signal handling for graceful shutdown.

src/teradata_mcp_server/tools/fs/fs_tools.py

@@ -21,6 +21,11 @@
 #       db_name - the database name to check for existenceAdd commentMore actions
 # #     Returns: True or False    
 def handle_fs_isFeatureStorePresent(conn: TeradataConnection, db_name: str, *args, **kwargs):
+    """ Check if a feature store is present in the specified database.
+    
+    Args:
+        db_name (str): The name of the database to check for the feature store.
+    """
 
     logger.info(f"Tool: handle_fs_isFeatureStorePresent: Args: db_name: {db_name}")
 
@@ -44,6 +49,10 @@ def handle_fs_isFeatureStorePresent(conn: TeradataConnection, db_name: str, *arg
 #       conn (TeradataConnection) - Teradata connection object for executing SQL queries
 # #     Returns: True or False    
 def handle_fs_getDataDomains(conn: TeradataConnection, fs_config, *args, **kwargs):
+    """ 
+    List the available data domains. Requires a configured `db_name`  in the feature store config. Use this to explore which entities can be used when building a dataset.
+    """
+
     db_name = fs_config.db_name
     logger.info(f"Tool: handle_fs_getDataDomains: Args: db_name: {db_name}")
 
@@ -86,6 +95,10 @@ def handle_fs_getDataDomains(conn: TeradataConnection, fs_config, *args, **kwarg
 #       conn (TeradataConnection) - Teradata connection object for executing SQL queries
 # #     Returns: True or False    
 def handle_fs_featureStoreContent(conn: TeradataConnection, fs_config, *args, **kwargs):
+    """ 
+    Returns a summary of the feature store content. Use this to understand what data is available in the feature store.
+    """
+
     db_name = fs_config.db_name
     logger.info(f"Tool: handle_fs_featureStoreContent: Args: db_name: {db_name}")
     metadata = {
@@ -125,6 +138,10 @@ def handle_fs_featureStoreContent(conn: TeradataConnection, fs_config, *args, **
 #       db_name - the database name to check for existence
 # #     Returns: the feature store schema, mainly the catalogs    
 def handle_fs_getFeatureDataModel(conn: TeradataConnection, fs_config, *args, **kwargs):
+    """ 
+    Returns the feature store data model, including the feature catalog, process catalog, and dataset catalog.
+    """
+
     db_name = fs_config.db_name
     logger.info(f"Tool: handle_fs_getFeatureDataModel: Args: db_name: {db_name}")
 
@@ -166,6 +183,9 @@ def handle_fs_getFeatureDataModel(conn: TeradataConnection, fs_config, *args, **
 #       db_name - the database name to check for existence
 # #     Returns: True or False    
 def handle_fs_getAvailableEntities(conn: TeradataConnection, fs_config, *args, **kwargs):
+    """ 
+    List the available entities for a given data domain. Requires a configured `db_name` and `data_domain` and  `entity` in the feature store config. Use this to explore which entities can be used when building a dataset.
+    """
     db_name = fs_config.db_name
     logger.info(f"Tool: handle_fs_getAvailableEntities: Args: db_name: {db_name}")
 
@@ -212,6 +232,9 @@ def handle_fs_getAvailableEntities(conn: TeradataConnection, fs_config, *args, *
 #       db_name - the database name to check for existence
 # #     Returns: True or False    
 def handle_fs_getAvailableDatasets(conn: TeradataConnection, fs_config, *args, **kwargs):
+    """ 
+    List the list of available datasets.Requires a configured `db_name` in the feature store config.Use this to explore the datasets that are available .
+    """
     db_name = fs_config.db_name
     logger.info(f"Tool: handle_fs_getAvailableDatasets: Args: db_name: {db_name}")
 
@@ -245,6 +268,10 @@ def handle_fs_getAvailableDatasets(conn: TeradataConnection, fs_config, *args, *
 #       db_name - the database name to check for existence
 # #     Returns: True or False    
 def handle_fs_getFeatures(conn: TeradataConnection, fs_config, *args, **kwargs):
+    """ 
+    List the list of features. Requires a configured `db_name` and  `data_domain` in the feature store config. Use this to explore the features available .
+    """
+
     db_name = fs_config.db_name
     logger.info(f"Tool: handle_fs_getFeatures: Args: db_name: {db_name}")
 
@@ -305,6 +332,15 @@ def handle_fs_getFeatures(conn: TeradataConnection, fs_config, *args, **kwargs):
 #       db_name - the database name to check for existence
 # #     Returns: True or False    
 def handle_fs_createDataset(conn: TeradataConnection, fs_config, entity_name: str, feature_selection: str, dataset_name: str, target_database: str, *args, **kwargs):
+    """ 
+    Create a dataset using selected features and an entity from the feature store. The dataset is created in the specified target database under the given name. Requires a configured feature store and data domain. Registers the dataset in the catalog automatically. Use this when you want to build and register a new dataset for analysis or modeling.
+    Args:
+        entity_name (str): Entity for which the dataset will be created. Available entities are reported in the feature catalog.
+        feature_selection (list[str]): List of features to include in the dataset. Available features are reported in the feature catalog.
+        dataset_name (str): The name of the dataset to create.
+        target_database (str): The database where the dataset will be created.
+    """
+
     db_name = fs_config.db_name
     logger.info(f"Tool: handle_fs_createDataset: Args: db_name: {db_name}")
 

src/teradata_mcp_server/tools/td_connect.py

@@ -21,11 +21,12 @@
 class TDConn:
     engine: Optional[Engine] = None
     connection_url: Optional[str] = None
+    tdml_context: Optional[bool] = False
 
     # Constructor
     #     It will read the connection URL from the environment variable DATABASE_URI
     #     It will parse the connection URL and create a SQLAlchemy engine connected to the database
-    def __init__(self, connection_url: Optional[str] = None):
+    def __init__(self, connection_url: Optional[str] = None, tdml_context: Optional[bool] = False):
         if connection_url is None and os.getenv("DATABASE_URI") is None:
             logger.warning("DATABASE_URI is not specified, database connection will not be established.")
             self.engine = None
@@ -64,9 +65,10 @@ def __init__(self, connection_url: Optional[str] = None):
                 logger.error(f"Error creating database engine: {e}")
                 self.engine = None
 
-            # Create the teradataml context 
-            if "<EVS or EFS enabled>":
-                import teradataml as tdml # import of the teradataml package
+            # Create the teradataml context
+            self.tdml_context = tdml_context
+            if self.tdml_context:
+                import teradataml as tdml  # import of the teradataml package
                 tdml.create_context(tdsqlengine=self.engine)
 
     # Destructor

test/EFS/efs_mcp_test.py

@@ -0,0 +1,58 @@
+import tdfs4ds
+from tdfs4ds.utils.lineage import crystallize_view
+from teradataml import create_context, DataFrame, execute_sql, db_list_tables
+from urllib.parse import urlparse
+import argparse
+import os
+
+parser = argparse.ArgumentParser(description="Teradata MCP Server")
+parser.add_argument('--database_uri', type=str, required=False, help='Database URI to connect to: teradata://username:password@host:1025/schemaname')
+parser.add_argument('--action', type=str, choices=['setup', 'cleanup'], required=True, help='Action to perform: setup or cleanup')
+# Extract known arguments and load them into the environment if provided
+args, unknown = parser.parse_known_args()
+
+database_name = 'demo_user'
+connection_url = args.database_uri or os.getenv("DATABASE_URI")
+
+if not connection_url:
+    raise ValueError("DATABASE_URI must be provided either as an argument or as an environment variable.")
+
+parsed_url = urlparse(connection_url)
+user = parsed_url.username
+password = parsed_url.password
+host = parsed_url.hostname
+port = parsed_url.port or 1025
+database = parsed_url.path.lstrip('/') or user
+
+eng = create_context(host = host, username=user, password = password)
+
+if args.action=='setup':
+    # Set up the feature store
+    tdfs4ds.setup(database=database_name)
+    tdfs4ds.connect(database=database_name)
+
+    # Define the feature store domain
+    tdfs4ds.DATA_DOMAIN='demo_dba'
+    tdfs4ds.VARCHAR_SIZE=50
+
+    # Create features (table space and skew)
+    df=DataFrame.from_query("SELECT databasename||'.'||tablename tablename, SUM(currentperm) currentperm, CAST((100-(AVG(currentperm)/MAX(currentperm)*100)) AS DECIMAL(5,2)) AS skew_pct FROM dbc.tablesizev GROUP BY 1;")
+    df = crystallize_view(df, view_name = 'efs_demo_dba_space', schema_name = database_name,output_view=True)
+
+    # upload the features in the physical feature store
+    tdfs4ds.upload_features(
+        df,
+        entity_id     = ['tablename'],
+        feature_names = df.columns[1::],
+        metadata      = {'project': 'dba'}
+    )
+
+    # Display our features
+    tdfs4ds.feature_catalog()
+
+elif args.action=='cleanup':
+    list_of_tables = db_list_tables()
+    [execute_sql(f"DROP VIEW {database_name}.{t}") for t in list_of_tables.TableName if t.startswith('FS_V')]
+    [execute_sql(f"DROP TABLE {database_name}.{t}") for t in list_of_tables.TableName if t.startswith('FS_T')]
+    [execute_sql(f"DROP TABLE {database_name}.{t}") for t in list_of_tables.TableName if t.startswith('FS_T')]
+    [execute_sql(f"DROP TABLE {database_name}.{t}") for t in list_of_tables.TableName if t.startswith('FS_T')]

test/EFS/requirements.txt

@@ -0,0 +1,5 @@
+tdfs4ds
+teradataml
+pandas
+json
+sqlalchemy