Coverage for mcpgateway/db.py: 74%

395 statements  

« 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. 

3 

4Copyright 2025 

5SPDX-License-Identifier: Apache-2.0 

6Authors: Mihai Criveti 

7 

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 

13 

14Updated to record server associations independently using many-to-many relationships, 

15and to record tool execution metrics. 

16""" 

17 

18# Standard 

19from datetime import datetime, timezone 

20import re 

21from typing import Any, Dict, List, Optional 

22import uuid 

23 

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 

55 

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 

61 

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" 

69 

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] = {} 

73 

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 ) 

86 

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 

93 

94# 4. Other backends (MySQL, MSSQL, etc.) leave `connect_args` empty. 

95 

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) 

108 

109 

110# --------------------------------------------------------------------------- 

111# 6. Function to return UTC timestamp 

112# --------------------------------------------------------------------------- 

113def utc_now() -> datetime: 

114 """Return the current Coordinated Universal Time (UTC). 

115 

116 Returns: 

117 datetime: A timezone-aware `datetime` whose `tzinfo` is 

118 `datetime.timezone.utc`. 

119 """ 

120 return datetime.now(timezone.utc) 

121 

122 

123# Session factory 

124SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) 

125 

126 

127class Base(DeclarativeBase): 

128 """Base class for all models.""" 

129 

130 

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# ) 

139 

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# ) 

147 

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# ) 

155 

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) 

163 

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) 

171 

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) 

179 

180 

181class ToolMetric(Base): 

182 """ 

183 ORM model for recording individual metrics for tool executions. 

184 

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. 

190 

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 """ 

195 

196 __tablename__ = "tool_metrics" 

197 

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) 

204 

205 # Relationship back to the Tool model. 

206 tool: Mapped["Tool"] = relationship("Tool", back_populates="metrics") 

207 

208 

209class ResourceMetric(Base): 

210 """ 

211 ORM model for recording metrics for resource invocations. 

212 

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 """ 

221 

222 __tablename__ = "resource_metrics" 

223 

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) 

230 

231 # Relationship back to the Resource model. 

232 resource: Mapped["Resource"] = relationship("Resource", back_populates="metrics") 

233 

234 

235class ServerMetric(Base): 

236 """ 

237 ORM model for recording metrics for server invocations. 

238 

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 """ 

247 

248 __tablename__ = "server_metrics" 

249 

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) 

256 

257 # Relationship back to the Server model. 

258 server: Mapped["Server"] = relationship("Server", back_populates="metrics") 

259 

260 

261class PromptMetric(Base): 

262 """ 

263 ORM model for recording metrics for prompt invocations. 

264 

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 """ 

273 

274 __tablename__ = "prompt_metrics" 

275 

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) 

282 

283 # Relationship back to the Prompt model. 

284 prompt: Mapped["Prompt"] = relationship("Prompt", back_populates="metrics") 

285 

286 

287class Tool(Base): 

288 """ 

289 ORM model for a registered Tool. 

290 

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 

295 

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. 

306 

307 The property `metrics_summary` returns a dictionary with these aggregated values. 

308 

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 """ 

318 

319 __tablename__ = "tools" 

320 

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="") 

336 

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) 

340 

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") 

346 

347 # Many-to-many relationship with Servers 

348 servers: Mapped[List["Server"]] = relationship("Server", secondary=server_tool_association, back_populates="tools") 

349 

350 # Relationship with ToolMetric records 

351 metrics: Mapped[List["ToolMetric"]] = relationship("ToolMetric", back_populates="tool", cascade="all, delete-orphan") 

352 

353 # @property 

354 # def gateway_slug(self) -> str: 

355 # return self.gateway.slug 

356 

357 _computed_name = Column("name", String, unique=True) # Stored column 

358 

359 @hybrid_property 

360 def name(self): 

361 """Return the display/lookup name. 

362 

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 

368 

369 original_slug = slugify(self.original_name) # pylint: disable=no-member 

370 

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}" 

375 

376 # No gateway → only the original name slug 

377 return original_slug 

378 

379 @name.setter 

380 def name(self, value): 

381 """Store an explicit value that overrides the calculated one. 

382 

383 Args: 

384 value (str): Value to set to _computed_name 

385 """ 

386 self._computed_name = value 

387 

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. 

394 

395 Returns: 

396 str: computed name for SQL use 

397 """ 

398 return cls._computed_name 

399 

400 __table_args__ = (UniqueConstraint("gateway_id", "original_name", name="uq_gateway_id__original_name"),) 

401 

402 @hybrid_property 

403 def gateway_slug(self): 

404 """Always returns the current slug from the related Gateway 

405 

406 Returns: 

407 str: slug for Python use 

408 """ 

409 return self.gateway.slug if self.gateway else None 

410 

411 @gateway_slug.expression 

412 def gateway_slug(cls): # pylint: disable=no-self-argument 

413 """For database queries - auto-joins to get current slug 

