Coverage for mcpgateway/db.py: 99%

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

2"""Location: ./mcpgateway/db.py

4SPDX-License-Identifier: Apache-2.0

5Authors: Mihai Criveti

7ContextForge Database Models.

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- Updated to record server associations independently using many-to-many relationships,

14- and to record tool execution metrics.

16Examples:

17 >>> from mcpgateway.db import connect_args

18 >>> isinstance(connect_args, dict)

19 True

20 >>> 'keepalives' in connect_args or 'check_same_thread' in connect_args or len(connect_args) == 0

21 True

22"""

24# Standard

25from contextlib import contextmanager

26from datetime import datetime, timedelta, timezone

27import logging

28import os

29from typing import Any, cast, Dict, Generator, List, Optional, TYPE_CHECKING

30import uuid

32# Third-Party

33import jsonschema

34from sqlalchemy import Boolean, Column, create_engine, DateTime, event, Float, ForeignKey, func, Index

35from sqlalchemy import inspect as sa_inspect

36from sqlalchemy import Integer, JSON, make_url, MetaData, select, String, Table, text, Text, UniqueConstraint

37from sqlalchemy.engine import Engine

38from sqlalchemy.event import listen

39from sqlalchemy.exc import OperationalError, ProgrammingError, SQLAlchemyError

40from sqlalchemy.ext.hybrid import hybrid_property

41from sqlalchemy.orm import DeclarativeBase, joinedload, Mapped, mapped_column, relationship, Session, sessionmaker

42from sqlalchemy.orm.attributes import get_history

43from sqlalchemy.pool import NullPool, QueuePool

44from sqlalchemy.types import TypeDecorator

46# First-Party

47from mcpgateway.common.validators import SecurityValidator

48from mcpgateway.config import settings

49from mcpgateway.utils.create_slug import slugify

50from mcpgateway.utils.db_isready import wait_for_db_ready

52logger = logging.getLogger(__name__)

54if TYPE_CHECKING:

55 # First-Party

56 from mcpgateway.common.models import ResourceContent

58# ResourceContent will be imported locally where needed to avoid circular imports

59# EmailUser models moved to this file to avoid circular imports

61# ---------------------------------------------------------------------------

62# 1. Parse the URL so we can inspect backend ("postgresql", "sqlite", ...)

63# and the specific driver ("psycopg", "asyncpg", empty string = default).

64# ---------------------------------------------------------------------------

65url = make_url(settings.database_url)

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

67driver = url.get_driver_name() or "default"

69# Start with an empty dict and add options only when the driver can accept

70# them; this prevents unexpected TypeError at connect time.

71connect_args: dict[str, object] = {}

73# ---------------------------------------------------------------------------

74# 2. PostgreSQL (synchronous psycopg3)

75# The keep-alive parameters below are recognised by libpq and let the

76# kernel detect broken network links quickly.

77#

78# Additionally, support PostgreSQL-specific options like search_path

79# via the 'options' query parameter in DATABASE_URL.

80# Example: postgresql+psycopg://user:pass@host/db?options=-c%20search_path=mcp_gateway

81#

82# IMPORTANT: Use postgresql+psycopg:// (not postgresql://) for psycopg3.

83# ---------------------------------------------------------------------------

84if backend == "postgresql" and driver in ("psycopg", "default", ""):

85 connect_args.update(

86 keepalives=1, # enable TCP keep-alive probes

87 keepalives_idle=30, # seconds of idleness before first probe

88 keepalives_interval=5, # seconds between probes

89 keepalives_count=5, # drop the link after N failed probes

90 # psycopg3: prepare_threshold controls automatic server-side prepared statements

91 # After N executions of the same query, psycopg3 prepares it server-side

92 # This significantly improves performance for frequently-executed queries

93 prepare_threshold=settings.db_prepare_threshold,

94 )

96 # Extract and apply PostgreSQL options from URL query parameters

97 # This allows users to specify search_path for custom schema support (Issue #1535)

98 url_options = url.query.get("options")

99 if url_options:

100 connect_args["options"] = url_options

101 logger.info(f"PostgreSQL connection options applied: {url_options}")

102

103 logger.info(f"psycopg3 prepare_threshold set to {settings.db_prepare_threshold}")

104

105# ---------------------------------------------------------------------------

106# 3. SQLite (optional) - only one extra flag and it is *SQLite-specific*.

107# ---------------------------------------------------------------------------

108elif backend == "sqlite":

109 # Allow pooled connections to hop across threads.

110 connect_args["check_same_thread"] = False

111

112# 4. Other backends leave `connect_args` empty.

113

114# ---------------------------------------------------------------------------

115# 5. Build the Engine with a single, clean connect_args mapping.

116# ---------------------------------------------------------------------------

117

118# Check for SQLALCHEMY_ECHO environment variable for query debugging

119# This is useful for N+1 detection and performance analysis

120_sqlalchemy_echo = os.getenv("SQLALCHEMY_ECHO", "").lower() in ("true", "1", "yes")

121

122

123def build_engine() -> Engine:

124 """Build the SQLAlchemy engine with appropriate settings.

125

126 This function constructs the SQLAlchemy engine using the database URL

127 and connection arguments determined by the backend type. It also configures

128 the connection pool size and timeout based on application settings.

129

130 Environment variables:

131 SQLALCHEMY_ECHO: Set to 'true' to log all SQL queries (useful for N+1 detection)

132

133 Returns:

134 SQLAlchemy Engine instance configured for the specified database.

135

136 Raises:

137 ValueError: If the database backend is not postgresql or sqlite.

138 """

139 if _sqlalchemy_echo:

140 logger.info("SQLALCHEMY_ECHO enabled - all SQL queries will be logged")

141

142 if backend == "sqlite":

143 # SQLite supports connection pooling with proper configuration

144 # For SQLite, we use a smaller pool size since it's file-based

145 sqlite_pool_size = min(settings.db_pool_size, 50) # Cap at 50 for SQLite

146 sqlite_max_overflow = min(settings.db_max_overflow, 20) # Cap at 20 for SQLite

147

148 logger.info("Configuring SQLite with pool_size=%s, max_overflow=%s", sqlite_pool_size, sqlite_max_overflow)

149

150 return create_engine(

151 settings.database_url,

152 pool_pre_ping=True, # quick liveness check per checkout

153 pool_size=sqlite_pool_size,

154 max_overflow=sqlite_max_overflow,

155 pool_timeout=settings.db_pool_timeout,

156 pool_recycle=settings.db_pool_recycle,

157 # SQLite specific optimizations

158 poolclass=QueuePool, # Explicit pool class

159 connect_args=connect_args,

160 # Log pool events in debug mode

161 echo_pool=settings.log_level == "DEBUG",

162 # Log all SQL queries when SQLALCHEMY_ECHO=true (useful for N+1 detection)

163 echo=_sqlalchemy_echo,

164 )

165

166 if backend != "postgresql":

167 raise ValueError(f"Unsupported database backend: '{backend}'. Only 'postgresql' and 'sqlite' are supported.")

168

169 # Determine if PgBouncer is in use (detected via URL or explicit config)

170 is_pgbouncer = "pgbouncer" in settings.database_url.lower()

171

172 # Determine pool class based on configuration

173 # - "auto": NullPool with PgBouncer (recommended), QueuePool otherwise

174 # - "null": Always NullPool (delegate pooling to PgBouncer/external pooler)

175 # - "queue": Always QueuePool (application-side pooling)

176 use_null_pool = False

177 if settings.db_pool_class == "null":

178 use_null_pool = True

179 logger.info("Using NullPool (explicit configuration)")

180 elif settings.db_pool_class == "auto" and is_pgbouncer:

181 use_null_pool = True

182 logger.info("PgBouncer detected - using NullPool (recommended: let PgBouncer handle pooling)")

183 elif settings.db_pool_class == "queue":

184 logger.info("Using QueuePool (explicit configuration)")

185 else:

186 logger.info("Using QueuePool with pool_size=%s, max_overflow=%s", settings.db_pool_size, settings.db_max_overflow)

187

188 # Determine pre_ping setting

189 # - "auto": Enabled for non-PgBouncer with QueuePool, disabled otherwise

190 # - "true": Always enable (validates connections, catches stale connections)

191 # - "false": Always disable

192 if settings.db_pool_pre_ping == "true":

193 use_pre_ping = True

194 logger.info("pool_pre_ping enabled (explicit configuration)")

195 elif settings.db_pool_pre_ping == "false":

196 use_pre_ping = False

197 logger.info("pool_pre_ping disabled (explicit configuration)")

198 else: # "auto"

199 # With NullPool, pre_ping is not needed (no pooled connections to validate)

200 # With QueuePool + PgBouncer, pre_ping helps detect stale connections

201 use_pre_ping = not use_null_pool and not is_pgbouncer

202 if is_pgbouncer and not use_null_pool:

203 logger.info("PgBouncer with QueuePool - consider enabling DB_POOL_PRE_PING=true to detect stale connections")

204

205 # Build engine with appropriate pool configuration

206 if use_null_pool:

207 return create_engine(

208 settings.database_url,

209 poolclass=NullPool,

210 connect_args=connect_args,

211 echo=_sqlalchemy_echo,

212 )

213

214 return create_engine(

215 settings.database_url,

216 pool_pre_ping=use_pre_ping,

217 pool_size=settings.db_pool_size,

218 max_overflow=settings.db_max_overflow,

219 pool_timeout=settings.db_pool_timeout,

220 pool_recycle=settings.db_pool_recycle,

221 connect_args=connect_args,

222 echo=_sqlalchemy_echo,

223 )

224

225

226engine = build_engine()

227

228# Initialize SQLAlchemy instrumentation for observability

229if settings.observability_enabled:

230 try:

231 # First-Party

232 from mcpgateway.instrumentation import instrument_sqlalchemy

233

234 instrument_sqlalchemy(engine)

235 logger.info("SQLAlchemy instrumentation enabled for observability")

236 except ImportError:

237 logger.warning("Failed to import SQLAlchemy instrumentation")

238

239

240# ---------------------------------------------------------------------------

241# 6. Function to return UTC timestamp

242# ---------------------------------------------------------------------------

243def utc_now() -> datetime:

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

245

246 Returns:

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

248 `datetime.timezone.utc`.

249

250 Examples:

251 >>> from mcpgateway.db import utc_now

252 >>> now = utc_now()

253 >>> now.tzinfo is not None

254 True

255 >>> str(now.tzinfo)

256 'UTC'

257 >>> isinstance(now, datetime)

258 True

259 """

260 return datetime.now(timezone.utc)

261

262

263class TokenEncryptionWriteError(ValueError):

264 """Raised when OAuth token encryption fails during DB write binding."""

265

266

267class EncryptedText(TypeDecorator): # pylint: disable=too-many-ancestors

268 """Text type that applies best-effort encryption/decryption at ORM boundary.

269

270 This preserves compatibility with service-layer encryption:

271 - Pre-encrypted values pass through unchanged.

272 - Plaintext values are encrypted when possible before persistence.

273 - On read, encrypted values are decrypted for runtime usage.

274 """

275

276 impl = Text

277 cache_ok = True

278

279 @property

280 def python_type(self):

281 """Return the Python type represented by this SQLAlchemy type.

282

283 Returns:

284 type: Python ``str`` type.

285 """

286 return str

287

288 @staticmethod

289 def _get_encryption():

290 """Resolve encryption service for column-level token protection.

291

292 Returns:

293 Optional[EncryptionService]: Encryption service instance when configured,

294 otherwise ``None``.

295 """

296 secret = getattr(settings, "auth_encryption_secret", None)

297 if not secret:

298 return None

299 try:

300 # First-Party

301 from mcpgateway.services.encryption_service import get_encryption_service # pylint: disable=import-outside-toplevel

302

303 return get_encryption_service(secret)

304 except Exception as exc:

305 logger.debug("Unable to initialize encryption service for EncryptedText: %s", exc)

306 return None

307

308 def process_literal_param(self, value, _dialect): # pylint: disable=unused-argument

309 """Render literal SQL parameter value via encrypted bind processing.

310

311 Args:

312 value (Any): Raw value from SQLAlchemy.

313 _dialect: SQLAlchemy dialect (unused).

314

315 Returns:

316 Any: Bound parameter value after encryption handling.

317 """

318 processed = self.process_bind_param(value, _dialect)

319 return processed

320

321 def process_bind_param(self, value, _dialect): # pylint: disable=unused-argument

322 """Encrypt plaintext values before persistence when encryption is available.

323

324 Args:

325 value (Any): Raw value from SQLAlchemy.

326 _dialect: SQLAlchemy dialect (unused).

327

328 Returns:

329 Any: Encrypted value for persistence or unchanged value when no

330 encryption is applied.

331

332 Raises:

333 TokenEncryptionWriteError: If encryption is configured and token

334 encryption fails.

335 """

336 if value in (None, "") or not isinstance(value, str):

337 return value

338

339 encryption = self._get_encryption()

340 if not encryption:

341 return value

342

343 try:

344 if encryption.is_encrypted(value):

345 return value

346 return encryption.encrypt_secret(value)

347 except Exception as exc:

348 logger.warning("EncryptedText bind encryption failed; rejecting token write")

349 logger.debug("EncryptedText bind encryption exception: %s", exc)

350 raise TokenEncryptionWriteError("OAuth token encryption failed during write") from exc

351

352 def process_result_value(self, value, _dialect): # pylint: disable=unused-argument

353 """Decrypt stored encrypted values when reading rows.

354

355 Args:

356 value (Any): Raw value loaded from database.

357 _dialect: SQLAlchemy dialect (unused).

358

359 Returns:

360 Any: Decrypted value when encrypted, otherwise unchanged.

361 """

362 if value in (None, "") or not isinstance(value, str):

363 return value

364

365 encryption = self._get_encryption()

366 if not encryption:

367 return value

368

369 try:

370 if not encryption.is_encrypted(value):

371 return value

372 decrypted = encryption.decrypt_secret_or_plaintext(value)

373 return decrypted if decrypted is not None else value

374 except Exception as exc:

375 logger.warning("EncryptedText result decryption failed, returning stored value")

376 logger.debug("EncryptedText result decryption exception: %s", exc)

377 return value

378

379

380# Configure SQLite for better concurrency if using SQLite

381if backend == "sqlite":

382

383 @event.listens_for(engine, "connect")

384 def set_sqlite_pragma(dbapi_conn, _connection_record):

385 """Set SQLite pragmas for better concurrency.

386

387 This is critical for running with multiple gunicorn workers.

388 WAL mode allows multiple readers and a single writer concurrently.

389

390 Args:

391 dbapi_conn: The raw DBAPI connection.

392 _connection_record: A SQLAlchemy-specific object that maintains

393 information about the connection's context.

394 """

395 cursor = dbapi_conn.cursor()

396 # Enable WAL mode for better concurrency

397 cursor.execute("PRAGMA journal_mode=WAL")

398 # Configure SQLite lock wait upper bound (ms) to prevent prolonged blocking under contention

399 cursor.execute(f"PRAGMA busy_timeout={settings.db_sqlite_busy_timeout}")

400 # Synchronous=NORMAL is safe with WAL mode and improves performance

401 cursor.execute("PRAGMA synchronous=NORMAL")

402 # Increase cache size for better performance (negative value = KB)

403 cursor.execute("PRAGMA cache_size=-64000") # 64MB cache

404 # Enable foreign key constraints for ON DELETE CASCADE support

405 cursor.execute("PRAGMA foreign_keys=ON")

406 cursor.close()

407

408

409# ---------------------------------------------------------------------------

410# Resilient Session class for graceful error recovery

411# ---------------------------------------------------------------------------

412class ResilientSession(Session):

413 """A Session subclass that auto-rollbacks on connection errors.

414

415 When a database operation fails due to a connection error (e.g., PgBouncer

416 query_wait_timeout), this session automatically rolls back to clear the

417 invalid transaction state. This prevents cascading PendingRollbackError

418 failures when multiple queries run within the same request.

419

420 Without this, the first failed query leaves the session in a "needs rollback"

421 state, and all subsequent queries fail with PendingRollbackError before

422 even attempting to use the database.

423 """

424

425 # Error types that indicate connection issues requiring rollback

426 _connection_error_patterns = (

427 "query_wait_timeout",

428 "server closed the connection unexpectedly",

429 "connection reset by peer",

430 "connection timed out",

431 "could not receive data from server",

432 "could not send data to server",

433 "terminating connection",

434 "no connection to the server",

435 )

436

437 def _is_connection_error(self, exception: Exception) -> bool:

438 """Check if an exception indicates a broken database connection.

439

440 Args:

441 exception: The exception to check.

442

443 Returns:

444 True if the exception indicates a connection error, False otherwise.

445 """

446 exc_name = type(exception).__name__

447 exc_msg = str(exception).lower()

448

449 # Check for known connection error types

450 if exc_name in ("ProtocolViolation", "OperationalError", "InterfaceError"):

451 return True

452

453 # Check for connection error patterns in message

454 for pattern in self._connection_error_patterns:

455 if pattern in exc_msg:

456 return True

457

458 return False

459

460 def _safe_rollback(self) -> None:

461 """Attempt to rollback, invalidating the session if rollback fails."""

462 try:

463 self.rollback()

464 except Exception:

465 try:

466 self.invalidate()

467 except Exception:

468 pass # nosec B110 - Best effort cleanup on connection failure

469

470 def execute(self, statement, params=None, **kw):

471 """Execute a statement with automatic rollback on connection errors.

472

473 Wraps the parent execute method to catch connection errors and

474 automatically rollback the session to prevent PendingRollbackError cascade.

475

476 Args:

477 statement: The SQL statement to execute.

478 params: Optional parameters for the statement.

479 **kw: Additional keyword arguments passed to Session.execute().

480

481 Returns:

482 The result of the execute operation.

483

484 Raises:

485 Exception: Re-raises any exception after rolling back on connection errors.

486 """

487 try:

488 return super().execute(statement, params, **kw)

489 except Exception as e:

490 if self._is_connection_error(e):

491 logger.warning(

492 "Connection error during execute, auto-rolling back session: %s",

493 type(e).__name__,

494 )

495 self._safe_rollback()

496 raise

497

498 def scalar(self, statement, params=None, **kw):

499 """Execute and return a scalar with automatic rollback on connection errors.

500

501 Wraps the parent scalar method to catch connection errors and

502 automatically rollback the session to prevent PendingRollbackError cascade.

503

504 Args:

505 statement: The SQL statement to execute.

506 params: Optional parameters for the statement.

507 **kw: Additional keyword arguments passed to Session.scalar().

508

509 Returns:

510 The scalar result of the query.

511

512 Raises:

513 Exception: Re-raises any exception after rolling back on connection errors.

514 """

515 try:

516 return super().scalar(statement, params, **kw)

517 except Exception as e:

518 if self._is_connection_error(e):

519 logger.warning(

520 "Connection error during scalar, auto-rolling back session: %s",

521 type(e).__name__,

522 )

523 self._safe_rollback()

524 raise

525

526 def scalars(self, statement, params=None, **kw):

527 """Execute and return scalars with automatic rollback on connection errors.

528

529 Wraps the parent scalars method to catch connection errors and

530 automatically rollback the session to prevent PendingRollbackError cascade.

531

532 Args:

533 statement: The SQL statement to execute.

534 params: Optional parameters for the statement.

535 **kw: Additional keyword arguments passed to Session.scalars().

536

537 Returns:

538 The scalars result of the query.

539

540 Raises:

541 Exception: Re-raises any exception after rolling back on connection errors.

542 """

543 try:

544 return super().scalars(statement, params, **kw)

545 except Exception as e:

546 if self._is_connection_error(e):

547 logger.warning(

548 "Connection error during scalars, auto-rolling back session: %s",

549 type(e).__name__,

550 )

551 self._safe_rollback()

552 raise

553

554

555# Session factory using ResilientSession

556# expire_on_commit=False prevents SQLAlchemy from expiring ORM objects after commit,

557# allowing continued access to attributes without re-querying the database.

558# This is essential when commits happen during read operations (e.g., to release transactions).

559SessionLocal = sessionmaker(class_=ResilientSession, autocommit=False, autoflush=False, expire_on_commit=False, bind=engine)

560

561

562@event.listens_for(ResilientSession, "after_transaction_end")

563def end_transaction_cleanup(_session, _transaction):

564 """Ensure connection is properly released after transaction ends.

565

566 This event fires after COMMIT or ROLLBACK, ensuring the connection

567 is returned to PgBouncer cleanly with no open transaction.

568

569 Args:

570 _session: The SQLAlchemy session that ended the transaction.

571 _transaction: The transaction that was ended.

572 """

573 # The transaction has already ended - nothing to do here

574 # This is just for monitoring/logging if needed

575

576

577@event.listens_for(ResilientSession, "before_commit")

578def before_commit_handler(session):

579 """Handler before commit to ensure transaction is in good state.

580

581 This is called before COMMIT, ensuring any pending work is flushed.

582 If the flush fails, the exception is propagated so the commit also fails

583 and the caller's error handling (e.g. get_db rollback) can clean up properly.

584

585 Args:

586 session: The SQLAlchemy session about to commit.

587 """

588 session.flush()

589

590

591# ---------------------------------------------------------------------------

592# Pool event listeners for connection resilience

593# These handlers ensure broken connections are properly invalidated and

594# discarded from the pool, preventing "poisoned" connections from causing

595# cascading failures (e.g., PendingRollbackError after PgBouncer timeout).

596#

597# Key issue: PgBouncer returns ProtocolViolation (SQL error 08P01) for

598# query_wait_timeout, but SQLAlchemy doesn't recognize this as a disconnect

599# by default. We must explicitly mark these errors as disconnects so the

600# connection pool properly invalidates these connections.

601#

602# References:

603# - https://github.com/zodb/relstorage/issues/412

604# - https://docs.sqlalchemy.org/en/20/core/pooling.html#custom-legacy-pessimistic-ping

605# ---------------------------------------------------------------------------

606@event.listens_for(engine, "handle_error")

607def handle_pool_error(exception_context):

608 """Mark PgBouncer and connection errors as disconnects for proper pool invalidation.

609

610 This event fires when an error occurs during query execution. By marking

611 certain errors as disconnects (is_disconnect=True), SQLAlchemy will:

612 1. Invalidate the current connection (discard from pool)

613 2. Invalidate all other pooled connections older than current time

614

615 Without this, PgBouncer errors like query_wait_timeout result in

616 ProtocolViolation which is classified as DatabaseError, not a disconnect.

617 The connection stays in the pool and causes PendingRollbackError on reuse.

618

619 Args:

620 exception_context: SQLAlchemy ExceptionContext with error details.

621 """

622 original = exception_context.original_exception

623 if original is None:

624 return

625

626 # Get the exception class name and message for pattern matching

627 exc_class = type(original).__name__

628 exc_msg = str(original).lower()

629

630 # List of error patterns that indicate the connection is broken

631 # and should be treated as a disconnect for pool invalidation

632 disconnect_patterns = [

633 # PgBouncer errors

634 "query_wait_timeout",

635 "server_login_retry",

636 "client_login_timeout",

637 "client_idle_timeout",

638 "idle_transaction_timeout",

639 "server closed the connection unexpectedly",

640 "connection reset by peer",

641 "connection timed out",

642 "no connection to the server",

643 "terminating connection",

644 "connection has been closed unexpectedly",

645 # PostgreSQL errors indicating dead connection

646 "could not receive data from server",

647 "could not send data to server",

648 "ssl connection has been closed unexpectedly",

649 "canceling statement due to conflict with recovery",

650 ]

651

652 # Check for ProtocolViolation or OperationalError with disconnect patterns

653 is_connection_error = exc_class in ("ProtocolViolation", "OperationalError", "InterfaceError", "DatabaseError")

654

655 if is_connection_error:

656 for pattern in disconnect_patterns:

657 if pattern in exc_msg:

658 exception_context.is_disconnect = True

659 logger.warning(

660 "Connection error detected, marking as disconnect for pool invalidation: %s: %s",

661 exc_class,

662 pattern,

663 )

664 return

665

666 # Also treat ProtocolViolation from PgBouncer as disconnect even without message match

667 # PgBouncer sends 08P01 PROTOCOL_VIOLATION for various connection issues

668 if exc_class == "ProtocolViolation":

669 exception_context.is_disconnect = True

670 logger.warning(

671 "ProtocolViolation detected (likely PgBouncer), marking as disconnect: %s",

672 exc_msg[:200],

673 )

674

675

676@event.listens_for(engine, "checkin")

677def reset_connection_on_checkin(dbapi_connection, _connection_record):

678 """Reset connection state when returned to pool.

679

680 This ensures transactions are properly closed before the connection

681 is returned to PgBouncer, preventing 'idle in transaction' buildup.

682 With PgBouncer in transaction mode, connections stays reserved until

683 the transaction ends - this rollback releases them immediately.

684

685 Args:

686 dbapi_connection: The raw DBAPI connection being checked in.

687 _connection_record: The connection record tracking this connection.

688 """

689 try:

690 # Issue a rollback to close any open transaction

691 # This is safe for both read and write operations:

692 # - For reads: rollback has no effect but closes the transaction

693 # - For writes: they should already be committed by the application

694 dbapi_connection.rollback()

695 except Exception as e:

696 # Connection may be invalid - log and try to force close

697 logger.debug("Connection checkin rollback failed: %s", e)

698 try:

699 # Try to close the raw connection to release it from PgBouncer

700 dbapi_connection.close()

701 except Exception: # nosec B110

702 pass # Nothing more we can do

703

704

705@event.listens_for(engine, "reset")

706def reset_connection_on_reset(dbapi_connection, _connection_record, _reset_state):

707 """Reset connection state when the pool resets a connection.

708

709 This handles the case where a connection is being reset before reuse.

710

711 Args:

712 dbapi_connection: The raw DBAPI connection being reset.

713 _connection_record: The connection record tracking this connection.

714 """

715 try:

716 dbapi_connection.rollback()

717 except Exception: # nosec B110

718 pass # Connection may be invalid

719

720

721def _refresh_gateway_slugs_batched(session: Session, batch_size: int) -> None:

722 """Refresh gateway slugs in small batches to reduce memory usage.

723

724 Args:

725 session: Active SQLAlchemy session.

726 batch_size: Maximum number of rows to process per batch.

727 """

728

729 last_id: Optional[str] = None

730

731 while True:

732 query = session.query(Gateway).order_by(Gateway.id)

733 if last_id is not None:

734 query = query.filter(Gateway.id > last_id)

735

736 gateways = query.limit(batch_size).all()

737 if not gateways:

738 break

739

740 updated = False

741 for gateway in gateways:

742 new_slug = slugify(gateway.name)

743 if gateway.slug != new_slug:

744 gateway.slug = new_slug

745 updated = True

746

747 if updated:

748 session.commit()

749

750 # Free ORM state from memory between batches

751 session.expire_all()

752 last_id = gateways[-1].id

753

754

755def _refresh_tool_names_batched(session: Session, batch_size: int) -> None:

756 """Refresh tool names in batches with eager-loaded gateways.

757

758 Uses joinedload(Tool.gateway) to avoid N+1 queries when accessing the

759 gateway relationship while regenerating tool names.

760

761 Args:

762 session: Active SQLAlchemy session.

763 batch_size: Maximum number of rows to process per batch.

764 """

765

766 last_id: Optional[str] = None

767 separator = settings.gateway_tool_name_separator

768

769 while True:

770 stmt = select(Tool).options(joinedload(Tool.gateway)).order_by(Tool.id).limit(batch_size)

771 if last_id is not None:

772 stmt = stmt.where(Tool.id > last_id)

773

774 tools = session.execute(stmt).scalars().all()

775 if not tools:

776 break

777

778 updated = False

779 for tool in tools:

780 # Prefer custom_name_slug when available; fall back to original_name

781 name_slug_source = getattr(tool, "custom_name_slug", None) or tool.original_name

782 name_slug = slugify(name_slug_source)

783

784 if tool.gateway:

785 gateway_slug = slugify(tool.gateway.name)

786 new_name = f"{gateway_slug}{separator}{name_slug}"

787 else:

788 new_name = name_slug

789

790 if tool.name != new_name:

791 tool.name = new_name

792 updated = True

793

794 if updated:

795 session.commit()

796

797 # Free ORM state from memory between batches

798 session.expire_all()

799 last_id = tools[-1].id

800

801

802def _refresh_prompt_names_batched(session: Session, batch_size: int) -> None:

803 """Refresh prompt names in batches with eager-loaded gateways.

804

805 Uses joinedload(Prompt.gateway) to avoid N+1 queries when accessing the

806 gateway relationship while regenerating prompt names.

807

808 Args:

809 session: Active SQLAlchemy session.

810 batch_size: Maximum number of rows to process per batch.

811 """

812 last_id: Optional[str] = None

813 separator = settings.gateway_tool_name_separator

814

815 while True:

816 stmt = select(Prompt).options(joinedload(Prompt.gateway)).order_by(Prompt.id).limit(batch_size)

817 if last_id is not None:

818 stmt = stmt.where(Prompt.id > last_id)

819

820 prompts = session.execute(stmt).scalars().all()

821 if not prompts:

822 break

823

824 updated = False

825 for prompt in prompts:

826 name_slug_source = getattr(prompt, "custom_name_slug", None) or prompt.original_name

827 name_slug = slugify(name_slug_source)

828

829 if prompt.gateway:

830 gateway_slug = slugify(prompt.gateway.name)

831 new_name = f"{gateway_slug}{separator}{name_slug}"

832 else:

833 new_name = name_slug

834

835 if prompt.name != new_name:

836 prompt.name = new_name

837 updated = True

838

839 if updated:

840 session.commit()

841

842 session.expire_all()

843 last_id = prompts[-1].id

844

845

846def refresh_slugs_on_startup(batch_size: Optional[int] = None) -> None:

847 """Refresh slugs for all gateways and tool names on startup.

848

849 This implementation avoids loading all rows into memory at once by

850 streaming through the tables in batches and eager-loading tool.gateway

851 relationships to prevent N+1 query patterns.

852

853 Args:

854 batch_size: Optional maximum number of rows to process per batch. If

855 not provided, the value is taken from

856 ``settings.slug_refresh_batch_size`` with a default of ``1000``.

857 """

858

859 effective_batch_size = batch_size or getattr(settings, "slug_refresh_batch_size", 1000)

860

861 try:

862 with cast(Any, SessionLocal)() as session:

863 # Skip if tables don't exist yet (fresh database)

864 try:

865 _refresh_gateway_slugs_batched(session, effective_batch_size)

866 except (OperationalError, ProgrammingError) as e:

867 # Table doesn't exist yet - expected on fresh database

868 logger.info("Gateway table not found, skipping slug refresh: %s", e)

869 return

870

871 try:

872 _refresh_tool_names_batched(session, effective_batch_size)

873 except (OperationalError, ProgrammingError) as e:

874 # Table doesn't exist yet - expected on fresh database

875 logger.info("Tool table not found, skipping tool name refresh: %s", e)

876

877 try:

878 _refresh_prompt_names_batched(session, effective_batch_size)

879 except (OperationalError, ProgrammingError) as e:

880 # Table doesn't exist yet - expected on fresh database

881 logger.info("Prompt table not found, skipping prompt name refresh: %s", e)

882

883 except SQLAlchemyError as e:

884 logger.warning("Failed to refresh slugs on startup (database error): %s", e)

885 except Exception as e:

886 logger.warning("Failed to refresh slugs on startup (unexpected error): %s", e)

887

888

889def _compute_metrics_summary(

890 raw_metrics: List[Any],

891 hourly_metrics: List[Any],

892 session: Optional[Any] = None,

893 entity_id: Optional[str] = None,

894 raw_metric_class: Optional[Any] = None,

895 hourly_metric_class: Optional[Any] = None,

896) -> Dict[str, Any]:

897 """Compute aggregated metrics from both raw and hourly tables without double-counting.

