Database Performance Observability Guide¶

# tests/helpers/query_counter.py
from contextlib import contextmanager
from sqlalchemy import event
from sqlalchemy.engine import Engine
import threading

class QueryCounter:
    """Thread-safe SQL query counter using SQLAlchemy events."""

    def __init__(self):
        self.count = 0
        self.queries = []
        self._lock = threading.Lock()

    def _before_execute(self, conn, cursor, statement, parameters, context, executemany):
        with self._lock:
            self.count += 1
            self.queries.append({
                'statement': statement,
                'parameters': parameters,
                'executemany': executemany
            })

    def reset(self):
        with self._lock:
            self.count = 0
            self.queries = []

@contextmanager
def count_queries(engine: Engine, print_queries: bool = False):
    """Context manager to count SQL queries.

    Usage:
        with count_queries(engine) as counter:
            # ... code that runs queries ...
        print(f"Executed {counter.count} queries")

    Args:
        engine: SQLAlchemy engine
        print_queries: If True, print each query as it's executed

    Yields:
        QueryCounter: Counter object with .count and .queries attributes
    """
    counter = QueryCounter()

    def before_execute(conn, cursor, statement, parameters, context, executemany):
        counter._before_execute(conn, cursor, statement, parameters, context, executemany)
        if print_queries:
            print(f"[Query #{counter.count}] {statement[:200]}...")

    event.listen(engine, "before_cursor_execute", before_execute)
    try:
        yield counter
    finally:
        event.remove(engine, "before_cursor_execute", before_execute)

# tests/conftest.py (add to existing)
import pytest
from tests.helpers.query_counter import count_queries

@pytest.fixture
def query_counter(db_engine):
    """Fixture to count database queries in tests.

    Usage:
        def test_something(query_counter, db_session):
            with query_counter() as counter:
                # do database operations
            assert counter.count <= 5, f"Too many queries: {counter.count}"
    """
    def _counter(print_queries=False):
        return count_queries(db_engine, print_queries=print_queries)
    return _counter

@pytest.fixture
def assert_max_queries(db_engine):
    """Fixture to assert maximum query count.

    Usage:
        def test_list_tools(assert_max_queries, db_session):
            with assert_max_queries(5):
                tools = tool_service.list_tools(db_session)
    """
    @contextmanager
    def _assert_max(max_count: int, message: str = None):
        with count_queries(db_engine) as counter:
            yield counter
        if counter.count > max_count:
            query_list = "\n".join(
                f"  {i+1}. {q['statement'][:100]}..."
                for i, q in enumerate(counter.queries)
            )
            msg = message or f"Expected at most {max_count} queries, got {counter.count}"
            raise AssertionError(f"{msg}\n\nQueries executed:\n{query_list}")
    return _assert_max

# BAD: N+1 pattern - 1 query for tools + N queries for each tool's gateway
tools = db.query(Tool).all()  # 1 query
for tool in tools:
    print(tool.gateway.name)  # N queries (one per tool!)

# GOOD: Eager loading - just 1-2 queries
from sqlalchemy.orm import joinedload
tools = db.query(Tool).options(joinedload(Tool.gateway)).all()
for tool in tools:
    print(tool.gateway.name)  # No additional queries

from sqlalchemy.orm import joinedload

# Before: N+1
tools = db.query(Tool).all()

# After: Single query with JOIN
tools = db.query(Tool).options(joinedload(Tool.gateway)).all()

from sqlalchemy.orm import selectinload

# Before: N+1
servers = db.query(Server).all()
for server in servers:
    for tool in server.tools:  # N additional queries!
        print(tool.name)

# After: 2 queries (servers + all tools)
servers = db.query(Server).options(selectinload(Server.tools)).all()

from sqlalchemy.orm import contains_eager

tools = (
    db.query(Tool)
    .join(Tool.gateway)
    .options(contains_eager(Tool.gateway))
    .filter(Gateway.status == 'active')
    .all()
)

# tests/performance/conftest.py
import pytest
from sqlalchemy import event
import time

@pytest.fixture
def performance_db(db_session, db_engine):
    """Database session with performance instrumentation."""
    queries = []

    def before_execute(conn, cursor, statement, parameters, context, executemany):
        conn.info['query_start'] = time.perf_counter()

    def after_execute(conn, cursor, statement, parameters, context, executemany):
        duration = (time.perf_counter() - conn.info.get('query_start', 0)) * 1000
        queries.append({
            'statement': statement,
            'duration_ms': duration,
            'parameters': parameters,
        })

    event.listen(db_engine, "before_cursor_execute", before_execute)
    event.listen(db_engine, "after_cursor_execute", after_execute)

    class PerfSession:
        session = db_session
        executed_queries = queries

        @property
        def query_count(self):
            return len(queries)

        @property
        def total_query_time_ms(self):
            return sum(q['duration_ms'] for q in queries)

        def get_slow_queries(self, threshold_ms=10):
            return [q for q in queries if q['duration_ms'] > threshold_ms]

        def print_summary(self):
            print(f"\n{'='*60}")
            print(f"Query Summary: {self.query_count} queries, {self.total_query_time_ms:.2f}ms total")
            print(f"{'='*60}")
            for i, q in enumerate(queries, 1):
                print(f"{i}. [{q['duration_ms']:.2f}ms] {q['statement'][:80]}...")

    yield PerfSession()

    event.remove(db_engine, "before_cursor_execute", before_execute)
    event.remove(db_engine, "after_cursor_execute", after_execute)