414 

415 Returns: 

416 str: slug for SQL use 

417 """ 

418 return select(Gateway.slug).where(Gateway.id == cls.gateway_id).scalar_subquery() 

419 

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. 

425 

426 Returns: 

427 int: The total count of tool executions. 

428 """ 

429 return len(self.metrics) 

430 

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. 

437 

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 

442 

443 @property 

444 def successful_executions(self) -> int: 

445 """ 

446 Returns the count of successful tool executions, 

447 computed from the associated ToolMetric records. 

448 

449 Returns: 

450 int: The count of successful tool executions. 

451 """ 

452 return sum(1 for m in self.metrics if m.is_success) 

453 

454 @property 

455 def failed_executions(self) -> int: 

456 """ 

457 Returns the count of failed tool executions, 

458 computed from the associated ToolMetric records. 

459 

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) 

464 

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. 

471 

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 

480 

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. 

486 

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 

492 

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. 

498 

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 

504 

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. 

510 

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 

516 

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. 

522 

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) 

529 

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). 

542 

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 } 

556 

557 

558class Resource(Base): 

559 """ 

560 ORM model for a registered Resource. 

561 

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 """ 

567 

568 __tablename__ = "resources" 

569 

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") 

581 

582 # Content storage - can be text or binary 

583 text_content: Mapped[Optional[str]] = mapped_column(Text) 

584 binary_content: Mapped[Optional[bytes]] 

585 

586 # Subscription tracking 

587 subscriptions: Mapped[List["ResourceSubscription"]] = relationship("ResourceSubscription", back_populates="resource", cascade="all, delete-orphan") 

588 

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") 

592 

593 # Many-to-many relationship with Servers 

594 servers: Mapped[List["Server"]] = relationship("Server", secondary=server_resource_association, back_populates="resources") 

595 

596 @property 

597 def content(self) -> ResourceContent: 

598 """ 

599 Returns the resource content in the appropriate format. 

600 

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. 

604 

605 Returns: 

606 ResourceContent: The resource content with appropriate format (text or blob). 

607 

608 Raises: 

609 ValueError: If the resource has no content available. 

610 """ 

611 

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") 

627 

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. 

633 

634 Returns: 

635 int: The total count of resource invocations. 

636 """ 

637 return len(self.metrics) 

638 

639 @property 

640 def successful_executions(self) -> int: 

641 """ 

642 Returns the count of successful resource invocations, 

643 computed from the associated ResourceMetric records. 

644 

645 Returns: 

646 int: The count of successful resource invocations. 

647 """ 

648 return sum(1 for m in self.metrics if m.is_success) 

649 

650 @property 

651 def failed_executions(self) -> int: 

652 """ 

653 Returns the count of failed resource invocations, 

654 computed from the associated ResourceMetric records. 

655 

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) 

660 

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. 

667 

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 

675 

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. 

681 

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 

687 

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. 

693 

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 

699 

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. 

705 

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 

711 

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. 

717 

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) 

724 

725 

726class ResourceSubscription(Base): 

727 """Tracks subscriptions to resource updates.""" 

728 

729 __tablename__ = "resource_subscriptions" 

730 

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) 

736 

737 resource: Mapped["Resource"] = relationship(back_populates="subscriptions") 

738 

739 

740class Prompt(Base): 

741 """ 

742 ORM model for a registered Prompt template. 

743 

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 """ 

757 

758 __tablename__ = "prompts" 

759 

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") 

769 

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") 

773 

774 # Many-to-many relationship with Servers 

775 servers: Mapped[List["Server"]] = relationship("Server", secondary=server_prompt_association, back_populates="prompts") 

776 

777 def validate_arguments(self, args: Dict[str, str]) -> None: 

778 """ 

779 Validate prompt arguments against the argument schema. 

780 

781 Args: 

782 args (Dict[str, str]): Dictionary of arguments to validate. 

783 

784 Raises: 

785 ValueError: If the arguments do not conform to the schema. 

786 

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)}") 

792 

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. 

798 

799 Returns: 

800 int: The total count of prompt invocations. 

801 """ 

802 return len(self.metrics) 

803 

804 @property 

805 def successful_executions(self) -> int: 

806 """ 

807 Returns the count of successful prompt invocations, 

808 computed from the associated PromptMetric records. 

809 

810 Returns: 

811 int: The count of successful prompt invocations. 

812 """ 

813 return sum(1 for m in self.metrics if m.is_success) 

814 

815 @property 

816 def failed_executions(self) -> int: 

817 """ 

818 Returns the count of failed prompt invocations, 

819 computed from the associated PromptMetric records. 

820 

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) 

825 

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. 

832 

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 

840 

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. 

846 

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 

852 

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. 

858 

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 

864 

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. 

870 

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 

876 

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. 

882 

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) 

889 

890 

891class Server(Base): 

892 """ 

893 ORM model for MCP Servers Catalog. 

894 

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 """ 

