Coverage for mcpgateway/db.py: 74%

1# -*- coding: utf-8 -*-

2"""MCP Gateway Database Models.

3

4 Copyright 2025

5 SPDX-License-Identifier: Apache-2.0

6 Authors: Mihai Criveti

7

8 This 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

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

15 and to record tool execution metrics.

16"""

17

18# Standard

19 from datetime import datetime, timezone

20 import re

21 from typing import Any, Dict, List, Optional

22 import uuid

23

24# Third-Party

25 import jsonschema

26 from 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)

44 from sqlalchemy.event import listen

45 from sqlalchemy.exc import SQLAlchemyError

46 from sqlalchemy.ext.hybrid import hybrid_property

47 from sqlalchemy.orm import (

48 DeclarativeBase,

49 Mapped,

50 mapped_column,

51 relationship,

52 sessionmaker,

53)

54 from sqlalchemy.orm.attributes import get_history

55

56# First-Party

57 from mcpgateway.config import settings

58 from mcpgateway.models import ResourceContent

59 from mcpgateway.utils.create_slug import slugify

60 from 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# ---------------------------------------------------------------------------

66 url = make_url(settings.database_url)

67 backend = url.get_backend_name() # e.g. 'postgresql', 'sqlite'

68 driver = 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.

72 connect_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# ---------------------------------------------------------------------------

79 if 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# ---------------------------------------------------------------------------

90 elif 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# ---------------------------------------------------------------------------

99 engine = 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# ---------------------------------------------------------------------------

113 def 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

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

125

126

127 class 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

157 server_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

165 server_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

173 server_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

181 class 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

209 class 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

235 class 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

261 class 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

287 class 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

558 class 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

726 class 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

740 class 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

891 class 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

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

1064 def 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

1101 class 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

1114 class 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

1129 def 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

1151 def 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

1171 def 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

1195 listen(Tool, "before_insert", validate_tool_schema)

1196 listen(Tool, "before_update", validate_tool_schema)

1197 listen(Tool, "before_insert", validate_tool_name)

1198 listen(Tool, "before_update", validate_tool_name)

1199 listen(Prompt, "before_insert", validate_prompt_schema)

1200 listen(Prompt, "before_update", validate_prompt_schema)

1201

1202

1203 def 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

1218 def 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

1232 if __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()