"""Generic query tool shared by all connectors."""

from pydantic import BaseModel, Field
from src.client import get_client


class QueryConnectorInput(BaseModel):
    connector_id: str = Field(
        description="Connector ID returned by the respective get_*_tables tool"
    )
    query: str = Field(description="Trino SQL query to execute")


def register_connector_tools(mcp):

    @mcp.tool(
        "query_connector",
        "Execute a Trino SQL query on any connector. "
        "Use the connector_id returned by any get_*_tables tool. "
        "Always call the corresponding get_*_schema tool first to know the table name and available columns. "
        "Always use double quotes for column names with special characters (/, -, etc).",
        QueryConnectorInput,
    )
    async def query_connector(args: QueryConnectorInput) -> str:
        return str(await get_client().query(args.connector_id, args.query))
