Coverage for mcpgateway/db.py: 67%

351 statements  

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

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 

18import re 

19from datetime import datetime, timezone 

20from typing import Any, Dict, List, Optional 

21 

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) 

49 

50from mcpgateway.config import settings 

51from mcpgateway.types import ResourceContent 

52 

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" 

60 

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

64 

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 ) 

77 

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 

84 

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

86 

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) 

99 

100# Session factory 

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

102 

103 

104class Base(DeclarativeBase): 

105 """Base class for all models.""" 

106 

107 

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

115 

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

123 

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

131 

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) 

139 

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) 

147 

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) 

155 

156 

157class ToolMetric(Base): 

158 """ 

159 ORM model for recording individual metrics for tool executions. 

160 

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. 

166 

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

171 

172 __tablename__ = "tool_metrics" 

173 

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) 

180 

181 # Relationship back to the Tool model. 

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

183 

184 

185class ResourceMetric(Base): 

186 """ 

187 ORM model for recording metrics for resource invocations. 

188 

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

197 

198 __tablename__ = "resource_metrics" 

199 

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) 

206 

207 # Relationship back to the Resource model. 

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

209 

210 

211class ServerMetric(Base): 

212 """ 

213 ORM model for recording metrics for server invocations. 

214 

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

223 

224 __tablename__ = "server_metrics" 

225 

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) 

232 

233 # Relationship back to the Server model. 

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

235 

236 

237class PromptMetric(Base): 

238 """ 

239 ORM model for recording metrics for prompt invocations. 

240 

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

249 

250 __tablename__ = "prompt_metrics" 

251 

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) 

258 

259 # Relationship back to the Prompt model. 

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

261 

262 

263class Tool(Base): 

264 """ 

265 ORM model for a registered Tool. 

266 

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 

271 

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. 

282 

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

284 

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

294 

295 __tablename__ = "tools" 

296 

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

309 

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) 

313 

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

318 

319 # Many-to-many relationship with Servers 

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

321 

322 # Relationship with ToolMetric records 

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

324 

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. 

330 

331 Returns: 

332 int: The total count of tool executions. 

333 """ 

334 return len(self.metrics) 

335 

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. 

342 

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 

347 

348 @property 

349 def successful_executions(self) -> int: 

350 """ 

351 Returns the count of successful tool executions, 

352 computed from the associated ToolMetric records. 

353 

354 Returns: 

355 int: The count of successful tool executions. 

356 """ 

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

358 

359 @property 

360 def failed_executions(self) -> int: 

361 """ 

362 Returns the count of failed tool executions, 

363 computed from the associated ToolMetric records. 

364 

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) 

369 

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. 

376 

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 

385 

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. 

391 

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 

397 

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. 

403 

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 

409 

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. 

415 

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 

421 

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. 

427 

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) 

434 

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

447 

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 } 

461 

462 

463class Resource(Base): 

464 """ 

465 ORM model for a registered Resource. 

466 

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

472 

473 __tablename__ = "resources" 

474 

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

486 

487 # Content storage - can be text or binary 

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

489 binary_content: Mapped[Optional[bytes]] 

490 

491 # Subscription tracking 

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

493 

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

497 

498 # Many-to-many relationship with Servers 

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

500 

501 @property 

502 def content(self) -> ResourceContent: 

503 """ 

504 Returns the resource content in the appropriate format. 

505 

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. 

509 

510 Returns: 

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

512 

513 Raises: 

514 ValueError: If the resource has no content available. 

515 """ 

516 

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

532 

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. 

538 

539 Returns: 

540 int: The total count of resource invocations. 

541 """ 

542 return len(self.metrics) 

543 

544 @property 

545 def successful_executions(self) -> int: 

546 """ 

547 Returns the count of successful resource invocations, 

548 computed from the associated ResourceMetric records. 

549 

550 Returns: 

551 int: The count of successful resource invocations. 

552 """ 

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

554 

555 @property 

556 def failed_executions(self) -> int: 

557 """ 

558 Returns the count of failed resource invocations, 

559 computed from the associated ResourceMetric records. 

560 

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) 

565 

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. 

572 

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 

580 

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. 

586 

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 

592 

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. 

598 

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 

604 

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. 

610 

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 

616 

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. 

622 

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) 

629 

630 

631class ResourceSubscription(Base): 

632 """Tracks subscriptions to resource updates.""" 

633 

634 __tablename__ = "resource_subscriptions" 

635 

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) 

641 

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

