AI_SQL / mcp_server.py
mgbam's picture
Update mcp_server.py
4cf3766 verified
raw
history blame
2.13 kB
"""
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 = FastMCP("EnterpriseData")
# ─── in-memory SQLite sample ─────────────────────────────────────────────
conn = sqlite3.connect(":memory:", check_same_thread=False)
cur = conn.cursor()
cur.execute("""
CREATE TABLE Customers (
CustomerID INTEGER PRIMARY KEY AUTOINCREMENT,
Name TEXT,
Region TEXT,
LastOrderDate TEXT
)
""")
cur.executemany(
"INSERT INTO Customers (Name, Region, LastOrderDate) VALUES (?,?,?)",
[
("Acme Corp", "Northeast", "2024-12-01"),
("Beta Inc", "West", "2025-06-01"),
("Gamma Co", "Northeast", "2023-09-15"),
("Delta LLC", "South", "2025-03-20"),
("Epsilon Ltd","Northeast", "2025-07-10"),
],
)
conn.commit()
# ─── SQLite tool ─────────────────────────────────────────────────────────
@mcp.tool()
def query_database(sql: str) -> str:
"""Run arbitrary SQL against the Customers SQLite table."""
try:
cur.execute(sql)
cols = [d[0] for d in cur.description or []]
rows = [dict(zip(cols, r)) for r in cur.fetchall()]
return json.dumps(rows)
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")