feat: add MCP tools for update data source, get/drop all data sources, run gsql, and list metadata

Author: xuanleilin

CHANGELOG.md

@@ -4,6 +4,8 @@
 ### 0.1.5
 - feat: add example using LangGraph LLM and agents to call TigerGraph-MCP
 - feat: upgrade TigerGraphX to v0.2.9 with new TigerGraphDatabase class integration
+- feat: add run algorithm subgraph to onboarding workflow
+- feat: add MCP tools for update data source, get/drop all data sources, run gsql, and list metadata
 
 ### 0.1.4
 - feat: enable multi-turn human-in-the-loop schema confirmation for iterative refinement

poetry.lock

@@ -10,7 +10,7 @@ requires-python = ">=3.10,<3.13"
 
 dependencies = [
     "mcp (>=1.9.2,<2.0.0)",
-    "tigergraphx (>=0.2.9,<0.3.0)",
+    "tigergraphx (>=0.2.10,<0.3.0)",
 ]
 
 [tool.poetry.group.dev.dependencies]

pyproject.toml

@@ -9,38 +9,67 @@
 
 class TestDataSourceTools(BaseFixture):
     @pytest.mark.asyncio
-    async def test_data_source_tools(self):
-        SAMPLE_PATH = "s3a://tigergraph-solution-kits/connected_customer/customer_360/data/Session.csv"  # Replace with actual path
+    async def test_data_source_tools_basic(self):
+        SAMPLE_PATH = "s3a://tigergraph-solution-kits/connected_customer/customer_360/data/Session.csv"
 
         async with stdio_client(self.server_params) as (read, write):
             async with ClientSession(read, write) as session:
                 await session.initialize()
 