898

899 This function prevents double-counting by including raw metrics only from hours

900 that have no corresponding hourly aggregate. This correctly handles:

901 - Normal operation (hourly rollup complete, raw data retained or deleted)

902 - Rollup lag (completed hour not yet rolled up)

903 - Rollup disabled or failed

904

905 The approach mirrors ``aggregate_metrics_combined()`` in metrics_query_service.py.

906

907 Args:

908 raw_metrics: List of raw metric objects loaded in memory (or None if using session)

909 hourly_metrics: List of hourly aggregate objects loaded in memory (or None if using session)

910 session: SQLAlchemy session for database queries (required if raw_metrics/hourly_metrics not loaded)

911 entity_id: ID of the entity (tool/resource/prompt/server/agent) for SQL query

912 raw_metric_class: ORM class for raw metrics (e.g., ToolMetric) for SQL query

913 hourly_metric_class: ORM class for hourly metrics (e.g., ToolMetricsHourly) for SQL query

914

915 Returns:

916 Dict with keys: total_executions, successful_executions, failed_executions,

917 failure_rate, min_response_time, max_response_time, avg_response_time,

918 last_execution_time

919

920 Raises:

921 ValueError: If both in-memory and SQL query parameters are incomplete

922 """

923 # Determine if we're using in-memory or SQL query path

924 use_memory = raw_metrics is not None and hourly_metrics is not None

925

926 if use_memory:

927 # ============================================================

928 # IN-MEMORY PATH: Iterate over loaded objects

929 # ============================================================

930

931 # Build set of hours already covered by hourly aggregates

932 covered_hours: set[datetime] = set()

933 for h in hourly_metrics:

934 hs = h.hour_start if h.hour_start.tzinfo is not None else h.hour_start.replace(tzinfo=timezone.utc)

935 covered_hours.add(hs)

936

937 total = 0

938 successful = 0

939 min_rt: Optional[float] = None

940 max_rt: Optional[float] = None

941 sum_rt = 0.0

942 last_time: Optional[datetime] = None

943

944 # Include raw metrics only from hours NOT covered by hourly aggregates

945 for m in raw_metrics:

946 metric_ts = m.timestamp if m.timestamp.tzinfo is not None else m.timestamp.replace(tzinfo=timezone.utc)

947 metric_hour = metric_ts.replace(minute=0, second=0, microsecond=0)

948 if metric_hour in covered_hours:

949 continue # Already counted in hourly aggregates

950

951 total += 1

952 if m.is_success:

953 successful += 1

954 rt = m.response_time

955 if min_rt is None or rt < min_rt:

956 min_rt = rt

957 if max_rt is None or rt > max_rt:

958 max_rt = rt

959 sum_rt += rt

960 if last_time is None or metric_ts > last_time:

961 last_time = metric_ts

962

963 # Process hourly aggregated metrics (completed hours)

964 for h in hourly_metrics:

965 total += h.total_count

966 successful += h.success_count

967 if h.min_response_time is not None:

968 if min_rt is None or h.min_response_time < min_rt:

969 min_rt = h.min_response_time

970 if h.max_response_time is not None:

971 if max_rt is None or h.max_response_time > max_rt:

972 max_rt = h.max_response_time

973 if h.avg_response_time is not None and h.total_count > 0:

974 sum_rt += h.avg_response_time * h.total_count

975 hs = h.hour_start if h.hour_start.tzinfo is not None else h.hour_start.replace(tzinfo=timezone.utc)

976 if last_time is None or hs > last_time:

977 last_time = hs

978

979 failed = total - successful

980 return {

981 "total_executions": total,

982 "successful_executions": successful,

983 "failed_executions": failed,

984 "failure_rate": failed / total if total > 0 else 0.0,

985 "min_response_time": min_rt,

986 "max_response_time": max_rt,

987 "avg_response_time": sum_rt / total if total > 0 else None,

988 "last_execution_time": last_time,

989 }

990

991 # ============================================================

992 # SQL QUERY PATH: hourly aggregates + uncovered raw metrics

993 # ============================================================

994 if session is None or entity_id is None or raw_metric_class is None or hourly_metric_class is None:

995 raise ValueError("For SQL query path, must provide: session, entity_id, raw_metric_class, hourly_metric_class")

996

997 # Third-Party

998 from sqlalchemy import case # pylint: disable=import-outside-toplevel

999

1000 # Determine the foreign key column name (tool_id, resource_id, etc.)

1001 class_name = raw_metric_class.__name__

1002 if class_name.endswith("Metric"):

1003 entity_type = class_name[:-6].lower() # ToolMetric -> tool

1004 fk_column_name = f"{entity_type}_id"

1005 else:

1006 raise ValueError(f"Cannot determine foreign key column for {class_name}")

1007

1008 fk_column_raw = getattr(raw_metric_class, fk_column_name)

1009 fk_column_hourly = getattr(hourly_metric_class, fk_column_name)

1010

1011 # Query 1: All hourly aggregates for this entity (includes max hour_start)

1012 hourly_result = (

1013 session.query(

1014 func.sum(hourly_metric_class.total_count), # pylint: disable=not-callable

1015 func.sum(hourly_metric_class.success_count), # pylint: disable=not-callable

1016 func.min(hourly_metric_class.min_response_time), # pylint: disable=not-callable

1017 func.max(hourly_metric_class.max_response_time), # pylint: disable=not-callable

1018 func.sum(hourly_metric_class.avg_response_time * hourly_metric_class.total_count), # weighted sum

1019 func.max(hourly_metric_class.hour_start), # pylint: disable=not-callable

1020 )

1021 .filter(fk_column_hourly == entity_id)

1022 .one()

1023 )

1024

1025 hourly_total = hourly_result[0] or 0

1026 hourly_successful = hourly_result[1] or 0

1027 hourly_min_rt = hourly_result[2]

1028 hourly_max_rt = hourly_result[3]

1029 hourly_weighted_sum_rt = hourly_result[4] or 0.0

1030 hourly_last_bucket = hourly_result[5]

1031

1032 # Query 2: Raw metrics from hours NOT covered by hourly aggregates.

1033 # Use max_hour_start to determine the boundary: hourly data covers

1034 # up to max_hour_start + 1h, raw data covers everything after that.

1035 # When no hourly data exists, all raw metrics are counted.

1036 raw_query = session.query(

1037 func.count(raw_metric_class.id), # pylint: disable=not-callable

1038 func.sum(case((raw_metric_class.is_success.is_(True), 1), else_=0)),

1039 func.min(raw_metric_class.response_time), # pylint: disable=not-callable

1040 func.max(raw_metric_class.response_time), # pylint: disable=not-callable

1041 func.sum(raw_metric_class.response_time), # pylint: disable=not-callable

1042 func.max(raw_metric_class.timestamp), # pylint: disable=not-callable

1043 ).filter(fk_column_raw == entity_id)

1044 if hourly_last_bucket is not None:

1045 # Only include raw metrics from after the last rolled-up hour

1046 hourly_coverage_end = hourly_last_bucket + timedelta(hours=1)

1047 raw_query = raw_query.filter(raw_metric_class.timestamp >= hourly_coverage_end)

1048

1049 raw_result = raw_query.one()

1050

1051 raw_total = raw_result[0] or 0

1052 raw_successful = raw_result[1] or 0

1053 raw_min_rt = raw_result[2]

1054 raw_max_rt = raw_result[3]

1055 raw_sum_rt = raw_result[4] or 0.0

1056 raw_last_time = raw_result[5]

1057

1058 # Aggregate totals

1059 total = hourly_total + raw_total

1060 successful = hourly_successful + raw_successful

1061 failed = total - successful

1062

1063 # Min/max across both sources

1064 min_rt = None

1065 if raw_min_rt is not None and hourly_min_rt is not None:

1066 min_rt = min(raw_min_rt, hourly_min_rt)

1067 elif raw_min_rt is not None:

1068 min_rt = raw_min_rt

1069 elif hourly_min_rt is not None:

1070 min_rt = hourly_min_rt

1071

1072 max_rt = None

1073 if raw_max_rt is not None and hourly_max_rt is not None:

1074 max_rt = max(raw_max_rt, hourly_max_rt)

1075 elif raw_max_rt is not None:

1076 max_rt = raw_max_rt

1077 elif hourly_max_rt is not None:

1078 max_rt = hourly_max_rt

1079

1080 # Weighted average response time

1081 avg_rt = None

1082 if total > 0:

1083 avg_rt = (hourly_weighted_sum_rt + raw_sum_rt) / total

1084

1085 # Last execution time: most recent (use hour_start for hourly, consistent with aggregate_metrics_combined)

1086 last_time = None

1087 if raw_last_time is not None and hourly_last_bucket is not None:

1088 last_time = max(raw_last_time, hourly_last_bucket)

1089 elif raw_last_time is not None:

1090 last_time = raw_last_time

1091 elif hourly_last_bucket is not None:

1092 last_time = hourly_last_bucket

1093

1094 return {

1095 "total_executions": total,

1096 "successful_executions": successful,

1097 "failed_executions": failed,

1098 "failure_rate": failed / total if total > 0 else 0.0,

1099 "min_response_time": min_rt,

1100 "max_response_time": max_rt,

1101 "avg_response_time": avg_rt,

1102 "last_execution_time": last_time,

1103 }

1104

1105

1106class Base(DeclarativeBase):

1107 """Base class for all models."""

1108

1109 # Naming convention for foreign keys

1110 metadata = MetaData(

1111 naming_convention={

1112 "fk": "fk_%(table_name)s_%(column_0_name)s",

1113 "pk": "pk_%(table_name)s",

1114 "ix": "ix_%(table_name)s_%(column_0_name)s",

1115 "uq": "uq_%(table_name)s_%(column_0_name)s",

1116 "ck": "ck_%(table_name)s_%(constraint_name)s",

1117 }

1118 )

1119

1120

1121# ---------------------------------------------------------------------------

1122# RBAC Models - SQLAlchemy Database Models

1123# ---------------------------------------------------------------------------

1124

1125

1126class Role(Base):

1127 """Role model for RBAC system."""

1128

1129 __tablename__ = "roles"

1130

1131 # Primary key

1132 id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()))

1133

1134 # Role metadata

1135 name: Mapped[str] = mapped_column(String(255), nullable=False)

1136 description: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

1137 scope: Mapped[str] = mapped_column(String(20), nullable=False) # 'global', 'team', 'personal'

1138

1139 # Permissions and inheritance

1140 permissions: Mapped[List[str]] = mapped_column(JSON, nullable=False, default=list)

1141 inherits_from: Mapped[Optional[str]] = mapped_column(String(36), ForeignKey("roles.id"), nullable=True)

1142

1143 # Metadata

1144 created_by: Mapped[str] = mapped_column(String(255), ForeignKey("email_users.email"), nullable=False)

1145 is_system_role: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)

1146 is_active: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True)

1147

1148 # Timestamps

1149 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False, default=utc_now)

1150 updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False, default=utc_now, onupdate=utc_now)

1151

1152 # Relationships

1153 parent_role: Mapped[Optional["Role"]] = relationship("Role", remote_side=[id], backref="child_roles")

1154 user_assignments: Mapped[List["UserRole"]] = relationship("UserRole", back_populates="role", cascade="all, delete-orphan")

1155

1156 def get_effective_permissions(self) -> List[str]:

1157 """Get all permissions including inherited ones.

1158

1159 Returns:

1160 List of permission strings including inherited permissions

1161 """

1162 effective_permissions = set(self.permissions)

1163 if self.parent_role:

1164 effective_permissions.update(self.parent_role.get_effective_permissions())

1165 return sorted(list(effective_permissions))

1166

1167

1168class UserRole(Base):

1169 """User role assignment model."""

1170

1171 __tablename__ = "user_roles"

1172

1173 # Primary key

1174 id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()))

1175

1176 # Assignment details

1177 user_email: Mapped[str] = mapped_column(String(255), ForeignKey("email_users.email"), nullable=False)

1178 role_id: Mapped[str] = mapped_column(String(36), ForeignKey("roles.id"), nullable=False)

1179 scope: Mapped[str] = mapped_column(String(20), nullable=False) # 'global', 'team', 'personal'

1180 scope_id: Mapped[Optional[str]] = mapped_column(String(36), nullable=True) # Team ID if team-scoped

1181

1182 # Grant metadata

1183 granted_by: Mapped[str] = mapped_column(String(255), ForeignKey("email_users.email"), nullable=False)

1184 granted_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False, default=utc_now)

1185 expires_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), nullable=True)

1186 is_active: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True)

1187 grant_source: Mapped[Optional[str]] = mapped_column(String(50), nullable=True, default=None)

1188

1189 # Relationships

1190 role: Mapped["Role"] = relationship("Role", back_populates="user_assignments")

1191

1192 def is_expired(self) -> bool:

1193 """Check if the role assignment has expired.

1194

1195 Returns:

1196 True if assignment has expired, False otherwise

1197 """

1198 if self.expires_at is None:

1199 return False

1200 expires_at = self.expires_at

1201 if expires_at.tzinfo is None:

1202 expires_at = expires_at.replace(tzinfo=timezone.utc)

1203

1204 return utc_now() > expires_at

1205

1206

1207class PermissionAuditLog(Base):

1208 """Permission audit log model."""

1209

1210 __tablename__ = "permission_audit_log"

1211

1212 # Primary key

1213 id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)

1214

1215 # Audit metadata

1216 timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False, default=utc_now)

1217 user_email: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

1218

1219 # Permission details

1220 permission: Mapped[str] = mapped_column(String(100), nullable=False)

1221 resource_type: Mapped[Optional[str]] = mapped_column(String(50), nullable=True)

1222 resource_id: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

1223 team_id: Mapped[Optional[str]] = mapped_column(String(36), nullable=True)

1224

1225 # Result

1226 granted: Mapped[bool] = mapped_column(Boolean, nullable=False)

1227 roles_checked: Mapped[Optional[Dict]] = mapped_column(JSON, nullable=True)

1228

1229 # Request metadata

1230 ip_address: Mapped[Optional[str]] = mapped_column(String(45), nullable=True) # IPv6 max length

1231 user_agent: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

1232

1233

1234# Permission constants for the system

1235class Permissions:

1236 """System permission constants."""

1237

1238 # User permissions

1239 USERS_CREATE = "users.create"

1240 USERS_READ = "users.read"

1241 USERS_UPDATE = "users.update"

1242 USERS_DELETE = "users.delete"

1243 USERS_INVITE = "users.invite"

1244

1245 # Team permissions

1246 TEAMS_CREATE = "teams.create"

1247 TEAMS_READ = "teams.read"

1248 TEAMS_UPDATE = "teams.update"

1249 TEAMS_DELETE = "teams.delete"

1250 TEAMS_JOIN = "teams.join"

1251 TEAMS_MANAGE_MEMBERS = "teams.manage_members"

1252

1253 # Tool permissions

1254 TOOLS_CREATE = "tools.create"

1255 TOOLS_READ = "tools.read"

1256 TOOLS_UPDATE = "tools.update"

1257 TOOLS_DELETE = "tools.delete"

1258 TOOLS_EXECUTE = "tools.execute"

1259

1260 # Resource permissions

1261 RESOURCES_CREATE = "resources.create"

1262 RESOURCES_READ = "resources.read"

1263 RESOURCES_UPDATE = "resources.update"

1264 RESOURCES_DELETE = "resources.delete"

1265 RESOURCES_SHARE = "resources.share"

1266

1267 # Gateway permissions

1268 GATEWAYS_CREATE = "gateways.create"

1269 GATEWAYS_READ = "gateways.read"

1270 GATEWAYS_UPDATE = "gateways.update"

1271 GATEWAYS_DELETE = "gateways.delete"

1272

1273 # Prompt permissions

1274 PROMPTS_CREATE = "prompts.create"

1275 PROMPTS_READ = "prompts.read"

1276 PROMPTS_UPDATE = "prompts.update"

1277 PROMPTS_DELETE = "prompts.delete"

1278 PROMPTS_EXECUTE = "prompts.execute"

1279

1280 # MCP method permission prefixes — used by token_catalog_service (generation-time)

1281 # and token_scoping middleware (runtime) to auto-grant servers.use transport access.

1282 MCP_METHOD_PREFIXES = ("tools.", "resources.", "prompts.")

1283

1284 # LLM proxy permissions

1285 LLM_READ = "llm.read"

1286 LLM_INVOKE = "llm.invoke"

1287

1288 # Server permissions

1289 SERVERS_CREATE = "servers.create"

1290 SERVERS_READ = "servers.read"

1291 SERVERS_USE = "servers.use"

1292 SERVERS_UPDATE = "servers.update"

1293 SERVERS_DELETE = "servers.delete"

1294 SERVERS_MANAGE = "servers.manage"

1295

1296 # Token permissions

1297 TOKENS_CREATE = "tokens.create"

1298 TOKENS_READ = "tokens.read"

1299 TOKENS_UPDATE = "tokens.update"

1300 TOKENS_REVOKE = "tokens.revoke"

1301

1302 # Admin permissions

1303 ADMIN_SYSTEM_CONFIG = "admin.system_config"

1304 ADMIN_USER_MANAGEMENT = "admin.user_management"

1305 ADMIN_SECURITY_AUDIT = "admin.security_audit"

1306 ADMIN_OVERVIEW = "admin.overview"

1307 ADMIN_DASHBOARD = "admin.dashboard"

1308 ADMIN_EVENTS = "admin.events"

1309 ADMIN_GRPC = "admin.grpc"

1310 ADMIN_PLUGINS = "admin.plugins"

1311 ADMIN_METRICS = "admin.metrics"

1312 ADMIN_EXPORT = "admin.export"

1313 ADMIN_IMPORT = "admin.import"

1314 ADMIN_SSO_PROVIDERS_CREATE = "admin.sso_providers:create"

1315 ADMIN_SSO_PROVIDERS_READ = "admin.sso_providers:read"

1316 ADMIN_SSO_PROVIDERS_UPDATE = "admin.sso_providers:update"

1317 ADMIN_SSO_PROVIDERS_DELETE = "admin.sso_providers:delete"

1318

1319 # Observability and audit read permissions

1320 LOGS_READ = "logs:read"

1321 METRICS_READ = "metrics:read"

1322 AUDIT_READ = "audit:read"

1323 SECURITY_READ = "security:read"

1324

1325 # A2A Agent permissions

1326 A2A_CREATE = "a2a.create"

1327 A2A_READ = "a2a.read"

1328 A2A_UPDATE = "a2a.update"

1329 A2A_DELETE = "a2a.delete"

1330 A2A_INVOKE = "a2a.invoke"

1331

1332 # Tag permissions

1333 TAGS_READ = "tags.read"

1334 TAGS_CREATE = "tags.create"

1335 TAGS_UPDATE = "tags.update"

1336 TAGS_DELETE = "tags.delete"

1337

1338 # Special permissions

1339 ALL_PERMISSIONS = "*" # Wildcard for all permissions

1340

1341 @classmethod

1342 def get_all_permissions(cls) -> List[str]:

1343 """Get list of all defined permissions.

1344

1345 Returns:

1346 List of all permission strings defined in the class

1347 """

1348 permissions = []

1349 for attr_name in dir(cls):

1350 if not attr_name.startswith("_") and attr_name.isupper() and attr_name != "ALL_PERMISSIONS":

1351 attr_value = getattr(cls, attr_name)

1352 if isinstance(attr_value, str):

1353 permissions.append(attr_value)

1354 return sorted(permissions)

1355

1356 @classmethod

1357 def get_permissions_by_resource(cls) -> Dict[str, List[str]]:

1358 """Get permissions organized by resource type.

1359

1360 Returns:

1361 Dictionary mapping resource types to their permissions

1362 """

1363 resource_permissions = {}

1364 for permission in cls.get_all_permissions():

1365 if "." in permission:

1366 resource_type = permission.split(".", 1)[0]

1367 elif ":" in permission:

1368 resource_type = permission.split(":", 1)[0]

1369 else:

1370 resource_type = permission

1371 if resource_type not in resource_permissions:

1372 resource_permissions[resource_type] = []

1373 resource_permissions[resource_type].append(permission)

1374 return resource_permissions

1375

1376

1377# ---------------------------------------------------------------------------

1378# Email-based User Authentication Models

1379# ---------------------------------------------------------------------------

1380

1381

1382class EmailUser(Base):

1383 """Email-based user model for authentication.

1384

1385 This model provides email-based authentication as the foundation

1386 for all multi-user features. Users are identified by email addresses

1387 instead of usernames.

1388

1389 Attributes:

1390 email (str): Primary key, unique email identifier

1391 password_hash (str): Argon2id hashed password

1392 full_name (str): Optional display name for professional appearance

1393 is_admin (bool): Admin privileges flag

1394 is_active (bool): Account status flag

1395 auth_provider (str): Authentication provider ('local', 'github', etc.)

1396 password_hash_type (str): Type of password hash used

1397 failed_login_attempts (int): Count of failed login attempts

1398 locked_until (datetime): Account lockout expiration

1399 created_at (datetime): Account creation timestamp

1400 updated_at (datetime): Last account update timestamp

1401 last_login (datetime): Last successful login timestamp

1402 email_verified_at (datetime): Email verification timestamp

1403

1404 Examples:

1405 >>> user = EmailUser(

1406 ... email="alice@example.com",

1407 ... password_hash="$argon2id$v=19$m=65536,t=3,p=1$...",

1408 ... full_name="Alice Smith",

1409 ... is_admin=False

1410 ... )

1411 >>> user.email

1412 'alice@example.com'

1413 >>> user.is_email_verified()

1414 False

1415 >>> user.is_account_locked()

1416 False

1417 """

1418

1419 __tablename__ = "email_users"

1420

1421 # Core identity fields

1422 email: Mapped[str] = mapped_column(String(255), primary_key=True, index=True)

1423 password_hash: Mapped[str] = mapped_column(String(255), nullable=False)

1424 full_name: Mapped[Optional[str]] = mapped_column(String(255), nullable=True, index=True)

1425 is_admin: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)

1426 # Track how admin status was granted: "sso" (synced from IdP), "manual" (Admin UI), "api" (API grant), or None (legacy)

1427 admin_origin: Mapped[Optional[str]] = mapped_column(String(20), nullable=True)

1428

1429 # Status fields

1430 is_active: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)

1431 email_verified_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), nullable=True)

1432

1433 # Security fields

1434 auth_provider: Mapped[str] = mapped_column(String(50), default="local", nullable=False)

1435 password_hash_type: Mapped[str] = mapped_column(String(20), default="argon2id", nullable=False)

1436 failed_login_attempts: Mapped[int] = mapped_column(Integer, default=0, nullable=False)

1437 locked_until: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), nullable=True)

1438 password_change_required: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)

1439 password_changed_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), default=utc_now, nullable=True)

1440

1441 # Timestamps

1442 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, nullable=False)

1443 updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, onupdate=utc_now, nullable=False)

1444 last_login: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), nullable=True)

1445

1446 def __repr__(self) -> str:

1447 """String representation of the user.

1448

1449 Returns:

1450 str: String representation of EmailUser instance

1451 """

1452 return f"<EmailUser(email='{self.email}', full_name='{self.full_name}', is_admin={self.is_admin})>"

1453

1454 def is_email_verified(self) -> bool:

1455 """Check if the user's email is verified.

1456

1457 Returns:

1458 bool: True if email is verified, False otherwise

1459

1460 Examples:

1461 >>> user = EmailUser(email="test@example.com")

1462 >>> user.is_email_verified()

1463 False

1464 >>> user.email_verified_at = utc_now()

1465 >>> user.is_email_verified()

1466 True

1467 """

1468 return self.email_verified_at is not None

1469

1470 def is_account_locked(self) -> bool:

1471 """Check if the account is currently locked.

1472

1473 Returns:

1474 bool: True if account is locked, False otherwise

1475

1476 Examples:

1477 >>> from datetime import timedelta

1478 >>> user = EmailUser(email="test@example.com")

1479 >>> user.is_account_locked()

1480 False

1481 >>> user.locked_until = utc_now() + timedelta(hours=1)

1482 >>> user.is_account_locked()

1483 True

1484 """

1485 if self.locked_until is None:

1486 return False

1487 locked_until = self.locked_until

1488 if locked_until.tzinfo is None:

1489 # Treat naive datetimes as UTC (SQLite strips timezone info)

1490 locked_until = locked_until.replace(tzinfo=timezone.utc)

1491 if utc_now() >= locked_until:

1492 # Lockout expired: reset counters so users get a fresh attempt window.

1493 self.failed_login_attempts = 0

1494 self.locked_until = None

1495 return False

1496 return True

1497

1498 def get_display_name(self) -> str:

1499 """Get the user's display name.

1500

1501 Returns the full_name if available, otherwise extracts

1502 the local part from the email address.

1503

1504 Returns:

1505 str: Display name for the user

1506

1507 Examples:

1508 >>> user = EmailUser(email="john@example.com", full_name="John Doe")

1509 >>> user.get_display_name()

1510 'John Doe'

1511 >>> user_no_name = EmailUser(email="jane@example.com")

1512 >>> user_no_name.get_display_name()

1513 'jane'

1514 """

1515 if self.full_name:

1516 return self.full_name

1517 return self.email.split("@")[0]

1518

1519 def reset_failed_attempts(self) -> None:

1520 """Reset failed login attempts counter.

1521

1522 Called after successful authentication to reset the

1523 failed attempts counter and clear any account lockout.

1524

1525 Examples:

1526 >>> user = EmailUser(email="test@example.com", failed_login_attempts=3)

1527 >>> user.reset_failed_attempts()

1528 >>> user.failed_login_attempts

1529 0

1530 >>> user.locked_until is None

1531 True

1532 """

1533 self.failed_login_attempts = 0

1534 self.locked_until = None

1535 self.last_login = utc_now()

1536

1537 def increment_failed_attempts(self, max_attempts: int = 5, lockout_duration_minutes: int = 30) -> bool:

1538 """Increment failed login attempts and potentially lock account.

1539

1540 Args:

1541 max_attempts: Maximum allowed failed attempts before lockout

1542 lockout_duration_minutes: Duration of lockout in minutes

1543

1544 Returns:

1545 bool: True if account is now locked, False otherwise

1546

1547 Examples:

1548 >>> user = EmailUser(email="test@example.com", password_hash="test", failed_login_attempts=0)

1549 >>> user.increment_failed_attempts(max_attempts=3)

1550 False

1551 >>> user.failed_login_attempts

1552 1

1553 >>> for _ in range(2):

1554 ... user.increment_failed_attempts(max_attempts=3)

1555 False

1556 True

1557 >>> user.is_account_locked()

1558 True

1559 """

1560 self.failed_login_attempts += 1

1561

1562 if self.failed_login_attempts >= max_attempts:

1563 self.locked_until = utc_now() + timedelta(minutes=lockout_duration_minutes)

1564 return True

1565

1566 return False

1567

1568 # Team relationships

1569 team_memberships: Mapped[List["EmailTeamMember"]] = relationship("EmailTeamMember", foreign_keys="EmailTeamMember.user_email", back_populates="user")

1570 created_teams: Mapped[List["EmailTeam"]] = relationship("EmailTeam", foreign_keys="EmailTeam.created_by", back_populates="creator")

1571 sent_invitations: Mapped[List["EmailTeamInvitation"]] = relationship("EmailTeamInvitation", foreign_keys="EmailTeamInvitation.invited_by", back_populates="inviter")

1572

1573 # API token relationships

1574 api_tokens: Mapped[List["EmailApiToken"]] = relationship("EmailApiToken", back_populates="user", cascade="all, delete-orphan")

1575

1576 def get_teams(self) -> List["EmailTeam"]:

1577 """Get all teams this user is a member of.

1578

1579 Returns:

1580 List[EmailTeam]: List of teams the user belongs to

1581

1582 Examples:

1583 >>> user = EmailUser(email="user@example.com")

1584 >>> teams = user.get_teams()

1585 >>> isinstance(teams, list)

1586 True

1587 """

1588 return [membership.team for membership in self.team_memberships if membership.is_active]

1589

1590 def get_personal_team(self) -> Optional["EmailTeam"]:

1591 """Get the user's personal team.

1592

1593 Returns:

1594 EmailTeam: The user's personal team or None if not found

1595

1596 Examples:

1597 >>> user = EmailUser(email="user@example.com")

1598 >>> personal_team = user.get_personal_team()

1599 """

1600 for team in self.created_teams:

1601 if team.is_personal and team.is_active:

1602 return team

1603 return None

1604

1605 def is_team_member(self, team_id: str) -> bool:

1606 """Check if user is a member of the specified team.

1607

1608 Args:

1609 team_id: ID of the team to check

1610

1611 Returns:

1612 bool: True if user is a member, False otherwise

1613

1614 Examples:

1615 >>> user = EmailUser(email="user@example.com")

1616 >>> user.is_team_member("team-123")

1617 False

1618 """

1619 return any(membership.team_id == team_id and membership.is_active for membership in self.team_memberships)

1620

1621 def get_team_role(self, team_id: str) -> Optional[str]:

1622 """Get user's role in a specific team.

1623

1624 Args:

1625 team_id: ID of the team to check

1626

1627 Returns:

1628 str: User's role or None if not a member

1629

1630 Examples:

1631 >>> user = EmailUser(email="user@example.com")

1632 >>> role = user.get_team_role("team-123")

1633 """

1634 for membership in self.team_memberships:

1635 if membership.team_id == team_id and membership.is_active:

1636 return membership.role

1637 return None

1638

1639

1640class EmailAuthEvent(Base):

1641 """Authentication event logging for email users.

1642

1643 This model tracks all authentication attempts for auditing,

1644 security monitoring, and compliance purposes.

1645

1646 Attributes:

1647 id (int): Primary key

1648 timestamp (datetime): Event timestamp

1649 user_email (str): Email of the user

1650 event_type (str): Type of authentication event

1651 success (bool): Whether the authentication was successful

1652 ip_address (str): Client IP address

1653 user_agent (str): Client user agent string

1654 failure_reason (str): Reason for authentication failure

1655 details (dict): Additional event details as JSON

1656

1657 Examples:

1658 >>> event = EmailAuthEvent(

1659 ... user_email="alice@example.com",

1660 ... event_type="login",

1661 ... success=True,

1662 ... ip_address="192.168.1.100"

1663 ... )

1664 >>> event.event_type

1665 'login'

1666 >>> event.success

1667 True

1668 """

1669

1670 __tablename__ = "email_auth_events"

1671

1672 # Primary key

1673 id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)

1674

1675 # Event details

1676 timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, nullable=False)

1677 user_email: Mapped[Optional[str]] = mapped_column(String(255), nullable=True, index=True)

1678 event_type: Mapped[str] = mapped_column(String(50), nullable=False)

1679 success: Mapped[bool] = mapped_column(Boolean, nullable=False)

1680

1681 # Client information

1682 ip_address: Mapped[Optional[str]] = mapped_column(String(45), nullable=True) # IPv6 compatible

1683 user_agent: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

1684

1685 # Failure information

1686 failure_reason: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

1687 details: Mapped[Optional[str]] = mapped_column(Text, nullable=True) # JSON string

1688

1689 def __repr__(self) -> str:

1690 """String representation of the auth event.

1691

1692 Returns:

1693 str: String representation of EmailAuthEvent instance

1694 """

1695 return f"<EmailAuthEvent(user_email='{self.user_email}', event_type='{self.event_type}', success={self.success})>"

1696

1697 @classmethod

1698 def create_login_attempt(

1699 cls,

1700 user_email: str,

1701 success: bool,

1702 ip_address: Optional[str] = None,

1703 user_agent: Optional[str] = None,

1704 failure_reason: Optional[str] = None,

1705 ) -> "EmailAuthEvent":

1706 """Create a login attempt event.

