moved EFS config setting method to the EFS tools and use the SQLalchemy connection

Author: remi-td

src/teradata_mcp_server/server.py

@@ -498,45 +498,18 @@ async def fs_setFeatureStoreConfig(
         data_domain: Optional[str] = None,
         db_name: Optional[str] = None,
         entity: Optional[str] = None,
-    ) -> ResponseType:
-        if db_name:
-            if tdfs4ds.connect(database=db_name):
-                logger.info(f"connected to the feature store of the {db_name} database")
-                # Reset data_domain if DB name changes
-                if not (fs_config.db_name and fs_config.db_name.upper() == db_name.upper()):
-                    fs_config.data_domain = None
-                
-                fs_config.db_name = db_name
-                logger.info(f"connected to the feature store of the {db_name} database")
-                fs_config.feature_catalog = f"{db_name}.{tdfs4ds.FEATURE_CATALOG_NAME_VIEW}"
-                logger.info(f"feature catalog {fs_config.feature_catalog}")
-                fs_config.process_catalog = f"{db_name}.{tdfs4ds.PROCESS_CATALOG_NAME_VIEW}"
-                logger.info(f"process catalog {fs_config.process_catalog}")
-                fs_config.dataset_catalog = f"{db_name}.FS_V_FS_DATASET_CATALOG"  # <- fixed line
-                logger.info(f"dataset catalog {fs_config.dataset_catalog}")
-
-        if fs_config.db_name is not None and data_domain is not None:
-            sql_query_ = f"SEL count(*) AS N FROM {fs_config.feature_catalog} WHERE UPPER(data_domain) = '{data_domain.upper()}'"
-            logger.info(f"{sql_query_}")
-            result = tdml.execute_sql(sql_query_)
-            logger.info(f"{result}")
-            if result.fetchall()[0][0] > 0:
-                fs_config.data_domain = data_domain
-            else:
-                fs_config.data_domain = None
-
-        if fs_config.db_name is not None and fs_config.data_domain is not None and entity is not None:
-            sql_query_ = f"SEL count(*) AS N FROM {fs_config.feature_catalog} WHERE UPPER(data_domain) = '{data_domain.upper()}' AND ENTITY_NAME = '{entity.upper()}'"
-            logger.info(f"{sql_query_}")
-            result = tdml.execute_sql(sql_query_)
-            logger.info(f"{result}")
-            if result.fetchall()[0][0] > 0:
-                fs_config.entity = entity
-        return format_text_response(f"Feature store config updated: {fs_config.dict(exclude_none=True)}")
+    ) -> td.FeatureStoreConfig:
+        with _tdconn.engine.connect() as conn:
+            return fs_config.fs_setFeatureStoreConfig(
+                conn=conn,
+                db_name=db_name,
+                data_domain=data_domain,
+                entity=entity,
+            )
 
     @mcp.tool(description="Display the current feature store configuration (database and data domain).")
     async def fs_getFeatureStoreConfig() -> ResponseType:
-        return format_text_response(f"Current feature store config: {fs_config.dict(exclude_none=True)}")
+        return format_text_response(f"Current feature store config: {fs_config.model_dump(exclude_none=True)}")
 
 #------------------ Main ------------------#
 # Main function to start the MCP server

src/teradata_mcp_server/tools/fs/fs_utils.py

@@ -1,5 +1,13 @@
 from typing import Any, List, Optional
+from typing import Optional
 from pydantic import Field, BaseModel
+import tdfs4ds
+import logging
+from sqlalchemy.engine import Connection
+from sqlalchemy import text
+
+
+logger = logging.getLogger("teradata_mcp_server")
 
 class FeatureStoreConfig(BaseModel):
     """
@@ -45,3 +53,52 @@ class FeatureStoreConfig(BaseModel):
             "Used to list and manage available datasets within the feature store."
         )
     )
+
+    def fs_setFeatureStoreConfig(
+        self,
+        conn: Connection,
+        db_name: Optional[str] = None,
+        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")
+                # Reset data_domain if DB name changes
+                if not (self.db_name and self.db_name.upper() == db_name.upper()):
+                    self.data_domain = None
+                
+                self.db_name = db_name
+                logger.info(f"connected to the feature store of the {db_name} database")
+                self.feature_catalog = f"{db_name}.{tdfs4ds.FEATURE_CATALOG_NAME_VIEW}"
+                logger.info(f"feature catalog {self.feature_catalog}")
+                self.process_catalog = f"{db_name}.{tdfs4ds.PROCESS_CATALOG_NAME_VIEW}"
+                logger.info(f"process catalog {self.process_catalog}")
+                self.dataset_catalog = f"{db_name}.FS_V_FS_DATASET_CATALOG"  # <- fixed line
+                logger.info(f"dataset catalog {self.dataset_catalog}")
+
+        if self.db_name is not None and data_domain is not None:
+            stmt = text(
+                f"SELECT COUNT(*) AS N FROM {self.feature_catalog} "
+                "WHERE UPPER(data_domain)=:domain"
+            )
+            result = conn.execute(stmt, {"domain": data_domain.upper()})
+            count = result.scalar_one_or_none() or 0
+            logger.info("Found %d matching data_domain rows", count)
+            if count > 0:
+                self.data_domain = data_domain
+            else:
+                self.data_domain = None
+
+        if self.db_name is not None and self.data_domain is not None and entity is not None:
+            stmt = text(
+                f"SELECT COUNT(*) AS N FROM {self.feature_catalog} "
+                "WHERE UPPER(data_domain)=:domain "
+                "AND ENTITY_NAME=:entity"
+            )
+            result = conn.execute(stmt, {"domain": self.data_domain.upper(), "entity": entity.upper()})
+            count = result.scalar_one_or_none() or 0
+            logger.info("Found %d matching entity rows", count)
+            if count > 0:
+                self.entity = entity
+        return self