Update mcp_server.py

mcp_server.py

@@ -1,68 +1,55 @@
 import os
 import json
 import sqlite3
-from mcp.server.fastmcp import FastMCP
 
-# HubSpot SDK
-from hubspot import HubSpot
-from hubspot.crm.contacts import ApiException as HSContactsError
+from mcp.server.fastmcp import FastMCP
+from connectors.hubspot_connector import query_hubspot
 
-# ─────────────────────────── MCP Server Init ─────────────────────────────
 mcp = FastMCP("EnterpriseData")
 
-# ───────────────────── In-Memory SQLite Sample ───────────────────────────
+# ─── 1) In-memory SQLite sample ─────────────────────────────────────────────
 conn = sqlite3.connect(":memory:", check_same_thread=False)
-cur = conn.cursor()
+cur  = conn.cursor()
 cur.execute("""
 CREATE TABLE Customers (
-    CustomerID INTEGER PRIMARY KEY AUTOINCREMENT,
-    Name TEXT,
-    Region TEXT,
-    LastOrderDate TEXT
+    CustomerID      INTEGER PRIMARY KEY AUTOINCREMENT,
+    Name            TEXT,
+    Region          TEXT,
+    LastOrderDate   TEXT
 )
 """)
 cur.executemany(
-    "INSERT INTO Customers (Name, Region, LastOrderDate) VALUES (?, ?, ?)",
+    "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"),
+        ("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()
 
 @mcp.tool()
 def query_database(sql: str) -> str:
-    """Run SQL query on Customers table and return results."""
+    """Run SQL against the in-memory Customers table and return JSON rows."""
     try:
         cur.execute(sql)
-        cols = [desc[0] for desc in cur.description or []]
-        rows = [dict(zip(cols, row)) for row in cur.fetchall()]
-        return json.dumps(rows, indent=2)
+        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 e:
         return json.dumps({"error": str(e)})
 
-# ───────────────────────── HubSpot Integration ───────────────────────────
-hs_token = os.getenv("HUBSPOT_TOKEN")
-hs_client = HubSpot(access_token=hs_token)
-
+# ─── 2) HubSpot tool ────────────────────────────────────────────────────────
 @mcp.tool()
-def query_hubspot_contacts(limit: int = 100) -> str:
+def hubspot_data(object_type: str, limit: int = 100) -> str:
     """
-    Fetch up to `limit` HubSpot contacts.
-    Example: query_hubspot_contacts(limit=50)
+    Fetch up to `limit` objects of type `contacts`, `companies`, or `deals`.
+    Example: object_type="contacts"
     """
-    try:
-        page = hs_client.crm.contacts.basic_api.get_page(limit=limit)
-        contacts = [contact.to_dict() for contact in page.results]
-        return json.dumps(contacts, indent=2)
-    except HSContactsError as e:
-        return json.dumps({"error": e.body})
-    except Exception as e:
-        return json.dumps({"error": str(e)})
+    return query_hubspot(object_type, limit)
 
-# ──────────────────────── Start the MCP Server ───────────────────────────
+# ─── 3) Start the MCP server ───────────────────────────────────────────────
 if __name__ == "__main__":
     mcp.run(transport="stdio")