added readme. Fixed handle_fs_createDataset

Author: remi-td

src/teradata_mcp_server/tools/fs/fs_tools.py

@@ -330,7 +330,7 @@ def handle_fs_getFeatures(conn: TeradataConnection, fs_config, *args, **kwargs):
 #       conn (TeradataConnection) - Teradata connection object for executing SQL queries
 #       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):
+def handle_fs_createDataset(conn: TeradataConnection, fs_config, entity_name: str, feature_selection: list[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:

test/EFS/Readme.md

@@ -0,0 +1,63 @@
+
+
+# EFS Test Script — Teradata Enterprise Feature Store
+
+## Purpose
+This script (`efs_mcp_test.py`) demonstrates and tests the Teradata MCP Server's functionality for the **Teradata Enterprise Feature Store** (EFS). 
+
+It runs through the following EFS functions:
+1. `fs_isFeatureStorePresent`
+2. `fs_setFeatureStoreConfig`
+3. `fs_getDataDomains`
+4. `fs_getAvailableDatasets`
+5. `fs_getAvailableEntities`
+6. `fs_getFeatures`
+7. `fs_createDataset`
+
+## Actions
+- **setup** – Create a demo feature store schema and load example features.  
+- **test** – Call the EFS MCP tools above in sequence to verify functionality.  
+- **cleanup** – Drop the demo feature store objects.
+
+
+## Requirements
+- Python 3.11+
+- Access to a Teradata system
+- MCP server exposing `fs_*` tools at `http://127.0.0.1:8001/mcp`
+- Python packages: `tdfs4ds`, `teradataml`, `langchain_mcp_adapters`
+
+## Setup
+Create and activate a virtual environment, then install the required packages:
+
+```bash
+python -m venv .venv
+source .venv/bin/activate  # On Windows use: .venv\Scripts\activate
+pip install --upgrade pip
+pip install -r requirements.txt
+```
+
+## Usage
+
+Make sure that you have the Teradata MCP server with EFS tools running on `127.0.0.1:8001/mcp` eg.
+
+```bash
+uv run teradata-mcp-server --profile dataScientist --mcp_transport streamable-http --mcp_port 8001
+```
+
+Use the test script
+
+```bash
+# Setup demo feature store in your schema
+python test/EFS/efs_mcp_test.py --action setup --database_uri "teradata://user:pass@host:1025/schema"
+
+# Run the tests
+python test/EFS/efs_mcp_test.py --action test
+
+# Clean up: drop the feature store
+python test/EFS/efs_mcp_test.py --action cleanup --database_uri "teradata://user:pass@host:1025/schema"
+```
+
+## Notes
+- Defaults: `database_name = demo_user`, `data_domain = demo_dba`.
+- Dataset created in test: `test_dataset`.
+- `feature_selection` must be a string (e.g. `"col1,col2"`), not a list.

test/EFS/efs_mcp_test.py

@@ -19,17 +19,18 @@ def main():
     data_domain = 'demo_dba'
     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.")
+    if args.action in ['setup', 'cleanup']:
+        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
+        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)
+        eng = create_context(host = host, username=user, password = password)
 
     if args.action=='setup':
         # Set up the feature store
@@ -64,18 +65,123 @@ def main():
                 "transport": "streamable_http"
             }
         })
-        async def _test():  # small async helper
+        async def runner():
             async with mcp_client.session("mcp_server") as mcp_session:
                 tools = await load_mcp_tools(mcp_session)
                 fs_tools = [t for t in tools if t.name.startswith('fs_')]
                 print("Available fs_ tools:", [t.name for t in fs_tools])
