Skip to content

Commit 2719c27

Browse files
antonbricksantonbricks
authored
Merge pull request #23 from antonbricks/streamlit-uc-connections
Streamlit uc connections
2 parents 011fccc + d838ac2 commit 2719c27

File tree

11 files changed

+699
-22
lines changed

11 files changed

+699
-22
lines changed

‎docs/docs/dash/external_services/_category_.json‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"label": "External services",
3+
"position": 8
4+
}

‎docs/docs/dash/external_services/external_connections.mdx‎

Lines changed: 215 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,215 @@
1+
---
2+
sidebar_position: 1
3+
---
4+
5+
# External connections
6+
7+
This recipe demonstrates how to use Unity Catalog-managed external HTTP connections for secure and governed access to MCP and non-MCP servers, for example, to GitHub, or Jira, and Slack.
8+
9+
## Code snippets
10+
11+
### Bearer Token Method
12+
13+
```python title="app.py"
14+
from databricks.sdk import WorkspaceClient
15+
from databricks.sdk.service.serving import ExternalFunctionRequestHttpMethod
16+
w = WorkspaceClient()
17+
18+
response = w.serving_endpoints.http_request(
19+
conn="github_connection",
20+
method=ExternalFunctionRequestHttpMethod.GET,
21+
path="/traffic/views",
22+
headers={"Accept": "application/vnd.github+json"},
23+
)
24+
25+
print(response.json())
26+
```
27+
28+
### Github Non-MCP API with Oauth On-Behalf-Of User
29+
30+
```python title="app.py"
31+
from databricks.sdk import WorkspaceClient
32+
from databricks.sdk.service.serving import ExternalFunctionRequestHttpMethod
33+
from flask import request
34+
token = request.headers.get("x-forwarded-access-token")
35+
w = WorkspaceClient(token=token, auth_type="pat")
36+
37+
response = w.serving_endpoints.http_request(
38+
conn="github_u2m",
39+
method=ExternalFunctionRequestHttpMethod.GET,
40+
path="/user",
41+
headers={"Accept": "application/vnd.github+json"},
42+
)
43+
44+
print(response.json())
45+
```
46+
47+
### Github API (Non-MCP) with OAuth User to Machine Per User (On-behalf-of-user)
48+
49+
```python title="app.py"
50+
import contextvars
51+
import os
52+
import uuid
53+
from pathlib import Path
54+
from mcp.server.fastmcp import FastMCP
55+
from fastapi import FastAPI, Request
56+
from fastapi.responses import FileResponse
57+
from typing import Dict, Any
58+
from databricks.sdk import WorkspaceClient
59+
from databricks.sdk.service.serving import ExternalFunctionRequestHttpMethod
60+
61+
62+
STATIC_DIR = Path(__file__).parent / "static"
63+
64+
# Create an MCP server
65+
mcp = FastMCP("Custom MCP Server on Databricks Apps")
66+
67+
# Context variable for headers
68+
header_store = contextvars.ContextVar("header_store")
69+
70+
def init_github_tools() -> Dict[str, Any]:
71+
headers = header_store.get({})
72+
token=headers["x-forwarded-access-token"]
73+
74+
json = {
75+
"jsonrpc": "2.0",
76+
"id": "init-1",
77+
"method": "initialize",
78+
"params": {}
79+
}
80+
81+
try:
82+
response = WorkspaceClient(token=token, auth_type="pat").serving_endpoints.http_request(
83+
conn="github_u2m_connection",
84+
method=ExternalFunctionRequestHttpMethod.POST,
85+
path="/",
86+
json=json,
87+
)
88+
except Exception as e:
89+
print(f"Error making the init request: {e}")
90+
return {"error": str(e)}
91+
92+
# Extract the Mcp-Session-Id from response headers
93+
print("MCP_HEADERS", response.headers)
94+
session_id = response.headers.get("mcp-session-id")
95+
if not session_id:
96+
print("No session ID returned by server.")
97+
98+
print(f"✅ Got MCP Session ID: {session_id}")
99+
return session_id
100+
101+
# Tool using header_store
102+
def list_github_tools() -> Dict[str, Any]:
103+
session_id = init_github_tools()
104+
headers = header_store.get({})
105+
token=headers["x-forwarded-access-token"]
106+
107+
json = {
108+
"jsonrpc": "2.0",
109+
"id": "list-1",
110+
"method": "tools/list",
111+
}
112+
113+
try:
114+
response = WorkspaceClient(token=token, auth_type="pat").serving_endpoints.http_request(
115+
conn="github_u2m_connection",
116+
method=ExternalFunctionRequestHttpMethod.POST,
117+
path="/",
118+
json=json,
119+
headers={
120+
"Mcp-Session-Id": session_id
121+
}
122+
)
123+
except Exception as e:
124+
print(f"Error: {e}")
125+
return {"error": str(e)}
126+
127+
print(f"Response list tools: {response.json()}")
128+
129+
return response.json()
130+
131+
132+
def call_github_tool(name: str, arguments: dict) -> Dict[str, Any]:
133+
session_id = init_github_tools()
134+
headers = header_store.get({})
135+
token=headers["x-forwarded-access-token"]
136+
json = {
137+
"jsonrpc": "2.0",
138+
"id": "call-1",
139+
"method": "tools/call",
140+
"params": {
141+
"name": name,
142+
"arguments": arguments
143+
}
144+
}
145+
try:
146+
response = WorkspaceClient(token=token, auth_type="pat").serving_endpoints.http_request(
147+
conn="github_u2m_connection",
148+
method=ExternalFunctionRequestHttpMethod.POST,
149+
path="/",
150+
json=json,
151+
headers={
152+
"Mcp-Session-Id": session_id
153+
}
154+
)
155+
except Exception as e:
156+
print(f"Error: {e}")
157+
return {"error": str(e)}
158+
159+
print(f"Response call tools: {response.json()}")
160+
return response.json()
161+
162+
@mcp._mcp_server.list_tools()
163+
async def list_tools():
164+
tools_dict = list_github_tools()
165+
return tools_dict.get("result", {}).get("tools", [])
166+
167+
@mcp._mcp_server.call_tool()
168+
async def call_tool(name: str, arguments: dict):
169+
response = call_github_tool(name, arguments)
170+
return response.get("result", {}).get("content", [])
171+
172+
mcp_app = mcp.streamable_http_app()
173+
174+
175+
app = FastAPI(
176+
lifespan=lambda _: mcp.session_manager.run(),
177+
)
178+
179+
@app.middleware("http")
180+
async def capture_headers(request: Request, call_next):
181+
header_store.set(dict(request.headers))
182+
return await call_next(request)
183+
184+
@app.get("/", include_in_schema=False)
185+
async def serve_index():
186+
return FileResponse(STATIC_DIR / "index.html")
187+
188+
189+
app.mount("/", mcp_app)
190+
```
191+
192+
## Resources
193+
194+
- [Unity Catalog HTTP Connection](https://docs.databricks.com/aws/en/query-federation/http), either MCP or non-MCP
195+
196+
## Permissions
197+
198+
Your [app service principal](https://docs.databricks.com/aws/en/dev-tools/databricks-apps/#how-does-databricks-apps-manage-authorization) needs the following permissions:
199+
200+
- `USE CONNECTION` permission on the HTTP Connection
201+
202+
When using OAuth User to Machine Per User (On-behalf-of-user), you need to configure [User authorization](https://docs.databricks.com/aws/en/dev-tools/databricks-apps/auth#user-authorization) by adding the Unity Catalog connection or other scopes.
203+
204+
## Dependencies
205+
206+
- [Databricks SDK for Python](https://pypi.org/project/databricks-sdk/) - `databricks-sdk`
207+
- [Streamlit](https://pypi.org/project/streamlit/) - `streamlit`
208+
209+
```python title="requirements.txt"
210+
databricks-sdk
211+
streamlit
212+
uvicorn
213+
mcp[cli]
214+
fastapi
215+
```

‎docs/docs/streamlit/external_services/_category_.json‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"label": "External services",
3+
"position": 8
4+
}

‎docs/docs/streamlit/external_services/external_connections.mdx‎

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
---
2+
sidebar_position: 1
3+
---
4+
5+
# External connections
6+
7+
This recipe demonstrates how to use Unity Catalog-managed external HTTP connections for secure and governed access to MCP and non-MCP servers, for example, to GitHub, or Jira, and Slack.
8+
9+
## Code snippets
10+
11+
### Bearer token
12+
13+
```python title="app.py"
14+
import streamlit as st
15+
from databricks.sdk import WorkspaceClient
16+
from databricks.sdk.service.serving import ExternalFunctionRequestHttpMethod
17+
18+
w = WorkspaceClient()
19+
20+
response = w.serving_endpoints.http_request(
21+
conn="github_connection",
22+
method=ExternalFunctionRequestHttpMethod.GET,
23+
path="/traffic/views",
24+
headers={"Accept": "application/vnd.github+json"},
25+
)
26+
27+
st.json(response.json())
28+
```
29+
30+
### Github MCP with OAuth User to Machine Per User (On-behalf-of-user)
31+
32+
```python title="app.py"
33+
import streamlit as st
34+
from databricks.sdk import WorkspaceClient
35+
from databricks.sdk.service.serving import ExternalFunctionRequestHttpMethod
36+
37+
token = st.context.headers.get("x-forwarded-access-token")
38+
w = WorkspaceClient(token=token, auth_type="pat")
39+
40+
response = w.serving_endpoints.http_request(
41+
conn="github_u2m",
42+
method=ExternalFunctionRequestHttpMethod.GET,
43+
path="/user",
44+
headers={"Accept": "application/vnd.github+json"},
45+
)
46+
47+
st.json(response.json())
48+
```
49+
50+
### Github API (Non-MCP) with OAuth User to Machine Per User (On-behalf-of-user)
51+
52+
```python title="app.py"
53+
import streamlit as st
54+
from databricks.sdk import WorkspaceClient
55+
from databricks.sdk.service.serving import ExternalFunctionRequestHttpMethod
56+
import json
57+
58+
59+
token = st.context.headers.get("x-forwarded-access-token")
60+
w = WorkspaceClient(token=token, auth_type="pat")
61+
62+
63+
def init_mcp_session(w: WorkspaceClient, connection_name: str):
64+
init_payload = {
65+
"jsonrpc": "2.0",
66+
"id": "init-1",
67+
"method": "initialize",
68+
"params": {}
69+
}
70+
response = w.serving_endpoints.http_request(
71+
conn=connection_name,
72+
method=ExternalFunctionRequestHttpMethod.POST,
73+
path="/",
74+
json=init_payload,
75+
)
76+
return response.headers.get("mcp-session-id")
77+
78+
79+
connection_name = "github_u2m_connection"
80+
http_method = ExternalFunctionRequestHttpMethod.POST
81+
path = "/"
82+
headers = {"Content-Type": "application/json"}
83+
payload = {"jsonrpc": "2.0", "id": "list-1", "method": "tools/list"}
84+
85+
if st.button("Run"):
86+
session_id = init_mcp_session(w, connection_name)
87+
headers["Mcp-Session-Id"] = session_id
88+
89+
response = w.serving_endpoints.http_request(
90+
conn=connection_name,
91+
method=http_method,
92+
path=path,
93+
headers=headers,
94+
json=payload,
95+
)
96+
st.json(response.json())
97+
```
98+
99+
## Resources
100+
101+
- [Unity Catalog HTTP Connection](https://docs.databricks.com/aws/en/query-federation/http), either MCP or non-MCP
102+
103+
## Permissions
104+
105+
Your [app service principal](https://docs.databricks.com/aws/en/dev-tools/databricks-apps/#how-does-databricks-apps-manage-authorization) needs the following permissions:
106+
107+
- `USE CONNECTION` permission on the HTTP Connection
108+
109+
When using OAuth User to Machine Per User (On-behalf-of-user), you need to configure [User authorization](https://docs.databricks.com/aws/en/dev-tools/databricks-apps/auth#user-authorization) by adding the Unity Catalog connection or other scopes.
110+
111+
## Dependencies
112+
113+
- [Databricks SDK for Python](https://pypi.org/project/databricks-sdk/) - `databricks-sdk`
114+
- [Streamlit](https://pypi.org/project/streamlit/) - `streamlit`
115+
116+
```python title="requirements.txt"
117+
databricks-sdk
118+
streamlit
119+
mcp[cli]
120+
```

0 commit comments

Comments
 (0)