Materialize/misc/mcp-materialize/tests/integration/test_mcp.py

# Copyright Materialize, Inc. and contributors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License in the LICENSE file at the
# root of this repository, or online at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os

import pytest
import pytest_asyncio
from mcp import Tool
from mcp.types import ToolAnnotations
from psycopg_pool import AsyncConnectionPool

from mcp_materialize.mz_client import MzClient


@pytest_asyncio.fixture(scope="function")
async def materialize_pool():
    conn = os.getenv("MZ_DSN", "postgres://materialize@localhost:6875/materialize")
    async with AsyncConnectionPool(
        conninfo=conn, min_size=1, max_size=10, open=False
    ) as pool:
        async with pool.connection() as conn:
            await conn.set_autocommit(True)
            async with conn.cursor() as cur:
                await cur.execute("DROP SCHEMA IF EXISTS materialize.tools CASCADE;")
                await cur.execute("CREATE SCHEMA materialize.tools;")

        yield pool


@pytest.mark.asyncio
async def test_basic_tool(materialize_pool):
    async with MzClient(pool=materialize_pool) as client:
        tools = await client.list_tools()
        assert len(tools) == 0

    async with materialize_pool.connection() as conn:
        await conn.set_autocommit(True)
        async with conn.cursor() as cur:
            await cur.execute(
                """CREATE OR REPLACE VIEW tools.my_tool AS
                SELECT 1 AS id, 'hello' AS result;"""
            )
            await cur.execute("CREATE INDEX my_tool_id_idx ON tools.my_tool (id);")
            await cur.execute("COMMENT ON VIEW tools.my_tool IS 'Get result from id';")

            await cur.execute(
                """CREATE OR REPLACE VIEW tools.missing_comment AS
                SELECT 1 AS id, 'goodbye' AS result;"""
            )
            await cur.execute(
                "CREATE INDEX missing_comment_id_idx ON tools.missing_comment (id);"
            )

            await cur.execute(
                """CREATE OR REPLACE VIEW tools.missing_idx AS
                SELECT 1 AS id, 'not it' AS result;"""
            )
            await cur.execute(
                "COMMENT ON VIEW tools.missing_idx IS 'Get result from id';"
            )

    async with MzClient(pool=materialize_pool) as client:
        tools = await client.list_tools()
        assert len(tools) == 1
        assert tools[0] == Tool(
            name="materialize_tools_my_tool_id_idx",
            description="Get result from id",
            inputSchema={
                "type": "object",
                "required": ["id"],
                "properties": {"id": {"type": "number"}},
            },
            outputSchema={
                "type": "object",
                "properties": {
                    "rows": {
                        "type": "array",
                        "items": {
                            "type": "object",
                            "required": ["result"],
                            "properties": {"result": {"type": "string"}},
                        },
                    }
                },
                "required": ["rows"],
            },
            annotations=ToolAnnotations(
                title="materialize::tools::my_tool(id)",
                readOnlyHint=True,
            ),
        )

        result = await client.call_tool("materialize_tools_my_tool_id_idx", {"id": 1})
        rows = result["rows"]
        assert len(rows) == 1
        assert rows[0] == {"result": "hello"}


@pytest.mark.asyncio
async def test_exists_tool(materialize_pool):
    async with materialize_pool.connection() as conn:
        await conn.set_autocommit(True)
        async with conn.cursor() as cur:
            await cur.execute("CREATE OR REPLACE VIEW tools.my_tool AS SELECT 1 AS id;")
            await cur.execute("CREATE INDEX my_tool_id_idx ON tools.my_tool (id);")
            await cur.execute("COMMENT ON VIEW tools.my_tool IS 'Check if id exists';")

    async with MzClient(pool=materialize_pool) as client:
        result = await client.call_tool("materialize_tools_my_tool_id_idx", {"id": 1})
        rows = result["rows"]
        assert len(rows) == 1
        assert rows[0] == {"exists": True}

        result = await client.call_tool("materialize_tools_my_tool_id_idx", {"id": 2})
        rows = result["rows"]
        assert len(rows) == 1
        assert rows[0] == {"exists": False}