-                # Create the data source
-                create_result = await session.call_tool(
-                    TigerGraphToolName.CREATE_DATA_SOURCE,
-                    arguments={
-                        "name": "data_source_1",
-                        "data_source_type": "s3",
-                        "access_key": "",
-                        "secret_key": "",
-                        "extra_config": {
-                            "file.reader.settings.fs.s3a.aws.credentials.provider": "org.apache.hadoop.fs.s3a.AnonymousAWSCredentialsProvider"
+                try:
+                    # Create the data source
+                    create_result = await session.call_tool(
+                        TigerGraphToolName.CREATE_DATA_SOURCE,
+                        arguments={
+                            "name": "data_source_1",
+                            "data_source_type": "s3",
+                            "access_key": "",
+                            "secret_key": "",
+                            "extra_config": {
+                                "file.reader.settings.fs.s3a.aws.credentials.provider": "org.apache.hadoop.fs.s3a.AnonymousAWSCredentialsProvider"
+                            },
                         },
-                    },
-                )
-                assert "✅ Successfully created data source" in str(create_result)
+                    )
+                    assert "✅ Successfully created data source" in str(create_result)
 
-                try:
                     # Get the data source
                     get_result = await session.call_tool(
                         TigerGraphToolName.GET_DATA_SOURCE,
+                        arguments={"name": "data_source_1"},
+                    )
+                    get_result_text = str(get_result)
+                    assert (
+                        "✅ Successfully retrieved configuration for data source"
+                        in get_result_text
+                    )
+
+                    # Update the data source
+                    update_result = await session.call_tool(
+                        TigerGraphToolName.UPDATE_DATA_SOURCE,
                         arguments={
                             "name": "data_source_1",
+                            "data_source_type": "s3",
+                            "access_key": "updated-key",
+                            "secret_key": "updated-secret",
+                            "extra_config": {
+                                "file.reader.settings.fs.s3a.aws.credentials.provider": "org.apache.hadoop.fs.s3a.AnonymousAWSCredentialsProvider"
+                            },
                         },
                     )
-                    get_result_text = str(get_result)
-                    assert "✅ Successfully retrieved configuration for data source" in get_result_text
+                    assert "✅ Successfully updated data source" in str(update_result)
+
+                    # Get the updated data source to verify
+                    get_updated_result = await session.call_tool(
+                        TigerGraphToolName.GET_DATA_SOURCE,
+                        arguments={"name": "data_source_1"},
+                    )
+                    get_updated_text = str(get_updated_result)
+                    assert (
+                        "✅ Successfully retrieved configuration for data source"
+                        in get_updated_text
+                    )
+                    assert "updated-key" in get_updated_text
+                    assert "updated-secret" in get_updated_text
 
                     # Preview sample data
                     preview_result = await session.call_tool(
@@ -66,3 +95,52 @@ async def test_data_source_tools(self):
                         arguments={"name": "data_source_1"},
                     )
                     assert "Successfully dropped data source" in str(drop_result)
+
+    @pytest.mark.asyncio
+    async def test_data_source_list_and_drop_all(self):
+        async with stdio_client(self.server_params) as (read, write):
+            async with ClientSession(read, write) as session:
+                await session.initialize()
+
+                # Create two data sources
+                try:
+                    for i in [1, 2]:
+                        result = await session.call_tool(
+                            TigerGraphToolName.CREATE_DATA_SOURCE,
+                            arguments={
+                                "name": f"tmp_data_source_{i}",
+                                "data_source_type": "s3",
+                                "access_key": "",
+                                "secret_key": "",
+                                "extra_config": {
+                                    "file.reader.settings.fs.s3a.aws.credentials.provider": "org.apache.hadoop.fs.s3a.AnonymousAWSCredentialsProvider"
+                                },
+                            },
+                        )
+                        assert (
+                            f"✅ Successfully created data source 'tmp_data_source_{i}'"
+                            in str(result)
+                        )
+
+                    # Get all data sources
+                    get_all_result = await session.call_tool(
+                        TigerGraphToolName.GET_ALL_DATA_SOURCES,
+                        arguments={},
+                    )
+                    get_all_text = str(get_all_result)
+                    assert (
+                        "✅ Successfully retrieved all data sources:"
+                        in get_all_text
+                    )
+                    assert "tmp_data_source_1" in get_all_text
+                    assert "tmp_data_source_2" in get_all_text
+
+                finally:
+                    # Drop all data sources
+                    drop_all_result = await session.call_tool(
+                        TigerGraphToolName.DROP_ALL_DATA_SOURCES,
+                        arguments={},
+                    )
+                    assert "✅ All data sources is dropped successfully." in str(
+                        drop_all_result
+                    )

tests/integration/tools/db/data_source_tools_test.py

@@ -0,0 +1,39 @@
+import pytest
+
+from mcp import ClientSession
+from mcp.client.stdio import stdio_client
+
+from tests.integration.base_graph_fixture import BaseFixture
+from tigergraph_mcp import TigerGraphToolName
+
+
+class TestGSQLTool(BaseFixture):
+    @pytest.mark.asyncio
+    async def test_gsql_tool(self):
+        async with stdio_client(self.server_params) as (read, write):
+            async with ClientSession(read, write) as session:
+                await session.initialize()
+
+                result = await session.call_tool(
+                    TigerGraphToolName.GSQL,
+                    arguments={"command": "LS"},
+                )
+
+                result_text = str(result)
+                assert "✅ GSQL Response" in result_text
+                assert "Graphs" in result_text or "Vertex Types" in result_text
+
+    @pytest.mark.asyncio
+    async def test_list_metadata_tool(self):
+        async with stdio_client(self.server_params) as (read, write):
+            async with ClientSession(read, write) as session:
+                await session.initialize()
+
+                # Test global metadata listing (no graph_name)
+                result = await session.call_tool(
+                    TigerGraphToolName.LIST_METADATA,
+                    arguments={},
+                )
+                result_text = str(result)
+                assert "✅ Successfully listed metadata" in result_text
+                assert "Graphs" in result_text or "Vertex Types" in result_text

tests/integration/tools/db/gsql_tools_test.py

@@ -45,9 +45,14 @@
     search,
     search_multi_vector_attributes,
     search_top_k_similar_nodes,
+    gsql,
+    list_metadata,
     create_data_source,
+    update_data_source,
     get_data_source,
     drop_data_source,
+    get_all_data_sources,
+    drop_all_data_sources,
     preview_sample_data,
 )
 
@@ -134,12 +139,22 @@ async def call_tool(name: str, arguments: Dict) -> List[TextContent]:
                     return await search_multi_vector_attributes(**arguments)
                 case TigerGraphToolName.SEARCH_TOP_K_SIMILAR_NODES:
                     return await search_top_k_similar_nodes(**arguments)
+                case TigerGraphToolName.GSQL:
+                    return await gsql(**arguments)
+                case TigerGraphToolName.LIST_METADATA:
+                    return await list_metadata(**arguments)
                 case TigerGraphToolName.CREATE_DATA_SOURCE:
                     return await create_data_source(**arguments)
+                case TigerGraphToolName.UPDATE_DATA_SOURCE:
+                    return await update_data_source(**arguments)
                 case TigerGraphToolName.GET_DATA_SOURCE:
                     return await get_data_source(**arguments)
                 case TigerGraphToolName.DROP_DATA_SOURCE:
                     return await drop_data_source(**arguments)