1707

1708 Args:

1709 user_email: Email address of the user

1710 success: Whether the login was successful

1711 ip_address: Client IP address

1712 user_agent: Client user agent

1713 failure_reason: Reason for failure (if applicable)

1714

1715 Returns:

1716 EmailAuthEvent: New authentication event

1717

1718 Examples:

1719 >>> event = EmailAuthEvent.create_login_attempt(

1720 ... user_email="user@example.com",

1721 ... success=True,

1722 ... ip_address="192.168.1.1"

1723 ... )

1724 >>> event.event_type

1725 'login'

1726 >>> event.success

1727 True

1728 """

1729 return cls(user_email=user_email, event_type="login", success=success, ip_address=ip_address, user_agent=user_agent, failure_reason=failure_reason)

1730

1731 @classmethod

1732 def create_registration_event(

1733 cls,

1734 user_email: str,

1735 success: bool,

1736 ip_address: Optional[str] = None,

1737 user_agent: Optional[str] = None,

1738 failure_reason: Optional[str] = None,

1739 ) -> "EmailAuthEvent":

1740 """Create a registration event.

1741

1742 Args:

1743 user_email: Email address of the user

1744 success: Whether the registration was successful

1745 ip_address: Client IP address

1746 user_agent: Client user agent

1747 failure_reason: Reason for failure (if applicable)

1748

1749 Returns:

1750 EmailAuthEvent: New authentication event

1751 """

1752 return cls(user_email=user_email, event_type="registration", success=success, ip_address=ip_address, user_agent=user_agent, failure_reason=failure_reason)

1753

1754 @classmethod

1755 def create_password_change_event(

1756 cls,

1757 user_email: str,

1758 success: bool,

1759 ip_address: Optional[str] = None,

1760 user_agent: Optional[str] = None,

1761 ) -> "EmailAuthEvent":

1762 """Create a password change event.

1763

1764 Args:

1765 user_email: Email address of the user

1766 success: Whether the password change was successful

1767 ip_address: Client IP address

1768 user_agent: Client user agent

1769

1770 Returns:

1771 EmailAuthEvent: New authentication event

1772 """

1773 return cls(user_email=user_email, event_type="password_change", success=success, ip_address=ip_address, user_agent=user_agent)

1774

1775

1776class PasswordResetToken(Base):

1777 """One-time password reset token record.

1778

1779 Stores only a SHA-256 hash of the user-facing token. Tokens are one-time use

1780 and expire after a configured duration.

1781 """

1782

1783 __tablename__ = "password_reset_tokens"

1784

1785 id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()))

1786 user_email: Mapped[str] = mapped_column(String(255), ForeignKey("email_users.email", ondelete="CASCADE"), nullable=False, index=True)

1787 token_hash: Mapped[str] = mapped_column(String(64), nullable=False, unique=True, index=True)

1788 expires_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)

1789 used_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), nullable=True)

1790 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, nullable=False)

1791 ip_address: Mapped[Optional[str]] = mapped_column(String(45), nullable=True)

1792 user_agent: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

1793

1794 user: Mapped["EmailUser"] = relationship("EmailUser")

1795

1796 __table_args__ = (Index("ix_password_reset_tokens_expires_at", "expires_at"),)

1797

1798 def is_expired(self) -> bool:

1799 """Return whether the reset token has expired.

1800

1801 Returns:

1802 bool: True when `expires_at` is in the past.

1803 """

1804 if self.expires_at is None:

1805 return False

1806 expires_at = self.expires_at

1807 if expires_at.tzinfo is None:

1808 expires_at = expires_at.replace(tzinfo=timezone.utc)

1809

1810 return expires_at <= utc_now()

1811

1812 def is_used(self) -> bool:

1813 """Return whether the reset token was already consumed.

1814

1815 Returns:

1816 bool: True when `used_at` is set.

1817 """

1818 return self.used_at is not None

1819

1820

1821class EmailTeam(Base):

1822 """Email-based team model for multi-team collaboration.

1823

1824 This model represents teams that users can belong to, with automatic

1825 personal team creation and role-based access control.

1826

1827 Attributes:

1828 id (str): Primary key UUID

1829 name (str): Team display name

1830 slug (str): URL-friendly team identifier

1831 description (str): Team description

1832 created_by (str): Email of the user who created the team

1833 is_personal (bool): Whether this is a personal team

1834 visibility (str): Team visibility (private, public)

1835 max_members (int): Maximum number of team members allowed

1836 created_at (datetime): Team creation timestamp

1837 updated_at (datetime): Last update timestamp

1838 is_active (bool): Whether the team is active

1839

1840 Examples:

1841 >>> team = EmailTeam(

1842 ... name="Engineering Team",

1843 ... slug="engineering-team",

1844 ... created_by="admin@example.com",

1845 ... is_personal=False

1846 ... )

1847 >>> team.name

1848 'Engineering Team'

1849 >>> team.is_personal

1850 False

1851 """

1852

1853 __tablename__ = "email_teams"

1854

1855 # Primary key

1856 id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: uuid.uuid4().hex)

1857

1858 # Basic team information

1859 name: Mapped[str] = mapped_column(String(255), nullable=False)

1860 slug: Mapped[str] = mapped_column(String(255), unique=True, nullable=False)

1861 description: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

1862 created_by: Mapped[str] = mapped_column(String(255), ForeignKey("email_users.email"), nullable=False)

1863

1864 # Team settings

1865 is_personal: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)

1866 visibility: Mapped[str] = mapped_column(String(20), default="public", nullable=False)

1867 max_members: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)

1868

1869 # Timestamps

1870 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, nullable=False)

1871 updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, onupdate=utc_now, nullable=False)

1872 is_active: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)

1873

1874 # Relationships

1875 members: Mapped[List["EmailTeamMember"]] = relationship("EmailTeamMember", back_populates="team", cascade="all, delete-orphan")

1876 invitations: Mapped[List["EmailTeamInvitation"]] = relationship("EmailTeamInvitation", back_populates="team", cascade="all, delete-orphan")

1877 api_tokens: Mapped[List["EmailApiToken"]] = relationship("EmailApiToken", back_populates="team", cascade="all, delete-orphan")

1878 creator: Mapped["EmailUser"] = relationship("EmailUser", foreign_keys=[created_by])

1879

1880 # Index for search and pagination performance

1881 __table_args__ = (Index("ix_email_teams_name_id", "name", "id"),)

1882

1883 def __repr__(self) -> str:

1884 """String representation of the team.

1885

1886 Returns:

1887 str: String representation of EmailTeam instance

1888 """

1889 return f"<EmailTeam(id='{self.id}', name='{self.name}', is_personal={self.is_personal})>"

1890

1891 def get_member_count(self) -> int:

1892 """Get the current number of team members.

1893

1894 Uses direct SQL COUNT to avoid loading all members into memory.

1895

1896 Returns:

1897 int: Number of active team members

1898

1899 Examples:

1900 >>> team = EmailTeam(name="Test Team", slug="test-team", created_by="admin@example.com")

1901 >>> team.get_member_count()

1902 0

1903 """

1904 # Third-Party

1905 from sqlalchemy.orm import object_session # pylint: disable=import-outside-toplevel

1906

1907 session = object_session(self)

1908 if session is None:

1909 # Fallback for detached objects (e.g., in doctests)

1910 return len([m for m in self.members if m.is_active])

1911

1912 count = session.query(func.count(EmailTeamMember.id)).filter(EmailTeamMember.team_id == self.id, EmailTeamMember.is_active.is_(True)).scalar() # pylint: disable=not-callable

1913 return count or 0

1914

1915 def is_member(self, user_email: str) -> bool:

1916 """Check if a user is a member of this team.

1917

1918 Uses direct SQL EXISTS to avoid loading all members into memory.

1919

1920 Args:

1921 user_email: Email address to check

1922

1923 Returns:

1924 bool: True if user is an active member, False otherwise

1925

1926 Examples:

1927 >>> team = EmailTeam(name="Test Team", slug="test-team", created_by="admin@example.com")

1928 >>> team.is_member("admin@example.com")

1929 False

1930 """

1931 # Third-Party

1932 from sqlalchemy.orm import object_session # pylint: disable=import-outside-toplevel

1933

1934 session = object_session(self)

1935 if session is None:

1936 # Fallback for detached objects (e.g., in doctests)

1937 return any(m.user_email == user_email and m.is_active for m in self.members)

1938

1939 exists = session.query(EmailTeamMember.id).filter(EmailTeamMember.team_id == self.id, EmailTeamMember.user_email == user_email, EmailTeamMember.is_active.is_(True)).first()

1940 return exists is not None

1941

1942 def get_member_role(self, user_email: str) -> Optional[str]:

1943 """Get the role of a user in this team.

1944

1945 Uses direct SQL query to avoid loading all members into memory.

1946

1947 Args:

1948 user_email: Email address to check

1949

1950 Returns:

1951 str: User's role or None if not a member

1952

1953 Examples:

1954 >>> team = EmailTeam(name="Test Team", slug="test-team", created_by="admin@example.com")

1955 >>> team.get_member_role("admin@example.com")

1956 """

1957 # Third-Party

1958 from sqlalchemy.orm import object_session # pylint: disable=import-outside-toplevel

1959

1960 session = object_session(self)

1961 if session is None:

1962 # Fallback for detached objects (e.g., in doctests)

1963 for member in self.members:

1964 if member.user_email == user_email and member.is_active:

1965 return member.role

1966 return None

1967

1968 member = session.query(EmailTeamMember.role).filter(EmailTeamMember.team_id == self.id, EmailTeamMember.user_email == user_email, EmailTeamMember.is_active.is_(True)).first()

1969 return member[0] if member else None

1970

1971

1972class EmailTeamMember(Base):

1973 """Team membership model linking users to teams with roles.

1974

1975 This model represents the many-to-many relationship between users and teams

1976 with additional role information and audit trails.

1977

1978 Attributes:

1979 id (str): Primary key UUID

1980 team_id (str): Foreign key to email_teams

1981 user_email (str): Foreign key to email_users

1982 role (str): Member role (owner, member)

1983 joined_at (datetime): When the user joined the team

1984 invited_by (str): Email of the user who invited this member

1985 is_active (bool): Whether the membership is active

1986 grant_source (str): Origin of the grant (e.g., 'sso', 'manual', 'bootstrap', 'auto')

1987

1988 Examples:

1989 >>> member = EmailTeamMember(

1990 ... team_id="team-123",

1991 ... user_email="user@example.com",

1992 ... role="member",

1993 ... invited_by="admin@example.com"

1994 ... )

1995 >>> member.role

1996 'member'

1997 """

1998

1999 __tablename__ = "email_team_members"

2000

2001 # Primary key

2002 id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: uuid.uuid4().hex)

2003

2004 # Foreign keys

2005 team_id: Mapped[str] = mapped_column(String(36), ForeignKey("email_teams.id", ondelete="CASCADE"), nullable=False)

2006 user_email: Mapped[str] = mapped_column(String(255), ForeignKey("email_users.email"), nullable=False)

2007

2008 # Membership details

2009 role: Mapped[str] = mapped_column(String(50), default="member", nullable=False)

2010 joined_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, nullable=False)

2011 invited_by: Mapped[Optional[str]] = mapped_column(String(255), ForeignKey("email_users.email"), nullable=True)

2012 is_active: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)

2013 grant_source: Mapped[Optional[str]] = mapped_column(String(50), nullable=True, default=None)

2014

2015 # Relationships

2016 team: Mapped["EmailTeam"] = relationship("EmailTeam", back_populates="members")

2017 user: Mapped["EmailUser"] = relationship("EmailUser", foreign_keys=[user_email])

2018 inviter: Mapped[Optional["EmailUser"]] = relationship("EmailUser", foreign_keys=[invited_by])

2019

2020 # Unique constraint to prevent duplicate memberships

2021 __table_args__ = (UniqueConstraint("team_id", "user_email", name="uq_team_member"),)

2022

2023 def __repr__(self) -> str:

2024 """String representation of the team member.

2025

2026 Returns:

2027 str: String representation of EmailTeamMember instance

2028 """

2029 return f"<EmailTeamMember(team_id='{self.team_id}', user_email='{self.user_email}', role='{self.role}')>"

2030

2031

2032# Team member history model

2033class EmailTeamMemberHistory(Base):

2034 """

2035 History of team member actions (add, remove, reactivate, role change).

2036

2037 This model records every membership-related event for audit and compliance.

2038 Each record tracks the team, user, role, action type, actor, and timestamp.

2039

2040 Attributes:

2041 id (str): Primary key UUID

2042 team_id (str): Foreign key to email_teams

2043 user_email (str): Foreign key to email_users

2044 role (str): Role at the time of action

2045 action (str): Action type ("added", "removed", "reactivated", "role_changed")

2046 action_by (str): Email of the user who performed the action

2047 action_timestamp (datetime): When the action occurred

2048

2049 Examples:

2050 >>> from mcpgateway.db import EmailTeamMemberHistory, utc_now

2051 >>> history = EmailTeamMemberHistory(

2052 ... team_id="team-123",

2053 ... user_email="user@example.com",

2054 ... role="member",

2055 ... action="added",

2056 ... action_by="admin@example.com",

2057 ... action_timestamp=utc_now()

2058 ... )

2059 >>> history.action

2060 'added'

2061 >>> history.role

2062 'member'

2063 >>> isinstance(history.action_timestamp, type(utc_now()))

2064 True

2065 """

2066

2067 __tablename__ = "email_team_member_history"

2068

2069 id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: uuid.uuid4().hex)

2070 team_member_id: Mapped[str] = mapped_column(String(36), ForeignKey("email_team_members.id", ondelete="CASCADE"), nullable=False)

2071 team_id: Mapped[str] = mapped_column(String(36), ForeignKey("email_teams.id"), nullable=False)

2072 user_email: Mapped[str] = mapped_column(String(255), ForeignKey("email_users.email"), nullable=False)

2073 role: Mapped[str] = mapped_column(String(50), default="member", nullable=False)

2074 action: Mapped[str] = mapped_column(String(50), nullable=False) # e.g. "added", "removed", "reactivated", "role_changed"

2075 action_by: Mapped[Optional[str]] = mapped_column(String(255), ForeignKey("email_users.email"), nullable=True)

2076 action_timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, nullable=False)

2077

2078 team_member: Mapped["EmailTeamMember"] = relationship("EmailTeamMember")

2079 team: Mapped["EmailTeam"] = relationship("EmailTeam")

2080 user: Mapped["EmailUser"] = relationship("EmailUser", foreign_keys=[user_email])

2081 actor: Mapped[Optional["EmailUser"]] = relationship("EmailUser", foreign_keys=[action_by])

2082

2083 def __repr__(self) -> str:

2084 """

2085 Return a string representation of the EmailTeamMemberHistory instance.

2086

2087 Returns:

2088 str: A string summarizing the team member history record.

2089

2090 Examples:

2091 >>> from mcpgateway.db import EmailTeamMemberHistory, utc_now

2092 >>> history = EmailTeamMemberHistory(

2093 ... team_member_id="tm-123",

2094 ... team_id="team-123",

2095 ... user_email="user@example.com",

2096 ... role="member",

2097 ... action="added",

2098 ... action_by="admin@example.com",

2099 ... action_timestamp=utc_now()

2100 ... )

2101 >>> isinstance(repr(history), str)

2102 True

2103 """

2104 return f"<EmailTeamMemberHistory(team_member_id='{self.team_member_id}', team_id='{self.team_id}', user_email='{self.user_email}', role='{self.role}', action='{self.action}', action_by='{self.action_by}', action_timestamp='{self.action_timestamp}')>"

2105

2106

2107class EmailTeamInvitation(Base):

2108 """Team invitation model for managing team member invitations.

2109

2110 This model tracks invitations sent to users to join teams, including

2111 expiration dates and invitation tokens.

2112

2113 Attributes:

2114 id (str): Primary key UUID

2115 team_id (str): Foreign key to email_teams

2116 email (str): Email address of the invited user

2117 role (str): Role the user will have when they accept

2118 invited_by (str): Email of the user who sent the invitation

2119 invited_at (datetime): When the invitation was sent

2120 expires_at (datetime): When the invitation expires

2121 token (str): Unique invitation token

2122 is_active (bool): Whether the invitation is still active

2123

2124 Examples:

2125 >>> invitation = EmailTeamInvitation(

2126 ... team_id="team-123",

2127 ... email="newuser@example.com",

2128 ... role="member",

2129 ... invited_by="admin@example.com"

2130 ... )

2131 >>> invitation.role

2132 'member'

2133 """

2134

2135 __tablename__ = "email_team_invitations"

2136

2137 # Primary key

2138 id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: uuid.uuid4().hex)

2139

2140 # Foreign keys

2141 team_id: Mapped[str] = mapped_column(String(36), ForeignKey("email_teams.id", ondelete="CASCADE"), nullable=False)

2142

2143 # Invitation details

2144 email: Mapped[str] = mapped_column(String(255), nullable=False)

2145 role: Mapped[str] = mapped_column(String(50), default="member", nullable=False)

2146 invited_by: Mapped[str] = mapped_column(String(255), ForeignKey("email_users.email"), nullable=False)

2147

2148 # Timing

2149 invited_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, nullable=False)

2150 expires_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)

2151

2152 # Security

2153 token: Mapped[str] = mapped_column(String(500), unique=True, nullable=False)

2154 is_active: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)

2155

2156 # Relationships

2157 team: Mapped["EmailTeam"] = relationship("EmailTeam", back_populates="invitations")

2158 inviter: Mapped["EmailUser"] = relationship("EmailUser", foreign_keys=[invited_by])

2159

2160 def __repr__(self) -> str:

2161 """String representation of the team invitation.

2162

2163 Returns:

2164 str: String representation of EmailTeamInvitation instance

2165 """

2166 return f"<EmailTeamInvitation(team_id='{self.team_id}', email='{self.email}', role='{self.role}')>"

2167

2168 def is_expired(self) -> bool:

2169 """Check if the invitation has expired.

2170

2171 Returns:

2172 bool: True if the invitation has expired, False otherwise

2173

2174 Examples:

2175 >>> from datetime import timedelta

2176 >>> invitation = EmailTeamInvitation(

2177 ... team_id="team-123",

2178 ... email="user@example.com",

2179 ... role="member",

2180 ... invited_by="admin@example.com",

2181 ... expires_at=utc_now() + timedelta(days=7)

2182 ... )

2183 >>> invitation.is_expired()

2184 False

2185 """

2186 now = utc_now()

2187 expires_at = self.expires_at

2188

2189 # Handle timezone awareness mismatch

2190 if now.tzinfo is not None and expires_at.tzinfo is None:

2191 expires_at = expires_at.replace(tzinfo=timezone.utc)

2192 elif now.tzinfo is None and expires_at.tzinfo is not None:

2193 now = now.replace(tzinfo=timezone.utc)

2194

2195 return now > expires_at

2196

2197 def is_valid(self) -> bool:

2198 """Check if the invitation is valid (active and not expired).

2199

2200 Returns:

2201 bool: True if the invitation is valid, False otherwise

2202

2203 Examples:

2204 >>> from datetime import timedelta

2205 >>> invitation = EmailTeamInvitation(

2206 ... team_id="team-123",

2207 ... email="user@example.com",

2208 ... role="member",

2209 ... invited_by="admin@example.com",

2210 ... expires_at=utc_now() + timedelta(days=7),

2211 ... is_active=True

2212 ... )

2213 >>> invitation.is_valid()

2214 True

2215 """

2216 return self.is_active and not self.is_expired()

2217

2218

2219class EmailTeamJoinRequest(Base):

2220 """Team join request model for managing public team join requests.

2221

2222 This model tracks user requests to join public teams, including

2223 approval workflow and expiration dates.

2224

2225 Attributes:

2226 id (str): Primary key UUID

2227 team_id (str): Foreign key to email_teams

2228 user_email (str): Email of the user requesting to join

2229 message (str): Optional message from the user

2230 status (str): Request status (pending, approved, rejected, expired)

2231 requested_at (datetime): When the request was made

2232 expires_at (datetime): When the request expires

2233 reviewed_at (datetime): When the request was reviewed

2234 reviewed_by (str): Email of user who reviewed the request

2235 notes (str): Optional admin notes

2236 """

2237

2238 __tablename__ = "email_team_join_requests"

2239

2240 # Primary key

2241 id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: uuid.uuid4().hex)

2242

2243 # Foreign keys

2244 team_id: Mapped[str] = mapped_column(String(36), ForeignKey("email_teams.id", ondelete="CASCADE"), nullable=False)

2245 user_email: Mapped[str] = mapped_column(String(255), ForeignKey("email_users.email"), nullable=False)

2246

2247 # Request details

2248 message: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

2249 status: Mapped[str] = mapped_column(String(20), default="pending", nullable=False)

2250

2251 # Timing

2252 requested_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, nullable=False)

2253 expires_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)

2254 reviewed_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), nullable=True)

2255 reviewed_by: Mapped[Optional[str]] = mapped_column(String(255), ForeignKey("email_users.email"), nullable=True)

2256 notes: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

2257

2258 # Relationships

2259 team: Mapped["EmailTeam"] = relationship("EmailTeam")

2260 user: Mapped["EmailUser"] = relationship("EmailUser", foreign_keys=[user_email])

2261 reviewer: Mapped[Optional["EmailUser"]] = relationship("EmailUser", foreign_keys=[reviewed_by])

2262

2263 # Unique constraint to prevent duplicate requests

2264 __table_args__ = (UniqueConstraint("team_id", "user_email", name="uq_team_join_request"),)

2265

2266 def __repr__(self) -> str:

2267 """String representation of the team join request.

2268

2269 Returns:

2270 str: String representation of the team join request.

2271 """

2272 return f"<EmailTeamJoinRequest(team_id='{self.team_id}', user_email='{self.user_email}', status='{self.status}')>"

2273

2274 def is_expired(self) -> bool:

2275 """Check if the join request has expired.

2276

2277 Returns:

2278 bool: True if the request has expired, False otherwise.

2279 """

2280 now = utc_now()

2281 expires_at = self.expires_at

2282

2283 # Handle timezone awareness mismatch

2284 if now.tzinfo is not None and expires_at.tzinfo is None:

2285 expires_at = expires_at.replace(tzinfo=timezone.utc)

2286 elif now.tzinfo is None and expires_at.tzinfo is not None:

2287 now = now.replace(tzinfo=timezone.utc)

2288

2289 return now > expires_at

2290

2291 def is_pending(self) -> bool:

2292 """Check if the join request is still pending.

2293

2294 Returns:

2295 bool: True if the request is pending and not expired, False otherwise.

2296 """

2297 return self.status == "pending" and not self.is_expired()

2298

2299

2300class PendingUserApproval(Base):

2301 """Model for pending SSO user registrations awaiting admin approval.

2302

2303 This model stores information about users who have authenticated via SSO

2304 but require admin approval before their account is fully activated.

2305

2306 Attributes:

2307 id (str): Primary key

2308 email (str): Email address of the pending user

2309 full_name (str): Full name from SSO provider

2310 auth_provider (str): SSO provider (github, google, etc.)

2311 sso_metadata (dict): Additional metadata from SSO provider

2312 requested_at (datetime): When the approval was requested

2313 expires_at (datetime): When the approval request expires

2314 approved_by (str): Email of admin who approved (if approved)

2315 approved_at (datetime): When the approval was granted

2316 status (str): Current status (pending, approved, rejected, expired)

2317 rejection_reason (str): Reason for rejection (if applicable)

2318 admin_notes (str): Notes from admin review

2319

2320 Examples:

2321 >>> from datetime import timedelta

2322 >>> approval = PendingUserApproval(

2323 ... email="newuser@example.com",

2324 ... full_name="New User",

2325 ... auth_provider="github",

2326 ... expires_at=utc_now() + timedelta(days=30),

2327 ... status="pending"

2328 ... )

2329 >>> approval.status

2330 'pending'

2331 """

2332

2333 __tablename__ = "pending_user_approvals"

2334

2335 # Primary key

2336 id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()))

2337

2338 # User details

2339 email: Mapped[str] = mapped_column(String(255), nullable=False, unique=True)

2340 full_name: Mapped[str] = mapped_column(String(255), nullable=False)

2341 auth_provider: Mapped[str] = mapped_column(String(50), nullable=False)

2342 sso_metadata: Mapped[Optional[Dict]] = mapped_column(JSON, nullable=True)

2343

2344 # Request details

2345 requested_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, nullable=False)

2346 expires_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)

2347

2348 # Approval details

2349 approved_by: Mapped[Optional[str]] = mapped_column(String(255), ForeignKey("email_users.email"), nullable=True)

2350 approved_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), nullable=True)

2351 status: Mapped[str] = mapped_column(String(20), default="pending", nullable=False) # pending, approved, rejected, expired

2352 rejection_reason: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

2353 admin_notes: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

2354

2355 # Relationships

2356 approver: Mapped[Optional["EmailUser"]] = relationship("EmailUser", foreign_keys=[approved_by])

2357

2358 def __repr__(self) -> str:

2359 """String representation of the pending approval.

2360

2361 Returns:

2362 str: String representation of PendingUserApproval instance

2363 """

2364 return f"<PendingUserApproval(email='{self.email}', status='{self.status}', provider='{self.auth_provider}')>"

2365

2366 def is_expired(self) -> bool:

2367 """Check if the approval request has expired.

2368

2369 Returns:

2370 bool: True if the approval request has expired

2371 """

2372 now = utc_now()

2373 expires_at = self.expires_at

2374

2375 # Handle timezone awareness mismatch

2376 if now.tzinfo is not None and expires_at.tzinfo is None:

2377 expires_at = expires_at.replace(tzinfo=timezone.utc)

2378 elif now.tzinfo is None and expires_at.tzinfo is not None:

2379 now = now.replace(tzinfo=timezone.utc)

2380

2381 return now > expires_at

2382

2383 def approve(self, admin_email: str, notes: Optional[str] = None) -> None:

2384 """Approve the user registration.

2385

2386 Args:

2387 admin_email: Email of the admin approving the request

2388 notes: Optional admin notes

2389 """

2390 self.status = "approved"

2391 self.approved_by = admin_email

2392 self.approved_at = utc_now()

2393 self.admin_notes = notes

2394

2395 def reject(self, admin_email: str, reason: str, notes: Optional[str] = None) -> None:

2396 """Reject the user registration.

2397

2398 Args:

2399 admin_email: Email of the admin rejecting the request

2400 reason: Reason for rejection

2401 notes: Optional admin notes

2402 """

2403 self.status = "rejected"

2404 self.approved_by = admin_email

2405 self.approved_at = utc_now()

2406 self.rejection_reason = reason

2407 self.admin_notes = notes

2408

2409

2410# Association table for servers and tools

2411server_tool_association = Table(

2412 "server_tool_association",

2413 Base.metadata,

2414 Column("server_id", String(36), ForeignKey("servers.id"), primary_key=True),

2415 Column("tool_id", String(36), ForeignKey("tools.id"), primary_key=True),

2416)

2417

2418# Association table for servers and resources

2419server_resource_association = Table(

2420 "server_resource_association",

2421 Base.metadata,

2422 Column("server_id", String(36), ForeignKey("servers.id"), primary_key=True),

2423 Column("resource_id", String(36), ForeignKey("resources.id"), primary_key=True),

2424)

2425

2426# Association table for servers and prompts

2427server_prompt_association = Table(

2428 "server_prompt_association",

2429 Base.metadata,

2430 Column("server_id", String(36), ForeignKey("servers.id"), primary_key=True),

2431 Column("prompt_id", String(36), ForeignKey("prompts.id"), primary_key=True),

2432)

2433

2434# Association table for servers and A2A agents

2435server_a2a_association = Table(

2436 "server_a2a_association",

2437 Base.metadata,

2438 Column("server_id", String(36), ForeignKey("servers.id"), primary_key=True),

2439 Column("a2a_agent_id", String(36), ForeignKey("a2a_agents.id"), primary_key=True),

2440)

2441

2442

2443class GlobalConfig(Base):

2444 """Global configuration settings.

2445

2446 Attributes:

2447 id (int): Primary key

2448 passthrough_headers (List[str]): List of headers allowed to be passed through globally

2449 """

2450

2451 __tablename__ = "global_config"

2452

2453 id: Mapped[int] = mapped_column(Integer, primary_key=True)

2454 passthrough_headers: Mapped[Optional[List[str]]] = mapped_column(JSON, nullable=True) # Store list of strings as JSON array

2455

2456

2457class ToolMetric(Base):

2458 """

2459 ORM model for recording individual metrics for tool executions.

2460

2461 Each record in this table corresponds to a single tool invocation and records:

2462 - timestamp (datetime): When the invocation occurred.

2463 - response_time (float): The execution time in seconds.

2464 - is_success (bool): True if the execution succeeded, False otherwise.

2465 - error_message (Optional[str]): Error message if the execution failed.

2466

2467 Aggregated metrics (such as total executions, successful/failed counts, failure rate,

2468 minimum, maximum, and average response times, and last execution time) should be computed

2469 on the fly using SQL aggregate functions over the rows in this table.

2470 """

2471

2472 __tablename__ = "tool_metrics"

2473

2474 id: Mapped[int] = mapped_column(primary_key=True)

2475 tool_id: Mapped[str] = mapped_column(String(36), ForeignKey("tools.id"), nullable=False, index=True)

2476 timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, index=True)

2477 response_time: Mapped[float] = mapped_column(Float, nullable=False)

2478 is_success: Mapped[bool] = mapped_column(Boolean, nullable=False)

2479 error_message: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

2480

2481 # Relationship back to the Tool model.

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

2483

2484

2485class ResourceMetric(Base):

2486 """

2487 ORM model for recording metrics for resource invocations.

2488

2489 Attributes:

2490 id (int): Primary key.

2491 resource_id (str): Foreign key linking to the resource.

2492 timestamp (datetime): The time when the invocation occurred.

2493 response_time (float): The response time in seconds.

2494 is_success (bool): True if the invocation succeeded, False otherwise.

2495 error_message (Optional[str]): Error message if the invocation failed.

2496 """

2497

2498 __tablename__ = "resource_metrics"

2499

2500 id: Mapped[int] = mapped_column(primary_key=True)

2501 resource_id: Mapped[str] = mapped_column(String(36), ForeignKey("resources.id"), nullable=False, index=True)

2502 timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, index=True)

2503 response_time: Mapped[float] = mapped_column(Float, nullable=False)

2504 is_success: Mapped[bool] = mapped_column(Boolean, nullable=False)

2505 error_message: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

2506

2507 # Relationship back to the Resource model.

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

2509

2510

2511class ServerMetric(Base):

2512 """

2513 ORM model for recording metrics for server invocations.

2514

2515 Attributes:

2516 id (int): Primary key.

2517 server_id (str): Foreign key linking to the server.

2518 timestamp (datetime): The time when the invocation occurred.

2519 response_time (float): The response time in seconds.

2520 is_success (bool): True if the invocation succeeded, False otherwise.

2521 error_message (Optional[str]): Error message if the invocation failed.

2522 """

2523

2524 __tablename__ = "server_metrics"

2525

2526 id: Mapped[int] = mapped_column(primary_key=True)

2527 server_id: Mapped[str] = mapped_column(String(36), ForeignKey("servers.id"), nullable=False, index=True)

2528 timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, index=True)

2529 response_time: Mapped[float] = mapped_column(Float, nullable=False)

2530 is_success: Mapped[bool] = mapped_column(Boolean, nullable=False)

2531 error_message: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

2532

2533 # Relationship back to the Server model.

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

2535

2536

2537class PromptMetric(Base):

2538 """

2539 ORM model for recording metrics for prompt invocations.

2540

2541 Attributes:

2542 id (int): Primary key.

