Coverage for mcpgateway/db.py: 67%
351 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-22 12:53 +0100
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-22 12:53 +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"""
18import re
19from datetime import datetime, timezone
20from typing import Any, Dict, List, Optional
22import jsonschema
23from sqlalchemy import (
24 JSON,
25 Boolean,
26 Column,
27 DateTime,
28 Float,
29 ForeignKey,
30 Integer,
31 String,
32 Table,
33 Text,
34 create_engine,
35 func,
36 make_url,
37 select,
38)
39from sqlalchemy.event import listen
40from sqlalchemy.exc import SQLAlchemyError
41from sqlalchemy.ext.hybrid import hybrid_property
42from sqlalchemy.orm import (
43 DeclarativeBase,
44 Mapped,
45 mapped_column,
46 relationship,
47 sessionmaker,
48)
50from mcpgateway.config import settings
51from mcpgateway.types import ResourceContent
53# ---------------------------------------------------------------------------
54# 1. Parse the URL so we can inspect backend ("postgresql", "sqlite", …)
55# and the specific driver ("psycopg2", "asyncpg", empty string = default).
56# ---------------------------------------------------------------------------
57url = make_url(settings.database_url)
58backend = url.get_backend_name() # e.g. 'postgresql', 'sqlite'
59driver = url.get_driver_name() or "default"
61# Start with an empty dict and add options only when the driver can accept
62# them; this prevents unexpected TypeError at connect time.
63connect_args: dict[str, object] = {}
65# ---------------------------------------------------------------------------
66# 2. PostgreSQL (synchronous psycopg2 only)
67# The keep-alive parameters below are recognised exclusively by libpq /
68# psycopg2 and let the kernel detect broken network links quickly.
69# ---------------------------------------------------------------------------
70if backend == "postgresql" and driver in ("psycopg2", "default", ""): 70 ↛ 71line 70 didn't jump to line 71 because the condition on line 70 was never true
71 connect_args.update(
72 keepalives=1, # enable TCP keep-alive probes
73 keepalives_idle=30, # seconds of idleness before first probe
74 keepalives_interval=5, # seconds between probes
75 keepalives_count=5, # drop the link after N failed probes
76 )
78# ---------------------------------------------------------------------------
79# 3. SQLite (optional) – only one extra flag and it is *SQLite-specific*.
80# ---------------------------------------------------------------------------
81elif backend == "sqlite": 81 ↛ 90line 81 didn't jump to line 90 because the condition on line 81 was always true
82 # Allow pooled connections to hop across threads.
83 connect_args["check_same_thread"] = False
85# 4. Other backends (MySQL, MSSQL, etc.) leave `connect_args` empty.
87# ---------------------------------------------------------------------------
88# 5. Build the Engine with a single, clean connect_args mapping.
89# ---------------------------------------------------------------------------
90engine = create_engine(
91 settings.database_url,
92 pool_pre_ping=True, # quick liveness check per checkout
93 pool_size=settings.db_pool_size,
94 max_overflow=settings.db_max_overflow,
95 pool_timeout=settings.db_pool_timeout,
96 pool_recycle=settings.db_pool_recycle,
97 connect_args=connect_args,
98)
100# Session factory
101SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
104class Base(DeclarativeBase):
105 """Base class for all models."""
108# # Association table for tools and gateways (federation)
109# tool_gateway_table = Table(
110# "tool_gateway_association",
111# Base.metadata,
112# Column("tool_id", Integer, ForeignKey("tools.id"), primary_key=True),
113# Column("gateway_id", Integer, ForeignKey("gateways.id"), primary_key=True),
114# )
116# # Association table for resources and gateways (federation)
117# resource_gateway_table = Table(
118# "resource_gateway_association",
119# Base.metadata,
120# Column("resource_id", Integer, ForeignKey("resources.id"), primary_key=True),
121# Column("gateway_id", Integer, ForeignKey("gateways.id"), primary_key=True),
122# )
124# # Association table for prompts and gateways (federation)
125# prompt_gateway_table = Table(
126# "prompt_gateway_association",
127# Base.metadata,
128# Column("prompt_id", Integer, ForeignKey("prompts.id"), primary_key=True),
129# Column("gateway_id", Integer, ForeignKey("gateways.id"), primary_key=True),
130# )
132# Association table for servers and tools
133server_tool_association = Table(
134 "server_tool_association",
135 Base.metadata,
136 Column("server_id", Integer, ForeignKey("servers.id"), primary_key=True),
137 Column("tool_id", Integer, ForeignKey("tools.id"), primary_key=True),
138)
140# Association table for servers and resources
141server_resource_association = Table(
142 "server_resource_association",
143 Base.metadata,
144 Column("server_id", Integer, ForeignKey("servers.id"), primary_key=True),
145 Column("resource_id", Integer, ForeignKey("resources.id"), primary_key=True),
146)
148# Association table for servers and prompts
149server_prompt_association = Table(
150 "server_prompt_association",
151 Base.metadata,
152 Column("server_id", Integer, ForeignKey("servers.id"), primary_key=True),
153 Column("prompt_id", Integer, ForeignKey("prompts.id"), primary_key=True),
154)
157class ToolMetric(Base):
158 """
159 ORM model for recording individual metrics for tool executions.
161 Each record in this table corresponds to a single tool invocation and records:
162 - timestamp (datetime): When the invocation occurred.
163 - response_time (float): The execution time in seconds.
164 - is_success (bool): True if the execution succeeded, False otherwise.
165 - error_message (Optional[str]): Error message if the execution failed.
167 Aggregated metrics (such as total executions, successful/failed counts, failure rate,
168 minimum, maximum, and average response times, and last execution time) should be computed
169 on the fly using SQL aggregate functions over the rows in this table.
170 """
172 __tablename__ = "tool_metrics"
174 id: Mapped[int] = mapped_column(primary_key=True)
175 tool_id: Mapped[int] = mapped_column(Integer, ForeignKey("tools.id"), nullable=False)
176 timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
177 response_time: Mapped[float] = mapped_column(Float, nullable=False)
178 is_success: Mapped[bool] = mapped_column(Boolean, nullable=False)
179 error_message: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
181 # Relationship back to the Tool model.
182 tool: Mapped["Tool"] = relationship("Tool", back_populates="metrics")
185class ResourceMetric(Base):
186 """
187 ORM model for recording metrics for resource invocations.
189 Attributes:
190 id (int): Primary key.
191 resource_id (int): Foreign key linking to the resource.
192 timestamp (datetime): The time when the invocation occurred.
193 response_time (float): The response time in seconds.
194 is_success (bool): True if the invocation succeeded, False otherwise.
195 error_message (Optional[str]): Error message if the invocation failed.
196 """
198 __tablename__ = "resource_metrics"
200 id: Mapped[int] = mapped_column(primary_key=True)
201 resource_id: Mapped[int] = mapped_column(Integer, ForeignKey("resources.id"), nullable=False)
202 timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
203 response_time: Mapped[float] = mapped_column(Float, nullable=False)
204 is_success: Mapped[bool] = mapped_column(Boolean, nullable=False)
205 error_message: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
207 # Relationship back to the Resource model.
208 resource: Mapped["Resource"] = relationship("Resource", back_populates="metrics")
211class ServerMetric(Base):
212 """
213 ORM model for recording metrics for server invocations.
215 Attributes:
216 id (int): Primary key.
217 server_id (int): Foreign key linking to the server.
218 timestamp (datetime): The time when the invocation occurred.
219 response_time (float): The response time in seconds.
220 is_success (bool): True if the invocation succeeded, False otherwise.
221 error_message (Optional[str]): Error message if the invocation failed.
222 """
224 __tablename__ = "server_metrics"
226 id: Mapped[int] = mapped_column(primary_key=True)
227 server_id: Mapped[int] = mapped_column(Integer, ForeignKey("servers.id"), nullable=False)
228 timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
229 response_time: Mapped[float] = mapped_column(Float, nullable=False)
230 is_success: Mapped[bool] = mapped_column(Boolean, nullable=False)
231 error_message: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
233 # Relationship back to the Server model.
234 server: Mapped["Server"] = relationship("Server", back_populates="metrics")
237class PromptMetric(Base):
238 """
239 ORM model for recording metrics for prompt invocations.
241 Attributes:
242 id (int): Primary key.
243 prompt_id (int): Foreign key linking to the prompt.
244 timestamp (datetime): The time when the invocation occurred.
245 response_time (float): The response time in seconds.
246 is_success (bool): True if the invocation succeeded, False otherwise.
247 error_message (Optional[str]): Error message if the invocation failed.
248 """
250 __tablename__ = "prompt_metrics"
252 id: Mapped[int] = mapped_column(primary_key=True)
253 prompt_id: Mapped[int] = mapped_column(Integer, ForeignKey("prompts.id"), nullable=False)
254 timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
255 response_time: Mapped[float] = mapped_column(Float, nullable=False)
256 is_success: Mapped[bool] = mapped_column(Boolean, nullable=False)
257 error_message: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
259 # Relationship back to the Prompt model.
260 prompt: Mapped["Prompt"] = relationship("Prompt", back_populates="metrics")
263class Tool(Base):
264 """
265 ORM model for a registered Tool.
267 Supports both local tools and federated tools from other gateways.
268 The integration_type field indicates the tool format:
269 - "MCP" for MCP-compliant tools (default)
270 - "REST" for REST tools
272 Additionally, this model provides computed properties for aggregated metrics based
273 on the associated ToolMetric records. These include:
274 - execution_count: Total number of invocations.
275 - successful_executions: Count of successful invocations.
276 - failed_executions: Count of failed invocations.
277 - failure_rate: Ratio of failed invocations to total invocations.
278 - min_response_time: Fastest recorded response time.
279 - max_response_time: Slowest recorded response time.
280 - avg_response_time: Mean response time.
281 - last_execution_time: Timestamp of the most recent invocation.
283 The property `metrics_summary` returns a dictionary with these aggregated values.
285 The following fields have been added to support tool invocation configuration:
286 - request_type: HTTP method to use when invoking the tool.
287 - auth_type: Type of authentication ("basic", "bearer", or None).
288 - auth_username: Username for basic authentication.
289 - auth_password: Password for basic authentication.
290 - auth_token: Token for bearer token authentication.
291 - auth_header_key: header key for authentication.
292 - auth_header_value: header value for authentication.
293 """
295 __tablename__ = "tools"
297 id: Mapped[int] = mapped_column(primary_key=True)
298 name: Mapped[str] = mapped_column(unique=True)
299 url: Mapped[str] = mapped_column(String, nullable=True)
300 description: Mapped[Optional[str]]
301 integration_type: Mapped[str] = mapped_column(default="MCP")
302 request_type: Mapped[str] = mapped_column(default="SSE")
303 headers: Mapped[Optional[Dict[str, str]]] = mapped_column(JSON)
304 input_schema: Mapped[Dict[str, Any]] = mapped_column(JSON)
305 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
306 updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc), onupdate=lambda: datetime.now(timezone.utc))
307 is_active: Mapped[bool] = mapped_column(default=True)
308 jsonpath_filter: Mapped[str] = mapped_column(default="")
310 # Request type and authentication fields
311 auth_type: Mapped[Optional[str]] = mapped_column(default=None) # "basic", "bearer", or None
312 auth_value: Mapped[Optional[str]] = mapped_column(default=None)
314 # Federation relationship with a local gateway
315 gateway_id: Mapped[Optional[int]] = mapped_column(ForeignKey("gateways.id"))
316 gateway: Mapped["Gateway"] = relationship("Gateway", back_populates="tools")
317 # federated_with = relationship("Gateway", secondary=tool_gateway_table, back_populates="federated_tools")
319 # Many-to-many relationship with Servers
320 servers: Mapped[List["Server"]] = relationship("Server", secondary=server_tool_association, back_populates="tools")
322 # Relationship with ToolMetric records
323 metrics: Mapped[List["ToolMetric"]] = relationship("ToolMetric", back_populates="tool", cascade="all, delete-orphan")
325 @hybrid_property
326 def execution_count(self) -> int:
327 """
328 Returns the number of times the tool has been executed,
329 calculated from the associated ToolMetric records.
331 Returns:
332 int: The total count of tool executions.
333 """
334 return len(self.metrics)
336 @execution_count.expression
337 # method is intentionally a class-level expression, so no `self`
338 # pylint: disable=no-self-argument
339 def execution_count(cls):
340 """
341 SQL expression to compute the execution count for the tool.
343 Returns:
344 int: Returns execution count of a given tool
345 """
346 return select(func.count(ToolMetric.id)).where(ToolMetric.tool_id == cls.id).label("execution_count") # pylint: disable=not-callable
348 @property
349 def successful_executions(self) -> int:
350 """
351 Returns the count of successful tool executions,
352 computed from the associated ToolMetric records.
354 Returns:
355 int: The count of successful tool executions.
356 """
357 return sum(1 for m in self.metrics if m.is_success)
359 @property
360 def failed_executions(self) -> int:
361 """
362 Returns the count of failed tool executions,
363 computed from the associated ToolMetric records.
365 Returns:
366 int: The count of failed tool executions.
367 """
368 return sum(1 for m in self.metrics if not m.is_success)
370 @property
371 def failure_rate(self) -> float:
372 """
373 Returns the failure rate (as a float between 0 and 1) computed as:
374 (failed executions) / (total executions).
375 Returns 0.0 if there are no executions.
377 Returns:
378 float: The failure rate as a value between 0 and 1.
379 """
380 total: int = self.execution_count
381 # execution_count is a @hybrid_property, not a callable here
382 if total == 0: # pylint: disable=comparison-with-callable
383 return 0.0
384 return self.failed_executions / total
386 @property
387 def min_response_time(self) -> Optional[float]:
388 """
389 Returns the minimum response time among all tool executions.
390 Returns None if no executions exist.
392 Returns:
393 Optional[float]: The minimum response time, or None if no executions exist.
394 """
395 times: List[float] = [m.response_time for m in self.metrics]
396 return min(times) if times else None
398 @property
399 def max_response_time(self) -> Optional[float]:
400 """
401 Returns the maximum response time among all tool executions.
402 Returns None if no executions exist.
404 Returns:
405 Optional[float]: The maximum response time, or None if no executions exist.
406 """
407 times: List[float] = [m.response_time for m in self.metrics]
408 return max(times) if times else None
410 @property
411 def avg_response_time(self) -> Optional[float]:
412 """
413 Returns the average response time among all tool executions.
414 Returns None if no executions exist.
416 Returns:
417 Optional[float]: The average response time, or None if no executions exist.
418 """
419 times: List[float] = [m.response_time for m in self.metrics]
420 return sum(times) / len(times) if times else None
422 @property
423 def last_execution_time(self) -> Optional[datetime]:
424 """
425 Returns the timestamp of the most recent tool execution.
426 Returns None if no executions exist.
428 Returns:
429 Optional[datetime]: The timestamp of the most recent execution, or None if no executions exist.
430 """
431 if not self.metrics:
432 return None
433 return max(m.timestamp for m in self.metrics)
435 @property
436 def metrics_summary(self) -> Dict[str, Any]:
437 """
438 Returns aggregated metrics for the tool as a dictionary with the following keys:
439 - total_executions: Total number of invocations.
440 - successful_executions: Number of successful invocations.
441 - failed_executions: Number of failed invocations.
442 - failure_rate: Failure rate (failed/total) or 0.0 if no invocations.
443 - min_response_time: Minimum response time (or None if no invocations).
444 - max_response_time: Maximum response time (or None if no invocations).
445 - avg_response_time: Average response time (or None if no invocations).
446 - last_execution_time: Timestamp of the most recent invocation (or None).
448 Returns:
449 Dict[str, Any]: Dictionary containing the aggregated metrics.
450 """
451 return {
452 "total_executions": self.execution_count,
453 "successful_executions": self.successful_executions,
454 "failed_executions": self.failed_executions,
455 "failure_rate": self.failure_rate,
456 "min_response_time": self.min_response_time,
457 "max_response_time": self.max_response_time,
458 "avg_response_time": self.avg_response_time,
459 "last_execution_time": self.last_execution_time,
460 }
463class Resource(Base):
464 """
465 ORM model for a registered Resource.
467 Resources represent content that can be read by clients.
468 Supports subscriptions for real-time updates.
469 Additionally, this model provides a relationship with ResourceMetric records
470 to capture invocation metrics (such as execution counts, response times, and failures).
471 """
473 __tablename__ = "resources"
475 id: Mapped[int] = mapped_column(primary_key=True)
476 uri: Mapped[str] = mapped_column(unique=True)
477 name: Mapped[str]
478 description: Mapped[Optional[str]]
479 mime_type: Mapped[Optional[str]]
480 size: Mapped[Optional[int]]
481 template: Mapped[Optional[str]] # URI template for parameterized resources
482 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
483 updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc), onupdate=lambda: datetime.now(timezone.utc))
484 is_active: Mapped[bool] = mapped_column(default=True)
485 metrics: Mapped[List["ResourceMetric"]] = relationship("ResourceMetric", back_populates="resource", cascade="all, delete-orphan")
487 # Content storage - can be text or binary
488 text_content: Mapped[Optional[str]] = mapped_column(Text)
489 binary_content: Mapped[Optional[bytes]]
491 # Subscription tracking
492 subscriptions: Mapped[List["ResourceSubscription"]] = relationship("ResourceSubscription", back_populates="resource", cascade="all, delete-orphan")
494 gateway_id: Mapped[Optional[int]] = mapped_column(ForeignKey("gateways.id"))
495 gateway: Mapped["Gateway"] = relationship("Gateway", back_populates="resources")
496 # federated_with = relationship("Gateway", secondary=resource_gateway_table, back_populates="federated_resources")
498 # Many-to-many relationship with Servers
499 servers: Mapped[List["Server"]] = relationship("Server", secondary=server_resource_association, back_populates="resources")
501 @property
502 def content(self) -> ResourceContent:
503 """
504 Returns the resource content in the appropriate format.
506 If text content exists, returns a ResourceContent with text.
507 Otherwise, if binary content exists, returns a ResourceContent with blob data.
508 Raises a ValueError if no content is available.
510 Returns:
511 ResourceContent: The resource content with appropriate format (text or blob).
513 Raises:
514 ValueError: If the resource has no content available.
515 """
517 if self.text_content is not None:
518 return ResourceContent(
519 type="resource",
520 uri=self.uri,
521 mime_type=self.mime_type,
522 text=self.text_content,
523 )
524 if self.binary_content is not None:
525 return ResourceContent(
526 type="resource",
527 uri=self.uri,
528 mime_type=self.mime_type or "application/octet-stream",
529 blob=self.binary_content,
530 )
531 raise ValueError("Resource has no content")
533 @property
534 def execution_count(self) -> int:
535 """
536 Returns the number of times the resource has been invoked,
537 calculated from the associated ResourceMetric records.
539 Returns:
540 int: The total count of resource invocations.
541 """
542 return len(self.metrics)
544 @property
545 def successful_executions(self) -> int:
546 """
547 Returns the count of successful resource invocations,
548 computed from the associated ResourceMetric records.
550 Returns:
551 int: The count of successful resource invocations.
552 """
553 return sum(1 for m in self.metrics if m.is_success)
555 @property
556 def failed_executions(self) -> int:
557 """
558 Returns the count of failed resource invocations,
559 computed from the associated ResourceMetric records.
561 Returns:
562 int: The count of failed resource invocations.
563 """
564 return sum(1 for m in self.metrics if not m.is_success)
566 @property
567 def failure_rate(self) -> float:
568 """
569 Returns the failure rate (as a float between 0 and 1) computed as:
570 (failed invocations) / (total invocations).
571 Returns 0.0 if there are no invocations.
573 Returns:
574 float: The failure rate as a value between 0 and 1.
575 """
576 total: int = self.execution_count
577 if total == 0:
578 return 0.0
579 return self.failed_executions / total
581 @property
582 def min_response_time(self) -> Optional[float]:
583 """
584 Returns the minimum response time among all resource invocations.
585 Returns None if no invocations exist.
587 Returns:
588 Optional[float]: The minimum response time, or None if no invocations exist.
589 """
590 times: List[float] = [m.response_time for m in self.metrics]
591 return min(times) if times else None
593 @property
594 def max_response_time(self) -> Optional[float]:
595 """
596 Returns the maximum response time among all resource invocations.
597 Returns None if no invocations exist.
599 Returns:
600 Optional[float]: The maximum response time, or None if no invocations exist.
601 """
602 times: List[float] = [m.response_time for m in self.metrics]
603 return max(times) if times else None
605 @property
606 def avg_response_time(self) -> Optional[float]:
607 """
608 Returns the average response time among all resource invocations.
609 Returns None if no invocations exist.
611 Returns:
612 Optional[float]: The average response time, or None if no invocations exist.
613 """
614 times: List[float] = [m.response_time for m in self.metrics]
615 return sum(times) / len(times) if times else None
617 @property
618 def last_execution_time(self) -> Optional[datetime]:
619 """
620 Returns the timestamp of the most recent resource invocation.
621 Returns None if no invocations exist.
623 Returns:
624 Optional[datetime]: The timestamp of the most recent invocation, or None if no invocations exist.
625 """
626 if not self.metrics:
627 return None
628 return max(m.timestamp for m in self.metrics)
631class ResourceSubscription(Base):
632 """Tracks subscriptions to resource updates."""
634 __tablename__ = "resource_subscriptions"
636 id: Mapped[int] = mapped_column(primary_key=True)
637 resource_id: Mapped[int] = mapped_column(ForeignKey("resources.id"))
638 subscriber_id: Mapped[str] # Client identifier
639 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
640 last_notification: Mapped[Optional[datetime]] = mapped_column(DateTime)
642 resource: Mapped["Resource"] = relationship(back_populates="subscriptions")
645class Prompt(Base):
646 """
647 ORM model for a registered Prompt template.
649 Represents a prompt template along with its argument schema.
650 Supports rendering and invocation of prompts.
651 Additionally, this model provides computed properties for aggregated metrics based
652 on the associated PromptMetric records. These include:
653 - execution_count: Total number of prompt invocations.
654 - successful_executions: Count of successful invocations.
655 - failed_executions: Count of failed invocations.
656 - failure_rate: Ratio of failed invocations to total invocations.
657 - min_response_time: Fastest recorded response time.
658 - max_response_time: Slowest recorded response time.
659 - avg_response_time: Mean response time.
660 - last_execution_time: Timestamp of the most recent invocation.
661 """
663 __tablename__ = "prompts"
665 id: Mapped[int] = mapped_column(primary_key=True)
666 name: Mapped[str] = mapped_column(unique=True)
667 description: Mapped[Optional[str]]
668 template: Mapped[str] = mapped_column(Text)
669 argument_schema: Mapped[Dict[str, Any]] = mapped_column(JSON)
670 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
671 updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc), onupdate=lambda: datetime.now(timezone.utc))
672 is_active: Mapped[bool] = mapped_column(default=True)
673 metrics: Mapped[List["PromptMetric"]] = relationship("PromptMetric", back_populates="prompt", cascade="all, delete-orphan")
675 gateway_id: Mapped[Optional[int]] = mapped_column(ForeignKey("gateways.id"))
676 gateway: Mapped["Gateway"] = relationship("Gateway", back_populates="prompts")
677 # federated_with = relationship("Gateway", secondary=prompt_gateway_table, back_populates="federated_prompts")
679 # Many-to-many relationship with Servers
680 servers: Mapped[List["Server"]] = relationship("Server", secondary=server_prompt_association, back_populates="prompts")
682 def validate_arguments(self, args: Dict[str, str]) -> None:
683 """
684 Validate prompt arguments against the argument schema.
686 Args:
687 args (Dict[str, str]): Dictionary of arguments to validate.
689 Raises:
690 ValueError: If the arguments do not conform to the schema.
692 """
693 try:
694 jsonschema.validate(args, self.argument_schema)
695 except jsonschema.exceptions.ValidationError as e:
696 raise ValueError(f"Invalid prompt arguments: {str(e)}")
698 @property
699 def execution_count(self) -> int:
700 """
701 Returns the number of times the prompt has been invoked,
702 calculated from the associated PromptMetric records.
704 Returns:
705 int: The total count of prompt invocations.
706 """
707 return len(self.metrics)
709 @property
710 def successful_executions(self) -> int:
711 """
712 Returns the count of successful prompt invocations,
713 computed from the associated PromptMetric records.
715 Returns:
716 int: The count of successful prompt invocations.
717 """
718 return sum(1 for m in self.metrics if m.is_success)
720 @property
721 def failed_executions(self) -> int:
722 """
723 Returns the count of failed prompt invocations,
724 computed from the associated PromptMetric records.
726 Returns:
727 int: The count of failed prompt invocations.
728 """
729 return sum(1 for m in self.metrics if not m.is_success)
731 @property
732 def failure_rate(self) -> float:
733 """
734 Returns the failure rate (as a float between 0 and 1) computed as:
735 (failed invocations) / (total invocations).
736 Returns 0.0 if there are no invocations.
738 Returns:
739 float: The failure rate as a value between 0 and 1.
740 """
741 total: int = self.execution_count
742 if total == 0:
743 return 0.0
744 return self.failed_executions / total
746 @property
747 def min_response_time(self) -> Optional[float]:
748 """
749 Returns the minimum response time among all prompt invocations.
750 Returns None if no invocations exist.
752 Returns:
753 Optional[float]: The minimum response time, or None if no invocations exist.
754 """
755 times: List[float] = [m.response_time for m in self.metrics]
756 return min(times) if times else None
758 @property
759 def max_response_time(self) -> Optional[float]:
760 """
761 Returns the maximum response time among all prompt invocations.
762 Returns None if no invocations exist.
764 Returns:
765 Optional[float]: The maximum response time, or None if no invocations exist.
766 """
767 times: List[float] = [m.response_time for m in self.metrics]
768 return max(times) if times else None
770 @property
771 def avg_response_time(self) -> Optional[float]:
772 """
773 Returns the average response time among all prompt invocations.
774 Returns None if no invocations exist.
776 Returns:
777 Optional[float]: The average response time, or None if no invocations exist.
778 """
779 times: List[float] = [m.response_time for m in self.metrics]
780 return sum(times) / len(times) if times else None
782 @property
783 def last_execution_time(self) -> Optional[datetime]:
784 """
785 Returns the timestamp of the most recent prompt invocation.
786 Returns None if no invocations exist.
788 Returns:
789 Optional[datetime]: The timestamp of the most recent invocation, or None if no invocations exist.
790 """
791 if not self.metrics:
792 return None
793 return max(m.timestamp for m in self.metrics)
796class Server(Base):
797 """
798 ORM model for MCP Servers Catalog.
800 Represents a server that composes catalog items (tools, resources, prompts).
801 Additionally, this model provides computed properties for aggregated metrics based
802 on the associated ServerMetric records. These include:
803 - execution_count: Total number of invocations.
804 - successful_executions: Count of successful invocations.
805 - failed_executions: Count of failed invocations.
806 - failure_rate: Ratio of failed invocations to total invocations.
807 - min_response_time: Fastest recorded response time.
808 - max_response_time: Slowest recorded response time.
809 - avg_response_time: Mean response time.
810 - last_execution_time: Timestamp of the most recent invocation.
811 """
813 __tablename__ = "servers"
815 id: Mapped[int] = mapped_column(primary_key=True)
816 name: Mapped[str] = mapped_column(unique=True)
817 description: Mapped[Optional[str]]
818 icon: Mapped[Optional[str]]
819 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
820 updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc), onupdate=lambda: datetime.now(timezone.utc))
821 is_active: Mapped[bool] = mapped_column(default=True)
822 metrics: Mapped[List["ServerMetric"]] = relationship("ServerMetric", back_populates="server", cascade="all, delete-orphan")
824 # Many-to-many relationships for associated items
825 tools: Mapped[List["Tool"]] = relationship("Tool", secondary=server_tool_association, back_populates="servers")
826 resources: Mapped[List["Resource"]] = relationship("Resource", secondary=server_resource_association, back_populates="servers")
827 prompts: Mapped[List["Prompt"]] = relationship("Prompt", secondary=server_prompt_association, back_populates="servers")
829 @property
830 def execution_count(self) -> int:
831 """
832 Returns the number of times the server has been invoked,
833 calculated from the associated ServerMetric records.
835 Returns:
836 int: The total count of server invocations.
837 """
838 return len(self.metrics)
840 @property
841 def successful_executions(self) -> int:
842 """
843 Returns the count of successful server invocations,
844 computed from the associated ServerMetric records.
846 Returns:
847 int: The count of successful server invocations.
848 """
849 return sum(1 for m in self.metrics if m.is_success)
851 @property
852 def failed_executions(self) -> int:
853 """
854 Returns the count of failed server invocations,
855 computed from the associated ServerMetric records.
857 Returns:
858 int: The count of failed server invocations.
859 """
860 return sum(1 for m in self.metrics if not m.is_success)
862 @property
863 def failure_rate(self) -> float:
864 """
865 Returns the failure rate (as a float between 0 and 1) computed as:
866 (failed invocations) / (total invocations).
867 Returns 0.0 if there are no invocations.
869 Returns:
870 float: The failure rate as a value between 0 and 1.
871 """
872 total: int = self.execution_count
873 if total == 0:
874 return 0.0
875 return self.failed_executions / total
877 @property
878 def min_response_time(self) -> Optional[float]:
879 """
880 Returns the minimum response time among all server invocations.
881 Returns None if no invocations exist.
883 Returns:
884 Optional[float]: The minimum response time, or None if no invocations exist.
885 """
886 times: List[float] = [m.response_time for m in self.metrics]
887 return min(times) if times else None
889 @property
890 def max_response_time(self) -> Optional[float]:
891 """
892 Returns the maximum response time among all server invocations.
893 Returns None if no invocations exist.
895 Returns:
896 Optional[float]: The maximum response time, or None if no invocations exist.
897 """
898 times: List[float] = [m.response_time for m in self.metrics]
899 return max(times) if times else None
901 @property
902 def avg_response_time(self) -> Optional[float]:
903 """
904 Returns the average response time among all server invocations.
905 Returns None if no invocations exist.
907 Returns:
908 Optional[float]: The average response time, or None if no invocations exist.
909 """
910 times: List[float] = [m.response_time for m in self.metrics]
911 return sum(times) / len(times) if times else None
913 @property
914 def last_execution_time(self) -> Optional[datetime]:
915 """
916 Returns the timestamp of the most recent server invocation.
917 Returns None if no invocations exist.
919 Returns:
920 Optional[datetime]: The timestamp of the most recent invocation, or None if no invocations exist.
921 """
922 if not self.metrics:
923 return None
924 return max(m.timestamp for m in self.metrics)
927class Gateway(Base):
928 """ORM model for a federated peer Gateway."""
930 __tablename__ = "gateways"
932 id: Mapped[int] = mapped_column(primary_key=True)
933 name: Mapped[str] = mapped_column(unique=True)
934 url: Mapped[str]
935 description: Mapped[Optional[str]]
936 transport: Mapped[str] = mapped_column(default="SSE")
937 capabilities: Mapped[Dict[str, Any]] = mapped_column(JSON)
938 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
939 updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc), onupdate=lambda: datetime.now(timezone.utc))
940 is_active: Mapped[bool] = mapped_column(default=True)
941 last_seen: Mapped[Optional[datetime]]
943 # Relationship with local tools this gateway provides
944 tools: Mapped[List["Tool"]] = relationship(back_populates="gateway", cascade="all, delete-orphan")
946 # Relationship with local prompts this gateway provides
947 prompts: Mapped[List["Prompt"]] = relationship(back_populates="gateway", cascade="all, delete-orphan")
949 # Relationship with local resources this gateway provides
950 resources: Mapped[List["Resource"]] = relationship(back_populates="gateway", cascade="all, delete-orphan")
952 # # Tools federated from this gateway
953 # federated_tools: Mapped[List["Tool"]] = relationship(secondary=tool_gateway_table, back_populates="federated_with")
955 # # Prompts federated from this resource
956 # federated_resources: Mapped[List["Resource"]] = relationship(secondary=resource_gateway_table, back_populates="federated_with")
958 # # Prompts federated from this gateway
959 # federated_prompts: Mapped[List["Prompt"]] = relationship(secondary=prompt_gateway_table, back_populates="federated_with")
961 # Authorizations
962 auth_type: Mapped[Optional[str]] = mapped_column(default=None) # "basic", "bearer", "headers" or None
963 auth_value: Mapped[Optional[Dict[str, str]]] = mapped_column(JSON)
966class SessionRecord(Base):
967 """ORM model for sessions from SSE client."""
969 __tablename__ = "mcp_sessions"
971 session_id: Mapped[str] = mapped_column(primary_key=True)
972 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc)) # pylint: disable=not-callable
973 last_accessed: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc), onupdate=lambda: datetime.now(timezone.utc)) # pylint: disable=not-callable
974 data: Mapped[str] = mapped_column(String, nullable=True)
976 messages: Mapped[List["SessionMessageRecord"]] = relationship("SessionMessageRecord", back_populates="session", cascade="all, delete-orphan")
979class SessionMessageRecord(Base):
980 """ORM model for messages from SSE client."""
982 __tablename__ = "mcp_messages"
984 id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
985 session_id: Mapped[str] = mapped_column(ForeignKey("mcp_sessions.session_id"))
986 message: Mapped[str] = mapped_column(String, nullable=True)
987 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc)) # pylint: disable=not-callable
988 last_accessed: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc), onupdate=lambda: datetime.now(timezone.utc)) # pylint: disable=not-callable
990 session: Mapped["SessionRecord"] = relationship("SessionRecord", back_populates="messages")
993# Event listeners for validation
994def validate_tool_schema(mapper, connection, target):
995 """
996 Validate tool schema before insert/update.
998 Args:
999 mapper: The mapper being used for the operation.
1000 connection: The database connection.
1001 target: The target object being validated.
1003 Raises:
1004 ValueError: If the tool input schema is invalid.
1005 """
1006 # You can use mapper and connection later, if required.
1007 _ = mapper
1008 _ = connection
1009 if hasattr(target, "input_schema"):
1010 try:
1011 jsonschema.Draft7Validator.check_schema(target.input_schema)
1012 except jsonschema.exceptions.SchemaError as e:
1013 raise ValueError(f"Invalid tool input schema: {str(e)}")
1016def validate_tool_name(mapper, connection, target):
1017 """
1018 Validate tool name before insert/update. Check if the name matches the required pattern.
1020 Args:
1021 mapper: The mapper being used for the operation.
1022 connection: The database connection.
1023 target: The target object being validated.
1025 Raises:
1026 ValueError: If the tool name contains invalid characters.
1027 """
1028 # You can use mapper and connection later, if required.
1029 _ = mapper
1030 _ = connection
1031 if hasattr(target, "name"):
1032 if not re.match(r"^[a-zA-Z0-9_-]+$", target.name):
1033 raise ValueError(f"Invalid tool name '{target.name}'. Only alphanumeric characters, hyphens, and underscores are allowed.")
1036def validate_prompt_schema(mapper, connection, target):
1037 """
1038 Validate prompt argument schema before insert/update.
1040 Args:
1041 mapper: The mapper being used for the operation.
1042 connection: The database connection.
1043 target: The target object being validated.
1045 Raises:
1046 ValueError: If the prompt argument schema is invalid.
1047 """
1048 # You can use mapper and connection later, if required.
1049 _ = mapper
1050 _ = connection
1051 if hasattr(target, "argument_schema"):
1052 try:
1053 jsonschema.Draft7Validator.check_schema(target.argument_schema)
1054 except jsonschema.exceptions.SchemaError as e:
1055 raise ValueError(f"Invalid prompt argument schema: {str(e)}")
1058# Register validation listeners
1060listen(Tool, "before_insert", validate_tool_schema)
1061listen(Tool, "before_update", validate_tool_schema)
1062listen(Tool, "before_insert", validate_tool_name)
1063listen(Tool, "before_update", validate_tool_name)
1064listen(Prompt, "before_insert", validate_prompt_schema)
1065listen(Prompt, "before_update", validate_prompt_schema)
1068def get_db():
1069 """
1070 Dependency to get database session.
1072 Yields:
1073 SessionLocal: A SQLAlchemy database session.
1074 """
1075 db = SessionLocal()
1076 try:
1077 yield db
1078 finally:
1079 db.close()
1082# Create all tables
1083def init_db():
1084 """
1085 Initialize database tables.
1087 Raises:
1088 Exception: If database initialization fails.
1089 """
1090 try:
1091 # Base.metadata.drop_all(bind=engine)
1092 Base.metadata.create_all(bind=engine)
1093 except SQLAlchemyError as e:
1094 raise Exception(f"Failed to initialize database: {str(e)}")
1097if __name__ == "__main__":
1098 init_db()