643 

644 

645class Prompt(Base): 

646 """ 

647 ORM model for a registered Prompt template. 

648 

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

662 

663 __tablename__ = "prompts" 

664 

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

674 

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

678 

679 # Many-to-many relationship with Servers 

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

681 

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

683 """ 

684 Validate prompt arguments against the argument schema. 

685 

686 Args: 

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

688 

689 Raises: 

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

691 

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

697 

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. 

703 

704 Returns: 

705 int: The total count of prompt invocations. 

706 """ 

707 return len(self.metrics) 

708 

709 @property 

710 def successful_executions(self) -> int: 

711 """ 

712 Returns the count of successful prompt invocations, 

713 computed from the associated PromptMetric records. 

714 

715 Returns: 

716 int: The count of successful prompt invocations. 

717 """ 

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

719 

720 @property 

721 def failed_executions(self) -> int: 

722 """ 

723 Returns the count of failed prompt invocations, 

724 computed from the associated PromptMetric records. 

725 

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) 

730 

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. 

737 

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 

745 

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. 

751 

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 

757 

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. 

763 

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 

769 

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. 

775 

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 

781 

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. 

787 

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) 

794 

795 

796class Server(Base): 

797 """ 

798 ORM model for MCP Servers Catalog. 

799 

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

812 

813 __tablename__ = "servers" 

814 

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

823 

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

828 

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. 

834 

835 Returns: 

836 int: The total count of server invocations. 

837 """ 

838 return len(self.metrics) 

839 

840 @property 

841 def successful_executions(self) -> int: 

842 """ 

843 Returns the count of successful server invocations, 

844 computed from the associated ServerMetric records. 

845 

846 Returns: 

847 int: The count of successful server invocations. 

848 """ 

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

850 

851 @property 

852 def failed_executions(self) -> int: 

853 """ 

854 Returns the count of failed server invocations, 

855 computed from the associated ServerMetric records. 

856 

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) 

861 

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. 

868 

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 

876 

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. 

882 

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 

888 

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. 

894 

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 

900 

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. 

906 

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 

912 

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. 

918 

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) 

925 

926 

927class Gateway(Base): 

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

929 

930 __tablename__ = "gateways" 

931 

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

942 

943 # Relationship with local tools this gateway provides 

944 tools: Mapped[List["Tool"]] = relationship(back_populates="gateway", cascade="all, delete-orphan") 

945 

946 # Relationship with local prompts this gateway provides 

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

948 

949 # Relationship with local resources this gateway provides 

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

951 

952 # # Tools federated from this gateway 

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

954 

955 # # Prompts federated from this resource 

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

957 

958 # # Prompts federated from this gateway 

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

960 

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) 

964 

965 

966class SessionRecord(Base): 

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

968 

969 __tablename__ = "mcp_sessions" 

970 

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) 

975 

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

977 

978 

979class SessionMessageRecord(Base): 

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

981 

982 __tablename__ = "mcp_messages" 

983 

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 

989 

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

991 

992 

993# Event listeners for validation 

994def validate_tool_schema(mapper, connection, target): 

995 """ 

996 Validate tool schema before insert/update. 

997 

998 Args: 

999 mapper: The mapper being used for the operation. 

1000 connection: The database connection. 

1001 target: The target object being validated. 

1002 

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

1014 

1015 

1016def validate_tool_name(mapper, connection, target): 

1017 """ 

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

1019 

1020 Args: 

1021 mapper: The mapper being used for the operation. 

1022 connection: The database connection. 

1023 target: The target object being validated. 

1024 

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

1034 

1035 

1036def validate_prompt_schema(mapper, connection, target): 

1037 """ 

1038 Validate prompt argument schema before insert/update. 

1039 

1040 Args: 

1041 mapper: The mapper being used for the operation. 

1042 connection: The database connection. 

1043 target: The target object being validated. 

1044 

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

1056 

1057 

1058# Register validation listeners 

1059 

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) 

1066 

1067 

1068def get_db(): 

1069 """ 

1070 Dependency to get database session. 

1071 

1072 Yields: 

1073 SessionLocal: A SQLAlchemy database session. 

1074 """ 

1075 db = SessionLocal() 

1076 try: 

1077 yield db 

1078 finally: 

1079 db.close() 

1080 

1081 

1082# Create all tables 

1083def init_db(): 

1084 """ 

1085 Initialize database tables. 

1086 

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

1095 

1096 

1097if __name__ == "__main__": 

1098 init_db()