2543 prompt_id (str): Foreign key linking to the prompt.

2544 timestamp (datetime): The time when the invocation occurred.

2545 response_time (float): The response time in seconds.

2546 is_success (bool): True if the invocation succeeded, False otherwise.

2547 error_message (Optional[str]): Error message if the invocation failed.

2548 """

2549

2550 __tablename__ = "prompt_metrics"

2551

2552 id: Mapped[int] = mapped_column(primary_key=True)

2553 prompt_id: Mapped[str] = mapped_column(String(36), ForeignKey("prompts.id"), nullable=False, index=True)

2554 timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, index=True)

2555 response_time: Mapped[float] = mapped_column(Float, nullable=False)

2556 is_success: Mapped[bool] = mapped_column(Boolean, nullable=False)

2557 error_message: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

2558

2559 # Relationship back to the Prompt model.

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

2561

2562

2563class A2AAgentMetric(Base):

2564 """

2565 ORM model for recording metrics for A2A agent interactions.

2566

2567 Attributes:

2568 id (int): Primary key.

2569 a2a_agent_id (str): Foreign key linking to the A2A agent.

2570 timestamp (datetime): The time when the interaction occurred.

2571 response_time (float): The response time in seconds.

2572 is_success (bool): True if the interaction succeeded, False otherwise.

2573 error_message (Optional[str]): Error message if the interaction failed.

2574 interaction_type (str): Type of interaction (invoke, query, etc.).

2575 """

2576

2577 __tablename__ = "a2a_agent_metrics"

2578

2579 id: Mapped[int] = mapped_column(primary_key=True)

2580 a2a_agent_id: Mapped[str] = mapped_column(String(36), ForeignKey("a2a_agents.id"), nullable=False, index=True)

2581 timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, index=True)

2582 response_time: Mapped[float] = mapped_column(Float, nullable=False)

2583 is_success: Mapped[bool] = mapped_column(Boolean, nullable=False)

2584 error_message: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

2585 interaction_type: Mapped[str] = mapped_column(String(50), nullable=False, default="invoke")

2586

2587 # Relationship back to the A2AAgent model.

2588 a2a_agent: Mapped["A2AAgent"] = relationship("A2AAgent", back_populates="metrics")

2589

2590

2591# ===================================

2592# Metrics Hourly Rollup Tables

2593# These tables store pre-aggregated hourly summaries for efficient historical queries.

2594# Raw metrics can be cleaned up after rollup, reducing storage while preserving trends.

2595# ===================================

2596

2597

2598class ToolMetricsHourly(Base):

2599 """

2600 Hourly rollup of tool metrics for efficient historical trend analysis.

2601

2602 This table stores pre-aggregated metrics per tool per hour, enabling fast

2603 queries for dashboards and reports without scanning millions of raw metrics.

2604

2605 Attributes:

2606 id: Primary key.

2607 tool_id: Foreign key to the tool (nullable for deleted tools).

2608 tool_name: Tool name snapshot (preserved even if tool is deleted).

2609 hour_start: Start of the aggregation hour (UTC).

2610 total_count: Total invocations during this hour.

2611 success_count: Successful invocations.

2612 failure_count: Failed invocations.

2613 min_response_time: Minimum response time in seconds.

2614 max_response_time: Maximum response time in seconds.

2615 avg_response_time: Average response time in seconds.

2616 p50_response_time: 50th percentile (median) response time.

2617 p95_response_time: 95th percentile response time.

2618 p99_response_time: 99th percentile response time.

2619 created_at: When this rollup was created.

2620 """

2621

2622 __tablename__ = "tool_metrics_hourly"

2623 __table_args__ = (

2624 UniqueConstraint("tool_id", "hour_start", name="uq_tool_metrics_hourly_tool_hour"),

2625 Index("ix_tool_metrics_hourly_hour_start", "hour_start"),

2626 )

2627

2628 id: Mapped[int] = mapped_column(primary_key=True)

2629 tool_id: Mapped[Optional[str]] = mapped_column(String(36), ForeignKey("tools.id", ondelete="SET NULL"), nullable=True, index=True)

2630 tool_name: Mapped[str] = mapped_column(String(255), nullable=False)

2631 hour_start: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)

2632 total_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)

2633 success_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)

2634 failure_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)

2635 min_response_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2636 max_response_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2637 avg_response_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2638 p50_response_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2639 p95_response_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2640 p99_response_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2641 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now)

2642

2643

2644class ResourceMetricsHourly(Base):

2645 """Hourly rollup of resource metrics for efficient historical trend analysis."""

2646

2647 __tablename__ = "resource_metrics_hourly"

2648 __table_args__ = (

2649 UniqueConstraint("resource_id", "hour_start", name="uq_resource_metrics_hourly_resource_hour"),

2650 Index("ix_resource_metrics_hourly_hour_start", "hour_start"),

2651 )

2652

2653 id: Mapped[int] = mapped_column(primary_key=True)

2654 resource_id: Mapped[Optional[str]] = mapped_column(String(36), ForeignKey("resources.id", ondelete="SET NULL"), nullable=True, index=True)

2655 resource_name: Mapped[str] = mapped_column(String(255), nullable=False)

2656 hour_start: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)

2657 total_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)

2658 success_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)

2659 failure_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)

2660 min_response_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2661 max_response_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2662 avg_response_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2663 p50_response_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2664 p95_response_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2665 p99_response_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2666 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now)

2667

2668

2669class PromptMetricsHourly(Base):

2670 """Hourly rollup of prompt metrics for efficient historical trend analysis."""

2671

2672 __tablename__ = "prompt_metrics_hourly"

2673 __table_args__ = (

2674 UniqueConstraint("prompt_id", "hour_start", name="uq_prompt_metrics_hourly_prompt_hour"),

2675 Index("ix_prompt_metrics_hourly_hour_start", "hour_start"),

2676 )

2677

2678 id: Mapped[int] = mapped_column(primary_key=True)

2679 prompt_id: Mapped[Optional[str]] = mapped_column(String(36), ForeignKey("prompts.id", ondelete="SET NULL"), nullable=True, index=True)

2680 prompt_name: Mapped[str] = mapped_column(String(255), nullable=False)

2681 hour_start: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)

2682 total_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)

2683 success_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)

2684 failure_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)

2685 min_response_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2686 max_response_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2687 avg_response_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2688 p50_response_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2689 p95_response_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2690 p99_response_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2691 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now)

2692

2693

2694class ServerMetricsHourly(Base):

2695 """Hourly rollup of server metrics for efficient historical trend analysis."""

2696

2697 __tablename__ = "server_metrics_hourly"

2698 __table_args__ = (

2699 UniqueConstraint("server_id", "hour_start", name="uq_server_metrics_hourly_server_hour"),

2700 Index("ix_server_metrics_hourly_hour_start", "hour_start"),

2701 )

2702

2703 id: Mapped[int] = mapped_column(primary_key=True)

2704 server_id: Mapped[Optional[str]] = mapped_column(String(36), ForeignKey("servers.id", ondelete="SET NULL"), nullable=True, index=True)

2705 server_name: Mapped[str] = mapped_column(String(255), nullable=False)

2706 hour_start: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)

2707 total_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)

2708 success_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)

2709 failure_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)

2710 min_response_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2711 max_response_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2712 avg_response_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2713 p50_response_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2714 p95_response_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2715 p99_response_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2716 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now)

2717

2718

2719class A2AAgentMetricsHourly(Base):

2720 """Hourly rollup of A2A agent metrics for efficient historical trend analysis."""

2721

2722 __tablename__ = "a2a_agent_metrics_hourly"

2723 __table_args__ = (

2724 UniqueConstraint("a2a_agent_id", "hour_start", "interaction_type", name="uq_a2a_agent_metrics_hourly_agent_hour_type"),

2725 Index("ix_a2a_agent_metrics_hourly_hour_start", "hour_start"),

2726 )

2727

2728 id: Mapped[int] = mapped_column(primary_key=True)

2729 a2a_agent_id: Mapped[Optional[str]] = mapped_column(String(36), ForeignKey("a2a_agents.id", ondelete="SET NULL"), nullable=True, index=True)

2730 agent_name: Mapped[str] = mapped_column(String(255), nullable=False)

2731 hour_start: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)

2732 interaction_type: Mapped[str] = mapped_column(String(50), nullable=False, default="invoke")

2733 total_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)

2734 success_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)

2735 failure_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)

2736 min_response_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2737 max_response_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2738 avg_response_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2739 p50_response_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2740 p95_response_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2741 p99_response_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2742 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now)

2743

2744

2745# ===================================

2746# Observability Models (OpenTelemetry-style traces, spans, events)

2747# ===================================

2748

2749

2750class ObservabilityTrace(Base):

2751 """

2752 ORM model for observability traces (similar to OpenTelemetry traces).

2753

2754 A trace represents a complete request flow through the system. It contains

2755 one or more spans representing individual operations.

2756

2757 Attributes:

2758 trace_id (str): Unique trace identifier (UUID or OpenTelemetry trace ID format).

2759 name (str): Human-readable name for the trace (e.g., "POST /tools/invoke").

2760 start_time (datetime): When the trace started.

2761 end_time (datetime): When the trace ended (optional, set when completed).

2762 duration_ms (float): Total duration in milliseconds.

2763 status (str): Trace status (success, error, timeout).

2764 status_message (str): Optional status message or error description.

2765 http_method (str): HTTP method for the request (GET, POST, etc.).

2766 http_url (str): Full URL of the request.

2767 http_status_code (int): HTTP response status code.

2768 user_email (str): User who initiated the request (if authenticated).

2769 user_agent (str): Client user agent string.

2770 ip_address (str): Client IP address.

2771 attributes (dict): Additional trace attributes (JSON).

2772 resource_attributes (dict): Resource attributes (service name, version, etc.).

2773 created_at (datetime): Trace creation timestamp.

2774 """

2775

2776 __tablename__ = "observability_traces"

2777

2778 # Primary key

2779 trace_id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()))

2780

2781 # Trace metadata

2782 name: Mapped[str] = mapped_column(String(255), nullable=False)

2783 start_time: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False, index=True)

2784 end_time: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), nullable=True)

2785 duration_ms: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2786 status: Mapped[str] = mapped_column(String(20), nullable=False, default="unset") # unset, ok, error

2787 status_message: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

2788

2789 # HTTP request context

2790 http_method: Mapped[Optional[str]] = mapped_column(String(10), nullable=True)

2791 http_url: Mapped[Optional[str]] = mapped_column(String(767), nullable=True)

2792 http_status_code: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)

2793

2794 # User context

2795 user_email: Mapped[Optional[str]] = mapped_column(String(255), nullable=True, index=True)

2796 user_agent: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

2797 ip_address: Mapped[Optional[str]] = mapped_column(String(45), nullable=True)

2798

2799 # Attributes (flexible key-value storage)

2800 attributes: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True, default=dict)

2801 resource_attributes: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True, default=dict)

2802

2803 # Timestamps

2804 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, nullable=False)

2805

2806 # Relationships

2807 spans: Mapped[List["ObservabilitySpan"]] = relationship("ObservabilitySpan", back_populates="trace", cascade="all, delete-orphan")

2808

2809 # Indexes for performance

2810 __table_args__ = (

2811 Index("idx_observability_traces_start_time", "start_time"),

2812 Index("idx_observability_traces_user_email", "user_email"),

2813 Index("idx_observability_traces_status", "status"),

2814 Index("idx_observability_traces_http_status_code", "http_status_code"),

2815 )

2816

2817

2818class ObservabilitySpan(Base):

2819 """

2820 ORM model for observability spans (similar to OpenTelemetry spans).

2821

2822 A span represents a single operation within a trace. Spans can be nested

2823 to represent hierarchical operations.

2824

2825 Attributes:

2826 span_id (str): Unique span identifier.

2827 trace_id (str): Parent trace ID.

2828 parent_span_id (str): Parent span ID (for nested spans).

2829 name (str): Span name (e.g., "database_query", "tool_invocation").

2830 kind (str): Span kind (internal, server, client, producer, consumer).

2831 start_time (datetime): When the span started.

2832 end_time (datetime): When the span ended.

2833 duration_ms (float): Span duration in milliseconds.

2834 status (str): Span status (success, error).

2835 status_message (str): Optional status message.

2836 attributes (dict): Span attributes (JSON).

2837 resource_name (str): Name of the resource being operated on.

2838 resource_type (str): Type of resource (tool, resource, prompt, gateway, etc.).

2839 resource_id (str): ID of the specific resource.

2840 created_at (datetime): Span creation timestamp.

2841 """

2842

2843 __tablename__ = "observability_spans"

2844

2845 # Primary key

2846 span_id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()))

2847

2848 # Trace relationship

2849 trace_id: Mapped[str] = mapped_column(String(36), ForeignKey("observability_traces.trace_id", ondelete="CASCADE"), nullable=False, index=True)

2850 parent_span_id: Mapped[Optional[str]] = mapped_column(String(36), ForeignKey("observability_spans.span_id", ondelete="CASCADE"), nullable=True, index=True)

2851

2852 # Span metadata

2853 name: Mapped[str] = mapped_column(String(255), nullable=False)

2854 kind: Mapped[str] = mapped_column(String(20), nullable=False, default="internal") # internal, server, client, producer, consumer

2855 start_time: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False, index=True)

2856 end_time: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), nullable=True)

2857 duration_ms: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

2858 status: Mapped[str] = mapped_column(String(20), nullable=False, default="unset")

2859 status_message: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

2860

2861 # Attributes

2862 attributes: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True, default=dict)

2863

2864 # Resource context

2865 resource_name: Mapped[Optional[str]] = mapped_column(String(255), nullable=True, index=True)

2866 resource_type: Mapped[Optional[str]] = mapped_column(String(50), nullable=True, index=True) # tool, resource, prompt, gateway, a2a_agent

2867 resource_id: Mapped[Optional[str]] = mapped_column(String(36), nullable=True, index=True)

2868

2869 # Timestamps

2870 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, nullable=False)

2871

2872 # Relationships

2873 trace: Mapped["ObservabilityTrace"] = relationship("ObservabilityTrace", back_populates="spans")

2874 parent_span: Mapped[Optional["ObservabilitySpan"]] = relationship("ObservabilitySpan", remote_side=[span_id], backref="child_spans")

2875 events: Mapped[List["ObservabilityEvent"]] = relationship("ObservabilityEvent", back_populates="span", cascade="all, delete-orphan")

2876

2877 # Indexes for performance

2878 __table_args__ = (

2879 Index("idx_observability_spans_trace_id", "trace_id"),

2880 Index("idx_observability_spans_parent_span_id", "parent_span_id"),

2881 Index("idx_observability_spans_start_time", "start_time"),

2882 Index("idx_observability_spans_resource_type", "resource_type"),

2883 Index("idx_observability_spans_resource_name", "resource_name"),

2884 )

2885

2886

2887class ObservabilityEvent(Base):

2888 """

2889 ORM model for observability events (logs within spans).

2890

2891 Events represent discrete occurrences within a span, such as log messages,

2892 exceptions, or state changes.

2893

2894 Attributes:

2895 id (int): Auto-incrementing primary key.

2896 span_id (str): Parent span ID.

2897 name (str): Event name (e.g., "exception", "log", "checkpoint").

2898 timestamp (datetime): When the event occurred.

2899 attributes (dict): Event attributes (JSON).

2900 severity (str): Log severity level (debug, info, warning, error, critical).

2901 message (str): Event message.

2902 exception_type (str): Exception class name (if event is an exception).

2903 exception_message (str): Exception message.

2904 exception_stacktrace (str): Exception stacktrace.

2905 created_at (datetime): Event creation timestamp.

2906 """

2907

2908 __tablename__ = "observability_events"

2909

2910 # Primary key

2911 id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)

2912

2913 # Span relationship

2914 span_id: Mapped[str] = mapped_column(String(36), ForeignKey("observability_spans.span_id", ondelete="CASCADE"), nullable=False, index=True)

2915

2916 # Event metadata

2917 name: Mapped[str] = mapped_column(String(255), nullable=False)

2918 timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False, default=utc_now, index=True)

2919 attributes: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True, default=dict)

2920

2921 # Log fields

2922 severity: Mapped[Optional[str]] = mapped_column(String(20), nullable=True, index=True) # debug, info, warning, error, critical

2923 message: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

2924

2925 # Exception fields

2926 exception_type: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

2927 exception_message: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

2928 exception_stacktrace: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

2929

2930 # Timestamps

2931 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, nullable=False)

2932

2933 # Relationships

2934 span: Mapped["ObservabilitySpan"] = relationship("ObservabilitySpan", back_populates="events")

2935

2936 # Indexes for performance

2937 __table_args__ = (

2938 Index("idx_observability_events_span_id", "span_id"),

2939 Index("idx_observability_events_timestamp", "timestamp"),

2940 Index("idx_observability_events_severity", "severity"),

2941 )

2942

2943

2944class ObservabilityMetric(Base):

2945 """

2946 ORM model for observability metrics (time-series numerical data).

2947

2948 Metrics represent numerical measurements over time, such as request rates,

2949 error rates, latencies, and custom business metrics.

2950

2951 Attributes:

2952 id (int): Auto-incrementing primary key.

2953 name (str): Metric name (e.g., "http.request.duration", "tool.invocation.count").

2954 metric_type (str): Metric type (counter, gauge, histogram).

2955 value (float): Metric value.

2956 timestamp (datetime): When the metric was recorded.

2957 unit (str): Metric unit (ms, count, bytes, etc.).

2958 attributes (dict): Metric attributes/labels (JSON).

2959 resource_type (str): Type of resource (tool, resource, prompt, etc.).

2960 resource_id (str): ID of the specific resource.

2961 trace_id (str): Associated trace ID (optional).

2962 created_at (datetime): Metric creation timestamp.

2963 """

2964

2965 __tablename__ = "observability_metrics"

2966

2967 # Primary key

2968 id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)

2969

2970 # Metric metadata

2971 name: Mapped[str] = mapped_column(String(255), nullable=False, index=True)

2972 metric_type: Mapped[str] = mapped_column(String(20), nullable=False) # counter, gauge, histogram

2973 value: Mapped[float] = mapped_column(Float, nullable=False)

2974 timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False, default=utc_now, index=True)

2975 unit: Mapped[Optional[str]] = mapped_column(String(20), nullable=True)

2976

2977 # Attributes/labels

2978 attributes: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True, default=dict)

2979

2980 # Resource context

2981 resource_type: Mapped[Optional[str]] = mapped_column(String(50), nullable=True, index=True)

2982 resource_id: Mapped[Optional[str]] = mapped_column(String(36), nullable=True, index=True)

2983

2984 # Trace association (optional)

2985 trace_id: Mapped[Optional[str]] = mapped_column(String(36), ForeignKey("observability_traces.trace_id", ondelete="SET NULL"), nullable=True, index=True)

2986

2987 # Timestamps

2988 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, nullable=False)

2989

2990 # Indexes for performance

2991 __table_args__ = (

2992 Index("idx_observability_metrics_name_timestamp", "name", "timestamp"),

2993 Index("idx_observability_metrics_resource_type", "resource_type"),

2994 Index("idx_observability_metrics_trace_id", "trace_id"),

2995 )

2996

2997

2998class ObservabilitySavedQuery(Base):

2999 """

3000 ORM model for saved observability queries (filter presets).

3001

3002 Allows users to save their filter configurations for quick access and

3003 historical query tracking. Queries can be personal or shared with the team.

3004

3005 Attributes:

3006 id (int): Auto-incrementing primary key.

3007 name (str): User-given name for the saved query.

3008 description (str): Optional description of what this query finds.

3009 user_email (str): Email of the user who created this query.

3010 filter_config (dict): JSON containing all filter values (time_range, status_filter, etc.).

3011 is_shared (bool): Whether this query is visible to other users.

3012 created_at (datetime): When the query was created.

3013 updated_at (datetime): When the query was last modified.

3014 last_used_at (datetime): When the query was last executed.

3015 use_count (int): How many times this query has been used.

3016 """

3017

3018 __tablename__ = "observability_saved_queries"

3019

3020 # Primary key

3021 id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)

3022

3023 # Query metadata

3024 name: Mapped[str] = mapped_column(String(255), nullable=False, index=True)

3025 description: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

3026 user_email: Mapped[str] = mapped_column(String(255), nullable=False, index=True)

3027

3028 # Filter configuration (stored as JSON)

3029 filter_config: Mapped[Dict[str, Any]] = mapped_column(JSON, nullable=False)

3030

3031 # Sharing settings

3032 is_shared: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)

3033

3034 # Timestamps and usage tracking

3035 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, nullable=False)

3036 updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, onupdate=utc_now, nullable=False)

3037 last_used_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), nullable=True)

3038 use_count: Mapped[int] = mapped_column(Integer, default=0, nullable=False)

3039

3040 # Indexes for performance

3041 __table_args__ = (

3042 Index("idx_observability_saved_queries_user_email", "user_email"),

3043 Index("idx_observability_saved_queries_is_shared", "is_shared"),

3044 Index("idx_observability_saved_queries_created_at", "created_at"),

3045 )

3046

3047

3048# ---------------------------------------------------------------------------

3049# Performance Monitoring Models

3050# ---------------------------------------------------------------------------

3051

3052

3053class PerformanceSnapshot(Base):

3054 """

3055 ORM model for point-in-time performance snapshots.

3056

3057 Stores comprehensive system, request, and worker metrics at regular intervals

3058 for historical analysis and trend detection.

3059

3060 Attributes:

3061 id (int): Auto-incrementing primary key.

3062 timestamp (datetime): When the snapshot was taken.

3063 host (str): Hostname of the machine.

3064 worker_id (str): Worker identifier (PID or UUID).

3065 metrics_json (dict): JSON blob containing all metrics data.

3066 created_at (datetime): Record creation timestamp.

3067 """

3068

3069 __tablename__ = "performance_snapshots"

3070

3071 # Primary key

3072 id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)

3073

3074 # Snapshot metadata

3075 timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, nullable=False, index=True)

3076 host: Mapped[str] = mapped_column(String(255), nullable=False, index=True)

3077 worker_id: Mapped[Optional[str]] = mapped_column(String(64), nullable=True, index=True)

3078

3079 # Metrics data (JSON blob)

3080 metrics_json: Mapped[Dict[str, Any]] = mapped_column(JSON, nullable=False)

3081

3082 # Timestamps

3083 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, nullable=False)

3084

3085 # Indexes for efficient querying

3086 __table_args__ = (

3087 Index("idx_performance_snapshots_timestamp", "timestamp"),

3088 Index("idx_performance_snapshots_host_timestamp", "host", "timestamp"),

3089 Index("idx_performance_snapshots_created_at", "created_at"),

3090 )

3091

3092

3093class PerformanceAggregate(Base):

3094 """

3095 ORM model for aggregated performance metrics.

3096

3097 Stores hourly and daily aggregations of performance data for efficient

3098 historical reporting and trend analysis.

3099

3100 Attributes:

3101 id (int): Auto-incrementing primary key.

3102 period_start (datetime): Start of the aggregation period.

3103 period_end (datetime): End of the aggregation period.

3104 period_type (str): Type of aggregation (hourly, daily).

3105 host (str): Hostname (None for cluster-wide aggregates).

3106 Various aggregate metrics for requests and resources.

3107 created_at (datetime): Record creation timestamp.

3108 """

3109

3110 __tablename__ = "performance_aggregates"

3111

3112 # Primary key

3113 id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)

3114

3115 # Period metadata

3116 period_start: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False, index=True)

3117 period_end: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)

3118 period_type: Mapped[str] = mapped_column(String(20), nullable=False, index=True) # hourly, daily

3119 host: Mapped[Optional[str]] = mapped_column(String(255), nullable=True, index=True)

3120

3121 # Request aggregates

3122 requests_total: Mapped[int] = mapped_column(Integer, default=0, nullable=False)

3123 requests_2xx: Mapped[int] = mapped_column(Integer, default=0, nullable=False)

3124 requests_4xx: Mapped[int] = mapped_column(Integer, default=0, nullable=False)

3125 requests_5xx: Mapped[int] = mapped_column(Integer, default=0, nullable=False)

3126 avg_response_time_ms: Mapped[float] = mapped_column(Float, default=0.0, nullable=False)

3127 p95_response_time_ms: Mapped[float] = mapped_column(Float, default=0.0, nullable=False)

3128 peak_requests_per_second: Mapped[float] = mapped_column(Float, default=0.0, nullable=False)

3129

3130 # Resource aggregates

3131 avg_cpu_percent: Mapped[float] = mapped_column(Float, default=0.0, nullable=False)

3132 avg_memory_percent: Mapped[float] = mapped_column(Float, default=0.0, nullable=False)

3133 peak_cpu_percent: Mapped[float] = mapped_column(Float, default=0.0, nullable=False)

3134 peak_memory_percent: Mapped[float] = mapped_column(Float, default=0.0, nullable=False)

3135

3136 # Timestamps

3137 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, nullable=False)

3138

3139 # Indexes and constraints

3140 __table_args__ = (

3141 Index("idx_performance_aggregates_period", "period_type", "period_start"),

3142 Index("idx_performance_aggregates_host_period", "host", "period_type", "period_start"),

3143 UniqueConstraint("period_type", "period_start", "host", name="uq_performance_aggregate_period_host"),

3144 )

3145

3146

3147class Tool(Base):

3148 """

3149 ORM model for a registered Tool.

3150

3151 Supports both local tools and federated tools from other gateways.

3152 The integration_type field indicates the tool format:

3153 - "MCP" for MCP-compliant tools (default)

3154 - "REST" for REST tools

3155

3156 Additionally, this model provides computed properties for aggregated metrics based

3157 on the associated ToolMetric records. These include:

3158 - execution_count: Total number of invocations.

3159 - successful_executions: Count of successful invocations.

3160 - failed_executions: Count of failed invocations.

3161 - failure_rate: Ratio of failed invocations to total invocations.

3162 - min_response_time: Fastest recorded response time.

3163 - max_response_time: Slowest recorded response time.

3164 - avg_response_time: Mean response time.

3165 - last_execution_time: Timestamp of the most recent invocation.

3166

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

3168

3169 Team association is handled via the `email_team` relationship (default lazy loading)

3170 which only includes active teams. For list operations, use explicit joinedload()

3171 to eager load team names. The `team` property provides convenient access to

3172 the team name:

3173 - team: Returns the team name if the tool belongs to an active team, otherwise None.

3174

3175 The following fields have been added to support tool invocation configuration:

3176 - request_type: HTTP method to use when invoking the tool.

3177 - auth_type: Type of authentication ("basic", "bearer", or None).

3178 - auth_username: Username for basic authentication.

3179 - auth_password: Password for basic authentication.

3180 - auth_token: Token for bearer token authentication.

3181 - auth_header_key: header key for authentication.

3182 - auth_header_value: header value for authentication.

3183 """

3184

3185 __tablename__ = "tools"

3186

3187 id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: uuid.uuid4().hex)

3188 original_name: Mapped[str] = mapped_column(String(255), nullable=False)

3189 url: Mapped[str] = mapped_column(String(767), nullable=True)

3190 original_description: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

3191 description: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

3192 integration_type: Mapped[str] = mapped_column(String(20), default="MCP")

3193 request_type: Mapped[str] = mapped_column(String(20), default="SSE")

3194 headers: Mapped[Optional[Dict[str, str]]] = mapped_column(JSON)

3195 input_schema: Mapped[Dict[str, Any]] = mapped_column(JSON)

3196 output_schema: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True)

3197 annotations: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, default=lambda: {})

3198 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now)

3199 updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, onupdate=utc_now)

3200 enabled: Mapped[bool] = mapped_column(default=True)

3201 reachable: Mapped[bool] = mapped_column(default=True)

3202 jsonpath_filter: Mapped[str] = mapped_column(Text, default="")

3203 tags: Mapped[List[str]] = mapped_column(JSON, default=list, nullable=False)

3204

3205 # Comprehensive metadata for audit tracking

3206 created_by: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

3207 created_from_ip: Mapped[Optional[str]] = mapped_column(String(45), nullable=True)

3208 created_via: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)

3209 created_user_agent: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

3210

3211 modified_by: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

3212 modified_from_ip: Mapped[Optional[str]] = mapped_column(String(45), nullable=True)

3213 modified_via: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)

3214 modified_user_agent: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

3215

3216 import_batch_id: Mapped[Optional[str]] = mapped_column(String(36), nullable=True)

3217 federation_source: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

3218 version: Mapped[int] = mapped_column(Integer, default=1, nullable=False)

3219

3220 # Request type and authentication fields

3221 auth_type: Mapped[Optional[str]] = mapped_column(String(20), default=None) # "basic", "bearer", or None

3222 auth_value: Mapped[Optional[str]] = mapped_column(Text, default=None)

3223

3224 # custom_name,custom_name_slug, display_name

3225 custom_name: Mapped[Optional[str]] = mapped_column(String(255), nullable=False)

3226 custom_name_slug: Mapped[Optional[str]] = mapped_column(String(255), nullable=False)

3227 display_name: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

3228 title: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

3229

3230 # Passthrough REST fields

3231 base_url: Mapped[Optional[str]] = mapped_column(String, nullable=True)

3232 path_template: Mapped[Optional[str]] = mapped_column(String, nullable=True)

3233 query_mapping: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True)

3234 header_mapping: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True)

3235 timeout_ms: Mapped[Optional[int]] = mapped_column(Integer, nullable=True, default=None)

3236 expose_passthrough: Mapped[bool] = mapped_column(Boolean, default=True)

3237 allowlist: Mapped[Optional[List[str]]] = mapped_column(JSON, nullable=True)

3238 plugin_chain_pre: Mapped[Optional[List[str]]] = mapped_column(JSON, nullable=True)

3239 plugin_chain_post: Mapped[Optional[List[str]]] = mapped_column(JSON, nullable=True)

3240

3241 # Federation relationship with a local gateway

3242 gateway_id: Mapped[Optional[str]] = mapped_column(ForeignKey("gateways.id", ondelete="CASCADE"))

3243 # gateway_slug: Mapped[Optional[str]] = mapped_column(ForeignKey("gateways.slug"))

3244 gateway: Mapped["Gateway"] = relationship("Gateway", primaryjoin="Tool.gateway_id == Gateway.id", foreign_keys=[gateway_id], back_populates="tools")

3245 # federated_with = relationship("Gateway", secondary=tool_gateway_table, back_populates="federated_tools")

3246

3247 # Many-to-many relationship with Servers

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

3249

3250 # Relationship with ToolMetric records

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

3252 metrics_hourly: Mapped[List["ToolMetricsHourly"]] = relationship(

3253 "ToolMetricsHourly",

3254 primaryjoin="Tool.id == foreign(ToolMetricsHourly.tool_id)",

3255 viewonly=True,

3256 )

3257

3258 # Team scoping fields for resource organization

3259 team_id: Mapped[Optional[str]] = mapped_column(String(36), ForeignKey("email_teams.id", ondelete="SET NULL"), nullable=True)

3260 owner_email: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

3261 visibility: Mapped[str] = mapped_column(String(20), nullable=False, default="public")

3262

3263 # Relationship for loading team names (only active teams)

3264 # Uses default lazy loading - team name is only loaded when accessed

3265 # For list/admin views, use explicit joinedload(DbTool.email_team) for single-query loading

3266 # This avoids adding overhead to hot paths like tool invocation that don't need team names

3267 email_team: Mapped[Optional["EmailTeam"]] = relationship(

3268 "EmailTeam",

3269 primaryjoin="and_(Tool.team_id == EmailTeam.id, EmailTeam.is_active == True)",

3270 foreign_keys=[team_id],

3271 )

3272

3273 @property

3274 def team(self) -> Optional[str]:

3275 """Return the team name from the eagerly-loaded email_team relationship.

