"""PostgreSQL-based persistent storage for API keys and roles. Provides production-ready storage backends with: - Connection pooling - Retry logic - Multi-instance support - Proper transaction handling - Domain event emission (CQRS compatible) Requires: psycopg2 (installed by the `postgres` extra) """ from collections.abc import Callable from datetime import datetime, timedelta, UTC import hmac import json import secrets import structlog from mcp_hangar.domain.contracts.authentication import ApiKeyMetadata, IApiKeyStore, IInitialAdminBootstrapStore from mcp_hangar.domain.contracts.authorization import IRoleStore from mcp_hangar.domain.events import ApiKeyCreated, ApiKeyRevoked, KeyRotated, RoleAssigned, RoleRevoked from mcp_hangar.domain.exceptions import ExpiredCredentialsError, RevokedCredentialsError from mcp_hangar.auth.roles import BUILTIN_ROLES from mcp_hangar.domain.value_objects import Permission, Principal, PrincipalId, PrincipalType, Role from mcp_hangar.domain.contracts.authorization import validate_role_scope logger = structlog.get_logger(__name__) # Dummy hash for constant-time comparison padding _DUMMY_HASH = "0" * 64 # SHA-256 length sentinel # SQL Schema for API Keys API_KEYS_SCHEMA = """ CREATE TABLE IF NOT EXISTS api_keys ( key_hash VARCHAR(64) PRIMARY KEY, key_id VARCHAR(32) NOT NULL UNIQUE, principal_id VARCHAR(256) NOT NULL, name VARCHAR(256) NOT NULL, tenant_id VARCHAR(256), groups JSONB DEFAULT '[]', created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(), expires_at TIMESTAMP WITH TIME ZONE, last_used_at TIMESTAMP WITH TIME ZONE, revoked BOOLEAN NOT NULL DEFAULT FALSE, revoked_at TIMESTAMP WITH TIME ZONE, metadata JSONB DEFAULT '{}', rotated_to_key_id VARCHAR(32), grace_until TIMESTAMP WITH TIME ZONE ); CREATE INDEX IF NOT EXISTS idx_api_keys_principal_id ON api_keys(principal_id); CREATE INDEX IF NOT EXISTS idx_api_keys_key_id ON api_keys(key_id); CREATE INDEX IF NOT EXISTS idx_api_keys_expires_at ON api_keys(expires_at) WHERE expires_at IS NOT NULL; -- Add rotation columns if they don't exist (migration) ALTER TABLE api_keys ADD COLUMN IF NOT EXISTS rotated_to_key_id VARCHAR(32); ALTER TABLE api_keys ADD COLUMN IF NOT EXISTS grace_until TIMESTAMP WITH TIME ZONE; -- Singleton durable claim for the initial API-key administrator. CREATE TABLE IF NOT EXISTS initial_admin_bootstrap ( singleton BOOLEAN PRIMARY KEY DEFAULT TRUE CHECK (singleton), principal_id VARCHAR(256) NOT NULL, key_id VARCHAR(32) NOT NULL, created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW() ); """ # SQL Schema for Roles ROLES_SCHEMA = """ CREATE TABLE IF NOT EXISTS roles ( name VARCHAR(128) PRIMARY KEY, description TEXT, permissions JSONB NOT NULL DEFAULT '[]', is_builtin BOOLEAN NOT NULL DEFAULT FALSE, created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(), updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW() ); CREATE TABLE IF NOT EXISTS role_assignments ( id SERIAL PRIMARY KEY, principal_id VARCHAR(256) NOT NULL, role_name VARCHAR(128) NOT NULL REFERENCES roles(name) ON DELETE CASCADE, scope VARCHAR(256) NOT NULL DEFAULT 'global', assigned_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(), assigned_by VARCHAR(256), UNIQUE(principal_id, role_name, scope) ); CREATE INDEX IF NOT EXISTS idx_role_assignments_principal_scope ON role_assignments(principal_id, scope); """ class PostgresApiKeyStore(IApiKeyStore, IInitialAdminBootstrapStore): """PostgreSQL-based API key store. Features: - Connection pooling via connection factory - Atomic operations with proper transactions - Optimistic locking for updates - Automatic last_used_at updates Multi-instance safe: Uses database-level locking. """ MAX_KEYS_PER_PRINCIPAL = 100 def __init__( self, connection_factory, table_prefix: str = "", event_publisher: Callable | None = None, ): """Initialize the PostgreSQL store. Args: connection_factory: An `IConnectionFactory` -- the shared port from `infrastructure.persistence.database_common`. This store knows SQL; it deliberately does not know psycopg2, pooling, or how a connection is obtained. One place holds that knowledge, and it is the factory (#779). table_prefix: Optional prefix for table names. event_publisher: Optional callback for publishing domain events. """ self._connections = connection_factory self._prefix = table_prefix self._table = f"{table_prefix}api_keys" if table_prefix else "api_keys" self._roles_table = f"{table_prefix}roles" if table_prefix else "roles" self._assignments_table = f"{table_prefix}role_assignments" if table_prefix else "role_assignments" self._bootstrap_table = f"{table_prefix}initial_admin_bootstrap" if table_prefix else "initial_admin_bootstrap" self._event_publisher = event_publisher def initialize(self) -> None: """Create tables if they don't exist.""" schema = API_KEYS_SCHEMA if self._prefix: schema = schema.replace("api_keys", self._table) schema = schema.replace("initial_admin_bootstrap", self._bootstrap_table) with self._connections.get_connection() as conn: with conn.cursor() as cur: cur.execute(schema) conn.commit() logger.info("postgres_api_key_store_initialized", table=self._table) def get_principal_for_key(self, key_hash: str) -> Principal | None: """Look up principal for an API key hash.""" with self._connections.get_connection() as conn, conn.cursor() as cur: cur.execute( f""" SELECT principal_id, tenant_id, groups, name, key_id, expires_at, revoked, metadata, rotated_to_key_id, grace_until FROM {self._table} WHERE key_hash = %s """, (key_hash,), ) row = cur.fetchone() if row is None: # Perform dummy comparison to equalize timing with found-key path hmac.compare_digest(key_hash.encode("utf-8"), _DUMMY_HASH.encode("utf-8")) # The SELECT above opened a transaction on the borrowed pooled # connection; the found-key path closes it via the last_used_at # UPDATE's commit, but this early return has nothing to write. # Commit so the connection is not handed back to the pool "idle # in transaction" (same pattern as PostgresMetricsHistoryStore # and PostgresEventStore). conn.commit() return None ( principal_id, tenant_id, groups, name, key_id, expires_at, revoked, metadata, rotated_to_key_id, grace_until, ) = row # Check revocation if revoked: raise RevokedCredentialsError( message="API key has been revoked", auth_method="api_key", ) # Check rotation and grace period if rotated_to_key_id is not None: # Key has been rotated - check grace period if grace_until: if datetime.now(UTC) >= grace_until: raise ExpiredCredentialsError( message="API key has been rotated and grace period has expired", auth_method="api_key", expired_at=grace_until.timestamp(), ) else: # No grace period set, reject immediately raise ExpiredCredentialsError( message="API key has been rotated and grace period has expired", auth_method="api_key", ) # Check expiration if expires_at and expires_at < datetime.now(UTC): raise ExpiredCredentialsError( message="API key has expired", auth_method="api_key", expired_at=expires_at.timestamp(), ) # Update last_used_at (fire and forget, don't fail auth on update error) try: cur.execute( f""" UPDATE {self._table} SET last_used_at = NOW() WHERE key_hash = %s """, (key_hash,), ) conn.commit() except Exception as e: # noqa: BLE001 -- infra-boundary: non-critical last_used_at update logger.warning("failed_to_update_last_used", error=str(e)) conn.rollback() # Parse groups from JSON if isinstance(groups, str): groups = json.loads(groups) return Principal( id=PrincipalId(principal_id), type=PrincipalType.SERVICE_ACCOUNT, tenant_id=tenant_id, groups=frozenset(groups or []), metadata={"key_id": key_id, "key_name": name, **(metadata or {})}, ) def create_key( self, principal_id: str, name: str, expires_at: datetime | None = None, groups: frozenset[str] | None = None, tenant_id: str | None = None, created_by: str | None = None, ) -> str: """Create a new API key. Emits: ApiKeyCreated event """ from .api_key_authenticator import ApiKeyAuthenticator with self._connections.get_connection() as conn, conn.cursor() as cur: # Check key count for principal cur.execute( f""" SELECT COUNT(*) FROM {self._table} WHERE principal_id = %s AND revoked = FALSE """, (principal_id,), ) count = cur.fetchone()[0] if count >= self.MAX_KEYS_PER_PRINCIPAL: raise ValueError( f"Principal {principal_id} has reached maximum API keys ({self.MAX_KEYS_PER_PRINCIPAL})" ) # Generate key raw_key = ApiKeyAuthenticator.generate_key() key_hash = ApiKeyAuthenticator._hash_key(raw_key) key_id = secrets.token_urlsafe(8) # Insert cur.execute( f""" INSERT INTO {self._table} (key_hash, key_id, principal_id, name, tenant_id, groups, expires_at) VALUES (%s, %s, %s, %s, %s, %s, %s) """, ( key_hash, key_id, principal_id, name, tenant_id, json.dumps(list(groups or [])), expires_at, ), ) conn.commit() logger.info( "api_key_created", key_id=key_id, principal_id=principal_id, name=name, expires_at=expires_at.isoformat() if expires_at else None, ) # Emit domain event if self._event_publisher: self._event_publisher( ApiKeyCreated( key_id=key_id, principal_id=principal_id, key_name=name, expires_at=expires_at.timestamp() if expires_at else None, created_by=created_by or "system", ) ) return raw_key def bootstrap_initial_admin( self, principal_id: str, key_name: str, groups: frozenset[str] | None = None, tenant_id: str | None = None, actor: str = "local-cli-bootstrap", ) -> tuple[str, str] | None: """Atomically create the first API-key administrator, if unclaimed.""" from .api_key_authenticator import ApiKeyAuthenticator with self._connections.get_connection() as conn, conn.cursor() as cur: try: cur.execute( f""" INSERT INTO {self._bootstrap_table} (singleton, principal_id, key_id) VALUES (TRUE, %s, %s) ON CONFLICT (singleton) DO NOTHING RETURNING singleton """, (principal_id, "pending"), ) if cur.fetchone() is None: conn.rollback() return None cur.execute(f"SELECT 1 FROM {self._roles_table} WHERE name = 'admin'") if cur.fetchone() is None: raise ValueError("Built-in admin role is not initialized") raw_key = ApiKeyAuthenticator.generate_key() key_hash = ApiKeyAuthenticator._hash_key(raw_key) key_id = secrets.token_urlsafe(8) cur.execute( f""" INSERT INTO {self._table} (key_hash, key_id, principal_id, name, tenant_id, groups) VALUES (%s, %s, %s, %s, %s, %s) """, (key_hash, key_id, principal_id, key_name, tenant_id, json.dumps(list(groups or []))), ) cur.execute( f""" INSERT INTO {self._assignments_table} (principal_id, role_name, scope, assigned_by) VALUES (%s, 'admin', 'global', %s) ON CONFLICT (principal_id, role_name, scope) DO NOTHING """, (principal_id, actor), ) cur.execute( f"UPDATE {self._bootstrap_table} SET key_id = %s WHERE singleton = TRUE", (key_id,), ) conn.commit() except Exception: conn.rollback() raise logger.info("initial_admin_bootstrapped", key_id=key_id, principal_id=principal_id) if self._event_publisher: self._event_publisher( ApiKeyCreated( key_id=key_id, principal_id=principal_id, key_name=key_name, expires_at=None, created_by=actor ) ) self._event_publisher( RoleAssigned(principal_id=principal_id, role_name="admin", scope="global", assigned_by=actor) ) return raw_key, key_id def is_initial_admin_bootstrapped(self) -> bool: """Whether the one-shot claim row exists (read-only, never spends it).""" with self._connections.get_connection() as conn, conn.cursor() as cur: cur.execute(f"SELECT 1 FROM {self._bootstrap_table} WHERE singleton = TRUE") exists = cur.fetchone() is not None # Read-only: close the transaction the SELECT opened so the pooled # connection is not returned "idle in transaction". conn.commit() return exists def revoke_key(self, key_id: str, revoked_by: str | None = None, reason: str | None = None) -> bool: """Revoke an API key. Emits: ApiKeyRevoked event """ with self._connections.get_connection() as conn, conn.cursor() as cur: # Get principal_id before revoking cur.execute( f""" SELECT principal_id FROM {self._table} WHERE key_id = %s AND revoked = FALSE """, (key_id,), ) row = cur.fetchone() principal_id = row[0] if row else None cur.execute( f""" UPDATE {self._table} SET revoked = TRUE, revoked_at = NOW() WHERE key_id = %s AND revoked = FALSE RETURNING key_id """, (key_id,), ) result = cur.fetchone() conn.commit() if result: logger.info("api_key_revoked", key_id=key_id) # Emit domain event if self._event_publisher and principal_id: self._event_publisher( ApiKeyRevoked( key_id=key_id, principal_id=principal_id, revoked_by=revoked_by or "system", reason=reason or "", ) ) return True return False def list_keys(self, principal_id: str) -> list[ApiKeyMetadata]: """List API keys for a principal.""" with self._connections.get_connection() as conn, conn.cursor() as cur: cur.execute( f""" SELECT key_id, name, principal_id, created_at, expires_at, last_used_at, revoked FROM {self._table} WHERE principal_id = %s ORDER BY created_at DESC """, (principal_id,), ) rows = cur.fetchall() # Read-only: close the transaction the SELECT opened so the pooled # connection is not returned "idle in transaction" (see # PostgresMetricsHistoryStore.query for the same pattern). conn.commit() return [ ApiKeyMetadata( key_id=row[0], name=row[1], principal_id=row[2], created_at=row[3], expires_at=row[4], last_used_at=row[5], revoked=row[6], ) for row in rows ] def count_keys(self, principal_id: str) -> int: """Count active keys for a principal.""" with self._connections.get_connection() as conn, conn.cursor() as cur: cur.execute( f""" SELECT COUNT(*) FROM {self._table} WHERE principal_id = %s AND revoked = FALSE """, (principal_id,), ) row = cur.fetchone() # Read-only: close the transaction the SELECT opened so the pooled # connection is not returned "idle in transaction". conn.commit() return int(row[0]) if row else 0 def rotate_key( self, key_id: str, grace_period_seconds: float = 86400, rotated_by: str = "system", ) -> str: """Rotate an API key with a grace period. Args: key_id: Unique identifier of the key to rotate. grace_period_seconds: How long the old key remains valid (default: 24h). rotated_by: Principal initiating the rotation. Returns: The new raw API key (only shown once!). Raises: ValueError: If key doesn't exist, is revoked, or already rotated. """ from .api_key_authenticator import ApiKeyAuthenticator with self._connections.get_connection() as conn, conn.cursor() as cur: # Look up existing key cur.execute( f""" SELECT key_hash, principal_id, name, tenant_id, groups, expires_at, revoked, rotated_to_key_id, grace_until FROM {self._table} WHERE key_id = %s """, (key_id,), ) row = cur.fetchone() if row is None: raise ValueError(f"API key not found: {key_id}") ( old_key_hash, principal_id, name, tenant_id, groups, expires_at, revoked, rotated_to_key_id, grace_until, ) = row if revoked: raise ValueError(f"Cannot rotate revoked key {key_id}") # Check if already rotated with active grace period if rotated_to_key_id is not None and grace_until: if datetime.now(UTC) < grace_until: raise ValueError(f"Key {key_id} already has pending rotation") # Generate new key raw_key = ApiKeyAuthenticator.generate_key() key_hash = ApiKeyAuthenticator._hash_key(raw_key) new_key_id = secrets.token_urlsafe(8) now = datetime.now(UTC) grace_until_dt = now + timedelta(seconds=grace_period_seconds) try: # Insert new key row cur.execute( f""" INSERT INTO {self._table} (key_hash, key_id, principal_id, name, tenant_id, groups, created_at, expires_at) VALUES (%s, %s, %s, %s, %s, %s, %s, %s) """, ( key_hash, new_key_id, principal_id, name, tenant_id, groups, now, expires_at, ), ) # Update old key to track rotation cur.execute( f""" UPDATE {self._table} SET rotated_to_key_id = %s, grace_until = %s WHERE key_id = %s """, (new_key_id, grace_until_dt, key_id), ) conn.commit() logger.info( "api_key_rotated", old_key_id=key_id, new_key_id=new_key_id, principal_id=principal_id, grace_until=grace_until_dt.isoformat(), rotated_by=rotated_by, ) # Emit domain event if self._event_publisher: self._event_publisher( KeyRotated( key_id=key_id, principal_id=principal_id, new_key_id=new_key_id, rotated_at=now.timestamp(), grace_until=grace_until_dt.timestamp(), rotated_by=rotated_by, ) ) return raw_key except Exception as e: # noqa: BLE001 -- infra-boundary: rotate_key error propagated after logging conn.rollback() logger.error("api_key_rotation_failed", key_id=key_id, error=str(e)) raise class PostgresRoleStore(IRoleStore): """PostgreSQL-based role store. Features: - Built-in roles seeded on init - Custom roles support """ def delete_role(self, role_name: str) -> None: """Delete a custom role.""" raise NotImplementedError("Not implemented in postgres store") def update_role(self, role_name: str, permissions: list[Permission], description: str | None = None) -> Role: """Update a custom role.""" raise NotImplementedError("Not implemented in postgres store") def list_all_roles(self) -> list[Role]: """List all roles.""" raise NotImplementedError("Not implemented in postgres store") """PostgreSQL-based role store. Features: - Built-in roles seeded on init - Custom roles support - Multi-scope assignments - Proper foreign key constraints - Domain event emission Multi-instance safe: Uses database-level constraints. Events emitted: - RoleAssigned: When a role is assigned - RoleRevoked: When a role is revoked """ def __init__( self, connection_factory, table_prefix: str = "", event_publisher: Callable | None = None, ): """Initialize the PostgreSQL store. Args: connection_factory: An `IConnectionFactory` -- the shared port from `infrastructure.persistence.database_common`. This store knows SQL; it deliberately does not know psycopg2, pooling, or how a connection is obtained. One place holds that knowledge, and it is the factory (#779). table_prefix: Optional prefix for table names. event_publisher: Optional callback for publishing domain events. """ self._connections = connection_factory self._prefix = table_prefix self._roles_table = f"{table_prefix}roles" if table_prefix else "roles" self._assignments_table = f"{table_prefix}role_assignments" if table_prefix else "role_assignments" self._event_publisher = event_publisher def initialize(self) -> None: """Create tables and seed built-in roles.""" schema = ROLES_SCHEMA if self._prefix: schema = schema.replace("roles", self._roles_table) schema = schema.replace("role_assignments", self._assignments_table) with self._connections.get_connection() as conn: with conn.cursor() as cur: cur.execute(schema) # Seed built-in roles for role_name, role in BUILTIN_ROLES.items(): permissions_json = json.dumps( [ {"resource_type": p.resource_type, "action": p.action, "resource_id": p.resource_id} for p in role.permissions ] ) cur.execute( f""" INSERT INTO {self._roles_table} (name, description, permissions, is_builtin) VALUES (%s, %s, %s, TRUE) ON CONFLICT (name) DO UPDATE SET description = EXCLUDED.description, permissions = EXCLUDED.permissions, updated_at = NOW() """, (role_name, role.description, permissions_json), ) conn.commit() logger.info("postgres_role_store_initialized", roles_table=self._roles_table) def get_role(self, role_name: str) -> Role | None: """Get role by name.""" with self._connections.get_connection() as conn, conn.cursor() as cur: cur.execute( f""" SELECT name, description, permissions FROM {self._roles_table} WHERE name = %s """, (role_name,), ) row = cur.fetchone() # Read-only: close the transaction the SELECT opened so the pooled # connection is not returned "idle in transaction" (covers both the # not-found early return and the row-found return below). conn.commit() if row is None: return None name, description, permissions_json = row if isinstance(permissions_json, str): permissions_json = json.loads(permissions_json) permissions = frozenset( Permission( resource_type=p["resource_type"], action=p["action"], resource_id=p.get("resource_id", "*"), ) for p in permissions_json ) return Role(name=name, description=description or "", permissions=permissions) def add_role(self, role: Role) -> None: """Add a custom role.""" with self._connections.get_connection() as conn: with conn.cursor() as cur: permissions_json = json.dumps( [ {"resource_type": p.resource_type, "action": p.action, "resource_id": p.resource_id} for p in role.permissions ] ) cur.execute( f""" INSERT INTO {self._roles_table} (name, description, permissions, is_builtin) VALUES (%s, %s, %s, FALSE) ON CONFLICT (name) DO UPDATE SET description = EXCLUDED.description, permissions = EXCLUDED.permissions, updated_at = NOW() """, (role.name, role.description, permissions_json), ) conn.commit() logger.info("role_created", role_name=role.name) def get_roles_for_principal( self, principal_id: str, scope: str = "*", ) -> list[Role]: """Get all roles assigned to a principal.""" with self._connections.get_connection() as conn, conn.cursor() as cur: if scope == "*": cur.execute( f""" SELECT r.name, r.description, r.permissions FROM {self._roles_table} r JOIN {self._assignments_table} a ON r.name = a.role_name WHERE a.principal_id = %s """, (principal_id,), ) else: cur.execute( f""" SELECT r.name, r.description, r.permissions FROM {self._roles_table} r JOIN {self._assignments_table} a ON r.name = a.role_name WHERE a.principal_id = %s AND (a.scope = %s OR a.scope = 'global') """, (principal_id, scope), ) rows = cur.fetchall() # Read-only: close the transaction the SELECT opened so the pooled # connection is not returned "idle in transaction". conn.commit() roles = [] for name, description, permissions_json in rows: if isinstance(permissions_json, str): permissions_json = json.loads(permissions_json) permissions = frozenset( Permission( resource_type=p["resource_type"], action=p["action"], resource_id=p.get("resource_id", "*"), ) for p in permissions_json ) roles.append(Role(name=name, description=description or "", permissions=permissions)) return roles def assign_role( self, principal_id: str, role_name: str, scope: str = "global", assigned_by: str | None = None, ) -> None: """Assign a role to a principal. Emits: RoleAssigned event """ validate_role_scope(scope) with self._connections.get_connection() as conn: with conn.cursor() as cur: # Verify role exists cur.execute(f"SELECT 1 FROM {self._roles_table} WHERE name = %s", (role_name,)) if cur.fetchone() is None: raise ValueError(f"Unknown role: {role_name}") cur.execute( f""" INSERT INTO {self._assignments_table} (principal_id, role_name, scope) VALUES (%s, %s, %s) ON CONFLICT (principal_id, role_name, scope) DO NOTHING RETURNING id """, (principal_id, role_name, scope), ) result = cur.fetchone() conn.commit() # Only emit event if actually inserted if result: logger.info("role_assigned", principal_id=principal_id, role_name=role_name, scope=scope) if self._event_publisher: self._event_publisher( RoleAssigned( principal_id=principal_id, role_name=role_name, scope=scope, assigned_by=assigned_by or "system", ) ) def revoke_role( self, principal_id: str, role_name: str, scope: str = "global", revoked_by: str | None = None, ) -> None: """Revoke a role from a principal. Emits: RoleRevoked event """ with self._connections.get_connection() as conn: with conn.cursor() as cur: cur.execute( f""" DELETE FROM {self._assignments_table} WHERE principal_id = %s AND role_name = %s AND scope = %s RETURNING id """, (principal_id, role_name, scope), ) result = cur.fetchone() conn.commit() if result: logger.info("role_revoked", principal_id=principal_id, role_name=role_name, scope=scope) if self._event_publisher: self._event_publisher( RoleRevoked( principal_id=principal_id, role_name=role_name, scope=scope, revoked_by=revoked_by or "system", ) )