Coverage for mcpgateway/db.py: 74%
395 statements
« prev ^ index » next coverage.py v7.9.2, created at 2025-07-09 11:03 +0100
« prev ^ index » next coverage.py v7.9.2, created at 2025-07-09 11:03 +0100
1# -*- coding: utf-8 -*-
2"""MCP Gateway Database Models.
4Copyright 2025
5SPDX-License-Identifier: Apache-2.0
6Authors: Mihai Criveti
8This module defines SQLAlchemy models for storing MCP entities including:
9- Tools with input schema validation
10- Resources with subscription tracking
11- Prompts with argument templates
12- Federated gateways with capability tracking
14Updated to record server associations independently using many-to-many relationships,
15and to record tool execution metrics.
16"""
18# Standard
19from datetime import datetime, timezone
20import re
21from typing import Any, Dict, List, Optional
22import uuid
24# Third-Party
25import jsonschema
26from sqlalchemy import (
27 Boolean,
28 Column,
29 create_engine,
30 DateTime,
31 event,
32 Float,
33 ForeignKey,
34 func,
35 Integer,
36 JSON,
37 make_url,
38 select,
39 String,
40 Table,
41 Text,
42 UniqueConstraint,
43)
44from sqlalchemy.event import listen
45from sqlalchemy.exc import SQLAlchemyError
46from sqlalchemy.ext.hybrid import hybrid_property
47from sqlalchemy.orm import (
48 DeclarativeBase,
49 Mapped,
50 mapped_column,
51 relationship,
52 sessionmaker,
53)
54from sqlalchemy.orm.attributes import get_history
56# First-Party
57from mcpgateway.config import settings
58from mcpgateway.models import ResourceContent
59from mcpgateway.utils.create_slug import slugify
60from mcpgateway.utils.db_isready import wait_for_db_ready
62# ---------------------------------------------------------------------------
63# 1. Parse the URL so we can inspect backend ("postgresql", "sqlite", ...)
64# and the specific driver ("psycopg2", "asyncpg", empty string = default).
65# ---------------------------------------------------------------------------
66url = make_url(settings.database_url)
67backend = url.get_backend_name() # e.g. 'postgresql', 'sqlite'
68driver = url.get_driver_name() or "default"
70# Start with an empty dict and add options only when the driver can accept
71# them; this prevents unexpected TypeError at connect time.
72connect_args: dict[str, object] = {}
74# ---------------------------------------------------------------------------
75# 2. PostgreSQL (synchronous psycopg2 only)
76# The keep-alive parameters below are recognised exclusively by libpq /
77# psycopg2 and let the kernel detect broken network links quickly.
78# ---------------------------------------------------------------------------
79if backend == "postgresql" and driver in ("psycopg2", "default", ""): 79 ↛ 80line 79 didn't jump to line 80 because the condition on line 79 was never true
80 connect_args.update(
81 keepalives=1, # enable TCP keep-alive probes
82 keepalives_idle=30, # seconds of idleness before first probe
83 keepalives_interval=5, # seconds between probes
84 keepalives_count=5, # drop the link after N failed probes
85 )
87# ---------------------------------------------------------------------------
88# 3. SQLite (optional) - only one extra flag and it is *SQLite-specific*.
89# ---------------------------------------------------------------------------
90elif backend == "sqlite": 90 ↛ 99line 90 didn't jump to line 99 because the condition on line 90 was always true
91 # Allow pooled connections to hop across threads.
92 connect_args["check_same_thread"] = False
94# 4. Other backends (MySQL, MSSQL, etc.) leave `connect_args` empty.
96# ---------------------------------------------------------------------------
97# 5. Build the Engine with a single, clean connect_args mapping.
98# ---------------------------------------------------------------------------
99engine = create_engine(
100 settings.database_url,
101 pool_pre_ping=True, # quick liveness check per checkout
102 pool_size=settings.db_pool_size,
103 max_overflow=settings.db_max_overflow,
104 pool_timeout=settings.db_pool_timeout,
105 pool_recycle=settings.db_pool_recycle,
106 connect_args=connect_args,
107)
110# ---------------------------------------------------------------------------
111# 6. Function to return UTC timestamp
112# ---------------------------------------------------------------------------
113def utc_now() -> datetime:
114 """Return the current Coordinated Universal Time (UTC).
116 Returns:
117 datetime: A timezone-aware `datetime` whose `tzinfo` is
118 `datetime.timezone.utc`.
119 """
120 return datetime.now(timezone.utc)
123# Session factory
124SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
127class Base(DeclarativeBase):
128 """Base class for all models."""
131# TODO: cleanup, not sure why this is commented out?
132# # Association table for tools and gateways (federation)
133# tool_gateway_table = Table(
134# "tool_gateway_association",
135# Base.metadata,
136# Column("tool_id", String, ForeignKey("tools.id"), primary_key=True),
137# Column("gateway_id", String, ForeignKey("gateways.id"), primary_key=True),
138# )
140# # Association table for resources and gateways (federation)
141# resource_gateway_table = Table(
142# "resource_gateway_association",
143# Base.metadata,
144# Column("resource_id", Integer, ForeignKey("resources.id"), primary_key=True),
145# Column("gateway_id", String, ForeignKey("gateways.id"), primary_key=True),
146# )
148# # Association table for prompts and gateways (federation)
149# prompt_gateway_table = Table(
150# "prompt_gateway_association",
151# Base.metadata,
152# Column("prompt_id", Integer, ForeignKey("prompts.id"), primary_key=True),
153# Column("gateway_id", String, ForeignKey("gateways.id"), primary_key=True),
154# )
156# Association table for servers and tools
157server_tool_association = Table(
158 "server_tool_association",
159 Base.metadata,
160 Column("server_id", String, ForeignKey("servers.id"), primary_key=True),
161 Column("tool_id", String, ForeignKey("tools.id"), primary_key=True),
162)
164# Association table for servers and resources
165server_resource_association = Table(
166 "server_resource_association",
167 Base.metadata,
168 Column("server_id", String, ForeignKey("servers.id"), primary_key=True),
169 Column("resource_id", Integer, ForeignKey("resources.id"), primary_key=True),
170)
172# Association table for servers and prompts
173server_prompt_association = Table(
174 "server_prompt_association",
175 Base.metadata,
176 Column("server_id", String, ForeignKey("servers.id"), primary_key=True),
177 Column("prompt_id", Integer, ForeignKey("prompts.id"), primary_key=True),
178)
181class ToolMetric(Base):
182 """
183 ORM model for recording individual metrics for tool executions.
185 Each record in this table corresponds to a single tool invocation and records:
186 - timestamp (datetime): When the invocation occurred.
187 - response_time (float): The execution time in seconds.
188 - is_success (bool): True if the execution succeeded, False otherwise.
189 - error_message (Optional[str]): Error message if the execution failed.
191 Aggregated metrics (such as total executions, successful/failed counts, failure rate,
192 minimum, maximum, and average response times, and last execution time) should be computed
193 on the fly using SQL aggregate functions over the rows in this table.
194 """
196 __tablename__ = "tool_metrics"
198 id: Mapped[int] = mapped_column(primary_key=True)
199 tool_id: Mapped[str] = mapped_column(String, ForeignKey("tools.id"), nullable=False)
200 timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now)
201 response_time: Mapped[float] = mapped_column(Float, nullable=False)
202 is_success: Mapped[bool] = mapped_column(Boolean, nullable=False)
203 error_message: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
205 # Relationship back to the Tool model.
206 tool: Mapped["Tool"] = relationship("Tool", back_populates="metrics")
209class ResourceMetric(Base):
210 """
211 ORM model for recording metrics for resource invocations.
213 Attributes:
214 id (int): Primary key.
215 resource_id (int): Foreign key linking to the resource.
216 timestamp (datetime): The time when the invocation occurred.
217 response_time (float): The response time in seconds.
218 is_success (bool): True if the invocation succeeded, False otherwise.
219 error_message (Optional[str]): Error message if the invocation failed.
220 """
222 __tablename__ = "resource_metrics"
224 id: Mapped[int] = mapped_column(primary_key=True)
225 resource_id: Mapped[int] = mapped_column(Integer, ForeignKey("resources.id"), nullable=False)
226 timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now)
227 response_time: Mapped[float] = mapped_column(Float, nullable=False)
228 is_success: Mapped[bool] = mapped_column(Boolean, nullable=False)
229 error_message: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
231 # Relationship back to the Resource model.
232 resource: Mapped["Resource"] = relationship("Resource", back_populates="metrics")
235class ServerMetric(Base):
236 """
237 ORM model for recording metrics for server invocations.
239 Attributes:
240 id (int): Primary key.
241 server_id (str): Foreign key linking to the server.
242 timestamp (datetime): The time when the invocation occurred.
243 response_time (float): The response time in seconds.
244 is_success (bool): True if the invocation succeeded, False otherwise.
245 error_message (Optional[str]): Error message if the invocation failed.
246 """
248 __tablename__ = "server_metrics"
250 id: Mapped[int] = mapped_column(primary_key=True)
251 server_id: Mapped[str] = mapped_column(String, ForeignKey("servers.id"), nullable=False)
252 timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now)
253 response_time: Mapped[float] = mapped_column(Float, nullable=False)
254 is_success: Mapped[bool] = mapped_column(Boolean, nullable=False)
255 error_message: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
257 # Relationship back to the Server model.
258 server: Mapped["Server"] = relationship("Server", back_populates="metrics")
261class PromptMetric(Base):
262 """
263 ORM model for recording metrics for prompt invocations.
265 Attributes:
266 id (int): Primary key.
267 prompt_id (int): Foreign key linking to the prompt.
268 timestamp (datetime): The time when the invocation occurred.
269 response_time (float): The response time in seconds.
270 is_success (bool): True if the invocation succeeded, False otherwise.
271 error_message (Optional[str]): Error message if the invocation failed.
272 """
274 __tablename__ = "prompt_metrics"
276 id: Mapped[int] = mapped_column(primary_key=True)
277 prompt_id: Mapped[int] = mapped_column(Integer, ForeignKey("prompts.id"), nullable=False)
278 timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now)
279 response_time: Mapped[float] = mapped_column(Float, nullable=False)
280 is_success: Mapped[bool] = mapped_column(Boolean, nullable=False)
281 error_message: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
283 # Relationship back to the Prompt model.
284 prompt: Mapped["Prompt"] = relationship("Prompt", back_populates="metrics")
287class Tool(Base):
288 """
289 ORM model for a registered Tool.
291 Supports both local tools and federated tools from other gateways.
292 The integration_type field indicates the tool format:
293 - "MCP" for MCP-compliant tools (default)
294 - "REST" for REST tools
296 Additionally, this model provides computed properties for aggregated metrics based
297 on the associated ToolMetric records. These include:
298 - execution_count: Total number of invocations.
299 - successful_executions: Count of successful invocations.
300 - failed_executions: Count of failed invocations.
301 - failure_rate: Ratio of failed invocations to total invocations.
302 - min_response_time: Fastest recorded response time.
303 - max_response_time: Slowest recorded response time.
304 - avg_response_time: Mean response time.
305 - last_execution_time: Timestamp of the most recent invocation.
307 The property `metrics_summary` returns a dictionary with these aggregated values.
309 The following fields have been added to support tool invocation configuration:
310 - request_type: HTTP method to use when invoking the tool.
311 - auth_type: Type of authentication ("basic", "bearer", or None).
312 - auth_username: Username for basic authentication.
313 - auth_password: Password for basic authentication.
314 - auth_token: Token for bearer token authentication.
315 - auth_header_key: header key for authentication.
316 - auth_header_value: header value for authentication.
317 """
319 __tablename__ = "tools"
321 id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: uuid.uuid4().hex)
322 original_name: Mapped[str] = mapped_column(String, nullable=False)
323 original_name_slug: Mapped[str] = mapped_column(String, nullable=False)
324 url: Mapped[str] = mapped_column(String, nullable=True)
325 description: Mapped[Optional[str]]
326 integration_type: Mapped[str] = mapped_column(default="MCP")
327 request_type: Mapped[str] = mapped_column(default="SSE")
328 headers: Mapped[Optional[Dict[str, str]]] = mapped_column(JSON)
329 input_schema: Mapped[Dict[str, Any]] = mapped_column(JSON)
330 annotations: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, default=lambda: {})
331 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now)
332 updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, onupdate=utc_now)
333 enabled: Mapped[bool] = mapped_column(default=True)
334 reachable: Mapped[bool] = mapped_column(default=True)
335 jsonpath_filter: Mapped[str] = mapped_column(default="")
337 # Request type and authentication fields
338 auth_type: Mapped[Optional[str]] = mapped_column(default=None) # "basic", "bearer", or None
339 auth_value: Mapped[Optional[str]] = mapped_column(default=None)
341 # Federation relationship with a local gateway
342 gateway_id: Mapped[Optional[str]] = mapped_column(ForeignKey("gateways.id"))
343 # gateway_slug: Mapped[Optional[str]] = mapped_column(ForeignKey("gateways.slug"))
344 gateway: Mapped["Gateway"] = relationship("Gateway", primaryjoin="Tool.gateway_id == Gateway.id", foreign_keys=[gateway_id], back_populates="tools")
345 # federated_with = relationship("Gateway", secondary=tool_gateway_table, back_populates="federated_tools")
347 # Many-to-many relationship with Servers
348 servers: Mapped[List["Server"]] = relationship("Server", secondary=server_tool_association, back_populates="tools")
350 # Relationship with ToolMetric records
351 metrics: Mapped[List["ToolMetric"]] = relationship("ToolMetric", back_populates="tool", cascade="all, delete-orphan")
353 # @property
354 # def gateway_slug(self) -> str:
355 # return self.gateway.slug
357 _computed_name = Column("name", String, unique=True) # Stored column
359 @hybrid_property
360 def name(self):
361 """Return the display/lookup name.
363 Returns:
364 str: Name to display
365 """
366 if self._computed_name: # pylint: disable=no-member 366 ↛ 367line 366 didn't jump to line 367 because the condition on line 366 was never true
367 return self._computed_name # orm column, resolved at runtime
369 original_slug = slugify(self.original_name) # pylint: disable=no-member
371 # Gateway present → prepend its slug and the configured separator
372 if self.gateway_id: # pylint: disable=no-member
373 gateway_slug = slugify(self.gateway.name) # pylint: disable=no-member
374 return f"{gateway_slug}{settings.gateway_tool_name_separator}{original_slug}"
376 # No gateway → only the original name slug
377 return original_slug
379 @name.setter
380 def name(self, value):
381 """Store an explicit value that overrides the calculated one.
383 Args:
384 value (str): Value to set to _computed_name
385 """
386 self._computed_name = value
388 @name.expression
389 def name(cls): # pylint: disable=no-self-argument
390 """
391 SQL expression used when the hybrid appears in a filter/order_by.
392 Simply forwards to the ``_computed_name`` column; the Python-side
393 reconstruction above is not needed on the SQL side.
395 Returns:
396 str: computed name for SQL use
397 """
398 return cls._computed_name
400 __table_args__ = (UniqueConstraint("gateway_id", "original_name", name="uq_gateway_id__original_name"),)
402 @hybrid_property
403 def gateway_slug(self):
404 """Always returns the current slug from the related Gateway
406 Returns:
407 str: slug for Python use
408 """
409 return self.gateway.slug if self.gateway else None
411 @gateway_slug.expression
412 def gateway_slug(cls): # pylint: disable=no-self-argument
413 """For database queries - auto-joins to get current slug
415 Returns:
416 str: slug for SQL use
417 """
418 return select(Gateway.slug).where(Gateway.id == cls.gateway_id).scalar_subquery()
420 @hybrid_property
421 def execution_count(self) -> int:
422 """
423 Returns the number of times the tool has been executed,
424 calculated from the associated ToolMetric records.
426 Returns:
427 int: The total count of tool executions.
428 """
429 return len(self.metrics)
431 @execution_count.expression
432 # method is intentionally a class-level expression, so no `self`
433 # pylint: disable=no-self-argument
434 def execution_count(cls):
435 """
436 SQL expression to compute the execution count for the tool.
438 Returns:
439 int: Returns execution count of a given tool
440 """
441 return select(func.count(ToolMetric.id)).where(ToolMetric.tool_id == cls.id).label("execution_count") # pylint: disable=not-callable
443 @property
444 def successful_executions(self) -> int:
445 """
446 Returns the count of successful tool executions,
447 computed from the associated ToolMetric records.
449 Returns:
450 int: The count of successful tool executions.
451 """
452 return sum(1 for m in self.metrics if m.is_success)
454 @property
455 def failed_executions(self) -> int:
456 """
457 Returns the count of failed tool executions,
458 computed from the associated ToolMetric records.
460 Returns:
461 int: The count of failed tool executions.
462 """
463 return sum(1 for m in self.metrics if not m.is_success)
465 @property
466 def failure_rate(self) -> float:
467 """
468 Returns the failure rate (as a float between 0 and 1) computed as:
469 (failed executions) / (total executions).
470 Returns 0.0 if there are no executions.
472 Returns:
473 float: The failure rate as a value between 0 and 1.
474 """
475 total: int = self.execution_count
476 # execution_count is a @hybrid_property, not a callable here
477 if total == 0: # pylint: disable=comparison-with-callable 477 ↛ 479line 477 didn't jump to line 479 because the condition on line 477 was always true
478 return 0.0
479 return self.failed_executions / total
481 @property
482 def min_response_time(self) -> Optional[float]:
483 """
484 Returns the minimum response time among all tool executions.
485 Returns None if no executions exist.
487 Returns:
488 Optional[float]: The minimum response time, or None if no executions exist.
489 """
490 times: List[float] = [m.response_time for m in self.metrics]
491 return min(times) if times else None
493 @property
494 def max_response_time(self) -> Optional[float]:
495 """
496 Returns the maximum response time among all tool executions.
497 Returns None if no executions exist.
499 Returns:
500 Optional[float]: The maximum response time, or None if no executions exist.
501 """
502 times: List[float] = [m.response_time for m in self.metrics]
503 return max(times) if times else None
505 @property
506 def avg_response_time(self) -> Optional[float]:
507 """
508 Returns the average response time among all tool executions.
509 Returns None if no executions exist.
511 Returns:
512 Optional[float]: The average response time, or None if no executions exist.
513 """
514 times: List[float] = [m.response_time for m in self.metrics]
515 return sum(times) / len(times) if times else None
517 @property
518 def last_execution_time(self) -> Optional[datetime]:
519 """
520 Returns the timestamp of the most recent tool execution.
521 Returns None if no executions exist.
523 Returns:
524 Optional[datetime]: The timestamp of the most recent execution, or None if no executions exist.
525 """
526 if not self.metrics: 526 ↛ 528line 526 didn't jump to line 528 because the condition on line 526 was always true
527 return None
528 return max(m.timestamp for m in self.metrics)
530 @property
531 def metrics_summary(self) -> Dict[str, Any]:
532 """
533 Returns aggregated metrics for the tool as a dictionary with the following keys:
534 - total_executions: Total number of invocations.
535 - successful_executions: Number of successful invocations.
536 - failed_executions: Number of failed invocations.
537 - failure_rate: Failure rate (failed/total) or 0.0 if no invocations.
538 - min_response_time: Minimum response time (or None if no invocations).
539 - max_response_time: Maximum response time (or None if no invocations).
540 - avg_response_time: Average response time (or None if no invocations).
541 - last_execution_time: Timestamp of the most recent invocation (or None).
543 Returns:
544 Dict[str, Any]: Dictionary containing the aggregated metrics.
545 """
546 return {
547 "total_executions": self.execution_count,
548 "successful_executions": self.successful_executions,
549 "failed_executions": self.failed_executions,
550 "failure_rate": self.failure_rate,
551 "min_response_time": self.min_response_time,
552 "max_response_time": self.max_response_time,
553 "avg_response_time": self.avg_response_time,
554 "last_execution_time": self.last_execution_time,
555 }
558class Resource(Base):
559 """
560 ORM model for a registered Resource.
562 Resources represent content that can be read by clients.
563 Supports subscriptions for real-time updates.
564 Additionally, this model provides a relationship with ResourceMetric records
565 to capture invocation metrics (such as execution counts, response times, and failures).
566 """
568 __tablename__ = "resources"
570 id: Mapped[int] = mapped_column(primary_key=True)
571 uri: Mapped[str] = mapped_column(unique=True)
572 name: Mapped[str]
573 description: Mapped[Optional[str]]
574 mime_type: Mapped[Optional[str]]
575 size: Mapped[Optional[int]]
576 template: Mapped[Optional[str]] # URI template for parameterized resources
577 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now)
578 updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, onupdate=utc_now)
579 is_active: Mapped[bool] = mapped_column(default=True)
580 metrics: Mapped[List["ResourceMetric"]] = relationship("ResourceMetric", back_populates="resource", cascade="all, delete-orphan")
582 # Content storage - can be text or binary
583 text_content: Mapped[Optional[str]] = mapped_column(Text)
584 binary_content: Mapped[Optional[bytes]]
586 # Subscription tracking
587 subscriptions: Mapped[List["ResourceSubscription"]] = relationship("ResourceSubscription", back_populates="resource", cascade="all, delete-orphan")
589 gateway_id: Mapped[Optional[str]] = mapped_column(ForeignKey("gateways.id"))
590 gateway: Mapped["Gateway"] = relationship("Gateway", back_populates="resources")
591 # federated_with = relationship("Gateway", secondary=resource_gateway_table, back_populates="federated_resources")
593 # Many-to-many relationship with Servers
594 servers: Mapped[List["Server"]] = relationship("Server", secondary=server_resource_association, back_populates="resources")
596 @property
597 def content(self) -> ResourceContent:
598 """
599 Returns the resource content in the appropriate format.
601 If text content exists, returns a ResourceContent with text.
602 Otherwise, if binary content exists, returns a ResourceContent with blob data.
603 Raises a ValueError if no content is available.
605 Returns:
606 ResourceContent: The resource content with appropriate format (text or blob).
608 Raises:
609 ValueError: If the resource has no content available.
610 """
612 if self.text_content is not None:
613 return ResourceContent(
614 type="resource",
615 uri=self.uri,
616 mime_type=self.mime_type,
617 text=self.text_content,
618 )
619 if self.binary_content is not None:
620 return ResourceContent(
621 type="resource",
622 uri=self.uri,
623 mime_type=self.mime_type or "application/octet-stream",
624 blob=self.binary_content,
625 )
626 raise ValueError("Resource has no content")
628 @property
629 def execution_count(self) -> int:
630 """
631 Returns the number of times the resource has been invoked,
632 calculated from the associated ResourceMetric records.
634 Returns:
635 int: The total count of resource invocations.
636 """
637 return len(self.metrics)
639 @property
640 def successful_executions(self) -> int:
641 """
642 Returns the count of successful resource invocations,
643 computed from the associated ResourceMetric records.
645 Returns:
646 int: The count of successful resource invocations.
647 """
648 return sum(1 for m in self.metrics if m.is_success)
650 @property
651 def failed_executions(self) -> int:
652 """
653 Returns the count of failed resource invocations,
654 computed from the associated ResourceMetric records.
656 Returns:
657 int: The count of failed resource invocations.
658 """
659 return sum(1 for m in self.metrics if not m.is_success)
661 @property
662 def failure_rate(self) -> float:
663 """
664 Returns the failure rate (as a float between 0 and 1) computed as:
665 (failed invocations) / (total invocations).
666 Returns 0.0 if there are no invocations.
668 Returns:
669 float: The failure rate as a value between 0 and 1.
670 """
671 total: int = self.execution_count
672 if total == 0:
673 return 0.0
674 return self.failed_executions / total
676 @property
677 def min_response_time(self) -> Optional[float]:
678 """
679 Returns the minimum response time among all resource invocations.
680 Returns None if no invocations exist.
682 Returns:
683 Optional[float]: The minimum response time, or None if no invocations exist.
684 """
685 times: List[float] = [m.response_time for m in self.metrics]
686 return min(times) if times else None
688 @property
689 def max_response_time(self) -> Optional[float]:
690 """
691 Returns the maximum response time among all resource invocations.
692 Returns None if no invocations exist.
694 Returns:
695 Optional[float]: The maximum response time, or None if no invocations exist.
696 """
697 times: List[float] = [m.response_time for m in self.metrics]
698 return max(times) if times else None
700 @property
701 def avg_response_time(self) -> Optional[float]:
702 """
703 Returns the average response time among all resource invocations.
704 Returns None if no invocations exist.
706 Returns:
707 Optional[float]: The average response time, or None if no invocations exist.
708 """
709 times: List[float] = [m.response_time for m in self.metrics]
710 return sum(times) / len(times) if times else None
712 @property
713 def last_execution_time(self) -> Optional[datetime]:
714 """
715 Returns the timestamp of the most recent resource invocation.
716 Returns None if no invocations exist.
718 Returns:
719 Optional[datetime]: The timestamp of the most recent invocation, or None if no invocations exist.
720 """
721 if not self.metrics:
722 return None
723 return max(m.timestamp for m in self.metrics)
726class ResourceSubscription(Base):
727 """Tracks subscriptions to resource updates."""
729 __tablename__ = "resource_subscriptions"
731 id: Mapped[int] = mapped_column(primary_key=True)
732 resource_id: Mapped[int] = mapped_column(ForeignKey("resources.id"))
733 subscriber_id: Mapped[str] # Client identifier
734 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now)
735 last_notification: Mapped[Optional[datetime]] = mapped_column(DateTime)
737 resource: Mapped["Resource"] = relationship(back_populates="subscriptions")
740class Prompt(Base):
741 """
742 ORM model for a registered Prompt template.
744 Represents a prompt template along with its argument schema.
745 Supports rendering and invocation of prompts.
746 Additionally, this model provides computed properties for aggregated metrics based
747 on the associated PromptMetric records. These include:
748 - execution_count: Total number of prompt invocations.
749 - successful_executions: Count of successful invocations.
750 - failed_executions: Count of failed invocations.
751 - failure_rate: Ratio of failed invocations to total invocations.
752 - min_response_time: Fastest recorded response time.
753 - max_response_time: Slowest recorded response time.
754 - avg_response_time: Mean response time.
755 - last_execution_time: Timestamp of the most recent invocation.
756 """
758 __tablename__ = "prompts"
760 id: Mapped[int] = mapped_column(primary_key=True)
761 name: Mapped[str] = mapped_column(unique=True)
762 description: Mapped[Optional[str]]
763 template: Mapped[str] = mapped_column(Text)
764 argument_schema: Mapped[Dict[str, Any]] = mapped_column(JSON)
765 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now)
766 updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, onupdate=utc_now)
767 is_active: Mapped[bool] = mapped_column(default=True)
768 metrics: Mapped[List["PromptMetric"]] = relationship("PromptMetric", back_populates="prompt", cascade="all, delete-orphan")
770 gateway_id: Mapped[Optional[str]] = mapped_column(ForeignKey("gateways.id"))
771 gateway: Mapped["Gateway"] = relationship("Gateway", back_populates="prompts")
772 # federated_with = relationship("Gateway", secondary=prompt_gateway_table, back_populates="federated_prompts")
774 # Many-to-many relationship with Servers
775 servers: Mapped[List["Server"]] = relationship("Server", secondary=server_prompt_association, back_populates="prompts")
777 def validate_arguments(self, args: Dict[str, str]) -> None:
778 """
779 Validate prompt arguments against the argument schema.
781 Args:
782 args (Dict[str, str]): Dictionary of arguments to validate.
784 Raises:
785 ValueError: If the arguments do not conform to the schema.
787 """
788 try:
789 jsonschema.validate(args, self.argument_schema)
790 except jsonschema.exceptions.ValidationError as e:
791 raise ValueError(f"Invalid prompt arguments: {str(e)}")
793 @property
794 def execution_count(self) -> int:
795 """
796 Returns the number of times the prompt has been invoked,
797 calculated from the associated PromptMetric records.
799 Returns:
800 int: The total count of prompt invocations.
801 """
802 return len(self.metrics)
804 @property
805 def successful_executions(self) -> int:
806 """
807 Returns the count of successful prompt invocations,
808 computed from the associated PromptMetric records.
810 Returns:
811 int: The count of successful prompt invocations.
812 """
813 return sum(1 for m in self.metrics if m.is_success)
815 @property
816 def failed_executions(self) -> int:
817 """
818 Returns the count of failed prompt invocations,
819 computed from the associated PromptMetric records.
821 Returns:
822 int: The count of failed prompt invocations.
823 """
824 return sum(1 for m in self.metrics if not m.is_success)
826 @property
827 def failure_rate(self) -> float:
828 """
829 Returns the failure rate (as a float between 0 and 1) computed as:
830 (failed invocations) / (total invocations).
831 Returns 0.0 if there are no invocations.
833 Returns:
834 float: The failure rate as a value between 0 and 1.
835 """
836 total: int = self.execution_count
837 if total == 0:
838 return 0.0
839 return self.failed_executions / total
841 @property
842 def min_response_time(self) -> Optional[float]:
843 """
844 Returns the minimum response time among all prompt invocations.
845 Returns None if no invocations exist.
847 Returns:
848 Optional[float]: The minimum response time, or None if no invocations exist.
849 """
850 times: List[float] = [m.response_time for m in self.metrics]
851 return min(times) if times else None
853 @property
854 def max_response_time(self) -> Optional[float]:
855 """
856 Returns the maximum response time among all prompt invocations.
857 Returns None if no invocations exist.
859 Returns:
860 Optional[float]: The maximum response time, or None if no invocations exist.
861 """
862 times: List[float] = [m.response_time for m in self.metrics]
863 return max(times) if times else None
865 @property
866 def avg_response_time(self) -> Optional[float]:
867 """
868 Returns the average response time among all prompt invocations.
869 Returns None if no invocations exist.
871 Returns:
872 Optional[float]: The average response time, or None if no invocations exist.
873 """
874 times: List[float] = [m.response_time for m in self.metrics]
875 return sum(times) / len(times) if times else None
877 @property
878 def last_execution_time(self) -> Optional[datetime]:
879 """
880 Returns the timestamp of the most recent prompt invocation.
881 Returns None if no invocations exist.
883 Returns:
884 Optional[datetime]: The timestamp of the most recent invocation, or None if no invocations exist.
885 """
886 if not self.metrics:
887 return None
888 return max(m.timestamp for m in self.metrics)
891class Server(Base):
892 """
893 ORM model for MCP Servers Catalog.
895 Represents a server that composes catalog items (tools, resources, prompts).
896 Additionally, this model provides computed properties for aggregated metrics based
897 on the associated ServerMetric records. These include:
898 - execution_count: Total number of invocations.
899 - successful_executions: Count of successful invocations.
900 - failed_executions: Count of failed invocations.
901 - failure_rate: Ratio of failed invocations to total invocations.
902 - min_response_time: Fastest recorded response time.
903 - max_response_time: Slowest recorded response time.
904 - avg_response_time: Mean response time.
905 - last_execution_time: Timestamp of the most recent invocation.
906 """
908 __tablename__ = "servers"
910 id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: uuid.uuid4().hex)
911 name: Mapped[str] = mapped_column(unique=True)
912 description: Mapped[Optional[str]]
913 icon: Mapped[Optional[str]]
914 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now)
915 updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, onupdate=utc_now)
916 is_active: Mapped[bool] = mapped_column(default=True)
917 metrics: Mapped[List["ServerMetric"]] = relationship("ServerMetric", back_populates="server", cascade="all, delete-orphan")
919 # Many-to-many relationships for associated items
920 tools: Mapped[List["Tool"]] = relationship("Tool", secondary=server_tool_association, back_populates="servers")
921 resources: Mapped[List["Resource"]] = relationship("Resource", secondary=server_resource_association, back_populates="servers")
922 prompts: Mapped[List["Prompt"]] = relationship("Prompt", secondary=server_prompt_association, back_populates="servers")
924 @property
925 def execution_count(self) -> int:
926 """
927 Returns the number of times the server has been invoked,
928 calculated from the associated ServerMetric records.
930 Returns:
931 int: The total count of server invocations.
932 """
933 return len(self.metrics)
935 @property
936 def successful_executions(self) -> int:
937 """
938 Returns the count of successful server invocations,
939 computed from the associated ServerMetric records.
941 Returns:
942 int: The count of successful server invocations.
943 """
944 return sum(1 for m in self.metrics if m.is_success)
946 @property
947 def failed_executions(self) -> int:
948 """
949 Returns the count of failed server invocations,
950 computed from the associated ServerMetric records.
952 Returns:
953 int: The count of failed server invocations.
954 """
955 return sum(1 for m in self.metrics if not m.is_success)
957 @property
958 def failure_rate(self) -> float:
959 """
960 Returns the failure rate (as a float between 0 and 1) computed as:
961 (failed invocations) / (total invocations).
962 Returns 0.0 if there are no invocations.
964 Returns:
965 float: The failure rate as a value between 0 and 1.
966 """
967 total: int = self.execution_count
968 if total == 0:
969 return 0.0
970 return self.failed_executions / total
972 @property
973 def min_response_time(self) -> Optional[float]:
974 """
975 Returns the minimum response time among all server invocations.
976 Returns None if no invocations exist.
978 Returns:
979 Optional[float]: The minimum response time, or None if no invocations exist.
980 """
981 times: List[float] = [m.response_time for m in self.metrics]
982 return min(times) if times else None
984 @property
985 def max_response_time(self) -> Optional[float]:
986 """
987 Returns the maximum response time among all server invocations.
988 Returns None if no invocations exist.
990 Returns:
991 Optional[float]: The maximum response time, or None if no invocations exist.
992 """
993 times: List[float] = [m.response_time for m in self.metrics]
994 return max(times) if times else None
996 @property
997 def avg_response_time(self) -> Optional[float]:
998 """
999 Returns the average response time among all server invocations.
1000 Returns None if no invocations exist.
1002 Returns:
1003 Optional[float]: The average response time, or None if no invocations exist.
1004 """
1005 times: List[float] = [m.response_time for m in self.metrics]
1006 return sum(times) / len(times) if times else None
1008 @property
1009 def last_execution_time(self) -> Optional[datetime]:
1010 """
1011 Returns the timestamp of the most recent server invocation.
1012 Returns None if no invocations exist.
1014 Returns:
1015 Optional[datetime]: The timestamp of the most recent invocation, or None if no invocations exist.
1016 """
1017 if not self.metrics:
1018 return None
1019 return max(m.timestamp for m in self.metrics)
1022class Gateway(Base):
1023 """ORM model for a federated peer Gateway."""
1025 __tablename__ = "gateways"
1027 id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: uuid.uuid4().hex)
1028 name: Mapped[str] = mapped_column(String, nullable=False)
1029 slug: Mapped[str] = mapped_column(String, nullable=False, unique=True)
1030 url: Mapped[str] = mapped_column(String, unique=True)
1031 description: Mapped[Optional[str]]
1032 transport: Mapped[str] = mapped_column(default="SSE")
1033 capabilities: Mapped[Dict[str, Any]] = mapped_column(JSON)
1034 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now)
1035 updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, onupdate=utc_now)
1036 enabled: Mapped[bool] = mapped_column(default=True)
1037 reachable: Mapped[bool] = mapped_column(default=True)
1038 last_seen: Mapped[Optional[datetime]]
1040 # Relationship with local tools this gateway provides
1041 tools: Mapped[List["Tool"]] = relationship(back_populates="gateway", foreign_keys="Tool.gateway_id", cascade="all, delete-orphan")
1043 # Relationship with local prompts this gateway provides
1044 prompts: Mapped[List["Prompt"]] = relationship(back_populates="gateway", cascade="all, delete-orphan")
1046 # Relationship with local resources this gateway provides
1047 resources: Mapped[List["Resource"]] = relationship(back_populates="gateway", cascade="all, delete-orphan")
1049 # # Tools federated from this gateway
1050 # federated_tools: Mapped[List["Tool"]] = relationship(secondary=tool_gateway_table, back_populates="federated_with")
1052 # # Prompts federated from this resource
1053 # federated_resources: Mapped[List["Resource"]] = relationship(secondary=resource_gateway_table, back_populates="federated_with")
1055 # # Prompts federated from this gateway
1056 # federated_prompts: Mapped[List["Prompt"]] = relationship(secondary=prompt_gateway_table, back_populates="federated_with")
1058 # Authorizations
1059 auth_type: Mapped[Optional[str]] = mapped_column(default=None) # "basic", "bearer", "headers" or None
1060 auth_value: Mapped[Optional[Dict[str, str]]] = mapped_column(JSON)
1063@event.listens_for(Gateway, "after_update")
1064def update_tool_names_on_gateway_update(_mapper, connection, target):
1065 """
1066 If a Gateway's name is updated, efficiently update all of its
1067 child Tools' names with a single SQL statement.
1069 Args:
1070 _mapper: Mapper
1071 connection: Connection
1072 target: Target
1073 """
1074 # 1. Check if the 'name' field was actually part of the update.
1075 # This is a concise way to see if the value has changed.
1076 if not get_history(target, "name").has_changes():
1077 return
1079 print(f"Gateway name changed for ID {target.id}. Issuing bulk update for tools.")
1081 # 2. Get a reference to the underlying database table for Tools
1082 tools_table = Tool.__table__
1084 # 3. Prepare the new values
1085 new_gateway_slug = slugify(target.name)
1086 separator = settings.gateway_tool_name_separator
1088 # 4. Construct a single, powerful UPDATE statement using SQLAlchemy Core.
1089 # This is highly efficient as it all happens in the database.
1090 stmt = (
1091 tools_table.update()
1092 .where(tools_table.c.gateway_id == target.id)
1093 .values(name=new_gateway_slug + separator + tools_table.c.original_name_slug)
1094 .execution_options(synchronize_session=False) # Important for bulk updates
1095 )
1097 # 5. Execute the statement using the connection from the ongoing transaction.
1098 connection.execute(stmt)
1101class SessionRecord(Base):
1102 """ORM model for sessions from SSE client."""
1104 __tablename__ = "mcp_sessions"
1106 session_id: Mapped[str] = mapped_column(primary_key=True)
1107 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now) # pylint: disable=not-callable
1108 last_accessed: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, onupdate=utc_now) # pylint: disable=not-callable
1109 data: Mapped[str] = mapped_column(String, nullable=True)
1111 messages: Mapped[List["SessionMessageRecord"]] = relationship("SessionMessageRecord", back_populates="session", cascade="all, delete-orphan")
1114class SessionMessageRecord(Base):
1115 """ORM model for messages from SSE client."""
1117 __tablename__ = "mcp_messages"
1119 id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
1120 session_id: Mapped[str] = mapped_column(ForeignKey("mcp_sessions.session_id"))
1121 message: Mapped[str] = mapped_column(String, nullable=True)
1122 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now) # pylint: disable=not-callable
1123 last_accessed: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, onupdate=utc_now) # pylint: disable=not-callable
1125 session: Mapped["SessionRecord"] = relationship("SessionRecord", back_populates="messages")
1128# Event listeners for validation
1129def validate_tool_schema(mapper, connection, target):
1130 """
1131 Validate tool schema before insert/update.
1133 Args:
1134 mapper: The mapper being used for the operation.
1135 connection: The database connection.
1136 target: The target object being validated.
1138 Raises:
1139 ValueError: If the tool input schema is invalid.
1140 """
1141 # You can use mapper and connection later, if required.
1142 _ = mapper
1143 _ = connection
1144 if hasattr(target, "input_schema"): 1144 ↛ exitline 1144 didn't return from function 'validate_tool_schema' because the condition on line 1144 was always true
1145 try:
1146 jsonschema.Draft7Validator.check_schema(target.input_schema)
1147 except jsonschema.exceptions.SchemaError as e:
1148 raise ValueError(f"Invalid tool input schema: {str(e)}")
1151def validate_tool_name(mapper, connection, target):
1152 """
1153 Validate tool name before insert/update. Check if the name matches the required pattern.
1155 Args:
1156 mapper: The mapper being used for the operation.
1157 connection: The database connection.
1158 target: The target object being validated.
1160 Raises:
1161 ValueError: If the tool name contains invalid characters.
1162 """
1163 # You can use mapper and connection later, if required.
1164 _ = mapper
1165 _ = connection
1166 if hasattr(target, "name"): 1166 ↛ exitline 1166 didn't return from function 'validate_tool_name' because the condition on line 1166 was always true
1167 if not re.match(r"^[a-zA-Z0-9_-]+$", target.name): 1167 ↛ 1168line 1167 didn't jump to line 1168 because the condition on line 1167 was never true
1168 raise ValueError(f"Invalid tool name '{target.name}'. Only alphanumeric characters, hyphens, and underscores are allowed.")
1171def validate_prompt_schema(mapper, connection, target):
1172 """
1173 Validate prompt argument schema before insert/update.
1175 Args:
1176 mapper: The mapper being used for the operation.
1177 connection: The database connection.
1178 target: The target object being validated.
1180 Raises:
1181 ValueError: If the prompt argument schema is invalid.
1182 """
1183 # You can use mapper and connection later, if required.
1184 _ = mapper
1185 _ = connection
1186 if hasattr(target, "argument_schema"):
1187 try:
1188 jsonschema.Draft7Validator.check_schema(target.argument_schema)
1189 except jsonschema.exceptions.SchemaError as e:
1190 raise ValueError(f"Invalid prompt argument schema: {str(e)}")
1193# Register validation listeners
1195listen(Tool, "before_insert", validate_tool_schema)
1196listen(Tool, "before_update", validate_tool_schema)
1197listen(Tool, "before_insert", validate_tool_name)
1198listen(Tool, "before_update", validate_tool_name)
1199listen(Prompt, "before_insert", validate_prompt_schema)
1200listen(Prompt, "before_update", validate_prompt_schema)
1203def get_db():
1204 """
1205 Dependency to get database session.
1207 Yields:
1208 SessionLocal: A SQLAlchemy database session.
1209 """
1210 db = SessionLocal()
1211 try:
1212 yield db
1213 finally:
1214 db.close()
1217# Create all tables
1218def init_db():
1219 """
1220 Initialize database tables.
1222 Raises:
1223 Exception: If database initialization fails.
1224 """
1225 try:
1226 # Base.metadata.drop_all(bind=engine)
1227 Base.metadata.create_all(bind=engine)
1228 except SQLAlchemyError as e:
1229 raise Exception(f"Failed to initialize database: {str(e)}")
1232if __name__ == "__main__":
1233 # Wait for database to be ready before initializing
1234 wait_for_db_ready(max_tries=int(settings.db_max_retries), interval=int(settings.db_retry_interval_ms) / 1000, sync=True) # Converting ms to s
1236 init_db()