3276

3277 Returns:

3278 Optional[str]: The team name if the tool belongs to an active team, otherwise None.

3279 """

3280 return self.email_team.name if self.email_team else None

3281

3282 # @property

3283 # def gateway_slug(self) -> str:

3284 # return self.gateway.slug

3285

3286 _computed_name: Mapped[str] = mapped_column("name", String(255), nullable=False) # Stored column

3287

3288 @hybrid_property

3289 def name(self) -> str:

3290 """Return the display/lookup name computed from gateway and custom slug.

3291

3292 Returns:

3293 str: Display/lookup name to use for this tool.

3294 """

3295 # Instance access resolves Column to Python value; cast ensures static acceptance

3296 if getattr(self, "_computed_name", None):

3297 return cast(str, getattr(self, "_computed_name"))

3298 custom_name_slug = slugify(getattr(self, "custom_name_slug"))

3299 if getattr(self, "gateway_id", None):

3300 gateway_slug = slugify(self.gateway.name) # type: ignore[attr-defined]

3301 return f"{gateway_slug}{settings.gateway_tool_name_separator}{custom_name_slug}"

3302 return custom_name_slug

3303

3304 @name.setter

3305 def name(self, value: str) -> None:

3306 """Setter for the stored name column.

3307

3308 Args:

3309 value: Explicit name to persist to the underlying column.

3310 """

3311 setattr(self, "_computed_name", value)

3312

3313 @name.expression

3314 @classmethod

3315 def name(cls) -> Any:

3316 """SQL expression for name used in queries (backs onto stored column).

3317

3318 Returns:

3319 Any: SQLAlchemy expression referencing the stored name column.

3320 """

3321 return cls._computed_name

3322

3323 __table_args__ = (

3324 UniqueConstraint("gateway_id", "original_name", name="uq_gateway_id__original_name"),

3325 UniqueConstraint("team_id", "owner_email", "name", name="uq_team_owner_email_name_tool"),

3326 Index("idx_tools_created_at_id", "created_at", "id"),

3327 )

3328

3329 @hybrid_property

3330 def gateway_slug(self) -> Optional[str]:

3331 """Python accessor returning the related gateway's slug if available.

3332

3333 Returns:

3334 Optional[str]: The gateway slug, or None if no gateway relation.

3335 """

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

3337

3338 @gateway_slug.expression

3339 @classmethod

3340 def gateway_slug(cls) -> Any:

3341 """SQL expression to select current gateway slug for this tool.

3342

3343 Returns:

3344 Any: SQLAlchemy scalar subquery selecting the gateway slug.

3345 """

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

3347

3348 def _metrics_loaded(self) -> bool:

3349 """Check if metrics relationship is loaded without triggering lazy load.

3350

3351 Returns:

3352 bool: True if metrics are loaded, False otherwise.

3353 """

3354 return "metrics" in sa_inspect(self).dict

3355

3356 def _get_metric_counts(self) -> tuple[int, int, int]:

3357 """Get total, successful, and failed metric counts in a single operation.

3358

3359 When metrics are already loaded, computes from memory in O(n).

3360 When not loaded, uses a single SQL query with conditional aggregation.

3361

3362 Note: For bulk operations, use metrics_summary which computes all fields

3363 in a single pass, or ensure metrics are preloaded via selectinload.

3364

3365 Returns:

3366 tuple[int, int, int]: (total, successful, failed) counts.

3367 """

3368 # If metrics are loaded, compute from memory in a single pass

3369 if self._metrics_loaded():

3370 total = 0

3371 successful = 0

3372 for m in self.metrics:

3373 total += 1

3374 if m.is_success:

3375 successful += 1

3376 return (total, successful, total - successful)

3377

3378 # Use single SQL query with conditional aggregation

3379 # Third-Party

3380 from sqlalchemy import case # pylint: disable=import-outside-toplevel

3381 from sqlalchemy.orm import object_session # pylint: disable=import-outside-toplevel

3382

3383 session = object_session(self)

3384 if session is None:

3385 return (0, 0, 0)

3386

3387 result = (

3388 session.query(

3389 func.count(ToolMetric.id), # pylint: disable=not-callable

3390 func.sum(case((ToolMetric.is_success.is_(True), 1), else_=0)),

3391 )

3392 .filter(ToolMetric.tool_id == self.id)

3393 .one()

3394 )

3395

3396 total = result[0] or 0

3397 successful = result[1] or 0

3398 return (total, successful, total - successful)

3399

3400 @hybrid_property

3401 def execution_count(self) -> int:

3402 """Number of ToolMetric records associated with this tool instance.

3403

3404 Note: Each property access may trigger a SQL query if metrics aren't loaded.

3405 For reading multiple metric fields, use metrics_summary or preload metrics.

3406

3407 Returns:

3408 int: Count of ToolMetric records for this tool.

3409 """

3410 return self._get_metric_counts()[0]

3411

3412 @execution_count.expression

3413 @classmethod

3414 def execution_count(cls) -> Any:

3415 """SQL expression that counts ToolMetric rows for this tool.

3416

3417 Returns:

3418 Any: SQLAlchemy labeled count expression for tool metrics.

3419 """

3420 return select(func.count(ToolMetric.id)).where(ToolMetric.tool_id == cls.id).correlate(cls).scalar_subquery().label("execution_count") # pylint: disable=not-callable

3421

3422 @property

3423 def successful_executions(self) -> int:

3424 """Count of successful tool executions.

3425

3426 Returns:

3427 int: The count of successful tool executions.

3428 """

3429 return self._get_metric_counts()[1]

3430

3431 @property

3432 def failed_executions(self) -> int:

3433 """Count of failed tool executions.

3434

3435 Returns:

3436 int: The count of failed tool executions.

3437 """

3438 return self._get_metric_counts()[2]

3439

3440 @property

3441 def failure_rate(self) -> float:

3442 """Failure rate as a float between 0 and 1.

3443

3444 Returns:

3445 float: The failure rate as a value between 0 and 1.

3446 """

3447 total, _, failed = self._get_metric_counts()

3448 return failed / total if total > 0 else 0.0

3449

3450 @property

3451 def min_response_time(self) -> Optional[float]:

3452 """Minimum response time among all tool executions.

3453

3454 Returns None if metrics are not loaded (use metrics_summary for SQL fallback).

3455

3456 Returns:

3457 Optional[float]: The minimum response time, or None.

3458 """

3459 if not self._metrics_loaded():

3460 return None

3461 times: List[float] = [m.response_time for m in self.metrics]

3462 return min(times) if times else None

3463

3464 @property

3465 def max_response_time(self) -> Optional[float]:

3466 """Maximum response time among all tool executions.

3467

3468 Returns None if metrics are not loaded (use metrics_summary for SQL fallback).

3469

3470 Returns:

3471 Optional[float]: The maximum response time, or None.

3472 """

3473 if not self._metrics_loaded():

3474 return None

3475 times: List[float] = [m.response_time for m in self.metrics]

3476 return max(times) if times else None

3477

3478 @property

3479 def avg_response_time(self) -> Optional[float]:

3480 """Average response time among all tool executions.

3481

3482 Returns None if metrics are not loaded (use metrics_summary for SQL fallback).

3483

3484 Returns:

3485 Optional[float]: The average response time, or None.

3486 """

3487 if not self._metrics_loaded():

3488 return None

3489 times: List[float] = [m.response_time for m in self.metrics]

3490 return sum(times) / len(times) if times else None

3491

3492 @property

3493 def last_execution_time(self) -> Optional[datetime]:

3494 """Timestamp of the most recent tool execution.

3495

3496 Returns None if metrics are not loaded (use metrics_summary for SQL fallback).

3497

3498 Returns:

3499 Optional[datetime]: The timestamp of the most recent execution, or None.

3500 """

3501 if not self._metrics_loaded():

3502 return None

3503 if not self.metrics:

3504 return None

3505 return max(m.timestamp for m in self.metrics)

3506

3507 @property

3508 def metrics_summary(self) -> Dict[str, Any]:

3509 """Aggregated metrics for the tool combining raw and hourly data without double-counting.

3510

3511 When metrics are loaded: computes from memory (raw + hourly)

3512 When not loaded: uses SQL queries with time partitioning

3513

3514 Returns:

3515 Dict[str, Any]: Dictionary containing aggregated metrics:

3516 - total_executions, successful_executions, failed_executions

3517 - failure_rate, min/max/avg_response_time, last_execution_time

3518 """

3519 # Try in-memory path first

3520 if self._metrics_loaded():

3521 try:

3522 hourly_metrics = self.metrics_hourly

3523 except AttributeError:

3524 hourly_metrics = [] # Relationship not loaded

3525 return _compute_metrics_summary(raw_metrics=self.metrics, hourly_metrics=hourly_metrics)

3526

3527 # SQL query path

3528 # Third-Party

3529 from sqlalchemy.orm import object_session # pylint: disable=import-outside-toplevel

3530

3531 session = object_session(self)

3532 if session is None:

3533 return {

3534 "total_executions": 0,

3535 "successful_executions": 0,

3536 "failed_executions": 0,

3537 "failure_rate": 0.0,

3538 "min_response_time": None,

3539 "max_response_time": None,

3540 "avg_response_time": None,

3541 "last_execution_time": None,

3542 }

3543

3544 return _compute_metrics_summary(

3545 raw_metrics=None,

3546 hourly_metrics=None,

3547 session=session,

3548 entity_id=self.id,

3549 raw_metric_class=ToolMetric,

3550 hourly_metric_class=ToolMetricsHourly,

3551 )

3552

3553

3554class Resource(Base):

3555 """

3556 ORM model for a registered Resource.

3557

3558 Resources represent content that can be read by clients.

3559 Supports subscriptions for real-time updates.

3560 Additionally, this model provides a relationship with ResourceMetric records

3561 to capture invocation metrics (such as execution counts, response times, and failures).

3562 """

3563

3564 __tablename__ = "resources"

3565

3566 id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: uuid.uuid4().hex)

3567 uri: Mapped[str] = mapped_column(String(767), nullable=False)

3568 name: Mapped[str] = mapped_column(String(255), nullable=False)

3569 description: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

3570 title: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

3571 mime_type: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

3572 size: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)

3573 uri_template: Mapped[Optional[str]] = mapped_column(Text, nullable=True) # URI template for parameterized resources

3574 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now)

3575 updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, onupdate=utc_now)

3576 # is_active: Mapped[bool] = mapped_column(default=True)

3577 enabled: Mapped[bool] = mapped_column(default=True)

3578 tags: Mapped[List[str]] = mapped_column(JSON, default=list, nullable=False)

3579

3580 # Comprehensive metadata for audit tracking

3581 created_by: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

3582 created_from_ip: Mapped[Optional[str]] = mapped_column(String(45), nullable=True)

3583 created_via: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)

3584 created_user_agent: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

3585

3586 modified_by: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

3587 modified_from_ip: Mapped[Optional[str]] = mapped_column(String(45), nullable=True)

3588 modified_via: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)

3589 modified_user_agent: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

3590

3591 import_batch_id: Mapped[Optional[str]] = mapped_column(String(36), nullable=True)

3592 federation_source: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

3593 version: Mapped[int] = mapped_column(Integer, default=1, nullable=False)

3594

3595 metrics: Mapped[List["ResourceMetric"]] = relationship("ResourceMetric", back_populates="resource", cascade="all, delete-orphan")

3596 metrics_hourly: Mapped[List["ResourceMetricsHourly"]] = relationship(

3597 "ResourceMetricsHourly",

3598 primaryjoin="Resource.id == foreign(ResourceMetricsHourly.resource_id)",

3599 viewonly=True,

3600 )

3601

3602 # Content storage - can be text or binary

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

3604 binary_content: Mapped[Optional[bytes]]

3605

3606 # Subscription tracking

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

3608

3609 gateway_id: Mapped[Optional[str]] = mapped_column(ForeignKey("gateways.id", ondelete="CASCADE"))

3610 gateway: Mapped["Gateway"] = relationship("Gateway", back_populates="resources")

3611 # federated_with = relationship("Gateway", secondary=resource_gateway_table, back_populates="federated_resources")

3612

3613 # Many-to-many relationship with Servers

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

3615 __table_args__ = (

3616 UniqueConstraint("team_id", "owner_email", "gateway_id", "uri", name="uq_team_owner_gateway_uri_resource"),

3617 Index("uq_team_owner_uri_resource_local", "team_id", "owner_email", "uri", unique=True, postgresql_where=text("gateway_id IS NULL"), sqlite_where=text("gateway_id IS NULL")),

3618 Index("idx_resources_created_at_id", "created_at", "id"),

3619 )

3620

3621 @property

3622 def content(self) -> "ResourceContent":

3623 """

3624 Returns the resource content in the appropriate format.

3625

3626 If text content exists, returns a ResourceContent with text.

3627 Otherwise, if binary content exists, returns a ResourceContent with blob data.

3628 Raises a ValueError if no content is available.

3629

3630 Returns:

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

3632

3633 Raises:

3634 ValueError: If the resource has no content available.

3635

3636 Examples:

3637 >>> resource = Resource(uri="test://example", name="test")

3638 >>> resource.text_content = "Hello, World!"

3639 >>> content = resource.content

3640 >>> content.text

3641 'Hello, World!'

3642 >>> content.type

3643 'resource'

3644

3645 >>> binary_resource = Resource(uri="test://binary", name="binary")

3646 >>> binary_resource.binary_content = b"\\x00\\x01\\x02"

3647 >>> binary_content = binary_resource.content

3648 >>> binary_content.blob

3649 b'\\x00\\x01\\x02'

3650

3651 >>> empty_resource = Resource(uri="test://empty", name="empty")

3652 >>> try:

3653 ... empty_resource.content

3654 ... except ValueError as e:

3655 ... str(e)

3656 'Resource has no content'

3657 """

3658

3659 # Local import to avoid circular import

3660 # First-Party

3661 from mcpgateway.common.models import ResourceContent # pylint: disable=import-outside-toplevel

3662

3663 if self.text_content is not None:

3664 return ResourceContent(

3665 type="resource",

3666 id=str(self.id),

3667 uri=self.uri,

3668 mime_type=self.mime_type,

3669 text=self.text_content,

3670 )

3671 if self.binary_content is not None:

3672 return ResourceContent(

3673 type="resource",

3674 id=str(self.id),

3675 uri=self.uri,

3676 mime_type=self.mime_type or "application/octet-stream",

3677 blob=self.binary_content,

3678 )

3679 raise ValueError("Resource has no content")

3680

3681 def _metrics_loaded(self) -> bool:

3682 """Check if metrics relationship is loaded without triggering lazy load.

3683

3684 Returns:

3685 bool: True if metrics are loaded, False otherwise.

3686 """

3687 return "metrics" in sa_inspect(self).dict

3688

3689 def _get_metric_counts(self) -> tuple[int, int, int]:

3690 """Get total, successful, and failed metric counts in a single operation.

3691

3692 When metrics are already loaded, computes from memory in O(n).

3693 When not loaded, uses a single SQL query with conditional aggregation.

3694

3695 Returns:

3696 tuple[int, int, int]: (total, successful, failed) counts.

3697 """

3698 if self._metrics_loaded():

3699 total = 0

3700 successful = 0

3701 for m in self.metrics:

3702 total += 1

3703 if m.is_success:

3704 successful += 1

3705 return (total, successful, total - successful)

3706

3707 # Third-Party

3708 from sqlalchemy import case # pylint: disable=import-outside-toplevel

3709 from sqlalchemy.orm import object_session # pylint: disable=import-outside-toplevel

3710

3711 session = object_session(self)

3712 if session is None:

3713 return (0, 0, 0)

3714

3715 result = (

3716 session.query(

3717 func.count(ResourceMetric.id), # pylint: disable=not-callable

3718 func.sum(case((ResourceMetric.is_success.is_(True), 1), else_=0)),

3719 )

3720 .filter(ResourceMetric.resource_id == self.id)

3721 .one()

3722 )

3723

3724 total = result[0] or 0

3725 successful = result[1] or 0

3726 return (total, successful, total - successful)

3727

3728 @hybrid_property

3729 def execution_count(self) -> int:

3730 """Number of ResourceMetric records associated with this resource instance.

3731

3732 Returns:

3733 int: Count of ResourceMetric records for this resource.

3734 """

3735 return self._get_metric_counts()[0]

3736

3737 @execution_count.expression

3738 @classmethod

3739 def execution_count(cls) -> Any:

3740 """SQL expression that counts ResourceMetric rows for this resource.

3741

3742 Returns:

3743 Any: SQLAlchemy labeled count expression for resource metrics.

3744 """

3745 return select(func.count(ResourceMetric.id)).where(ResourceMetric.resource_id == cls.id).correlate(cls).scalar_subquery().label("execution_count") # pylint: disable=not-callable

3746

3747 @property

3748 def successful_executions(self) -> int:

3749 """Count of successful resource invocations.

3750

3751 Returns:

3752 int: The count of successful resource invocations.

3753 """

3754 return self._get_metric_counts()[1]

3755

3756 @property

3757 def failed_executions(self) -> int:

3758 """Count of failed resource invocations.

3759

3760 Returns:

3761 int: The count of failed resource invocations.

3762 """

3763 return self._get_metric_counts()[2]

3764

3765 @property

3766 def failure_rate(self) -> float:

3767 """Failure rate as a float between 0 and 1.

3768

3769 Returns:

3770 float: The failure rate as a value between 0 and 1.

3771 """

3772 total, _, failed = self._get_metric_counts()

3773 return failed / total if total > 0 else 0.0

3774

3775 @property

3776 def min_response_time(self) -> Optional[float]:

3777 """Minimum response time among all resource invocations.

3778

3779 Returns None if metrics are not loaded. Note: counts may be non-zero

3780 (via SQL) while timing is None. Use service layer converters for

3781 consistent metrics, or preload metrics via selectinload.

3782

3783 Returns:

3784 Optional[float]: The minimum response time, or None.

3785 """

3786 if not self._metrics_loaded():

3787 return None

3788 times: List[float] = [m.response_time for m in self.metrics]

3789 return min(times) if times else None

3790

3791 @property

3792 def max_response_time(self) -> Optional[float]:

3793 """Maximum response time among all resource invocations.

3794

3795 Returns None if metrics are not loaded. Note: counts may be non-zero

3796 (via SQL) while timing is None. Use service layer converters for

3797 consistent metrics, or preload metrics via selectinload.

3798

3799 Returns:

3800 Optional[float]: The maximum response time, or None.

3801 """

3802 if not self._metrics_loaded():

3803 return None

3804 times: List[float] = [m.response_time for m in self.metrics]

3805 return max(times) if times else None

3806

3807 @property

3808 def avg_response_time(self) -> Optional[float]:

3809 """Average response time among all resource invocations.

3810

3811 Returns None if metrics are not loaded. Note: counts may be non-zero

3812 (via SQL) while timing is None. Use service layer converters for

3813 consistent metrics, or preload metrics via selectinload.

3814

3815 Returns:

3816 Optional[float]: The average response time, or None.

3817 """

3818 if not self._metrics_loaded():

3819 return None

3820 times: List[float] = [m.response_time for m in self.metrics]

3821 return sum(times) / len(times) if times else None

3822

3823 @property

3824 def last_execution_time(self) -> Optional[datetime]:

3825 """Timestamp of the most recent resource invocation.

3826

3827 Returns None if metrics are not loaded. Note: counts may be non-zero

3828 (via SQL) while timing is None. Use service layer converters for

3829 consistent metrics, or preload metrics via selectinload.

3830

3831 Returns:

3832 Optional[datetime]: The timestamp of the most recent invocation, or None.

3833 """

3834 if not self._metrics_loaded():

3835 return None

3836 if not self.metrics:

3837 return None

3838 return max(m.timestamp for m in self.metrics)

3839

3840 @property

3841 def metrics_summary(self) -> Dict[str, Any]:

3842 """Aggregated metrics for the resource combining raw and hourly data without double-counting.

3843

3844 When metrics are loaded: computes from memory (raw + hourly)

3845 When not loaded: uses SQL queries with time partitioning

3846

3847 Returns:

3848 Dict[str, Any]: Dictionary containing aggregated metrics:

3849 - total_executions, successful_executions, failed_executions

3850 - failure_rate, min/max/avg_response_time, last_execution_time

3851 """

3852 # Try in-memory path first

3853 if self._metrics_loaded():

3854 try:

3855 hourly_metrics = self.metrics_hourly

3856 except AttributeError:

3857 hourly_metrics = []

3858 return _compute_metrics_summary(raw_metrics=self.metrics, hourly_metrics=hourly_metrics)

3859

3860 # SQL query path

3861 # Third-Party

3862 from sqlalchemy.orm import object_session # pylint: disable=import-outside-toplevel

3863

3864 session = object_session(self)

3865 if session is None:

3866 return {

3867 "total_executions": 0,

3868 "successful_executions": 0,

3869 "failed_executions": 0,

3870 "failure_rate": 0.0,

3871 "min_response_time": None,

3872 "max_response_time": None,

3873 "avg_response_time": None,

3874 "last_execution_time": None,

3875 }

3876

3877 return _compute_metrics_summary(

3878 raw_metrics=None,

3879 hourly_metrics=None,

3880 session=session,

3881 entity_id=self.id,

3882 raw_metric_class=ResourceMetric,

3883 hourly_metric_class=ResourceMetricsHourly,

3884 )

3885

3886 # Team scoping fields for resource organization

3887 team_id: Mapped[Optional[str]] = mapped_column(String(36), ForeignKey("email_teams.id", ondelete="SET NULL"), nullable=True)

3888 owner_email: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

3889 visibility: Mapped[str] = mapped_column(String(20), nullable=False, default="public")

3890

3891

3892class ResourceSubscription(Base):

3893 """Tracks subscriptions to resource updates."""

3894

3895 __tablename__ = "resource_subscriptions"

3896

3897 id: Mapped[int] = mapped_column(primary_key=True)

3898 resource_id: Mapped[str] = mapped_column(ForeignKey("resources.id"))

3899 subscriber_id: Mapped[str] = mapped_column(String(255), nullable=False) # Client identifier

3900 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now)

3901 last_notification: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), nullable=True)

3902

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

3904

3905

3906class ToolOpsTestCases(Base):

3907 """

3908 ORM model for a registered Tool test cases.

3909

3910 Represents a tool and the generated test cases.

3911 Includes:

3912 - tool_id: unique tool identifier

3913 - test_cases: generated test cases.

3914 - run_status: status of test case generation

3915 """

3916

3917 __tablename__ = "toolops_test_cases"

3918

3919 tool_id: Mapped[str] = mapped_column(String(255), primary_key=True)

3920 test_cases: Mapped[Dict[str, Any]] = mapped_column(JSON)

3921 run_status: Mapped[str] = mapped_column(String(255), nullable=False)

3922

3923

3924class Prompt(Base):

3925 """

3926 ORM model for a registered Prompt template.

3927

3928 Represents a prompt template along with its argument schema.

3929 Supports rendering and invocation of prompts.

3930 Additionally, this model provides computed properties for aggregated metrics based

3931 on the associated PromptMetric records. These include:

3932 - execution_count: Total number of prompt invocations.

3933 - successful_executions: Count of successful invocations.

3934 - failed_executions: Count of failed invocations.

3935 - failure_rate: Ratio of failed invocations to total invocations.

3936 - min_response_time: Fastest recorded response time.

3937 - max_response_time: Slowest recorded response time.

3938 - avg_response_time: Mean response time.

3939 - last_execution_time: Timestamp of the most recent invocation.

3940 """

3941

3942 __tablename__ = "prompts"

3943

3944 id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: uuid.uuid4().hex)

3945 original_name: Mapped[str] = mapped_column(String(255), nullable=False)

3946 custom_name: Mapped[str] = mapped_column(String(255), nullable=False)

3947 custom_name_slug: Mapped[str] = mapped_column(String(255), nullable=False)

3948 display_name: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

3949 title: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

3950 name: Mapped[str] = mapped_column(String(255), nullable=False)

3951 description: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

3952 template: Mapped[str] = mapped_column(Text)

3953 argument_schema: Mapped[Dict[str, Any]] = mapped_column(JSON)

3954 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now)

3955 updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, onupdate=utc_now)

3956 # is_active: Mapped[bool] = mapped_column(default=True)

3957 enabled: Mapped[bool] = mapped_column(default=True)

3958 tags: Mapped[List[str]] = mapped_column(JSON, default=list, nullable=False)

3959

3960 # Comprehensive metadata for audit tracking

3961 created_by: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

3962 created_from_ip: Mapped[Optional[str]] = mapped_column(String(45), nullable=True)

3963 created_via: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)

3964 created_user_agent: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

3965

3966 modified_by: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

3967 modified_from_ip: Mapped[Optional[str]] = mapped_column(String(45), nullable=True)

3968 modified_via: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)

3969 modified_user_agent: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

3970

3971 import_batch_id: Mapped[Optional[str]] = mapped_column(String(36), nullable=True)

3972 federation_source: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

3973 version: Mapped[int] = mapped_column(Integer, default=1, nullable=False)

3974

3975 metrics: Mapped[List["PromptMetric"]] = relationship("PromptMetric", back_populates="prompt", cascade="all, delete-orphan")

3976 metrics_hourly: Mapped[List["PromptMetricsHourly"]] = relationship(

3977 "PromptMetricsHourly",

3978 primaryjoin="Prompt.id == foreign(PromptMetricsHourly.prompt_id)",

3979 viewonly=True,

3980 )

3981

3982 gateway_id: Mapped[Optional[str]] = mapped_column(ForeignKey("gateways.id", ondelete="CASCADE"))

3983 gateway: Mapped["Gateway"] = relationship("Gateway", back_populates="prompts")

3984 # federated_with = relationship("Gateway", secondary=prompt_gateway_table, back_populates="federated_prompts")

3985

3986 # Many-to-many relationship with Servers

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

3988

3989 # Team scoping fields for resource organization

3990 team_id: Mapped[Optional[str]] = mapped_column(String(36), ForeignKey("email_teams.id", ondelete="SET NULL"), nullable=True)

3991 owner_email: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

3992 visibility: Mapped[str] = mapped_column(String(20), nullable=False, default="public")

3993

3994 __table_args__ = (

3995 UniqueConstraint("team_id", "owner_email", "gateway_id", "name", name="uq_team_owner_gateway_name_prompt"),

3996 UniqueConstraint("gateway_id", "original_name", name="uq_gateway_id__original_name_prompt"),

3997 Index("uq_team_owner_name_prompt_local", "team_id", "owner_email", "name", unique=True, postgresql_where=text("gateway_id IS NULL"), sqlite_where=text("gateway_id IS NULL")),

3998 Index("idx_prompts_created_at_id", "created_at", "id"),

3999 )

4000

4001 @hybrid_property

4002 def gateway_slug(self) -> Optional[str]:

4003 """Return the related gateway's slug if available.

4004

4005 Returns:

4006 Optional[str]: Gateway slug or None when no gateway is attached.

4007 """

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

4009

4010 @gateway_slug.expression

4011 @classmethod

4012 def gateway_slug(cls) -> Any:

4013 """SQL expression to select current gateway slug for this prompt.

4014

4015 Returns:

4016 Any: SQLAlchemy scalar subquery selecting the gateway slug.

4017 """

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

4019

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

4021 """

4022 Validate prompt arguments against the argument schema.

4023

4024 Args:

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

4026

4027 Raises:

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

4029

4030 Examples:

4031 >>> prompt = Prompt(

4032 ... name="test_prompt",

4033 ... template="Hello {name}",

4034 ... argument_schema={

4035 ... "type": "object",

4036 ... "properties": {

4037 ... "name": {"type": "string"}

4038 ... },

4039 ... "required": ["name"]

4040 ... }

4041 ... )

4042 >>> prompt.validate_arguments({"name": "Alice"}) # No exception

4043 >>> try:

4044 ... prompt.validate_arguments({"age": 25}) # Missing required field

4045 ... except ValueError as e:

4046 ... "name" in str(e)

4047 True

4048 """

4049 try:

4050 jsonschema.validate(args, self.argument_schema)

4051 except jsonschema.exceptions.ValidationError as e:

4052 raise ValueError(f"Invalid prompt arguments: {str(e)}") from e

4053

4054 def _metrics_loaded(self) -> bool:

4055 """Check if metrics relationship is loaded without triggering lazy load.

4056

4057 Returns:

4058 bool: True if metrics are loaded, False otherwise.

4059 """

4060 return "metrics" in sa_inspect(self).dict

4061

4062 def _get_metric_counts(self) -> tuple[int, int, int]:

4063 """Get total, successful, and failed metric counts in a single operation.

4064

4065 When metrics are already loaded, computes from memory in O(n).

4066 When not loaded, uses a single SQL query with conditional aggregation.

4067

4068 Returns:

4069 tuple[int, int, int]: (total, successful, failed) counts.

4070 """

4071 if self._metrics_loaded():

4072 total = 0

4073 successful = 0

4074 for m in self.metrics:

4075 total += 1

4076 if m.is_success:

4077 successful += 1

4078 return (total, successful, total - successful)

4079

4080 # Third-Party

4081 from sqlalchemy import case # pylint: disable=import-outside-toplevel

4082 from sqlalchemy.orm import object_session # pylint: disable=import-outside-toplevel

4083

4084 session = object_session(self)

4085 if session is None:

4086 return (0, 0, 0)

4087

4088 result = (

4089 session.query(

4090 func.count(PromptMetric.id), # pylint: disable=not-callable

4091 func.sum(case((PromptMetric.is_success.is_(True), 1), else_=0)),

4092 )

4093 .filter(PromptMetric.prompt_id == self.id)

4094 .one()

4095 )

4096

4097 total = result[0] or 0

4098 successful = result[1] or 0

4099 return (total, successful, total - successful)

4100

4101 @hybrid_property

4102 def execution_count(self) -> int:

4103 """Number of PromptMetric records associated with this prompt instance.

4104

4105 Returns:

4106 int: Count of PromptMetric records for this prompt.

4107 """

4108 return self._get_metric_counts()[0]

4109

4110 @execution_count.expression

4111 @classmethod

4112 def execution_count(cls) -> Any:

4113 """SQL expression that counts PromptMetric rows for this prompt.

4114

4115 Returns:

4116 Any: SQLAlchemy labeled count expression for prompt metrics.

4117 """

4118 return select(func.count(PromptMetric.id)).where(PromptMetric.prompt_id == cls.id).correlate(cls).scalar_subquery().label("execution_count") # pylint: disable=not-callable

4119

4120 @property

4121 def successful_executions(self) -> int:

4122 """Count of successful prompt invocations.

4123

4124 Returns:

4125 int: The count of successful prompt invocations.

4126 """

4127 return self._get_metric_counts()[1]

4128

4129 @property

4130 def failed_executions(self) -> int:

4131 """Count of failed prompt invocations.

4132

4133 Returns:

4134 int: The count of failed prompt invocations.

4135 """

4136 return self._get_metric_counts()[2]

4137

4138 @property

4139 def failure_rate(self) -> float:

4140 """Failure rate as a float between 0 and 1.

4141

4142 Returns:

4143 float: The failure rate as a value between 0 and 1.

4144 """

4145 total, _, failed = self._get_metric_counts()

4146 return failed / total if total > 0 else 0.0

4147

4148 @property

4149 def min_response_time(self) -> Optional[float]:

4150 """Minimum response time among all prompt invocations.

4151

4152 Returns None if metrics are not loaded. Note: counts may be non-zero

4153 (via SQL) while timing is None. Use service layer converters for

4154 consistent metrics, or preload metrics via selectinload.

4155

4156 Returns:

4157 Optional[float]: The minimum response time, or None.

