removed teradataml context from TDConn

Author: remi-td

src/teradata_mcp_server/server.py

@@ -11,7 +11,6 @@
 from mcp.server.fastmcp.prompts.base import UserMessage, TextContent
 from dotenv import load_dotenv
 import tdfs4ds
-import teradataml as tdml
 import inspect
 from sqlalchemy.engine import Connection
 import argparse
@@ -67,8 +66,6 @@
 else:
     _enableEVS = False
 
-_requireTdmlContext = _enableEFS or _enableEVS
-
 # Set up logging
 os.makedirs("logs", exist_ok=True)
 logging.basicConfig(
@@ -88,11 +85,13 @@
 
 # Connect to the Teradata server
 # Initiate base connection
-_tdconn = td.TDConn(tdml_context=_requireTdmlContext)
+_tdconn = td.TDConn()
 
 # If the feature store is enabled, set it up
-if any(re.match(pattern, 'fs_*') for pattern in config.get('tool',[])):
+if _enableEFS:
     fs_config = td.FeatureStoreConfig() 
+    import teradataml as tdml  # import of the teradataml package
+    tdml.create_context(tdsqlengine=_tdconn.engine)
 
 # If the enterprise vector store is enabled, set it up
 if _enableEVS and (len(os.getenv("VS_NAME", "").strip()) > 0):
@@ -142,7 +141,12 @@ def execute_db_tool(tool, *args, **kwargs):
     # (Re)initialize if needed
     if not getattr(_tdconn, "engine", None):
         logger.info("Reinitializing TDConn")
-        _tdconn = td.TDConn(tdml_context=_requireTdmlContext)
+        _tdconn = td.TDConn()
+        if _enableEFS:
+            fs_config = td.FeatureStoreConfig() 
+            import teradataml as tdml  # import of the teradataml package
+            tdml.create_context(tdsqlengine=_tdconn.engine)
+
 
     # Check is the first argument of the tool is a SQLAlchemy Connection
     sig = inspect.signature(tool)

src/teradata_mcp_server/tools/fs/fs_tools.py

@@ -9,6 +9,7 @@
 
 
 logger = logging.getLogger("teradata_mcp_server")
+
 from teradata_mcp_server.tools.utils import serialize_teradata_types, rows_to_json, create_response
 
 #------------------ Do not make changes above  ------------------#
@@ -235,11 +236,12 @@ def handle_fs_getAvailableDatasets(conn: TeradataConnection, fs_config, *args, *
     """ 
     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}")
 
     is_a_feature_store = False
-    
+
     try:
         is_a_feature_store = tdfs4ds.connect(database=db_name)
     except Exception as e:
@@ -248,13 +250,13 @@ def handle_fs_getAvailableDatasets(conn: TeradataConnection, fs_config, *args, *
 
     if not is_a_feature_store:
         return create_response({"error": f"There is no feature store in {db_name}"}, {"tool_name": "handle_fs_getAvailableDatasets", "db_name": db_name})
-    
+
     try:
         data = tdfs4ds.dataset_catalog().to_pandas()
     except Exception as e:
         logger.error(f"Error retrieving available datasets: {e}")
         return create_response({"error": str(e)}, {"tool_name": "handle_fs_getAvailableDatasets", "db_name": db_name})
-        
+
     metadata = {
         "tool_name": "fs_getAvailableDatasets",
         "db_name": db_name,

src/teradata_mcp_server/tools/fs/fs_utils.py

@@ -6,7 +6,6 @@
 from sqlalchemy.engine import Connection
 from sqlalchemy import text
 
-
 logger = logging.getLogger("teradata_mcp_server")
 
 class FeatureStoreConfig(BaseModel):
@@ -61,6 +60,7 @@ def fs_setFeatureStoreConfig(
         data_domain: Optional[str] = None,
         entity: Optional[str] = None,
     ) -> "FeatureStoreConfig":
+
         if db_name:
             if tdfs4ds.connect(database=db_name):
                 logger.info(f"connected to the feature store of the {db_name} database")

src/teradata_mcp_server/tools/td_connect.py

@@ -21,7 +21,6 @@
 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
@@ -67,9 +66,6 @@ def __init__(self, connection_url: Optional[str] = None, tdml_context: Optional[
 
             # 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
     #     It will close the SQLAlchemy connection and engine