+                case TigerGraphToolName.GET_ALL_DATA_SOURCES:
+                    return await get_all_data_sources(**arguments)
+                case TigerGraphToolName.DROP_ALL_DATA_SOURCES:
+                    return await drop_all_data_sources(**arguments)
                 case TigerGraphToolName.PREVIEW_SAMPLE_DATA:
                     return await preview_sample_data(**arguments)
                 case _:

tigergraph_mcp/server.py

@@ -41,9 +41,14 @@
     search_top_k_similar_nodes,
 )
 from .db import (
+    gsql,
+    list_metadata,
     create_data_source,
+    update_data_source,
     get_data_source,
     drop_data_source,
+    get_all_data_sources,
+    drop_all_data_sources,
     preview_sample_data,
 )
 
@@ -91,9 +96,15 @@
     "search",
     "search_multi_vector_attributes",
     "search_top_k_similar_nodes",
+    # Tools for GSQL Operations
+    "gsql",
+    "list_metadata",
     # Tools for Data Source Operations
     "create_data_source",
+    "update_data_source",
     "get_data_source",
     "drop_data_source",
+    "get_all_data_sources",
+    "drop_all_data_sources",
     "preview_sample_data",
 ]

tigergraph_mcp/tools/__init__.py

@@ -5,17 +5,27 @@
 # Permission is granted to use, copy, modify, and distribute this software
 # under the License. The software is provided "AS IS", without warranty.
 
+from .gsql import gsql, list_metadata
 from .data_source import (
     create_data_source,
+    update_data_source,
     get_data_source,
     drop_data_source,
+    get_all_data_sources,
+    drop_all_data_sources,
     preview_sample_data,
 )
 
 __all__ = [
+    # Tools for GSQL Operations
+    "gsql",
+    "list_metadata",
     # Tools for Data Source Operations
     "create_data_source",
+    "update_data_source",
     "get_data_source",
     "drop_data_source",
+    "get_all_data_sources",
+    "drop_all_data_sources",
     "preview_sample_data",
 ]

tigergraph_mcp/tools/db/__init__.py

@@ -6,14 +6,20 @@
 # under the License. The software is provided "AS IS", without warranty.
 
 from .create_data_source_tool import create_data_source
+from .update_data_source_tool import update_data_source
 from .get_data_source_tool import get_data_source
 from .drop_data_source_tool import drop_data_source
+from .get_all_data_sources_tool import get_all_data_sources
+from .drop_all_data_sources_tool import drop_all_data_sources
 from .preview_sample_data_tool import preview_sample_data
 
 
 __all__ = [
     "create_data_source",
+    "update_data_source",
     "get_data_source",
     "drop_data_source",
+    "get_all_data_sources",
+    "drop_all_data_sources",
     "preview_sample_data",
 ]

tigergraph_mcp/tools/db/data_source/__init__.py

@@ -0,0 +1,56 @@
+# Copyright 2025 TigerGraph Inc.
+# Licensed under the Apache License, Version 2.0.
+# See the LICENSE file or https://www.apache.org/licenses/LICENSE-2.0
+#
+# Permission is granted to use, copy, modify, and distribute this software
+# under the License. The software is provided "AS IS", without warranty.
+
+from typing import Optional, List
+from pydantic import Field, BaseModel
+from mcp.types import Tool, TextContent
+
+from tigergraphx import TigerGraphDatabase
+from tigergraph_mcp.tools import TigerGraphToolName
+
+
+class DropAllDataSourcesToolInput(BaseModel):
+    """Input schema for dropping all data sources."""
+
+    graph_name: Optional[str] = Field(
+        default=None,
+        description="Optional graph name to scope the data source removal.",
+    )
+
+
+tools = [
+    Tool(
+        name=TigerGraphToolName.DROP_ALL_DATA_SOURCES,
+        description="""Drops all data sources in TigerGraph using TigerGraphX.
+
+You may optionally scope the drop operation to a specific graph.
+
+Example:
+```python
+graph_name = "MyGraph"
+```""",
+        inputSchema=DropAllDataSourcesToolInput.model_json_schema(),
+    )
+]
+
+
+async def drop_all_data_sources(
+    graph_name: Optional[str] = None,
+) -> List[TextContent]:
+    try:
+        db = TigerGraphDatabase()
+        result = db.drop_all_data_sources(graph_name=graph_name)
+
+        if "is dropped" in result or "is dropped successfully" in result:
+            message = f"✅ {result}"
+        else:
+            message = f"⚠️ Attempted to drop all data sources but got unexpected response:\n{result}"
+
+    except Exception as e:
+        message = f"❌ Error dropping all data sources: {str(e)}"
+
+    return [TextContent(type="text", text=message)]