-                fs_set_tool = next((t for t in fs_tools if t.name == 'fs_setFeatureStoreConfig'), None)
-                if not fs_set_tool:
-                    raise RuntimeError('fs_setFeatureStoreConfig tool not found')
-                response = await fs_set_tool.arun({"data_domain": data_domain, "db_name": database_name})
-                print("fs_setFeatureStoreConfig response:", response)
-        import asyncio
-        asyncio.run(_test())
+
+                # Map tool names for quick access
+                tool_by_name = {t.name: t for t in fs_tools}
+
+                import json as _json
+
+                async def _call(name: str, payload: dict | None = None):
+                    if name not in tool_by_name:
+                        raise RuntimeError(f"Tool {name} not found")
+                    tool = tool_by_name[name]
+                    tool_input = payload or {}
+                    # StructuredTool expects a single positional/named argument: tool_input
+                    resp = await tool.arun(tool_input=tool_input)
+                    # Try to parse JSON text if needed
+                    if isinstance(resp, str):
+                        try:
+                            return _json.loads(resp)
+                        except Exception:
+                            return resp
+                    return resp
+
+                # 1) fs_isFeatureStorePresent
+                print("\n[1/8] fs_isFeatureStorePresent…")
+                r1 = await _call('fs_isFeatureStorePresent', {"db_name": database_name})
+                print("fs_isFeatureStorePresent →", r1)
+
+                # 2) fs_setFeatureStoreConfig
+                print("\n[2/8] fs_setFeatureStoreConfig…")
+                r_set = await _call('fs_setFeatureStoreConfig', {"db_name": database_name, "data_domain": data_domain})
+                print("fs_setFeatureStoreConfig →", r_set)
+
+                # 3) fs_getDataDomains
+                print("\n[3/8] fs_getDataDomains…")
+                r2 = await _call('fs_getDataDomains')
+                print("fs_getDataDomains →", r2)
+
+                # 4) fs_getAvailableDatasets
+                print("\n[4/8] fs_getAvailableDatasets…")
+                r3 = await _call('fs_getAvailableDatasets')
+                print("fs_getAvailableDatasets →", r3)
+
+                # 5) fs_getAvailableEntities
+                print("\n[5/8] fs_getAvailableEntities…")
+                r4 = await _call('fs_getAvailableEntities')
+                print("fs_getAvailableEntities →", r4)
+
+                # 6) fs_setFeatureStoreConfig (entity)
+                print("\n[6/8] fs_setFeatureStoreConfig (entity)…")
+                def _extract_entity_name(payload):
+                    res = payload.get("results") if isinstance(payload, dict) else payload
+                    if isinstance(res, list) and res and isinstance(res[0], dict):
+                        for key in ("ENTITY_NAME", "entity_name", "entity", "name"):
+                            if key in res[0] and res[0][key]:
+                                return res[0][key]
+                    if isinstance(res, str):
+                        lines = [ln.strip() for ln in res.splitlines() if ln.strip()]
+                        if lines:
+                            parts = lines[-1].split()
+                            if parts:
+                                return parts[-1]
+                    return None
+                entity_name = _extract_entity_name(r4) or "tablename"
+                r_set_entity = await _call('fs_setFeatureStoreConfig', {"entity": entity_name})
+                print("fs_setFeatureStoreConfig (entity) →", r_set_entity)
+
+                # 7) fs_getFeatures
+                print("\n[7/8] fs_getFeatures…")
+                r5 = await _call('fs_getFeatures')
+                print("fs_getFeatures →", r5)
+
+                # Extract feature names from r5
+                def _extract_feature_names(payload):
+                    # Accept either {"results": [...]} or a raw list
+                    items = payload.get("results") if isinstance(payload, dict) else payload
+                    if not isinstance(items, list):
+                        return []
+                    names = []
+                    for row in items:
+                        if not isinstance(row, dict):
+                            continue
+                        for key in ("feature_name", "FEATURE_NAME", "name", "FEATURE", "feature"):
+                            if key in row and row[key] is not None:
+                                names.append(row[key])
+                                break
+                    return names
+
+                feature_selection = _extract_feature_names(r5)
+                print(f"Extracted {len(feature_selection)} feature names for dataset creation")
+
+                # 8) fs_createDataset
+                print("\n[8/8] fs_createDataset…")
+                create_payload = {
+                    "entity_name": entity_name,
+                    "feature_selection": feature_selection,
+                    "dataset_name": "test_efs_dataset",
+                    "target_database": database_name,
+                }
+                print("fs_createDataset payload:", create_payload)
+                r6 = await _call('fs_createDataset', create_payload)
+
+                # If tool returned an error payload (not exception), also retry with CSV features
+                if isinstance(r6, dict) and isinstance(r6.get("results"), dict):
+                    err = r6["results"].get("error")
+                    if isinstance(err, str) and ("NoneType" in err or "string" in err):
+                        create_payload_retry = dict(create_payload)
+                        create_payload_retry["feature_selection"] = ",".join(feature_selection)
+                        print("Retrying fs_createDataset with CSV features (error payload):", create_payload_retry)
+                        r6 = await _call('fs_createDataset', create_payload_retry)
+
+                print("fs_createDataset →", r6)
+
+        asyncio.run(runner())
 
     elif args.action=='cleanup':
         list_of_tables = db_list_tables()

test/EFS/requirements.txt

@@ -1,5 +1,7 @@
 tdfs4ds
 teradataml
 pandas
-json
-sqlalchemy
+sqlalchemy
+langchain_mcp_adapters
+mcp
+anyio