907 

908 __tablename__ = "servers" 

909 

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") 

918 

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") 

923 

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. 

929 

930 Returns: 

931 int: The total count of server invocations. 

932 """ 

933 return len(self.metrics) 

934 

935 @property 

936 def successful_executions(self) -> int: 

937 """ 

938 Returns the count of successful server invocations, 

939 computed from the associated ServerMetric records. 

940 

941 Returns: 

942 int: The count of successful server invocations. 

943 """ 

944 return sum(1 for m in self.metrics if m.is_success) 

945 

946 @property 

947 def failed_executions(self) -> int: 

948 """ 

949 Returns the count of failed server invocations, 

950 computed from the associated ServerMetric records. 

951 

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) 

956 

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. 

963 

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 

971 

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. 

977 

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 

983 

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. 

989 

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 

995 

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. 

1001 

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 

1007 

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. 

1013 

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) 

1020 

1021 

1022class Gateway(Base): 

1023 """ORM model for a federated peer Gateway.""" 

1024 

1025 __tablename__ = "gateways" 

1026 

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]] 

1039 

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") 

1042 

1043 # Relationship with local prompts this gateway provides 

1044 prompts: Mapped[List["Prompt"]] = relationship(back_populates="gateway", cascade="all, delete-orphan") 

1045 

1046 # Relationship with local resources this gateway provides 

1047 resources: Mapped[List["Resource"]] = relationship(back_populates="gateway", cascade="all, delete-orphan") 

1048 

1049 # # Tools federated from this gateway 

1050 # federated_tools: Mapped[List["Tool"]] = relationship(secondary=tool_gateway_table, back_populates="federated_with") 

1051 

1052 # # Prompts federated from this resource 

1053 # federated_resources: Mapped[List["Resource"]] = relationship(secondary=resource_gateway_table, back_populates="federated_with") 

1054 

1055 # # Prompts federated from this gateway 

1056 # federated_prompts: Mapped[List["Prompt"]] = relationship(secondary=prompt_gateway_table, back_populates="federated_with") 

1057 

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) 

1061 

1062 

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. 

1068 

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 

1078 

1079 print(f"Gateway name changed for ID {target.id}. Issuing bulk update for tools.") 

1080 

1081 # 2. Get a reference to the underlying database table for Tools 

1082 tools_table = Tool.__table__ 

1083 

1084 # 3. Prepare the new values 

1085 new_gateway_slug = slugify(target.name) 

1086 separator = settings.gateway_tool_name_separator 

1087 

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 ) 

1096 

1097 # 5. Execute the statement using the connection from the ongoing transaction. 

1098 connection.execute(stmt) 

1099 

1100 

1101class SessionRecord(Base): 

1102 """ORM model for sessions from SSE client.""" 

1103 

1104 __tablename__ = "mcp_sessions" 

1105 

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) 

1110 

1111 messages: Mapped[List["SessionMessageRecord"]] = relationship("SessionMessageRecord", back_populates="session", cascade="all, delete-orphan") 

1112 

1113 

1114class SessionMessageRecord(Base): 

1115 """ORM model for messages from SSE client.""" 

1116 

1117 __tablename__ = "mcp_messages" 

1118 

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 

1124 

1125 session: Mapped["SessionRecord"] = relationship("SessionRecord", back_populates="messages") 

1126 

1127 

1128# Event listeners for validation 

1129def validate_tool_schema(mapper, connection, target): 

1130 """ 

1131 Validate tool schema before insert/update. 

1132 

1133 Args: 

1134 mapper: The mapper being used for the operation. 

1135 connection: The database connection. 

1136 target: The target object being validated. 

1137 

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)}") 

1149 

1150 

1151def validate_tool_name(mapper, connection, target): 

1152 """ 

1153 Validate tool name before insert/update. Check if the name matches the required pattern. 

1154 

1155 Args: 

1156 mapper: The mapper being used for the operation. 

1157 connection: The database connection. 

1158 target: The target object being validated. 

1159 

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.") 

1169 

1170 

1171def validate_prompt_schema(mapper, connection, target): 

1172 """ 

1173 Validate prompt argument schema before insert/update. 

1174 

1175 Args: 

1176 mapper: The mapper being used for the operation. 

1177 connection: The database connection. 

1178 target: The target object being validated. 

1179 

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)}") 

1191 

1192 

1193# Register validation listeners 

1194 

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) 

1201 

1202 

1203def get_db(): 

1204 """ 

1205 Dependency to get database session. 

1206 

1207 Yields: 

1208 SessionLocal: A SQLAlchemy database session. 

1209 """ 

1210 db = SessionLocal() 

1211 try: 

1212 yield db 

1213 finally: 

1214 db.close() 

1215 

1216 

1217# Create all tables 

1218def init_db(): 

1219 """ 

1220 Initialize database tables. 

1221 

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)}") 

1230 

1231 

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 

1235 

1236 init_db()