4158 """

4159 if not self._metrics_loaded():

4160 return None

4161 times: List[float] = [m.response_time for m in self.metrics]

4162 return min(times) if times else None

4163

4164 @property

4165 def max_response_time(self) -> Optional[float]:

4166 """Maximum response time among all prompt invocations.

4167

4168 Returns None if metrics are not loaded. Note: counts may be non-zero

4169 (via SQL) while timing is None. Use service layer converters for

4170 consistent metrics, or preload metrics via selectinload.

4171

4172 Returns:

4173 Optional[float]: The maximum response time, or None.

4174 """

4175 if not self._metrics_loaded():

4176 return None

4177 times: List[float] = [m.response_time for m in self.metrics]

4178 return max(times) if times else None

4179

4180 @property

4181 def avg_response_time(self) -> Optional[float]:

4182 """Average response time among all prompt invocations.

4183

4184 Returns None if metrics are not loaded. Note: counts may be non-zero

4185 (via SQL) while timing is None. Use service layer converters for

4186 consistent metrics, or preload metrics via selectinload.

4187

4188 Returns:

4189 Optional[float]: The average response time, or None.

4190 """

4191 if not self._metrics_loaded():

4192 return None

4193 times: List[float] = [m.response_time for m in self.metrics]

4194 return sum(times) / len(times) if times else None

4195

4196 @property

4197 def last_execution_time(self) -> Optional[datetime]:

4198 """Timestamp of the most recent prompt invocation.

4199

4200 Returns None if metrics are not loaded. Note: counts may be non-zero

4201 (via SQL) while timing is None. Use service layer converters for

4202 consistent metrics, or preload metrics via selectinload.

4203

4204 Returns:

4205 Optional[datetime]: The timestamp of the most recent invocation, or None if no invocations exist.

4206 """

4207 if not self._metrics_loaded():

4208 return None

4209 if not self.metrics:

4210 return None

4211 return max(m.timestamp for m in self.metrics)

4212

4213 @property

4214 def metrics_summary(self) -> Dict[str, Any]:

4215 """Aggregated metrics for the prompt combining raw and hourly data without double-counting.

4216

4217 When metrics are loaded: computes from memory (raw + hourly)

4218 When not loaded: uses SQL queries with time partitioning

4219

4220 Returns:

4221 Dict[str, Any]: Dictionary containing aggregated metrics:

4222 - total_executions, successful_executions, failed_executions

4223 - failure_rate, min/max/avg_response_time, last_execution_time

4224 """

4225 # Try in-memory path first

4226 if self._metrics_loaded():

4227 try:

4228 hourly_metrics = self.metrics_hourly

4229 except AttributeError:

4230 hourly_metrics = []

4231 return _compute_metrics_summary(raw_metrics=self.metrics, hourly_metrics=hourly_metrics)

4232

4233 # SQL query path

4234 # Third-Party

4235 from sqlalchemy.orm import object_session # pylint: disable=import-outside-toplevel

4236

4237 session = object_session(self)

4238 if session is None:

4239 return {

4240 "total_executions": 0,

4241 "successful_executions": 0,

4242 "failed_executions": 0,

4243 "failure_rate": 0.0,

4244 "min_response_time": None,

4245 "max_response_time": None,

4246 "avg_response_time": None,

4247 "last_execution_time": None,

4248 }

4249

4250 return _compute_metrics_summary(

4251 raw_metrics=None,

4252 hourly_metrics=None,

4253 session=session,

4254 entity_id=self.id,

4255 raw_metric_class=PromptMetric,

4256 hourly_metric_class=PromptMetricsHourly,

4257 )

4258

4259

4260class Server(Base):

4261 """

4262 ORM model for MCP Servers Catalog.

4263

4264 Represents a server that composes catalog items (tools, resources, prompts).

4265 Additionally, this model provides computed properties for aggregated metrics based

4266 on the associated ServerMetric records. These include:

4267 - execution_count: Total number of invocations.

4268 - successful_executions: Count of successful invocations.

4269 - failed_executions: Count of failed invocations.

4270 - failure_rate: Ratio of failed invocations to total invocations.

4271 - min_response_time: Fastest recorded response time.

4272 - max_response_time: Slowest recorded response time.

4273 - avg_response_time: Mean response time.

4274 - last_execution_time: Timestamp of the most recent invocation.

4275 """

4276

4277 __tablename__ = "servers"

4278

4279 id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: uuid.uuid4().hex)

4280 name: Mapped[str] = mapped_column(String(255), nullable=False)

4281 description: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

4282 icon: Mapped[Optional[str]] = mapped_column(String(767), nullable=True)

4283 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now)

4284 updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, onupdate=utc_now)

4285 # is_active: Mapped[bool] = mapped_column(default=True)

4286 enabled: Mapped[bool] = mapped_column(default=True)

4287 tags: Mapped[List[str]] = mapped_column(JSON, default=list, nullable=False)

4288

4289 # Comprehensive metadata for audit tracking

4290 created_by: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

4291 created_from_ip: Mapped[Optional[str]] = mapped_column(String(45), nullable=True)

4292 created_via: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)

4293 created_user_agent: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

4294

4295 modified_by: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

4296 modified_from_ip: Mapped[Optional[str]] = mapped_column(String(45), nullable=True)

4297 modified_via: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)

4298 modified_user_agent: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

4299

4300 import_batch_id: Mapped[Optional[str]] = mapped_column(String(36), nullable=True)

4301 federation_source: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

4302 version: Mapped[int] = mapped_column(Integer, default=1, nullable=False)

4303

4304 metrics: Mapped[List["ServerMetric"]] = relationship("ServerMetric", back_populates="server", cascade="all, delete-orphan")

4305 metrics_hourly: Mapped[List["ServerMetricsHourly"]] = relationship(

4306 "ServerMetricsHourly",

4307 primaryjoin="Server.id == foreign(ServerMetricsHourly.server_id)",

4308 viewonly=True,

4309 )

4310

4311 # Many-to-many relationships for associated items

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

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

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

4315 a2a_agents: Mapped[List["A2AAgent"]] = relationship("A2AAgent", secondary=server_a2a_association, back_populates="servers")

4316

4317 # API token relationships

4318 scoped_tokens: Mapped[List["EmailApiToken"]] = relationship("EmailApiToken", back_populates="server")

4319

4320 def _metrics_loaded(self) -> bool:

4321 """Check if metrics relationship is loaded without triggering lazy load.

4322

4323 Returns:

4324 bool: True if metrics are loaded, False otherwise.

4325 """

4326 return "metrics" in sa_inspect(self).dict

4327

4328 def _get_metric_counts(self) -> tuple[int, int, int]:

4329 """Get total, successful, and failed metric counts in a single operation.

4330

4331 When metrics are already loaded, computes from memory in O(n).

4332 When not loaded, uses a single SQL query with conditional aggregation.

4333

4334 Returns:

4335 tuple[int, int, int]: (total, successful, failed) counts.

4336 """

4337 if self._metrics_loaded():

4338 total = 0

4339 successful = 0

4340 for m in self.metrics:

4341 total += 1

4342 if m.is_success:

4343 successful += 1

4344 return (total, successful, total - successful)

4345

4346 # Third-Party

4347 from sqlalchemy import case # pylint: disable=import-outside-toplevel

4348 from sqlalchemy.orm import object_session # pylint: disable=import-outside-toplevel

4349

4350 session = object_session(self)

4351 if session is None:

4352 return (0, 0, 0)

4353

4354 result = (

4355 session.query(

4356 func.count(ServerMetric.id), # pylint: disable=not-callable

4357 func.sum(case((ServerMetric.is_success.is_(True), 1), else_=0)),

4358 )

4359 .filter(ServerMetric.server_id == self.id)

4360 .one()

4361 )

4362

4363 total = result[0] or 0

4364 successful = result[1] or 0

4365 return (total, successful, total - successful)

4366

4367 @hybrid_property

4368 def execution_count(self) -> int:

4369 """Number of ServerMetric records associated with this server instance.

4370

4371 Returns:

4372 int: Count of ServerMetric records for this server.

4373 """

4374 return self._get_metric_counts()[0]

4375

4376 @execution_count.expression

4377 @classmethod

4378 def execution_count(cls) -> Any:

4379 """SQL expression that counts ServerMetric rows for this server.

4380

4381 Returns:

4382 Any: SQLAlchemy labeled count expression for server metrics.

4383 """

4384 return select(func.count(ServerMetric.id)).where(ServerMetric.server_id == cls.id).correlate(cls).scalar_subquery().label("execution_count") # pylint: disable=not-callable

4385

4386 @property

4387 def successful_executions(self) -> int:

4388 """Count of successful server invocations.

4389

4390 Returns:

4391 int: The count of successful server invocations.

4392 """

4393 return self._get_metric_counts()[1]

4394

4395 @property

4396 def failed_executions(self) -> int:

4397 """Count of failed server invocations.

4398

4399 Returns:

4400 int: The count of failed server invocations.

4401 """

4402 return self._get_metric_counts()[2]

4403

4404 @property

4405 def failure_rate(self) -> float:

4406 """Failure rate as a float between 0 and 1.

4407

4408 Returns:

4409 float: The failure rate as a value between 0 and 1.

4410 """

4411 total, _, failed = self._get_metric_counts()

4412 return failed / total if total > 0 else 0.0

4413

4414 @property

4415 def min_response_time(self) -> Optional[float]:

4416 """Minimum response time among all server invocations.

4417

4418 Returns None if metrics are not loaded. Note: counts may be non-zero

4419 (via SQL) while timing is None. Use service layer converters for

4420 consistent metrics, or preload metrics via selectinload.

4421

4422 Returns:

4423 Optional[float]: The minimum response time, or None.

4424 """

4425 if not self._metrics_loaded():

4426 return None

4427 times: List[float] = [m.response_time for m in self.metrics]

4428 return min(times) if times else None

4429

4430 @property

4431 def max_response_time(self) -> Optional[float]:

4432 """Maximum response time among all server invocations.

4433

4434 Returns None if metrics are not loaded. Note: counts may be non-zero

4435 (via SQL) while timing is None. Use service layer converters for

4436 consistent metrics, or preload metrics via selectinload.

4437

4438 Returns:

4439 Optional[float]: The maximum response time, or None.

4440 """

4441 if not self._metrics_loaded():

4442 return None

4443 times: List[float] = [m.response_time for m in self.metrics]

4444 return max(times) if times else None

4445

4446 @property

4447 def avg_response_time(self) -> Optional[float]:

4448 """Average response time among all server invocations.

4449

4450 Returns None if metrics are not loaded. Note: counts may be non-zero

4451 (via SQL) while timing is None. Use service layer converters for

4452 consistent metrics, or preload metrics via selectinload.

4453

4454 Returns:

4455 Optional[float]: The average response time, or None.

4456 """

4457 if not self._metrics_loaded():

4458 return None

4459 times: List[float] = [m.response_time for m in self.metrics]

4460 return sum(times) / len(times) if times else None

4461

4462 @property

4463 def last_execution_time(self) -> Optional[datetime]:

4464 """Timestamp of the most recent server invocation.

4465

4466 Returns None if metrics are not loaded. Note: counts may be non-zero

4467 (via SQL) while timing is None. Use service layer converters for

4468 consistent metrics, or preload metrics via selectinload.

4469

4470 Returns:

4471 Optional[datetime]: The timestamp of the most recent invocation, or None.

4472 """

4473 if not self._metrics_loaded():

4474 return None

4475 if not self.metrics:

4476 return None

4477 return max(m.timestamp for m in self.metrics)

4478

4479 @property

4480 def metrics_summary(self) -> Dict[str, Any]:

4481 """Aggregated metrics for the server combining raw and hourly data without double-counting.

4482

4483 When metrics are loaded: computes from memory (raw + hourly)

4484 When not loaded: uses SQL queries with time partitioning

4485

4486 Returns:

4487 Dict[str, Any]: Dictionary containing aggregated metrics:

4488 - total_executions, successful_executions, failed_executions

4489 - failure_rate, min/max/avg_response_time, last_execution_time

4490 """

4491 # Try in-memory path first

4492 if self._metrics_loaded():

4493 try:

4494 hourly_metrics = self.metrics_hourly

4495 except AttributeError:

4496 hourly_metrics = []

4497 return _compute_metrics_summary(raw_metrics=self.metrics, hourly_metrics=hourly_metrics)

4498

4499 # SQL query path

4500 # Third-Party

4501 from sqlalchemy.orm import object_session # pylint: disable=import-outside-toplevel

4502

4503 session = object_session(self)

4504 if session is None:

4505 return {

4506 "total_executions": 0,

4507 "successful_executions": 0,

4508 "failed_executions": 0,

4509 "failure_rate": 0.0,

4510 "min_response_time": None,

4511 "max_response_time": None,

4512 "avg_response_time": None,

4513 "last_execution_time": None,

4514 }

4515

4516 return _compute_metrics_summary(

4517 raw_metrics=None,

4518 hourly_metrics=None,

4519 session=session,

4520 entity_id=self.id,

4521 raw_metric_class=ServerMetric,

4522 hourly_metric_class=ServerMetricsHourly,

4523 )

4524

4525 # Team scoping fields for resource organization

4526 team_id: Mapped[Optional[str]] = mapped_column(String(36), ForeignKey("email_teams.id", ondelete="SET NULL"), nullable=True)

4527 owner_email: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

4528 visibility: Mapped[str] = mapped_column(String(20), nullable=False, default="public")

4529

4530 # OAuth 2.0 configuration for RFC 9728 Protected Resource Metadata

4531 # When enabled, MCP clients can authenticate using OAuth with browser-based IDP SSO

4532 oauth_enabled: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)

4533 oauth_config: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True)

4534

4535 # Relationship for loading team names (only active teams)

4536 # Uses default lazy loading - team name is only loaded when accessed

4537 # For list/admin views, use explicit joinedload(DbServer.email_team) for single-query loading

4538 # This avoids adding overhead to hot paths that don't need team names

4539 email_team: Mapped[Optional["EmailTeam"]] = relationship(

4540 "EmailTeam",

4541 primaryjoin="and_(Server.team_id == EmailTeam.id, EmailTeam.is_active == True)",

4542 foreign_keys=[team_id],

4543 )

4544

4545 @property

4546 def team(self) -> Optional[str]:

4547 """Return the team name from the `email_team` relationship.

4548

4549 Returns:

4550 Optional[str]: The team name if the server belongs to an active team, otherwise None.

4551 """

4552 return self.email_team.name if self.email_team else None

4553

4554 __table_args__ = (

4555 UniqueConstraint("team_id", "owner_email", "name", name="uq_team_owner_name_server"),

4556 Index("idx_servers_created_at_id", "created_at", "id"),

4557 )

4558

4559

4560class Gateway(Base):

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

4562

4563 __tablename__ = "gateways"

4564

4565 id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: uuid.uuid4().hex)

4566 name: Mapped[str] = mapped_column(String(255), nullable=False)

4567 slug: Mapped[str] = mapped_column(String(255), nullable=False)

4568 url: Mapped[str] = mapped_column(String(767), nullable=False)

4569 description: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

4570 transport: Mapped[str] = mapped_column(String(20), default="SSE")

4571 capabilities: Mapped[Dict[str, Any]] = mapped_column(JSON)

4572 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now)

4573 updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, onupdate=utc_now)

4574 enabled: Mapped[bool] = mapped_column(default=True)

4575 reachable: Mapped[bool] = mapped_column(default=True)

4576 last_seen: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), nullable=True)

4577 tags: Mapped[List[Dict[str, str]]] = mapped_column(JSON, default=list, nullable=False)

4578

4579 # Comprehensive metadata for audit tracking

4580 created_by: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

4581 created_from_ip: Mapped[Optional[str]] = mapped_column(String(45), nullable=True)

4582 created_via: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)

4583 created_user_agent: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

4584

4585 modified_by: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

4586 modified_from_ip: Mapped[Optional[str]] = mapped_column(String(45), nullable=True)

4587 modified_via: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)

4588 modified_user_agent: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

4589

4590 import_batch_id: Mapped[Optional[str]] = mapped_column(String(36), nullable=True)

4591 federation_source: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

4592 version: Mapped[int] = mapped_column(Integer, default=1, nullable=False)

4593

4594 # Header passthrough configuration

4595 passthrough_headers: Mapped[Optional[List[str]]] = mapped_column(JSON, nullable=True) # Store list of strings as JSON array

4596

4597 # CA certificate

4598 ca_certificate: Mapped[Optional[bytes]] = mapped_column(Text, nullable=True)

4599 ca_certificate_sig: Mapped[Optional[str]] = mapped_column(String(64), nullable=True)

4600 signing_algorithm: Mapped[Optional[str]] = mapped_column(String(20), nullable=True, default="ed25519") # e.g., "sha256"

4601

4602 # mTLS client certificate/key for upstream gateway authentication

4603 client_cert: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

4604 client_key: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

4605

4606 # Relationship with local tools this gateway provides

4607 tools: Mapped[List["Tool"]] = relationship(back_populates="gateway", foreign_keys="Tool.gateway_id", cascade="all, delete-orphan", passive_deletes=True)

4608

4609 # Relationship with local prompts this gateway provides

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

4611

4612 # Relationship with local resources this gateway provides

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

4614

4615 # # Tools federated from this gateway

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

4617

4618 # # Prompts federated from this resource

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

4620

4621 # # Prompts federated from this gateway

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

4623

4624 # Authorizations

4625 auth_type: Mapped[Optional[str]] = mapped_column(String(20), default=None) # "basic", "bearer", "authheaders", "oauth", "query_param" or None

4626 auth_value: Mapped[Optional[Dict[str, str]]] = mapped_column(JSON)

4627 auth_query_params: Mapped[Optional[Dict[str, str]]] = mapped_column(

4628 JSON,

4629 nullable=True,

4630 default=None,

4631 comment="Encrypted query parameters for auth. Format: {'param_name': 'encrypted_value'}",

4632 )

4633

4634 # OAuth configuration

4635 oauth_config: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True, comment="OAuth 2.0 configuration including grant_type, client_id, encrypted client_secret, URLs, and scopes")

4636

4637 # Team scoping fields for resource organization

4638 team_id: Mapped[Optional[str]] = mapped_column(String(36), ForeignKey("email_teams.id", ondelete="SET NULL"), nullable=True)

4639 owner_email: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

4640 visibility: Mapped[str] = mapped_column(String(20), nullable=False, default="public")

4641

4642 # Relationship for loading team names (only active teams)

4643 # Uses default lazy loading - team name is only loaded when accessed

4644 # For list/admin views, use explicit joinedload(DbGateway.email_team) for single-query loading

4645 # This avoids adding overhead to hot paths that don't need team names

4646 email_team: Mapped[Optional["EmailTeam"]] = relationship(

4647 "EmailTeam",

4648 primaryjoin="and_(Gateway.team_id == EmailTeam.id, EmailTeam.is_active == True)",

4649 foreign_keys=[team_id],

4650 )

4651

4652 @property

4653 def team(self) -> Optional[str]:

4654 """Return the team name from the `email_team` relationship.

4655

4656 Returns:

4657 Optional[str]: The team name if the gateway belongs to an active team, otherwise None.

4658 """

4659 return self.email_team.name if self.email_team else None

4660

4661 # Per-gateway refresh configuration

4662 refresh_interval_seconds: Mapped[Optional[int]] = mapped_column(Integer, nullable=True, comment="Per-gateway refresh interval in seconds; NULL uses global default")

4663 last_refresh_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), nullable=True, comment="Timestamp of the last successful tools/resources/prompts refresh")

4664

4665 # Gateway mode: 'cache' (default) or 'direct_proxy'

4666 # - 'cache': Tools/resources/prompts are cached in database upon gateway registration (current behavior)

4667 # - 'direct_proxy': All RPC calls are proxied directly to remote MCP server with no database caching

4668 gateway_mode: Mapped[str] = mapped_column(String(20), nullable=False, default="cache", comment="Gateway mode: 'cache' (database caching) or 'direct_proxy' (pass-through mode)")

4669

4670 # Relationship with OAuth tokens

4671 oauth_tokens: Mapped[List["OAuthToken"]] = relationship("OAuthToken", back_populates="gateway", cascade="all, delete-orphan")

4672

4673 # Relationship with registered OAuth clients (DCR)

4674

4675 registered_oauth_clients: Mapped[List["RegisteredOAuthClient"]] = relationship("RegisteredOAuthClient", back_populates="gateway", cascade="all, delete-orphan")

4676

4677 __table_args__ = (

4678 UniqueConstraint("team_id", "owner_email", "slug", name="uq_team_owner_slug_gateway"),

4679 Index("idx_gateways_created_at_id", "created_at", "id"),

4680 )

4681

4682

4683@event.listens_for(Gateway, "after_update")

4684def update_tool_names_on_gateway_update(_mapper, connection, target):

4685 """

4686 If a Gateway's name is updated, efficiently update all of its

4687 child Tools' names with a single SQL statement.

4688

4689 Args:

4690 _mapper: Mapper

4691 connection: Connection

4692 target: Target

4693 """

4694 # 1. Check if the 'name' field was actually part of the update.

4695 # This is a concise way to see if the value has changed.

4696 if not get_history(target, "name").has_changes():

4697 return

4698

4699 logger.info("Gateway name changed for ID %s. Issuing bulk update for tools.", target.id)

4700

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

4702 tools_table = Tool.__table__

4703

4704 # 3. Prepare the new values

4705 new_gateway_slug = slugify(target.name)

4706 separator = settings.gateway_tool_name_separator

4707

4708 # 4. Construct a single, powerful UPDATE statement using SQLAlchemy Core.

4709 # This is highly efficient as it all happens in the database.

4710 stmt = (

4711 cast(Any, tools_table)

4712 .update()

4713 .where(tools_table.c.gateway_id == target.id)

4714 .values(name=new_gateway_slug + separator + tools_table.c.custom_name_slug)

4715 .execution_options(synchronize_session=False)

4716 )

4717

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

4719 connection.execute(stmt)

4720

4721

4722@event.listens_for(Gateway, "after_update")

4723def update_prompt_names_on_gateway_update(_mapper, connection, target):

4724 """Update prompt names when a gateway name changes.

4725

4726 Args:

4727 _mapper: SQLAlchemy mapper for the Gateway model.

4728 connection: Database connection for the update transaction.

4729 target: Gateway instance being updated.

4730 """

4731 if not get_history(target, "name").has_changes():

4732 return

4733

4734 logger.info("Gateway name changed for ID %s. Issuing bulk update for prompts.", target.id)

4735

4736 prompts_table = Prompt.__table__

4737 new_gateway_slug = slugify(target.name)

4738 separator = settings.gateway_tool_name_separator

4739

4740 stmt = (

4741 cast(Any, prompts_table)

4742 .update()

4743 .where(prompts_table.c.gateway_id == target.id)

4744 .values(name=new_gateway_slug + separator + prompts_table.c.custom_name_slug)

4745 .execution_options(synchronize_session=False)

4746 )

4747

4748 connection.execute(stmt)

4749

4750

4751class A2AAgent(Base):

4752 """

4753 ORM model for A2A (Agent-to-Agent) compatible agents.

4754

4755 A2A agents represent external AI agents that can be integrated into the gateway

4756 and exposed as tools within virtual servers. They support standardized

4757 Agent-to-Agent communication protocols for interoperability.

4758 """

4759

4760 __tablename__ = "a2a_agents"

4761

4762 id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: uuid.uuid4().hex)

4763 name: Mapped[str] = mapped_column(String(255), nullable=False)

4764 slug: Mapped[str] = mapped_column(String(255), nullable=False)

4765 description: Mapped[Optional[str]] = mapped_column(Text)

4766 endpoint_url: Mapped[str] = mapped_column(String(767), nullable=False)

4767 agent_type: Mapped[str] = mapped_column(String(50), nullable=False, default="generic") # e.g., "openai", "anthropic", "custom"

4768 protocol_version: Mapped[str] = mapped_column(String(10), nullable=False, default="1.0")

4769 capabilities: Mapped[Dict[str, Any]] = mapped_column(JSON, default=dict)

4770 # Configuration

4771 config: Mapped[Dict[str, Any]] = mapped_column(JSON, default=dict)

4772

4773 # Authorizations

4774 auth_type: Mapped[Optional[str]] = mapped_column(String(20), default=None) # "basic", "bearer", "authheaders", "oauth", "query_param" or None

4775 auth_value: Mapped[Optional[str]] = mapped_column(Text)

4776 auth_query_params: Mapped[Optional[Dict[str, str]]] = mapped_column(

4777 JSON,

4778 nullable=True,

4779 default=None,

4780 comment="Encrypted query parameters for auth. Format: {'param_name': 'encrypted_value'}",

4781 )

4782

4783 # OAuth configuration

4784 oauth_config: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True, comment="OAuth 2.0 configuration including grant_type, client_id, encrypted client_secret, URLs, and scopes")

4785

4786 # Header passthrough configuration

4787 passthrough_headers: Mapped[Optional[List[str]]] = mapped_column(JSON, nullable=True) # Store list of strings as JSON array

4788

4789 # Status and metadata

4790 enabled: Mapped[bool] = mapped_column(Boolean, default=True)

4791 reachable: Mapped[bool] = mapped_column(Boolean, default=True)

4792 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now)

4793 updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, onupdate=utc_now)

4794 last_interaction: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True))

4795

4796 # Tags for categorization

4797 tags: Mapped[List[str]] = mapped_column(JSON, default=list, nullable=False)

4798

4799 # Comprehensive metadata for audit tracking

4800 created_by: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

4801 created_from_ip: Mapped[Optional[str]] = mapped_column(String(45), nullable=True)

4802 created_via: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)

4803 created_user_agent: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

4804

4805 modified_by: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

4806 modified_from_ip: Mapped[Optional[str]] = mapped_column(String(45), nullable=True)

4807 modified_via: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)

4808 modified_user_agent: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

4809

4810 import_batch_id: Mapped[Optional[str]] = mapped_column(String(36), nullable=True)

4811 federation_source: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

4812 version: Mapped[int] = mapped_column(Integer, default=1, nullable=False)

4813

4814 # Team scoping fields for resource organization

4815 team_id: Mapped[Optional[str]] = mapped_column(String(36), ForeignKey("email_teams.id", ondelete="SET NULL"), nullable=True)

4816 owner_email: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

4817 visibility: Mapped[str] = mapped_column(String(20), nullable=False, default="public")

4818

4819 # Associated tool ID (A2A agents are automatically registered as tools)

4820 tool_id: Mapped[Optional[str]] = mapped_column(String(36), ForeignKey("tools.id", ondelete="SET NULL"), nullable=True)

4821

4822 # Relationships

4823 servers: Mapped[List["Server"]] = relationship("Server", secondary=server_a2a_association, back_populates="a2a_agents")

4824 tool: Mapped[Optional["Tool"]] = relationship("Tool", foreign_keys=[tool_id])

4825 metrics: Mapped[List["A2AAgentMetric"]] = relationship("A2AAgentMetric", back_populates="a2a_agent", cascade="all, delete-orphan")

4826 __table_args__ = (

4827 UniqueConstraint("team_id", "owner_email", "slug", name="uq_team_owner_slug_a2a_agent"),

4828 Index("idx_a2a_agents_created_at_id", "created_at", "id"),

4829 Index("idx_a2a_agents_tool_id", "tool_id"),

4830 )

4831

4832 # Relationship with OAuth tokens

4833 # oauth_tokens: Mapped[List["OAuthToken"]] = relationship("OAuthToken", back_populates="gateway", cascade="all, delete-orphan")

4834

4835 # Relationship with registered OAuth clients (DCR)

4836 # registered_oauth_clients: Mapped[List["RegisteredOAuthClient"]] = relationship("RegisteredOAuthClient", back_populates="gateway", cascade="all, delete-orphan")

4837

4838 def _metrics_loaded(self) -> bool:

4839 """Check if metrics relationship is loaded without triggering lazy load.

4840

4841 Returns:

4842 bool: True if metrics are loaded, False otherwise.

4843 """

4844 return "metrics" in sa_inspect(self).dict

4845

4846 @property

4847 def execution_count(self) -> int:

4848 """Total number of interactions with this agent.

4849 Returns 0 if metrics are not loaded (avoids lazy loading).

4850

4851 Returns:

4852 int: The total count of interactions.

4853 """

4854 if not self._metrics_loaded():

4855 return 0

4856 return len(self.metrics)

4857

4858 @property

4859 def successful_executions(self) -> int:

4860 """Number of successful interactions.

4861 Returns 0 if metrics are not loaded (avoids lazy loading).

4862

4863 Returns:

4864 int: The count of successful interactions.

4865 """

4866 if not self._metrics_loaded():

4867 return 0

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

4869

4870 @property

4871 def failed_executions(self) -> int:

4872 """Number of failed interactions.

4873 Returns 0 if metrics are not loaded (avoids lazy loading).

4874

4875 Returns:

4876 int: The count of failed interactions.

4877 """

4878 if not self._metrics_loaded():

4879 return 0

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

4881

4882 @property

4883 def failure_rate(self) -> float:

4884 """Failure rate as a percentage.

4885 Returns 0.0 if metrics are not loaded (avoids lazy loading).

4886

4887 Returns:

4888 float: The failure rate percentage.

4889 """

4890 if not self._metrics_loaded():

4891 return 0.0

4892 if not self.metrics:

4893 return 0.0

4894 return (self.failed_executions / len(self.metrics)) * 100

4895

4896 @property

4897 def avg_response_time(self) -> Optional[float]:

4898 """Average response time in seconds.

4899 Returns None if metrics are not loaded (avoids lazy loading).

4900

4901 Returns:

4902 Optional[float]: The average response time, or None if no metrics.

4903 """

4904 if not self._metrics_loaded():

4905 return None

4906 if not self.metrics:

4907 return None

4908 return sum(m.response_time for m in self.metrics) / len(self.metrics)

4909

4910 @property

4911 def last_execution_time(self) -> Optional[datetime]:

4912 """Timestamp of the most recent interaction.

4913 Returns None if metrics are not loaded (avoids lazy loading).

4914

4915 Returns:

4916 Optional[datetime]: The timestamp of the last interaction, or None if no metrics.

4917 """

4918 if not self._metrics_loaded():

4919 return None

4920 if not self.metrics:

4921 return None

4922 return max(m.timestamp for m in self.metrics)

4923

4924 def __repr__(self) -> str:

4925 """Return a string representation of the A2AAgent instance.

4926

4927 Returns:

4928 str: A formatted string containing the agent's ID, name, and type.

4929

4930 Examples:

4931 >>> agent = A2AAgent(id='123', name='test-agent', agent_type='custom')

4932 >>> repr(agent)

4933 "<A2AAgent(id='123', name='test-agent', agent_type='custom')>"

4934 """

4935 return f"<A2AAgent(id='{self.id}', name='{self.name}', agent_type='{self.agent_type}')>"

4936

4937

4938class GrpcService(Base):

4939 """

4940 ORM model for gRPC services with reflection-based discovery.

4941

4942 gRPC services represent external gRPC servers that can be automatically discovered

4943 via server reflection and exposed as MCP tools. The gateway translates between

4944 gRPC/Protobuf and MCP/JSON protocols.

