Update mcp_server.py

mcp_server.py

@@ -1,10 +1,17 @@
+"""
+mcp_server.py – MCP tool server exposing:
+  • query_database   – SQLite sample
+  • query_salesforce – Salesforce via CData JDBC
+Add more connectors under connectors/ and register the same way.
+"""
+
 from mcp.server.fastmcp import FastMCP
 import sqlite3, json
+from connectors.salesforce_connector import sf_query   # ← new connector
 
-# ─── MCP instance ───────────────────────────────────────────────────────
 mcp = FastMCP("EnterpriseData")
 
-# ─── In-memory SQLite sample DB ─────────────────────────────────────────
+# ─── in-memory SQLite sample ─────────────────────────────────────────────
 conn = sqlite3.connect(":memory:", check_same_thread=False)
 cur  = conn.cursor()
 
@@ -16,7 +23,6 @@ CREATE TABLE Customers (
     LastOrderDate   TEXT
 )
 """)
-
 cur.executemany(
     "INSERT INTO Customers (Name, Region, LastOrderDate) VALUES (?,?,?)",
     [
@@ -25,17 +31,14 @@ cur.executemany(
         ("Gamma Co",   "Northeast", "2023-09-15"),
         ("Delta LLC",  "South",     "2025-03-20"),
         ("Epsilon Ltd","Northeast", "2025-07-10"),
-    ]
+    ],
 )
 conn.commit()
 
-# ─── SQL tool exposed to the agent ──────────────────────────────────────
+# ─── SQLite tool ─────────────────────────────────────────────────────────
 @mcp.tool()
 def query_database(sql: str) -> str:
-    """
-    Execute raw SQL and return rows as JSON.
-    Example:  SELECT * FROM Customers WHERE Region='West';
-    """
+    """Run arbitrary SQL against the Customers SQLite table."""
     try:
         cur.execute(sql)
         cols = [d[0] for d in cur.description or []]
@@ -44,5 +47,11 @@ def query_database(sql: str) -> str:
     except Exception as exc:
         return json.dumps({"error": str(exc)})
 
+# ─── Salesforce tool (premium) ───────────────────────────────────────────
+@mcp.tool()
+def query_salesforce(sql: str) -> str:
+    """Run SOQL-as-SQL against Salesforce objects (requires Pro plan)."""
+    return sf_query(sql)
+
 if __name__ == "__main__":
     mcp.run(transport="stdio")