@pytest.fixture
def seed_performance_data(db_session):
    """Seed database with realistic data volumes for performance testing."""
    from mcpgateway.db import Gateway, Tool, Resource, Prompt

    # Create gateways
    gateways = []
    for i in range(10):
        gw = Gateway(
            name=f"gateway-{i}",
            url=f"http://gateway-{i}.local:8000",
            status="active"
        )
        db_session.add(gw)
        gateways.append(gw)
    db_session.flush()

    # Create tools (100 per gateway = 1000 total)
    for gw in gateways:
        for j in range(100):
            tool = Tool(
                name=f"tool-{gw.id}-{j}",
                description=f"Test tool {j}",
                gateway_id=gw.id,
                input_schema={"type": "object"}
            )
            db_session.add(tool)

    db_session.commit()
    return {"gateways": 10, "tools": 1000}

# tests/performance/test_n_plus_one.py
import pytest
from mcpgateway.services.tool_service import ToolService

class TestN1Detection:
    """Tests to detect and prevent N+1 query patterns."""

    def test_list_tools_no_n_plus_one(self, performance_db, seed_performance_data):
        """Listing tools should not cause N+1 queries for gateways."""
        service = ToolService()

        # List all tools
        tools = service.list_tools(performance_db.session)

        # Access gateway for each tool (this would trigger N+1 if not eager loaded)
        for tool in tools:
            _ = tool.gateway.name if tool.gateway else None

        performance_db.print_summary()

        # Should be ~2-3 queries max, not 1000+
        assert performance_db.query_count < 10, (
            f"Potential N+1: {performance_db.query_count} queries for {len(tools)} tools"
        )

    def test_get_tool_with_relations(self, performance_db, seed_performance_data):
        """Getting a single tool should load relations efficiently."""
        service = ToolService()

        tool = service.get_tool(performance_db.session, tool_id="tool-1-0")

        # Access related objects
        _ = tool.gateway
        _ = tool.servers

        # Should be 1-3 queries, not more
        assert performance_db.query_count <= 3


class TestQueryCountRegression:
    """Regression tests for query counts."""

    QUERY_BUDGETS = {
        'list_tools': 5,
        'list_gateways': 3,
        'list_servers': 5,
        'get_tool_detail': 3,
    }

    def test_list_tools_query_budget(self, assert_max_queries, db_session):
        """Tool listing should stay within query budget."""
        from mcpgateway.services.tool_service import ToolService

        with assert_max_queries(self.QUERY_BUDGETS['list_tools']):
            ToolService().list_tools(db_session)

    def test_list_gateways_query_budget(self, assert_max_queries, db_session):
        """Gateway listing should stay within query budget."""
        from mcpgateway.services.gateway_service import GatewayService

        with assert_max_queries(self.QUERY_BUDGETS['list_gateways']):
            GatewayService().list_gateways(db_session)

# mcpgateway/middleware/query_logging.py
import logging
import time
from contextvars import ContextVar
from starlette.middleware.base import BaseHTTPMiddleware
from sqlalchemy import event

logger = logging.getLogger(__name__)

# Context var to track queries per request
request_queries: ContextVar[list] = ContextVar('request_queries', default=[])

class QueryLoggingMiddleware(BaseHTTPMiddleware):
    """Middleware to log query counts and timing per request."""

    async def dispatch(self, request, call_next):
        queries = []
        request_queries.set(queries)

        start_time = time.perf_counter()
        response = await call_next(request)
        total_time = (time.perf_counter() - start_time) * 1000

        query_count = len(queries)
        query_time = sum(q.get('duration_ms', 0) for q in queries)

        # Log warning for high query counts
        if query_count > 20:
            logger.warning(
                f"HIGH QUERY COUNT: {request.method} {request.url.path} "
                f"- {query_count} queries in {query_time:.2f}ms "
                f"(total: {total_time:.2f}ms)"
            )
        elif query_count > 10:
            logger.info(
                f"Query summary: {request.method} {request.url.path} "
                f"- {query_count} queries in {query_time:.2f}ms"
            )

        # Add query info to response headers (development only)
        response.headers['X-Query-Count'] = str(query_count)
        response.headers['X-Query-Time-Ms'] = f"{query_time:.2f}"

        return response


def setup_query_tracking(engine):
    """Setup event listeners to track queries per request."""

    @event.listens_for(engine, "before_cursor_execute")
    def before_execute(conn, cursor, statement, parameters, context, executemany):
        conn.info['query_start'] = time.perf_counter()

    @event.listens_for(engine, "after_cursor_execute")
    def after_execute(conn, cursor, statement, parameters, context, executemany):
        duration = (time.perf_counter() - conn.info.get('query_start', 0)) * 1000
        try:
            queries = request_queries.get()
            queries.append({
                'statement': statement,
                'duration_ms': duration,
            })
        except LookupError:
            pass  # No request context

# In a test or debug script
from tests.helpers.query_counter import count_queries
from mcpgateway.db import engine, SessionLocal

with count_queries(engine, print_queries=True) as counter:
    db = SessionLocal()
    # Your code here
    db.close()

print(f"\nTotal: {counter.count} queries")