4945 """

4946

4947 __tablename__ = "grpc_services"

4948

4949 id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: uuid.uuid4().hex)

4950 name: Mapped[str] = mapped_column(String(255), nullable=False, unique=True)

4951 slug: Mapped[str] = mapped_column(String(255), nullable=False, unique=True)

4952 description: Mapped[Optional[str]] = mapped_column(Text)

4953 target: Mapped[str] = mapped_column(String(767), nullable=False) # host:port format

4954

4955 # Configuration

4956 reflection_enabled: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)

4957 tls_enabled: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)

4958 tls_cert_path: Mapped[Optional[str]] = mapped_column(String(767))

4959 tls_key_path: Mapped[Optional[str]] = mapped_column(String(767))

4960 grpc_metadata: Mapped[Dict[str, str]] = mapped_column(JSON, default=dict) # gRPC metadata headers

4961

4962 # Status

4963 enabled: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)

4964 reachable: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)

4965

4966 # Discovery results from reflection

4967 service_count: Mapped[int] = mapped_column(Integer, default=0, nullable=False)

4968 method_count: Mapped[int] = mapped_column(Integer, default=0, nullable=False)

4969 discovered_services: Mapped[Dict[str, Any]] = mapped_column(JSON, default=dict) # Service descriptors

4970 last_reflection: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True))

4971

4972 # Tags for categorization

4973 tags: Mapped[List[str]] = mapped_column(JSON, default=list, nullable=False)

4974

4975 # Timestamps

4976 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now)

4977 updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, onupdate=utc_now)

4978

4979 # Comprehensive metadata for audit tracking

4980 created_by: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

4981 created_from_ip: Mapped[Optional[str]] = mapped_column(String(45), nullable=True)

4982 created_via: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)

4983 created_user_agent: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

4984

4985 modified_by: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

4986 modified_from_ip: Mapped[Optional[str]] = mapped_column(String(45), nullable=True)

4987 modified_via: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)

4988 modified_user_agent: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

4989

4990 import_batch_id: Mapped[Optional[str]] = mapped_column(String(36), nullable=True)

4991 federation_source: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

4992 version: Mapped[int] = mapped_column(Integer, default=1, nullable=False)

4993

4994 # Team scoping fields for resource organization

4995 team_id: Mapped[Optional[str]] = mapped_column(String(36), ForeignKey("email_teams.id", ondelete="SET NULL"), nullable=True)

4996 owner_email: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

4997 visibility: Mapped[str] = mapped_column(String(20), nullable=False, default="public")

4998

4999 def __repr__(self) -> str:

5000 """Return a string representation of the GrpcService instance.

5001

5002 Returns:

5003 str: A formatted string containing the service's ID, name, and target.

5004 """

5005 return f"<GrpcService(id='{self.id}', name='{self.name}', target='{self.target}')>"

5006

5007

5008class SessionRecord(Base):

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

5010

5011 __tablename__ = "mcp_sessions"

5012

5013 session_id: Mapped[str] = mapped_column(String(255), primary_key=True)

5014 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now) # pylint: disable=not-callable

5015 last_accessed: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, onupdate=utc_now) # pylint: disable=not-callable

5016 data: Mapped[str] = mapped_column(Text, nullable=True)

5017

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

5019

5020

5021class SessionMessageRecord(Base):

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

5023

5024 __tablename__ = "mcp_messages"

5025

5026 id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)

5027 session_id: Mapped[str] = mapped_column(String(255), ForeignKey("mcp_sessions.session_id"))

5028 message: Mapped[str] = mapped_column(Text, nullable=True)

5029 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now) # pylint: disable=not-callable

5030 last_accessed: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, onupdate=utc_now) # pylint: disable=not-callable

5031

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

5033

5034

5035class OAuthToken(Base):

5036 """ORM model for OAuth access and refresh tokens with user association."""

5037

5038 __tablename__ = "oauth_tokens"

5039

5040 id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: uuid.uuid4().hex)

5041 gateway_id: Mapped[str] = mapped_column(String(36), ForeignKey("gateways.id", ondelete="CASCADE"), nullable=False)

5042 user_id: Mapped[str] = mapped_column(String(255), nullable=False) # OAuth provider's user ID

5043 app_user_email: Mapped[str] = mapped_column(String(255), ForeignKey("email_users.email", ondelete="CASCADE"), nullable=False) # ContextForge user

5044 access_token: Mapped[str] = mapped_column(EncryptedText(), nullable=False)

5045 refresh_token: Mapped[Optional[str]] = mapped_column(EncryptedText(), nullable=True)

5046 token_type: Mapped[str] = mapped_column(String(50), default="Bearer")

5047 expires_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), nullable=True)

5048 scopes: Mapped[Optional[List[str]]] = mapped_column(JSON, nullable=True)

5049 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now)

5050 updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, onupdate=utc_now)

5051

5052 # Relationships

5053 gateway: Mapped["Gateway"] = relationship("Gateway", back_populates="oauth_tokens")

5054 app_user: Mapped["EmailUser"] = relationship("EmailUser", foreign_keys=[app_user_email])

5055

5056 # Unique constraint: one token per user per gateway

5057 __table_args__ = (UniqueConstraint("gateway_id", "app_user_email", name="uq_oauth_gateway_user"),)

5058

5059

5060class OAuthState(Base):

5061 """ORM model for OAuth authorization states with TTL for CSRF protection."""

5062

5063 __tablename__ = "oauth_states"

5064

5065 id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: uuid.uuid4().hex)

5066 gateway_id: Mapped[str] = mapped_column(String(36), ForeignKey("gateways.id", ondelete="CASCADE"), nullable=False)

5067 state: Mapped[str] = mapped_column(String(500), nullable=False, unique=True) # The state parameter

5068 code_verifier: Mapped[Optional[str]] = mapped_column(String(128), nullable=True) # PKCE code verifier (RFC 7636)

5069 app_user_email: Mapped[Optional[str]] = mapped_column(String(255), nullable=True) # Requesting user context for token association

5070 expires_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)

5071 used: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)

5072 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now)

5073

5074 # Relationships

5075 gateway: Mapped["Gateway"] = relationship("Gateway")

5076

5077 # Index for efficient lookups

5078 __table_args__ = (Index("idx_oauth_state_lookup", "gateway_id", "state"),)

5079

5080

5081class RegisteredOAuthClient(Base):

5082 """Stores dynamically registered OAuth clients (RFC 7591 client mode).

5083

5084 This model maintains client credentials obtained through Dynamic Client

5085 Registration with upstream Authorization Servers.

5086 """

5087

5088 __tablename__ = "registered_oauth_clients"

5089

5090 id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()))

5091 gateway_id: Mapped[str] = mapped_column(String(36), ForeignKey("gateways.id", ondelete="CASCADE"), nullable=False, index=True)

5092

5093 # Registration details

5094 issuer: Mapped[str] = mapped_column(String(500), nullable=False) # AS issuer URL

5095 client_id: Mapped[str] = mapped_column(String(500), nullable=False)

5096 client_secret_encrypted: Mapped[Optional[str]] = mapped_column(Text, nullable=True) # Encrypted

5097

5098 # RFC 7591 fields

5099 redirect_uris: Mapped[str] = mapped_column(Text, nullable=False) # JSON array

5100 grant_types: Mapped[str] = mapped_column(Text, nullable=False) # JSON array

5101 response_types: Mapped[Optional[str]] = mapped_column(Text, nullable=True) # JSON array

5102 scope: Mapped[Optional[str]] = mapped_column(String(1000), nullable=True)

5103 token_endpoint_auth_method: Mapped[str] = mapped_column(String(50), default="client_secret_basic")

5104

5105 # Registration management (RFC 7591 section 4)

5106 registration_client_uri: Mapped[Optional[str]] = mapped_column(String(500), nullable=True)

5107 registration_access_token_encrypted: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

5108

5109 # Metadata

5110 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now)

5111 expires_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), nullable=True)

5112 is_active: Mapped[bool] = mapped_column(Boolean, default=True)

5113

5114 # Relationships

5115 gateway: Mapped["Gateway"] = relationship("Gateway", back_populates="registered_oauth_clients")

5116

5117 # Unique constraint: one registration per gateway+issuer

5118 __table_args__ = (Index("idx_gateway_issuer", "gateway_id", "issuer", unique=True),)

5119

5120

5121class EmailApiToken(Base):

5122 """Email user API token model for token catalog management.

5123

5124 This model provides comprehensive API token management with scoping,

5125 revocation, and usage tracking for email-based users.

5126

5127 Attributes:

5128 id (str): Unique token identifier

5129 user_email (str): Owner's email address

5130 team_id (str): Team the token is associated with (required for team-based access)

5131 name (str): Human-readable token name

5132 jti (str): JWT ID for revocation checking

5133 token_hash (str): Hashed token value for security

5134 server_id (str): Optional server scope limitation

5135 resource_scopes (List[str]): Permission scopes like ['tools.read']

5136 ip_restrictions (List[str]): IP address/CIDR restrictions

5137 time_restrictions (dict): Time-based access restrictions

5138 usage_limits (dict): Rate limiting and usage quotas

5139 created_at (datetime): Token creation timestamp

5140 expires_at (datetime): Optional expiry timestamp

5141 last_used (datetime): Last usage timestamp

5142 is_active (bool): Active status flag

5143 description (str): Token description

5144 tags (List[str]): Organizational tags

5145

5146 Examples:

5147 >>> token = EmailApiToken(

5148 ... user_email="alice@example.com",

5149 ... name="Production API Access",

5150 ... server_id="prod-server-123",

5151 ... resource_scopes=["tools.read", "resources.read"],

5152 ... description="Read-only access to production tools"

5153 ... )

5154 >>> token.is_scoped_to_server("prod-server-123")

5155 True

5156 >>> token.has_permission("tools.read")

5157 True

5158 """

5159

5160 __tablename__ = "email_api_tokens"

5161

5162 # Core identity fields

5163 id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()))

5164 user_email: Mapped[str] = mapped_column(String(255), ForeignKey("email_users.email", ondelete="CASCADE"), nullable=False, index=True)

5165 team_id: Mapped[Optional[str]] = mapped_column(String(36), ForeignKey("email_teams.id", ondelete="SET NULL"), nullable=True, index=True)

5166 name: Mapped[str] = mapped_column(String(255), nullable=False)

5167 jti: Mapped[str] = mapped_column(String(36), unique=True, nullable=False, default=lambda: str(uuid.uuid4()))

5168 token_hash: Mapped[str] = mapped_column(String(255), nullable=False)

5169

5170 # Scoping fields

5171 server_id: Mapped[Optional[str]] = mapped_column(String(36), ForeignKey("servers.id", ondelete="CASCADE"), nullable=True)

5172 resource_scopes: Mapped[Optional[List[str]]] = mapped_column(JSON, nullable=True, default=list)

5173 ip_restrictions: Mapped[Optional[List[str]]] = mapped_column(JSON, nullable=True, default=list)

5174 time_restrictions: Mapped[Optional[dict]] = mapped_column(JSON, nullable=True, default=dict)

5175 usage_limits: Mapped[Optional[dict]] = mapped_column(JSON, nullable=True, default=dict)

5176

5177 # Lifecycle fields

5178 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, nullable=False)

5179 expires_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), nullable=True)

5180 last_used: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), nullable=True)

5181 is_active: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)

5182

5183 # Metadata fields

5184 description: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

5185 tags: Mapped[Optional[List[str]]] = mapped_column(JSON, nullable=True, default=list)

5186

5187 # Unique constraint for user+name+team_id combination (per-team scope).

5188 # The composite UniqueConstraint handles non-NULL team_id rows. SQL NULL != NULL

5189 # semantics mean it cannot protect global-scope tokens (team_id IS NULL), so we add

5190 # a partial unique index for that case — matching the pattern used by resources/prompts.

5191 __table_args__ = (

5192 UniqueConstraint("user_email", "name", "team_id", name="uq_email_api_tokens_user_name_team"),

5193 Index("uq_email_api_tokens_user_name_global", "user_email", "name", unique=True, postgresql_where=text("team_id IS NULL"), sqlite_where=text("team_id IS NULL")),

5194 Index("idx_email_api_tokens_user_email", "user_email"),

5195 Index("idx_email_api_tokens_jti", "jti"),

5196 Index("idx_email_api_tokens_expires_at", "expires_at"),

5197 Index("idx_email_api_tokens_is_active", "is_active"),

5198 )

5199

5200 # Relationships

5201 user: Mapped["EmailUser"] = relationship("EmailUser", back_populates="api_tokens")

5202 team: Mapped[Optional["EmailTeam"]] = relationship("EmailTeam", back_populates="api_tokens")

5203 server: Mapped[Optional["Server"]] = relationship("Server", back_populates="scoped_tokens")

5204

5205 def is_scoped_to_server(self, server_id: str) -> bool:

5206 """Check if token is scoped to a specific server.

5207

5208 Args:

5209 server_id: Server ID to check against.

5210

5211 Returns:

5212 bool: True if token is scoped to the server, False otherwise.

5213 """

5214 return self.server_id == server_id if self.server_id else False

5215

5216 def has_permission(self, permission: str) -> bool:

5217 """Check if token has a specific permission.

5218

5219 Args:

5220 permission: Permission string to check for.

5221

5222 Returns:

5223 bool: True if token has the permission, False otherwise.

5224 """

5225 return permission in (self.resource_scopes or [])

5226

5227 def is_team_token(self) -> bool:

5228 """Check if this is a team-based token.

5229

5230 Returns:

5231 bool: True if token is associated with a team, False otherwise.

5232 """

5233 return self.team_id is not None

5234

5235 def get_effective_permissions(self) -> List[str]:

5236 """Get effective permissions for this token.

5237

5238 For team tokens, this should inherit team permissions.

5239 For personal tokens, this uses the resource_scopes.

5240

5241 Returns:

5242 List[str]: List of effective permissions for this token.

5243 """

5244 if self.is_team_token() and self.team:

5245 # For team tokens, we would inherit team permissions

5246 # This would need to be implemented based on your RBAC system

5247 return self.resource_scopes or []

5248 return self.resource_scopes or []

5249

5250 def is_expired(self) -> bool:

5251 """Check if token is expired.

5252

5253 Returns:

5254 bool: True if token is expired, False otherwise.

5255 """

5256 if not self.expires_at:

5257 return False

5258 expires_at = self.expires_at

5259 if expires_at.tzinfo is None:

5260 expires_at = expires_at.replace(tzinfo=timezone.utc)

5261 return utc_now() > expires_at

5262

5263 def is_valid(self) -> bool:

5264 """Check if token is valid (active and not expired).

5265

5266 Returns:

5267 bool: True if token is valid, False otherwise.

5268 """

5269 return self.is_active and not self.is_expired()

5270

5271

5272class TokenUsageLog(Base):

5273 """Token usage logging for analytics and security monitoring.

5274

5275 This model tracks every API request made with email API tokens

5276 for security auditing and usage analytics.

5277

5278 Attributes:

5279 id (int): Auto-incrementing log ID

5280 token_jti (str): Token JWT ID reference

5281 user_email (str): Token owner's email

5282 timestamp (datetime): Request timestamp

5283 endpoint (str): API endpoint accessed

5284 method (str): HTTP method used

5285 ip_address (str): Client IP address

5286 user_agent (str): Client user agent

5287 status_code (int): HTTP response status

5288 response_time_ms (int): Response time in milliseconds

5289 blocked (bool): Whether request was blocked

5290 block_reason (str): Reason for blocking if applicable

5291

5292 Examples:

5293 >>> log = TokenUsageLog(

5294 ... token_jti="token-uuid-123",

5295 ... user_email="alice@example.com",

5296 ... endpoint="/tools",

5297 ... method="GET",

5298 ... ip_address="192.168.1.100",

5299 ... status_code=200,

5300 ... response_time_ms=45

5301 ... )

5302 """

5303

5304 __tablename__ = "token_usage_logs"

5305

5306 # Primary key

5307 id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)

5308

5309 # Token reference

5310 token_jti: Mapped[str] = mapped_column(String(36), nullable=False, index=True)

5311 user_email: Mapped[str] = mapped_column(String(255), nullable=False, index=True)

5312

5313 # Timestamp

5314 timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, nullable=False, index=True)

5315

5316 # Request details

5317 endpoint: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

5318 method: Mapped[Optional[str]] = mapped_column(String(10), nullable=True)

5319 ip_address: Mapped[Optional[str]] = mapped_column(String(45), nullable=True) # IPv6 max length

5320 user_agent: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

5321

5322 # Response details

5323 status_code: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)

5324 response_time_ms: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)

5325

5326 # Security fields

5327 blocked: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)

5328 block_reason: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

5329

5330 # Indexes for performance

5331 __table_args__ = (

5332 Index("idx_token_usage_logs_token_jti_timestamp", "token_jti", "timestamp"),

5333 Index("idx_token_usage_logs_user_email_timestamp", "user_email", "timestamp"),

5334 )

5335

5336

5337class TokenRevocation(Base):

5338 """Token revocation blacklist for immediate token invalidation.

5339

5340 This model maintains a blacklist of revoked JWT tokens to provide

5341 immediate token invalidation capabilities.

5342

5343 Attributes:

5344 jti (str): JWT ID (primary key)

5345 revoked_at (datetime): Revocation timestamp

5346 revoked_by (str): Email of user who revoked the token

5347 reason (str): Optional reason for revocation

5348

5349 Examples:

5350 >>> revocation = TokenRevocation(

5351 ... jti="token-uuid-123",

5352 ... revoked_by="admin@example.com",

5353 ... reason="Security compromise"

5354 ... )

5355 """

5356

5357 __tablename__ = "token_revocations"

5358

5359 # JWT ID as primary key

5360 jti: Mapped[str] = mapped_column(String(36), primary_key=True)

5361

5362 # Revocation details

5363 revoked_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, nullable=False)

5364 revoked_by: Mapped[str] = mapped_column(String(255), ForeignKey("email_users.email"), nullable=False)

5365 reason: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

5366

5367 # Relationship

5368 revoker: Mapped["EmailUser"] = relationship("EmailUser")

5369

5370

5371class SSOProvider(Base):

5372 """SSO identity provider configuration for OAuth2/OIDC authentication.

5373

5374 Stores configuration and credentials for external identity providers

5375 like GitHub, Google, IBM Security Verify, Okta, Microsoft Entra ID,

5376 and any generic OIDC-compliant provider (Keycloak, Auth0, Authentik, etc.).

5377

5378 Attributes:

5379 id (str): Unique provider ID (e.g., 'github', 'google', 'ibm_verify')

5380 name (str): Human-readable provider name

5381 display_name (str): Display name for UI

5382 provider_type (str): Protocol type ('oauth2', 'oidc')

5383 is_enabled (bool): Whether provider is active

5384 client_id (str): OAuth client ID

5385 client_secret_encrypted (str): Encrypted client secret

5386 authorization_url (str): OAuth authorization endpoint

5387 token_url (str): OAuth token endpoint

5388 userinfo_url (str): User info endpoint

5389 issuer (str): OIDC issuer (optional)

5390 jwks_uri (str): OIDC JWKS endpoint for token signature verification (optional)

5391 trusted_domains (List[str]): Auto-approved email domains

5392 scope (str): OAuth scope string

5393 auto_create_users (bool): Auto-create users on first login

5394 team_mapping (dict): Organization/domain to team mapping rules

5395 created_at (datetime): Provider creation timestamp

5396 updated_at (datetime): Last configuration update

5397

5398 Examples:

5399 >>> provider = SSOProvider(

5400 ... id="github",

5401 ... name="github",

5402 ... display_name="GitHub",

5403 ... provider_type="oauth2",

5404 ... client_id="gh_client_123",

5405 ... authorization_url="https://github.com/login/oauth/authorize",

5406 ... token_url="https://github.com/login/oauth/access_token",

5407 ... userinfo_url="https://api.github.com/user",

5408 ... scope="user:email"

5409 ... )

5410 """

5411

5412 __tablename__ = "sso_providers"

5413

5414 # Provider identification

5415 id: Mapped[str] = mapped_column(String(50), primary_key=True) # github, google, ibm_verify, okta, keycloak, entra, or any custom ID

5416 name: Mapped[str] = mapped_column(String(100), nullable=False, unique=True)

5417 display_name: Mapped[str] = mapped_column(String(100), nullable=False)

5418 provider_type: Mapped[str] = mapped_column(String(20), nullable=False) # oauth2, oidc

5419 is_enabled: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)

5420

5421 # OAuth2/OIDC Configuration

5422 client_id: Mapped[str] = mapped_column(String(255), nullable=False)

5423 client_secret_encrypted: Mapped[str] = mapped_column(Text, nullable=False) # Encrypted storage

5424 authorization_url: Mapped[str] = mapped_column(String(500), nullable=False)

5425 token_url: Mapped[str] = mapped_column(String(500), nullable=False)

5426 userinfo_url: Mapped[str] = mapped_column(String(500), nullable=False)

5427 issuer: Mapped[Optional[str]] = mapped_column(String(500), nullable=True) # For OIDC

5428 jwks_uri: Mapped[Optional[str]] = mapped_column(String(500), nullable=True) # OIDC JWKS endpoint for token signature verification

5429

5430 # Provider Settings

5431 trusted_domains: Mapped[List[str]] = mapped_column(JSON, default=list, nullable=False)

5432 scope: Mapped[str] = mapped_column(String(200), default="openid profile email", nullable=False)

5433 auto_create_users: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)

5434 team_mapping: Mapped[dict] = mapped_column(JSON, default=dict, nullable=False)

5435

5436 # Provider-specific metadata (e.g., role mappings, claim configurations)

5437 provider_metadata: Mapped[dict] = mapped_column(JSON, default=dict, nullable=False)

5438

5439 # Timestamps

5440 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, nullable=False)

5441 updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, onupdate=utc_now, nullable=False)

5442

5443 def __repr__(self):

5444 """String representation of SSO provider.

5445

5446 Returns:

5447 String representation of the SSO provider instance

5448 """

5449 return f"<SSOProvider(id='{self.id}', name='{self.name}', enabled={self.is_enabled})>"

5450

5451

5452class SSOAuthSession(Base):

5453 """Tracks SSO authentication sessions and state.

5454

5455 Maintains OAuth state parameters and callback information during

5456 the SSO authentication flow for security and session management.

5457

5458 Attributes:

5459 id (str): Unique session ID (UUID)

5460 provider_id (str): Reference to SSO provider

5461 state (str): OAuth state parameter for CSRF protection

5462 code_verifier (str): PKCE code verifier (for OAuth 2.1)

5463 nonce (str): OIDC nonce parameter

5464 redirect_uri (str): OAuth callback URI

5465 expires_at (datetime): Session expiration time

5466 user_email (str): User email after successful auth (optional)

5467 created_at (datetime): Session creation timestamp

5468

5469 Examples:

5470 >>> session = SSOAuthSession(

5471 ... provider_id="github",

5472 ... state="csrf-state-token",

5473 ... redirect_uri="https://gateway.example.com/auth/sso-callback/github"

5474 ... )

5475 """

5476

5477 __tablename__ = "sso_auth_sessions"

5478

5479 # Session identification

5480 id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()))

5481 provider_id: Mapped[str] = mapped_column(String(50), ForeignKey("sso_providers.id"), nullable=False)

5482

5483 # OAuth/OIDC parameters

5484 state: Mapped[str] = mapped_column(String(128), nullable=False, unique=True) # CSRF protection

5485 code_verifier: Mapped[Optional[str]] = mapped_column(String(128), nullable=True) # PKCE

5486 nonce: Mapped[Optional[str]] = mapped_column(String(128), nullable=True) # OIDC

5487 redirect_uri: Mapped[str] = mapped_column(String(500), nullable=False)

5488

5489 # Session lifecycle

5490 expires_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: utc_now() + timedelta(minutes=10), nullable=False) # 10-minute expiration

5491 user_email: Mapped[Optional[str]] = mapped_column(String(255), ForeignKey("email_users.email"), nullable=True)

5492 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, nullable=False)

5493

5494 # Relationships

5495 provider: Mapped["SSOProvider"] = relationship("SSOProvider")

5496 user: Mapped[Optional["EmailUser"]] = relationship("EmailUser")

5497

5498 @property

5499 def is_expired(self) -> bool:

5500 """Check if SSO auth session has expired.

5501

5502 Returns:

5503 True if the session has expired, False otherwise

5504 """

5505 now = utc_now()

5506 expires = self.expires_at

5507

5508 # Handle timezone mismatch by converting naive datetime to UTC if needed

5509 if expires.tzinfo is None:

5510 # expires_at is timezone-naive, assume it's UTC

5511 expires = expires.replace(tzinfo=timezone.utc)

5512 elif now.tzinfo is None:

5513 # now is timezone-naive (shouldn't happen with utc_now, but just in case)

5514 now = now.replace(tzinfo=timezone.utc)

5515

5516 return now > expires

5517

5518 def __repr__(self):

5519 """String representation of SSO auth session.

5520

5521 Returns:

5522 str: String representation of the session object

5523 """

5524 return f"<SSOAuthSession(id='{self.id}', provider='{self.provider_id}', expired={self.is_expired})>"

5525

5526

5527# Event listeners for validation

5528def validate_tool_schema(mapper, connection, target):

5529 """

5530 Validate tool schema before insert/update.

5531

5532 Args:

5533 mapper: The mapper being used for the operation.

5534 connection: The database connection.

5535 target: The target object being validated.

5536

5537 Raises:

5538 ValueError: If the tool input schema is invalid.

5539

5540 """

5541 # You can use mapper and connection later, if required.

5542 _ = mapper

5543 _ = connection

5544

5545 allowed_validator_names = {

5546 "Draft4Validator",

5547 "Draft6Validator",

5548 "Draft7Validator",

5549 "Draft201909Validator",

5550 "Draft202012Validator",

5551 }

5552

5553 if hasattr(target, "input_schema"):

5554 schema = target.input_schema

5555 if schema is None:

5556 return

5557

5558 try:

5559 # If $schema is missing, default to Draft 2020-12 as per MCP spec.

5560 if schema.get("$schema") is None:

5561 validator_cls = jsonschema.Draft202012Validator

5562 else:

5563 validator_cls = jsonschema.validators.validator_for(schema)

5564

5565 if validator_cls.__name__ not in allowed_validator_names:

5566 logger.warning(f"Unsupported JSON Schema draft: {validator_cls.__name__}")

5567

5568 validator_cls.check_schema(schema)

5569 except jsonschema.exceptions.SchemaError as e:

5570 logger.warning(f"Invalid tool input schema: {str(e)}")

5571 if settings.json_schema_validation_strict:

5572 raise ValueError(f"Invalid tool input schema: {str(e)}") from e

5573

5574

5575def validate_tool_name(mapper, connection, target):

5576 """

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

5578

5579 Args:

5580 mapper: The mapper being used for the operation.

5581 connection: The database connection.

5582 target: The target object being validated.

5583

5584 Raises:

5585 ValueError: If the tool name contains invalid characters.

5586 """

5587 # You can use mapper and connection later, if required.

5588 _ = mapper

5589 _ = connection

5590 if hasattr(target, "name"):

5591 try:

5592 SecurityValidator.validate_tool_name(target.name)

5593 except ValueError as e:

5594 raise ValueError(f"Invalid tool name: {str(e)}") from e

5595

5596

5597def validate_prompt_schema(mapper, connection, target):

5598 """

5599 Validate prompt argument schema before insert/update.

5600

5601 Args:

5602 mapper: The mapper being used for the operation.

5603 connection: The database connection.

5604 target: The target object being validated.

5605

5606 Raises:

5607 ValueError: If the prompt argument schema is invalid.

5608 """

5609 # You can use mapper and connection later, if required.

5610 _ = mapper

5611 _ = connection

5612

5613 allowed_validator_names = {

5614 "Draft4Validator",

5615 "Draft6Validator",

5616 "Draft7Validator",

5617 "Draft201909Validator",

5618 "Draft202012Validator",

5619 }

5620

5621 if hasattr(target, "argument_schema"):

5622 schema = target.argument_schema

5623 if schema is None:

5624 return

5625

5626 try:

5627 # If $schema is missing, default to Draft 2020-12 as per MCP spec.

5628 if schema.get("$schema") is None:

5629 validator_cls = jsonschema.Draft202012Validator

5630 else:

5631 validator_cls = jsonschema.validators.validator_for(schema)

5632

5633 if validator_cls.__name__ not in allowed_validator_names:

5634 logger.warning(f"Unsupported JSON Schema draft: {validator_cls.__name__}")

5635

5636 validator_cls.check_schema(schema)

5637 except jsonschema.exceptions.SchemaError as e:

5638 logger.warning(f"Invalid prompt argument schema: {str(e)}")

5639 if settings.json_schema_validation_strict:

5640 raise ValueError(f"Invalid prompt argument schema: {str(e)}") from e

5641

5642

5643# Register validation listeners

5644

5645listen(Tool, "before_insert", validate_tool_schema)

5646listen(Tool, "before_update", validate_tool_schema)

5647listen(Tool, "before_insert", validate_tool_name)

5648listen(Tool, "before_update", validate_tool_name)

5649listen(Prompt, "before_insert", validate_prompt_schema)

5650listen(Prompt, "before_update", validate_prompt_schema)

5651

5652

5653def get_db() -> Generator[Session, Any, None]:

5654 """

5655 Dependency to get database session.

5656

5657 Commits the transaction on successful completion to avoid implicit rollbacks

5658 for read-only operations. Rolls back explicitly on exception.

5659

5660 Yields:

5661 SessionLocal: A SQLAlchemy database session.

5662

5663 Raises:

5664 Exception: Re-raises any exception after rolling back the transaction.

5665

5666 Examples:

5667 >>> from mcpgateway.db import get_db

5668 >>> gen = get_db()

5669 >>> db = next(gen)

5670 >>> hasattr(db, 'query')

5671 True

5672 >>> hasattr(db, 'commit')

5673 True

5674 >>> gen.close()

5675 """

5676 db = SessionLocal()

5677 try:

5678 yield db

5679 db.commit()

5680 except Exception:

5681 try:

5682 db.rollback()

5683 except Exception:

5684 try:

5685 db.invalidate()

5686 except Exception:

5687 pass # nosec B110 - Best effort cleanup on connection failure

5688 raise

5689 finally:

5690 db.close()

5691

5692

5693def get_for_update(

5694 db: Session,

5695 model,

5696 entity_id=None,

5697 where: Optional[Any] = None,

5698 skip_locked: bool = False,

5699 nowait: bool = False,

5700 lock_timeout_ms: Optional[int] = None,

5701 options: Optional[List] = None,

5702):

5703 """Get entity with row lock for update operations.

5704

5705 Args:

5706 db: SQLAlchemy Session

5707 model: ORM model class

5708 entity_id: Primary key value (optional if `where` provided)

5709 where: Optional SQLAlchemy WHERE clause to locate rows for conflict detection

5710 skip_locked: If False (default), wait for locked rows. If True, skip locked

5711 rows (returns None if row is locked). Use False for conflict checks and

5712 entity updates to ensure consistency. Use True only for job-queue patterns.

5713 nowait: If True, fail immediately if row is locked (raises OperationalError).

5714 Use this for operations that should not block. Default False.

5715 lock_timeout_ms: Optional lock timeout in milliseconds for PostgreSQL.

5716 If set, the query will wait at most this long for locks before failing.

5717 Only applies to PostgreSQL. Default None (use database default).

5718 options: Optional list of loader options (e.g., selectinload(...))

5719

5720 Returns:

5721 The model instance or None

5722

5723 Raises:

5724 sqlalchemy.exc.OperationalError: If nowait=True and row is locked, or if

5725 lock_timeout_ms is exceeded.

5726

5727 Notes:

5728 - On PostgreSQL this acquires a FOR UPDATE row lock.