@pytest.mark.asyncio
async def test_type_handling_keys(materialize_pool):
    async with materialize_pool.connection() as conn:
        await conn.set_autocommit(True)
        async with conn.cursor() as cur:
            await cur.execute(
                """
            CREATE OR REPLACE VIEW tools.all_types AS
            SELECT
                1::smallint                                  AS smallint_col,
                2::integer                                   AS integer_col,
                3::bigint                                    AS bigint_col,
                2::uint4                                     AS uint2_col,
                4::uint4                                     AS uint4_col,
                8::uint4                                     AS uint8_col,
                4.5::real                                    AS real_col,
                6.7::double precision                        AS double_col,
                1.23::numeric                                AS numeric_col,
                true                                         AS boolean_col,
                'a'::char                                    AS char_col,
                'abc'::varchar                               AS varchar_col,
                'abc'::text                                  AS text_col,
                '2024-01-01'::date                           AS date_col,
                '12:34:56'::time                             AS time_col,
                '2024-01-01 12:34:56'::timestamp             AS timestamp_col,
                '2024-01-01 12:34:56+00'::timestamptz        AS timestamptz_col,
                decode('DEADBEEF', 'hex')::bytea             AS bytea_col,
                '{"a": 1, "b": [1, 2, 3]}'::jsonb            AS jsonb_col,
                '550e8400-e29b-41d4-a716-446655440000'::uuid AS uuid_col;
                """
            )
            await cur.execute("CREATE DEFAULT INDEX all_types_idx ON tools.all_types;")
            await cur.execute("COMMENT ON VIEW tools.all_types IS 'All types';")

    async with MzClient(pool=materialize_pool) as client:
        tools = await client.list_tools()
        assert len(tools) == 1

        assert tools[0].name == "materialize_tools_all_types_idx"
        assert tools[0].description == "All types"
        assert sorted(tools[0].inputSchema["required"]) == sorted(
            [
                "bigint_col",
                "boolean_col",
                "bytea_col",
                "char_col",
                "date_col",
                "double_col",
                "integer_col",
                "jsonb_col",
                "numeric_col",
                "real_col",
                "smallint_col",
                "text_col",
                "time_col",
                "timestamp_col",
                "timestamptz_col",
                "uint2_col",
                "uint4_col",
                "uint8_col",
                "uuid_col",
                "varchar_col",
            ]
        )

    assert tools[0].inputSchema["properties"] == {
        "bigint_col": {"type": "number"},
        "boolean_col": {"type": "boolean"},
        "bytea_col": {
            "type": "string",
            "contentEncoding": "base64",
            "contentMediaType": "application/octet-stream",
        },
        "char_col": {"type": "string"},
        "smallint_col": {"type": "number"},
        "double_col": {"type": "number"},
        "text_col": {"type": "string"},
        "integer_col": {"type": "number"},
        "uint2_col": {"type": "number"},
        "uint4_col": {"type": "number"},
        "uint8_col": {"type": "number"},
        "date_col": {"type": "string", "format": "date"},
        "time_col": {"type": "string", "format": "time"},
        "timestamp_col": {"type": "string", "format": "date-time"},
        "timestamptz_col": {"type": "string", "format": "date-time"},
        "jsonb_col": {"type": "object"},
        "numeric_col": {"type": "number"},
        "real_col": {"type": "number"},
        "varchar_col": {"type": "string"},
        "uuid_col": {"type": "string", "format": "uuid"},
    }


@pytest.mark.asyncio
async def test_type_handling_values(materialize_pool):
    async with materialize_pool.connection() as conn:
        await conn.set_autocommit(True)
        async with conn.cursor() as cur:
            await cur.execute(
                """
            CREATE OR REPLACE VIEW tools.all_types AS
            SELECT
                1                                            AS id,
                1::smallint                                  AS smallint_col,
                2::integer                                   AS integer_col,
                3::bigint                                    AS bigint_col,
                2::uint4                                     AS uint2_col,
                4::uint4                                     AS uint4_col,
                8::uint4                                     AS uint8_col,
                4.5::real                                    AS real_col,
                6.7::double precision                        AS double_col,
                1.23::numeric                                AS numeric_col,
                true                                         AS boolean_col,
                'a'::char                                    AS char_col,
                'abc'::varchar                               AS varchar_col,
                'abc'::text                                  AS text_col,
                '2024-01-01'::date                           AS date_col,
                '12:34:56'::time                             AS time_col,
                '2024-01-01 12:34:56'::timestamp             AS timestamp_col,
                '2024-01-01 12:34:56+00'::timestamptz        AS timestamptz_col,
                decode('DEADBEEF', 'hex')::bytea             AS bytea_col,
                '{"a": 1, "b": [1, 2, 3]}'::jsonb            AS jsonb_col,
                '550e8400-e29b-41d4-a716-446655440000'::uuid AS uuid_col;
                """
            )
            await cur.execute("CREATE INDEX all_types_idx ON tools.all_types (id);")
            await cur.execute("COMMENT ON VIEW tools.all_types IS 'All types';")

    async with MzClient(pool=materialize_pool) as client:
        results = await client.call_tool("materialize_tools_all_types_idx", {"id": 1})
        assert len(results) == 1