feat: add read tools (list_tables, describe_table, get_records, sql_query)
This commit is contained in:
78
src/grist_mcp/tools/read.py
Normal file
78
src/grist_mcp/tools/read.py
Normal file
@@ -0,0 +1,78 @@
|
||||
"""Read tools - query tables and records."""
|
||||
|
||||
from grist_mcp.auth import Agent, Authenticator, Permission
|
||||
from grist_mcp.grist_client import GristClient
|
||||
|
||||
|
||||
async def list_tables(
|
||||
agent: Agent,
|
||||
auth: Authenticator,
|
||||
document: str,
|
||||
client: GristClient | None = None,
|
||||
) -> dict:
|
||||
"""List all tables in a document."""
|
||||
auth.authorize(agent, document, Permission.READ)
|
||||
|
||||
if client is None:
|
||||
doc = auth.get_document(document)
|
||||
client = GristClient(doc)
|
||||
|
||||
tables = await client.list_tables()
|
||||
return {"tables": tables}
|
||||
|
||||
|
||||
async def describe_table(
|
||||
agent: Agent,
|
||||
auth: Authenticator,
|
||||
document: str,
|
||||
table: str,
|
||||
client: GristClient | None = None,
|
||||
) -> dict:
|
||||
"""Get column information for a table."""
|
||||
auth.authorize(agent, document, Permission.READ)
|
||||
|
||||
if client is None:
|
||||
doc = auth.get_document(document)
|
||||
client = GristClient(doc)
|
||||
|
||||
columns = await client.describe_table(table)
|
||||
return {"table": table, "columns": columns}
|
||||
|
||||
|
||||
async def get_records(
|
||||
agent: Agent,
|
||||
auth: Authenticator,
|
||||
document: str,
|
||||
table: str,
|
||||
filter: dict | None = None,
|
||||
sort: str | None = None,
|
||||
limit: int | None = None,
|
||||
client: GristClient | None = None,
|
||||
) -> dict:
|
||||
"""Fetch records from a table."""
|
||||
auth.authorize(agent, document, Permission.READ)
|
||||
|
||||
if client is None:
|
||||
doc = auth.get_document(document)
|
||||
client = GristClient(doc)
|
||||
|
||||
records = await client.get_records(table, filter=filter, sort=sort, limit=limit)
|
||||
return {"records": records}
|
||||
|
||||
|
||||
async def sql_query(
|
||||
agent: Agent,
|
||||
auth: Authenticator,
|
||||
document: str,
|
||||
query: str,
|
||||
client: GristClient | None = None,
|
||||
) -> dict:
|
||||
"""Run a read-only SQL query."""
|
||||
auth.authorize(agent, document, Permission.READ)
|
||||
|
||||
if client is None:
|
||||
doc = auth.get_document(document)
|
||||
client = GristClient(doc)
|
||||
|
||||
records = await client.sql_query(query)
|
||||
return {"records": records}
|
||||
82
tests/test_tools_read.py
Normal file
82
tests/test_tools_read.py
Normal file
@@ -0,0 +1,82 @@
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from grist_mcp.tools.read import list_tables, describe_table, get_records, sql_query
|
||||
from grist_mcp.auth import Authenticator, Agent, Permission
|
||||
from grist_mcp.config import Config, Document, Token, TokenScope
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def config():
|
||||
return Config(
|
||||
documents={
|
||||
"budget": Document(
|
||||
url="https://grist.example.com",
|
||||
doc_id="abc123",
|
||||
api_key="key",
|
||||
),
|
||||
},
|
||||
tokens=[
|
||||
Token(
|
||||
token="test-token",
|
||||
name="test-agent",
|
||||
scope=[TokenScope(document="budget", permissions=["read"])],
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def auth(config):
|
||||
return Authenticator(config)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def agent(auth):
|
||||
return auth.authenticate("test-token")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_client():
|
||||
client = AsyncMock()
|
||||
client.list_tables.return_value = ["Table1", "Table2"]
|
||||
client.describe_table.return_value = [
|
||||
{"id": "Name", "type": "Text", "formula": ""},
|
||||
]
|
||||
client.get_records.return_value = [
|
||||
{"id": 1, "Name": "Alice"},
|
||||
]
|
||||
client.sql_query.return_value = [{"Name": "Alice"}]
|
||||
return client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_tables(agent, auth, mock_client):
|
||||
result = await list_tables(agent, auth, "budget", client=mock_client)
|
||||
|
||||
assert result == {"tables": ["Table1", "Table2"]}
|
||||
mock_client.list_tables.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_describe_table(agent, auth, mock_client):
|
||||
result = await describe_table(agent, auth, "budget", "Table1", client=mock_client)
|
||||
|
||||
assert result == {
|
||||
"table": "Table1",
|
||||
"columns": [{"id": "Name", "type": "Text", "formula": ""}],
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_records(agent, auth, mock_client):
|
||||
result = await get_records(agent, auth, "budget", "Table1", client=mock_client)
|
||||
|
||||
assert result == {"records": [{"id": 1, "Name": "Alice"}]}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sql_query(agent, auth, mock_client):
|
||||
result = await sql_query(agent, auth, "budget", "SELECT * FROM Table1", client=mock_client)
|
||||
|
||||
assert result == {"records": [{"Name": "Alice"}]}
|
||||
Reference in New Issue
Block a user