5729 - On SQLite (or other backends that don't support FOR UPDATE) it

5730 falls back to a regular select; when ``options`` is None it uses

5731 ``db.get`` for efficiency, otherwise it executes a select with

5732 the provided loader options.

5733 """

5734 dialect = ""

5735 try:

5736 dialect = db.bind.dialect.name

5737 except Exception:

5738 dialect = ""

5739

5740 # Build base select statement. Prefer `where` when provided, otherwise use primary key `entity_id`.

5741 if where is not None:

5742 stmt = select(model).where(where)

5743 elif entity_id is not None:

5744 stmt = select(model).where(model.id == entity_id)

5745 else:

5746 return None

5747

5748 if options:

5749 stmt = stmt.options(*options)

5750

5751 if dialect != "postgresql":

5752 # SQLite and others: no FOR UPDATE support

5753 # Use db.get optimization only when querying by primary key without loader options

5754 if not options and where is None and entity_id is not None:

5755 return db.get(model, entity_id)

5756 return db.execute(stmt).scalar_one_or_none()

5757

5758 # PostgreSQL: set lock timeout if specified

5759 if lock_timeout_ms is not None:

5760 db.execute(text(f"SET LOCAL lock_timeout = '{lock_timeout_ms}ms'"))

5761

5762 # PostgreSQL: apply FOR UPDATE with optional nowait

5763 stmt = stmt.with_for_update(skip_locked=skip_locked, nowait=nowait)

5764 return db.execute(stmt).scalar_one_or_none()

5765

5766

5767# Using the existing get_db generator to create a context manager for fresh sessions

5768fresh_db_session = contextmanager(get_db) # type: ignore

5769

5770

5771def extract_json_field(column, json_path: str, dialect_name: Optional[str] = None):

5772 """Extract a JSON field in a database-agnostic way.

5773

5774 This function provides cross-database compatibility for JSON field extraction,

5775 supporting both SQLite and PostgreSQL backends.

5776

5777 Args:

5778 column: SQLAlchemy column containing JSON data

5779 json_path: JSON path in SQLite format (e.g., '$.\"tool.name\"')

5780 dialect_name: Optional database dialect name to override global backend.

5781 If not provided, uses the global backend from DATABASE_URL.

5782 Use this when querying a different database than the default.

5783

5784 Returns:

5785 SQLAlchemy expression for extracting the JSON field as text

5786

5787 Note:

5788 - For SQLite: Uses json_extract(column, '$.\"key\"')

5789 - For PostgreSQL: Uses column ->> 'key' operator

5790 - Backend-specific behavior is tested via unit tests in test_db.py

5791 """

5792 effective_backend = dialect_name if dialect_name is not None else backend

5793

5794 if effective_backend == "postgresql":

5795 # PostgreSQL uses ->> operator for text extraction

5796 # Convert $.\"key\" or $.\"nested.key\" format to just the key

5797 # Handle both simple keys and nested keys with dots

5798 path_key = json_path.replace('$."', "").replace('"', "")

5799 return column.op("->>")(path_key)

5800

5801 # SQLite and other databases use json_extract function

5802 # Keep the original $.\"key\" format

5803 return func.json_extract(column, json_path)

5804

5805

5806# Create all tables

5807def init_db():

5808 """

5809 Initialize database tables.

5810

5811 Raises:

5812 Exception: If database initialization fails.

5813 """

5814 try:

5815 # Base.metadata.drop_all(bind=engine)

5816 Base.metadata.create_all(bind=engine)

5817 except SQLAlchemyError as e:

5818 raise Exception(f"Failed to initialize database: {str(e)}")

5819

5820

5821# ============================================================================

5822# Structured Logging Models

5823# ============================================================================

5824

5825

5826class StructuredLogEntry(Base):

5827 """Structured log entry for comprehensive logging and analysis.

5828

5829 Stores all log entries with correlation IDs, performance metrics,

5830 and security context for advanced search and analytics.

5831 """

5832

5833 __tablename__ = "structured_log_entries"

5834

5835 # Primary key

5836 id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: uuid.uuid4().hex)

5837

5838 # Timestamps

5839 timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False, index=True, default=utc_now)

5840

5841 # Correlation and request tracking

5842 correlation_id: Mapped[Optional[str]] = mapped_column(String(64), index=True, nullable=True)

5843 request_id: Mapped[Optional[str]] = mapped_column(String(64), index=True, nullable=True)

5844

5845 # Log metadata

5846 level: Mapped[str] = mapped_column(String(20), nullable=False, index=True) # DEBUG, INFO, WARNING, ERROR, CRITICAL

5847 component: Mapped[str] = mapped_column(String(100), nullable=False, index=True)

5848 message: Mapped[str] = mapped_column(Text, nullable=False)

5849 logger: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

5850

5851 # User and request context

5852 user_id: Mapped[Optional[str]] = mapped_column(String(255), index=True, nullable=True)

5853 user_email: Mapped[Optional[str]] = mapped_column(String(255), index=True, nullable=True)

5854 client_ip: Mapped[Optional[str]] = mapped_column(String(45), nullable=True) # IPv6 max length

5855 user_agent: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

5856 request_path: Mapped[Optional[str]] = mapped_column(String(500), nullable=True)

5857 request_method: Mapped[Optional[str]] = mapped_column(String(10), nullable=True)

5858

5859 # Performance data

5860 duration_ms: Mapped[Optional[float]] = mapped_column(Float, nullable=True)

5861 operation_type: Mapped[Optional[str]] = mapped_column(String(100), index=True, nullable=True)

5862

5863 # Security context

5864 is_security_event: Mapped[bool] = mapped_column(Boolean, default=False, index=True, nullable=False)

5865 security_severity: Mapped[Optional[str]] = mapped_column(String(20), index=True, nullable=True) # LOW, MEDIUM, HIGH, CRITICAL

5866 threat_indicators: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True)

5867

5868 # Structured context data

5869 context: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True)

5870 error_details: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True)

5871 performance_metrics: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True)

5872

5873 # System information

5874 hostname: Mapped[str] = mapped_column(String(255), nullable=False)

5875 process_id: Mapped[int] = mapped_column(Integer, nullable=False)

5876 thread_id: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)

5877 version: Mapped[str] = mapped_column(String(50), nullable=False)

5878 environment: Mapped[str] = mapped_column(String(50), nullable=False, default="production")

5879

5880 # OpenTelemetry trace context

5881 trace_id: Mapped[Optional[str]] = mapped_column(String(32), index=True, nullable=True)

5882 span_id: Mapped[Optional[str]] = mapped_column(String(16), nullable=True)

5883

5884 # Indexes for performance

5885 __table_args__ = (

5886 Index("idx_log_correlation_time", "correlation_id", "timestamp"),

5887 Index("idx_log_user_time", "user_id", "timestamp"),

5888 Index("idx_log_level_time", "level", "timestamp"),

5889 Index("idx_log_component_time", "component", "timestamp"),

5890 Index("idx_log_security", "is_security_event", "security_severity", "timestamp"),

5891 Index("idx_log_operation", "operation_type", "timestamp"),

5892 Index("idx_log_trace", "trace_id", "timestamp"),

5893 )

5894

5895

5896class PerformanceMetric(Base):

5897 """Aggregated performance metrics from log analysis.

5898

5899 Stores time-windowed aggregations of operation performance

5900 for analytics and trend analysis.

5901 """

5902

5903 __tablename__ = "performance_metrics"

5904

5905 # Primary key

5906 id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: uuid.uuid4().hex)

5907

5908 # Timestamp

5909 timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False, index=True, default=utc_now)

5910

5911 # Metric identification

5912 operation_type: Mapped[str] = mapped_column(String(100), nullable=False, index=True)

5913 component: Mapped[str] = mapped_column(String(100), nullable=False, index=True)

5914

5915 # Aggregated metrics

5916 request_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)

5917 error_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)

5918 error_rate: Mapped[float] = mapped_column(Float, nullable=False, default=0.0)

5919

5920 # Duration metrics (in milliseconds)

5921 avg_duration_ms: Mapped[float] = mapped_column(Float, nullable=False)

5922 min_duration_ms: Mapped[float] = mapped_column(Float, nullable=False)

5923 max_duration_ms: Mapped[float] = mapped_column(Float, nullable=False)

5924 p50_duration_ms: Mapped[float] = mapped_column(Float, nullable=False)

5925 p95_duration_ms: Mapped[float] = mapped_column(Float, nullable=False)

5926 p99_duration_ms: Mapped[float] = mapped_column(Float, nullable=False)

5927

5928 # Time window

5929 window_start: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False, index=True)

5930 window_end: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)

5931 window_duration_seconds: Mapped[int] = mapped_column(Integer, nullable=False)

5932

5933 # Additional context

5934 metric_metadata: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True)

5935

5936 __table_args__ = (

5937 Index("idx_perf_operation_time", "operation_type", "window_start"),

5938 Index("idx_perf_component_time", "component", "window_start"),

5939 Index("idx_perf_window", "window_start", "window_end"),

5940 )

5941

5942

5943class SecurityEvent(Base):

5944 """Security event logging for threat detection and audit trails.

5945

5946 Specialized table for security events with enhanced context

5947 and threat analysis capabilities.

5948 """

5949

5950 __tablename__ = "security_events"

5951

5952 # Primary key

5953 id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: uuid.uuid4().hex)

5954

5955 # Timestamps

5956 timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False, index=True, default=utc_now)

5957 detected_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False, default=utc_now)

5958

5959 # Correlation tracking

5960 correlation_id: Mapped[Optional[str]] = mapped_column(String(64), index=True, nullable=True)

5961 log_entry_id: Mapped[Optional[str]] = mapped_column(String(36), ForeignKey("structured_log_entries.id"), index=True, nullable=True)

5962

5963 # Event classification

5964 event_type: Mapped[str] = mapped_column(String(100), nullable=False, index=True) # auth_failure, suspicious_activity, rate_limit, etc.

5965 severity: Mapped[str] = mapped_column(String(20), nullable=False, index=True) # LOW, MEDIUM, HIGH, CRITICAL

5966 category: Mapped[str] = mapped_column(String(50), nullable=False, index=True) # authentication, authorization, data_access, etc.

5967

5968 # User and request context

5969 user_id: Mapped[Optional[str]] = mapped_column(String(255), index=True, nullable=True)

5970 user_email: Mapped[Optional[str]] = mapped_column(String(255), index=True, nullable=True)

5971 client_ip: Mapped[str] = mapped_column(String(45), nullable=False, index=True)

5972 user_agent: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

5973

5974 # Event details

5975 description: Mapped[str] = mapped_column(Text, nullable=False)

5976 action_taken: Mapped[Optional[str]] = mapped_column(String(100), nullable=True) # blocked, allowed, flagged, etc.

5977

5978 # Threat analysis

5979 threat_score: Mapped[float] = mapped_column(Float, nullable=False, default=0.0) # 0.0-1.0

5980 threat_indicators: Mapped[Dict[str, Any]] = mapped_column(JSON, nullable=False, default=dict)

5981 failed_attempts_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)

5982

5983 # Resolution tracking

5984 resolved: Mapped[bool] = mapped_column(Boolean, default=False, index=True, nullable=False)

5985 resolved_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), nullable=True)

5986 resolved_by: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

5987 resolution_notes: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

5988

5989 # Alert tracking

5990 alert_sent: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)

5991 alert_sent_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), nullable=True)

5992 alert_recipients: Mapped[Optional[List[str]]] = mapped_column(JSON, nullable=True)

5993

5994 # Additional context

5995 context: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True)

5996

5997 __table_args__ = (

5998 Index("idx_security_type_time", "event_type", "timestamp"),

5999 Index("idx_security_severity_time", "severity", "timestamp"),

6000 Index("idx_security_user_time", "user_id", "timestamp"),

6001 Index("idx_security_ip_time", "client_ip", "timestamp"),

6002 Index("idx_security_unresolved", "resolved", "severity", "timestamp"),

6003 )

6004

6005

6006# ---------------------------------------------------------------------------

6007# LLM Provider Configuration Models

6008# ---------------------------------------------------------------------------

6009

6010

6011class LLMProviderType:

6012 """Constants for LLM provider types."""

6013

6014 OPENAI = "openai"

6015 AZURE_OPENAI = "azure_openai"

6016 ANTHROPIC = "anthropic"

6017 BEDROCK = "bedrock"

6018 GOOGLE_VERTEX = "google_vertex"

6019 WATSONX = "watsonx"

6020 OLLAMA = "ollama"

6021 OPENAI_COMPATIBLE = "openai_compatible"

6022 COHERE = "cohere"

6023 MISTRAL = "mistral"

6024 GROQ = "groq"

6025 TOGETHER = "together"

6026

6027 @classmethod

6028 def get_all_types(cls) -> List[str]:

6029 """Get list of all supported provider types.

6030

6031 Returns:

6032 List of provider type strings.

6033 """

6034 return [

6035 cls.OPENAI,

6036 cls.AZURE_OPENAI,

6037 cls.ANTHROPIC,

6038 cls.BEDROCK,

6039 cls.GOOGLE_VERTEX,

6040 cls.WATSONX,

6041 cls.OLLAMA,

6042 cls.OPENAI_COMPATIBLE,

6043 cls.COHERE,

6044 cls.MISTRAL,

6045 cls.GROQ,

6046 cls.TOGETHER,

6047 ]

6048

6049 @classmethod

6050 def get_provider_defaults(cls) -> Dict[str, Dict[str, Any]]:

6051 """Get default configuration for each provider type.

6052

6053 Returns:

6054 Dictionary mapping provider type to default config.

6055 """

6056 return {

6057 cls.OPENAI: {

6058 "api_base": "https://api.openai.com/v1",

6059 "default_model": "gpt-4o",

6060 "supports_model_list": True,

6061 "models_endpoint": "/models",

6062 "requires_api_key": True,

6063 "description": "OpenAI GPT models (GPT-4, GPT-4o, etc.)",

6064 },

6065 cls.AZURE_OPENAI: {

6066 "api_base": "https://{resource}.openai.azure.com/openai/deployments/{deployment}",

6067 "default_model": "",

6068 "supports_model_list": False,

6069 "requires_api_key": True,

6070 "description": "Azure OpenAI Service",

6071 },

6072 cls.ANTHROPIC: {

6073 "api_base": "https://api.anthropic.com",

6074 "default_model": "claude-sonnet-4-20250514",

6075 "supports_model_list": False,

6076 "requires_api_key": True,

6077 "description": "Anthropic Claude models",

6078 },

6079 cls.OLLAMA: {

6080 "api_base": "http://localhost:11434",

6081 "default_model": "llama3.2",

6082 "supports_model_list": True,

6083 "models_endpoint": "/api/tags",

6084 "requires_api_key": False,

6085 "description": "Local Ollama server",

6086 },

6087 cls.OPENAI_COMPATIBLE: {

6088 "api_base": "http://localhost:8080/v1",

6089 "default_model": "",

6090 "supports_model_list": True,

6091 "models_endpoint": "/models",

6092 "requires_api_key": False,

6093 "description": "Any OpenAI-compatible API server",

6094 },

6095 cls.COHERE: {

6096 "api_base": "https://api.cohere.ai/v1",

6097 "default_model": "command-r-plus",

6098 "supports_model_list": True,

6099 "models_endpoint": "/models",

6100 "requires_api_key": True,

6101 "description": "Cohere Command models",

6102 },

6103 cls.MISTRAL: {

6104 "api_base": "https://api.mistral.ai/v1",

6105 "default_model": "mistral-large-latest",

6106 "supports_model_list": True,

6107 "models_endpoint": "/models",

6108 "requires_api_key": True,

6109 "description": "Mistral AI models",

6110 },

6111 cls.GROQ: {

6112 "api_base": "https://api.groq.com/openai/v1",

6113 "default_model": "llama-3.3-70b-versatile",

6114 "supports_model_list": True,

6115 "models_endpoint": "/models",

6116 "requires_api_key": True,

6117 "description": "Groq high-speed inference",

6118 },

6119 cls.TOGETHER: {

6120 "api_base": "https://api.together.xyz/v1",

6121 "default_model": "meta-llama/Llama-3.3-70B-Instruct-Turbo",

6122 "supports_model_list": True,

6123 "models_endpoint": "/models",

6124 "requires_api_key": True,

6125 "description": "Together AI inference",

6126 },

6127 cls.BEDROCK: {

6128 "api_base": "",

6129 "default_model": "anthropic.claude-3-sonnet-20240229-v1:0",

6130 "supports_model_list": False,

6131 "requires_api_key": False,

6132 "description": "AWS Bedrock (uses IAM credentials)",

6133 },

6134 cls.GOOGLE_VERTEX: {

6135 "api_base": "",

6136 "default_model": "gemini-1.5-pro",

6137 "supports_model_list": False,

6138 "requires_api_key": False,

6139 "description": "Google Vertex AI (uses service account)",

6140 },

6141 cls.WATSONX: {

6142 "api_base": "https://us-south.ml.cloud.ibm.com",

6143 "default_model": "ibm/granite-13b-chat-v2",

6144 "supports_model_list": False,

6145 "requires_api_key": True,

6146 "description": "IBM watsonx.ai",

6147 },

6148 }

6149

6150

6151class LLMProvider(Base):

6152 """ORM model for LLM provider configurations.

6153

6154 Stores credentials and settings for external LLM providers

6155 used by the internal LLM Chat feature.

6156

6157 Attributes:

6158 id: Unique identifier (UUID)

6159 name: Display name (unique)

6160 slug: URL-safe identifier (unique)

6161 provider_type: Provider type (openai, anthropic, etc.)

6162 api_key: Encrypted API key

6163 api_base: Base URL for API requests

6164 api_version: API version (for Azure OpenAI)

6165 config: Provider-specific settings (JSON)

6166 default_model: Default model ID

6167 default_temperature: Default temperature (0.0-2.0)

6168 default_max_tokens: Default max tokens

6169 enabled: Whether provider is enabled

6170 health_status: Current health status (healthy/unhealthy/unknown)

6171 last_health_check: Last health check timestamp

6172 plugin_ids: Attached plugin IDs (JSON)

6173 """

6174

6175 __tablename__ = "llm_providers"

6176

6177 # Primary key

6178 id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: uuid.uuid4().hex)

6179

6180 # Basic info

6181 name: Mapped[str] = mapped_column(String(255), nullable=False)

6182 slug: Mapped[str] = mapped_column(String(255), nullable=False)

6183 description: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

6184

6185 # Provider type

6186 provider_type: Mapped[str] = mapped_column(String(50), nullable=False)

6187

6188 # Credentials (encrypted)

6189 api_key: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

6190 api_base: Mapped[Optional[str]] = mapped_column(String(512), nullable=True)

6191 api_version: Mapped[Optional[str]] = mapped_column(String(50), nullable=True)

6192

6193 # Provider-specific configuration

6194 config: Mapped[Dict[str, Any]] = mapped_column(JSON, default=dict, nullable=False)

6195

6196 # Default settings

6197 default_model: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

6198 default_temperature: Mapped[float] = mapped_column(Float, default=0.7, nullable=False)

6199 default_max_tokens: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)

6200

6201 # Status

6202 enabled: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)

6203 health_status: Mapped[str] = mapped_column(String(20), default="unknown", nullable=False)

6204 last_health_check: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), nullable=True)

6205

6206 # Plugin integration

6207 plugin_ids: Mapped[List[str]] = mapped_column(JSON, default=list, nullable=False)

6208

6209 # Timestamps

6210 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, nullable=False)

6211 updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, onupdate=utc_now, nullable=False)

6212

6213 # Audit fields

6214 created_by: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

6215 modified_by: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

6216

6217 # Relationships

6218 models: Mapped[List["LLMModel"]] = relationship("LLMModel", back_populates="provider", cascade="all, delete-orphan")

6219

6220 __table_args__ = (

6221 UniqueConstraint("name", name="uq_llm_providers_name"),

6222 UniqueConstraint("slug", name="uq_llm_providers_slug"),

6223 Index("idx_llm_providers_enabled", "enabled"),

6224 Index("idx_llm_providers_type", "provider_type"),

6225 Index("idx_llm_providers_health", "health_status"),

6226 )

6227

6228 def __repr__(self) -> str:

6229 """Return string representation.

6230

6231 Returns:

6232 String representation of the provider.

6233 """

6234 return f"<LLMProvider(id='{self.id}', name='{self.name}', type='{self.provider_type}')>"

6235

6236

6237class LLMModel(Base):

6238 """ORM model for LLM model definitions.

6239

6240 Stores model metadata and capabilities for each provider.

6241

6242 Attributes:

6243 id: Unique identifier (UUID)

6244 provider_id: Foreign key to llm_providers

6245 model_id: Provider's model ID (e.g., gpt-4o)

6246 model_name: Display name

6247 model_alias: Optional routing alias

6248 supports_chat: Whether model supports chat completions

6249 supports_streaming: Whether model supports streaming

6250 supports_function_calling: Whether model supports function/tool calling

6251 supports_vision: Whether model supports vision/images

6252 context_window: Maximum context tokens

6253 max_output_tokens: Maximum output tokens

6254 enabled: Whether model is enabled

6255 deprecated: Whether model is deprecated

6256 """

6257

6258 __tablename__ = "llm_models"

6259

6260 # Primary key

6261 id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: uuid.uuid4().hex)

6262

6263 # Provider relationship

6264 provider_id: Mapped[str] = mapped_column(String(36), ForeignKey("llm_providers.id", ondelete="CASCADE"), nullable=False)

6265

6266 # Model identification

6267 model_id: Mapped[str] = mapped_column(String(255), nullable=False)

6268 model_name: Mapped[str] = mapped_column(String(255), nullable=False)

6269 model_alias: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)

6270 description: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

6271

6272 # Capabilities

6273 supports_chat: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)

6274 supports_streaming: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)

6275 supports_function_calling: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)

6276 supports_vision: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)

6277

6278 # Limits

6279 context_window: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)

6280 max_output_tokens: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)

6281

6282 # Status

6283 enabled: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)

6284 deprecated: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)

6285

6286 # Timestamps

6287 created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, nullable=False)

6288 updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utc_now, onupdate=utc_now, nullable=False)

6289

6290 # Relationship

6291 provider: Mapped["LLMProvider"] = relationship("LLMProvider", back_populates="models")

6292

6293 __table_args__ = (

6294 UniqueConstraint("provider_id", "model_id", name="uq_llm_models_provider_model"),

6295 Index("idx_llm_models_provider", "provider_id"),

6296 Index("idx_llm_models_enabled", "enabled"),

6297 Index("idx_llm_models_deprecated", "deprecated"),

6298 )

6299

6300 def __repr__(self) -> str:

6301 """Return string representation.

6302

6303 Returns:

6304 String representation of the model.

6305 """

6306 return f"<LLMModel(id='{self.id}', model_id='{self.model_id}', provider_id='{self.provider_id}')>"

6307

6308

6309class AuditTrail(Base):

6310 """Comprehensive audit trail for data access and changes.

6311

6312 Tracks all significant system changes and data access for

6313 compliance and security auditing.

6314 """

6315

6316 __tablename__ = "audit_trails"

6317

6318 # Primary key

6319 id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: uuid.uuid4().hex)

6320

6321 # Timestamps

6322 timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False, index=True, default=utc_now)

6323

6324 # Correlation tracking

6325 correlation_id: Mapped[Optional[str]] = mapped_column(String(64), index=True, nullable=True)

6326 request_id: Mapped[Optional[str]] = mapped_column(String(64), index=True, nullable=True)

6327

6328 # Action details

6329 action: Mapped[str] = mapped_column(String(100), nullable=False, index=True) # create, read, update, delete, execute, etc.

6330 resource_type: Mapped[str] = mapped_column(String(100), nullable=False, index=True) # tool, resource, prompt, user, etc.

6331 resource_id: Mapped[Optional[str]] = mapped_column(String(255), index=True, nullable=True)

6332 resource_name: Mapped[Optional[str]] = mapped_column(String(500), nullable=True)

6333

6334 # User context

6335 user_id: Mapped[str] = mapped_column(String(255), nullable=False, index=True)

6336 user_email: Mapped[Optional[str]] = mapped_column(String(255), index=True, nullable=True)

6337 team_id: Mapped[Optional[str]] = mapped_column(String(36), index=True, nullable=True)

6338

6339 # Request context

6340 client_ip: Mapped[Optional[str]] = mapped_column(String(45), nullable=True)

6341 user_agent: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

6342 request_path: Mapped[Optional[str]] = mapped_column(String(500), nullable=True)

6343 request_method: Mapped[Optional[str]] = mapped_column(String(10), nullable=True)

6344

6345 # Change tracking

6346 old_values: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True)

6347 new_values: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True)

6348 changes: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True)

6349

6350 # Data classification

6351 data_classification: Mapped[Optional[str]] = mapped_column(String(50), index=True, nullable=True) # public, internal, confidential, restricted

6352 requires_review: Mapped[bool] = mapped_column(Boolean, default=False, index=True, nullable=False)

6353

6354 # Result

6355 success: Mapped[bool] = mapped_column(Boolean, nullable=False, index=True)

6356 error_message: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

6357

6358 # Additional context

6359 context: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True)

6360

6361 __table_args__ = (

6362 Index("idx_audit_action_time", "action", "timestamp"),

6363 Index("idx_audit_resource_time", "resource_type", "resource_id", "timestamp"),

6364 Index("idx_audit_user_time", "user_id", "timestamp"),

6365 Index("idx_audit_classification", "data_classification", "timestamp"),

6366 Index("idx_audit_review", "requires_review", "timestamp"),

6367 )

6368

6369

6370if __name__ == "__main__":

6371 # Wait for database to be ready before initializing

6372 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

6373

6374 init_db()

6375

6376

6377@event.listens_for(Gateway, "before_insert")

6378def set_gateway_slug(_mapper, _conn, target):

6379 """Set the slug for a Gateway before insert.

6380

6381 Args:

6382 _mapper: Mapper

6383 _conn: Connection

6384 target: Target Gateway instance

6385 """

6386

6387 target.slug = slugify(target.name)

6388

6389

6390@event.listens_for(A2AAgent, "before_insert")

6391def set_a2a_agent_slug(_mapper, _conn, target):

6392 """Set the slug for an A2AAgent before insert.

6393

6394 Args:

6395 _mapper: Mapper

6396 _conn: Connection

6397 target: Target A2AAgent instance

6398 """

6399 target.slug = slugify(target.name)

6400

6401

6402@event.listens_for(GrpcService, "before_insert")

6403def set_grpc_service_slug(_mapper, _conn, target):

6404 """Set the slug for a GrpcService before insert.

6405

6406 Args:

6407 _mapper: Mapper

6408 _conn: Connection

6409 target: Target GrpcService instance

6410 """

6411 target.slug = slugify(target.name)

6412

6413

6414@event.listens_for(LLMProvider, "before_insert")

6415def set_llm_provider_slug(_mapper, _conn, target):

6416 """Set the slug for an LLMProvider before insert.

6417

6418 Args:

6419 _mapper: Mapper

6420 _conn: Connection

6421 target: Target LLMProvider instance

6422 """

6423 target.slug = slugify(target.name)

6424

6425

6426@event.listens_for(EmailTeam, "before_insert")

6427def set_email_team_slug(_mapper, _conn, target):

6428 """Set the slug for an EmailTeam before insert.

6429

6430 Args:

6431 _mapper: Mapper

6432 _conn: Connection

6433 target: Target EmailTeam instance

6434 """

6435 if not target.slug:

6436 target.slug = slugify(target.name)

6437

6438

6439@event.listens_for(Tool, "before_insert")

6440@event.listens_for(Tool, "before_update")

6441def set_custom_name_and_slug(mapper, connection, target): # pylint: disable=unused-argument

6442 """

6443 Event listener to set custom_name, custom_name_slug, and name for Tool before insert/update.

6444

6445 - Sets custom_name to original_name if not provided.

6446 - Calculates custom_name_slug from custom_name using slugify.

6447 - Updates name to gateway_slug + separator + custom_name_slug.

6448 - Sets display_name to custom_name if not provided.

6449

6450 Note: The gateway relationship must be explicitly set (via target.gateway = gateway_obj)

6451 before adding the tool to the session if gateway namespacing is needed. If only

6452 gateway_id is set without the relationship, we look up the gateway name via a direct

6453 SQL query.

6454

6455 Args:

6456 mapper: SQLAlchemy mapper for the Tool model.

6457 connection: Database connection.

6458 target: The Tool instance being inserted or updated.

6459 """

6460 # Set custom_name to original_name if not provided

6461 if not target.custom_name:

6462 target.custom_name = target.original_name

6463 # Set display_name to custom_name if not provided

6464 if not target.display_name:

6465 target.display_name = target.custom_name

6466 # Always update custom_name_slug from custom_name

6467 target.custom_name_slug = slugify(target.custom_name)

6468

6469 # Get gateway_slug - check for explicitly set gateway relationship first

6470 gateway_slug = ""

6471 if target.gateway:

6472 # Gateway relationship is already loaded

6473 gateway_slug = slugify(target.gateway.name)

6474 elif target.gateway_id:

6475 # Gateway relationship not loaded but gateway_id is set

6476 # Use a cached gateway name if available from gateway_name_cache attribute

6477 if hasattr(target, "gateway_name_cache") and target.gateway_name_cache:

6478 gateway_slug = slugify(target.gateway_name_cache)

6479 else:

6480 # Fall back to querying the database

6481 try:

6482 result = connection.execute(text("SELECT name FROM gateways WHERE id = :gw_id"), {"gw_id": target.gateway_id})

6483 row = result.fetchone()

6484 if row:

6485 gateway_slug = slugify(row[0])

6486 except Exception: # nosec B110 - intentionally proceed without prefix on failure

6487 pass

6488

6489 if gateway_slug:

6490 sep = settings.gateway_tool_name_separator

6491 target.name = f"{gateway_slug}{sep}{target.custom_name_slug}"

6492 else:

6493 target.name = target.custom_name_slug

6494

6495

6496@event.listens_for(Prompt, "before_insert")

6497@event.listens_for(Prompt, "before_update")

6498def set_prompt_name_and_slug(mapper, connection, target): # pylint: disable=unused-argument

6499 """Set name fields for Prompt before insert/update.

6500

6501 - Sets original_name from name if missing (legacy compatibility).

6502 - Sets custom_name to original_name if not provided.

6503 - Sets display_name to custom_name if not provided.

6504 - Calculates custom_name_slug from custom_name.

6505 - Updates name to gateway_slug + separator + custom_name_slug.

6506

6507 Note: The gateway relationship must be explicitly set (via target.gateway = gateway_obj)

6508 before adding the prompt to the session if gateway namespacing is needed. If only

6509 gateway_id is set without the relationship, we look up the gateway name via a direct

6510 SQL query.

6511

6512 Args:

6513 mapper: SQLAlchemy mapper for the Prompt model.

6514 connection: Database connection for the insert/update.

6515 target: Prompt instance being inserted or updated.

6516 """

6517 if not target.original_name:

6518 target.original_name = target.name

6519 if not target.custom_name:

6520 target.custom_name = target.original_name

6521 if not target.display_name:

6522 target.display_name = target.custom_name

6523 target.custom_name_slug = slugify(target.custom_name)

6524

6525 # Get gateway_slug - check for explicitly set gateway relationship first

6526 gateway_slug = ""

6527 if target.gateway:

6528 # Gateway relationship is already loaded

6529 gateway_slug = slugify(target.gateway.name)

6530 elif target.gateway_id:

6531 # Gateway relationship not loaded but gateway_id is set

6532 # Use a cached gateway name if available from gateway_name_cache attribute

6533 if hasattr(target, "gateway_name_cache") and target.gateway_name_cache:

6534 gateway_slug = slugify(target.gateway_name_cache)

6535 else:

6536 # Fall back to querying the database

6537 try:

6538 result = connection.execute(text("SELECT name FROM gateways WHERE id = :gw_id"), {"gw_id": target.gateway_id})

6539 row = result.fetchone()

6540 if row:

6541 gateway_slug = slugify(row[0])

6542 except Exception: # nosec B110 - intentionally proceed without prefix on failure

6543 pass

6544

6545 if gateway_slug:

6546 sep = settings.gateway_tool_name_separator

6547 target.name = f"{gateway_slug}{sep}{target.custom_name_slug}"

6548 else:

6549 target.name = target.custom_name_slug

Coverage for mcpgateway / db.py: 99%

2291 statements