# Implementation Plan: Hydrate Workstation Roles and Account Feature Controls

## Overview

Build an in-memory database at startup that maps resource types to their Cerbos policy requirements. Use this database
to determine when resources need hydration with `workstation_roles` or `account_feature_controls` before calling Cerbos.

## Problem Statement

The authorization service needs to know when a resource requires additional data from external services:

- **workstation_roles**: Requires fetching workstation roles for the user from another service
- **account_feature_controls**: Requires fetching feature controls enabled for the tenant from another service

Currently, the `Resource` class has stub methods for this purpose, but they are not implemented.

## Solution Architecture Overview

The solution has three main components:

1. **Policy Metadata Database** - Parse Cerbos policies at startup to know which resources need which data
2. **Principal Hydration** - Fetch and inject workstation roles into the principal before calling Cerbos
3. **Resource Hydration** - Fetch and inject account feature controls into resources before calling Cerbos

### Key Insight: Two Different Hydration Targets

- **workstation_roles** → Affects the **Principal** (user's derived roles)
    - Imported via `importDerivedRoles` in policy files
    - Must be added to principal's tenant roles via `_build_principal`
    - Example: User with `workstation_catalog` role can perform certain actions

- **account_feature_controls** → Affects the **Resource** (tenant's feature flags)
    - Imported via `variables.import` in policy files
    - Must be added to `resource.attributes.tenant.account_feature_controls` array
    - Example: Tenant with feature 43 enabled allows `bulk_create` action

### 1. Policy Metadata Database

> **Architecture note**: The policy directory is crawled **once** in CI (Jenkinsfile), not at each
> Fargate task startup. The resulting `PolicyMetadataDatabase` is serialized to JSON and written to
> Redis. At startup, each task reads the entry from Redis, deserializes it, and injects the instance
> via the existing datasources/dependency-injection pattern. This mirrors the existing
> `bludgeon_cached_cerbos_decisions` CI pattern.

Create an in-memory database that maps `resource_type` → policy metadata, including:

- Whether the resource policy imports `workstation_roles` (from `importDerivedRoles`)
- Whether the resource policy imports `account_feature_controls` (from `variables.import`)

This database will be built by parsing YAML policy files from `cerbos/policies/` directory at startup.

### 2. Database Structure

```python
class ResourcePolicyMetadata(BaseModel):
    """Metadata about a resource policy's requirements."""
    resource_type: str
    requires_workstation_roles: bool
    requires_account_feature_controls: bool
    # Could add more metadata in the future:
    # - other imported derived roles
    # - other imported variables
    # - policy file path


class PolicyMetadataDatabase:
    """In-memory database of resource policy metadata."""

    def __init__(self):
        self._policies: Dict[str, ResourcePolicyMetadata] = {}

    def add_policy(self, metadata: ResourcePolicyMetadata) -> None:
        """Add a resource policy's metadata."""
        self._policies[metadata.resource_type] = metadata

    def get_policy(self, resource_type: str) -> Optional[ResourcePolicyMetadata]:
        """Get metadata for a resource type."""
        return self._policies.get(resource_type)

    def requires_workstation_roles(self, resource_type: str) -> bool:
        """Check if resource type requires workstation roles."""
        policy = self.get_policy(resource_type)
        return policy.requires_workstation_roles if policy else False

    def requires_account_feature_controls(self, resource_type: str) -> bool:
        """Check if resource type requires account feature controls."""
        policy = self.get_policy(resource_type)
        return policy.requires_account_feature_controls if policy else False
```

## Implementation Steps

### Step 1: Create Policy Parser Module

**File**: `pdp/connectors/cerbos_policy_parser.py`

Create a module to parse Cerbos policy YAML files and extract metadata:

```python
import logging
import yaml
from pathlib import Path
from typing import List, Optional

logger = logging.getLogger(__name__)


class CerbosPolicyParser:
    """Parser for Cerbos policy YAML files."""

    def __init__(self, policies_dir: str = "cerbos/policies"):
        self.policies_dir = Path(policies_dir)

    def parse_resource_policy(self, policy_file: Path) -> Optional[ResourcePolicyMetadata]:
        """Parse a single resource policy YAML file."""
        try:
            with open(policy_file, 'r') as f:
                policy_data = yaml.safe_load(f)

            # Check if this is a resourcePolicy file
            if 'resourcePolicy' not in policy_data:
                return None

            resource_policy = policy_data['resourcePolicy']
            resource_type = resource_policy.get('resource')

            if not resource_type:
                logger.warning(f"No resource type found in {policy_file}")
                return None

            # Check for workstation_roles in importDerivedRoles
            derived_roles = resource_policy.get('importDerivedRoles', [])
            requires_workstation_roles = 'workstation_roles' in derived_roles

            # Check for account_feature_controls in variables.import
            variables_import = resource_policy.get('variables', {}).get('import', [])
            requires_account_feature_controls = 'account_feature_controls' in variables_import

            return ResourcePolicyMetadata(
                resource_type=resource_type,
                requires_workstation_roles=requires_workstation_roles,
                requires_account_feature_controls=requires_account_feature_controls,
            )
        except Exception as e:
            logger.error(f"Error parsing policy file {policy_file}: {e}")
            return None

    def build_database(self) -> PolicyMetadataDatabase:
        """Build the policy metadata database by parsing all policy files."""
        database = PolicyMetadataDatabase()

        # Find all YAML files in policies directory (recursively)
        policy_files = list(self.policies_dir.rglob("*.yaml")) + list(self.policies_dir.rglob("*.yml"))

        logger.info(f"Found {len(policy_files)} policy files to parse")

        for policy_file in policy_files:
            metadata = self.parse_resource_policy(policy_file)
            if metadata:
                database.add_policy(metadata)
                logger.debug(
                    f"Added policy for {metadata.resource_type}: "
                    f"workstation_roles={metadata.requires_workstation_roles}, "
                    f"feature_controls={metadata.requires_account_feature_controls}"
                )

        logger.info(f"Built policy database with {len(database._policies)} resource types")
        return database
```

**Key decisions**:

- Use PyYAML to parse policy files
- Recursively search `cerbos/policies/` for all `.yaml` and `.yml` files
- Only parse files that contain `resourcePolicy` (skip derived_roles and variables files)
- Log warnings for unparseable files but don't fail the CI step

Add `model_dump()` / `model_validate()` support to `PolicyMetadataDatabase` so it can be round-tripped through JSON:

```python
class PolicyMetadataDatabase:
    """In-memory database of resource policy metadata."""

    # ... existing methods ...

    def to_json(self) -> str:
        """Serialize the database to a JSON string for caching."""
        data = {
            resource_type: meta.model_dump()
            for resource_type, meta in self._policies.items()
        }
        return json.dumps(data)

    @classmethod
    def from_json(cls, raw: str) -> "PolicyMetadataDatabase":
        """Deserialize a PolicyMetadataDatabase from a cached JSON string."""
        db = cls()
        data = json.loads(raw)
        for resource_type, meta_dict in data.items():
            db.add_policy(ResourcePolicyMetadata(**meta_dict))
        return db
```

### Step 2: Add Cache Key Constant

**File**: `pdp/constants/constants.py`

Add a new cache entry constant alongside the existing ones:

```python
CACHE_ENTRY_CERBOS_POLICY_METADATA = "cerbos_policy_metadata"

class CacheEntryType(str, Enum):
    ...
    CACHE_ENTRY_CERBOS_POLICY_METADATA = CACHE_ENTRY_CERBOS_POLICY_METADATA
```

The cache key used at runtime will be the bare string `"cerbos_policy_metadata"` (a single,
non-identity-scoped entry — no UUID suffix).

### Step 3: Create CLI Command to Seed the Cache

**File**: `pdp/cli/commands/cerbos.py` *(new file)*

Modelled on `pdp/cli/commands/redis.py`. This command is invoked by Jenkins (not at task startup)
to crawl the policy directory and write the serialized database to Redis.

```python
"""Commands to manage Cerbos policy metadata."""

import asyncio

import typer

from pdp import config
from pdp.connectors.cerbos_policy_parser import CerbosPolicyParser
from pdp.connectors.redis_client import RedisConnector
from pdp.constants.constants import CACHE_ENTRY_CERBOS_POLICY_METADATA

cli: typer.Typer = typer.Typer(
    short_help="Commands to manage Cerbos policy metadata", no_args_is_help=True
)


async def _seed_policy_metadata_cache(policies_dir: str) -> None:
    """Crawl policy dir, build database, write to Redis."""
    redis_connector = RedisConnector(
        config.REDIS_URL,
        use_redis_cache=config.CACHE_USE_REDIS,
    )

    parser = CerbosPolicyParser(policies_dir=policies_dir)
    db = parser.build_database()
    serialized = db.to_json()

    await redis_connector.set(
        key=CACHE_ENTRY_CERBOS_POLICY_METADATA,
        value=serialized,
    )
    typer.secho(
        f"Wrote policy metadata database ({len(db._policies)} resource types) to cache "
        f"key '{CACHE_ENTRY_CERBOS_POLICY_METADATA}'.",
        fg="green",
    )


@cli.command("seed_policy_metadata_cache", short_help="Build and cache the policy metadata DB.")
def seed_policy_metadata_cache(
    policies_dir: str = typer.Option(
        "cerbos/policies",
        help="Path to the Cerbos policies directory.",
    ),
) -> None:
    """Crawl cerbos/policies, build PolicyMetadataDatabase, write to Redis."""
    asyncio.run(_seed_policy_metadata_cache(policies_dir))
```

Register the new CLI group in `pdp/cli/main.py` (or wherever other CLI groups are added):

```python
from pdp.cli.commands import cerbos as cerbos_commands
app.add_typer(cerbos_commands.cli, name="cerbos")
```

### Step 4: Initialize Database at Startup (Redis Read)

**File**: `pdp/fastapi/datasources.py`

At startup, each Fargate task reads the pre-seeded entry from Redis instead of crawling the
filesystem. Fail loudly if the entry is missing so a misconfigured deployment is caught immediately.

```python
# Add to imports
from pdp.connectors.cerbos_policy_parser import PolicyMetadataDatabase
from pdp.constants.constants import CACHE_ENTRY_CERBOS_POLICY_METADATA

# Add constant for data sources key
POLICY_METADATA_DATABASE_KEY = "POLICY_METADATA_DATABASE"


# In datasources_lifespan function, after redis_connector is available:
async def datasources_lifespan(app: FastAPI) -> AsyncIterator[Dict[str, Any]]:
    """Fastapi Lifespan function to create shared external data source connections."""
    # ... existing initialization code (redis_connector must be ready first) ...

    raw = await redis_connector.get(key=CACHE_ENTRY_CERBOS_POLICY_METADATA)
    if raw is None:
        raise RuntimeError(
            f"Cache entry '{CACHE_ENTRY_CERBOS_POLICY_METADATA}' not found. "
            "Run 'make seed_policy_metadata_cache' (or the Jenkins stage) to populate it."
        )
    policy_metadata_db = PolicyMetadataDatabase.from_json(raw)
    logger.info(
        "[lifespan] Loaded policy metadata database from cache "
        f"({len(policy_metadata_db._policies)} resource types)"
    )

    # ... existing DATA_SOURCES assignments ...
    DATA_SOURCES[POLICY_METADATA_DATABASE_KEY] = policy_metadata_db

    # ... rest of function ...


# Add dependency injector
def get_policy_metadata_database() -> PolicyMetadataDatabase:
    """Dependency injector method for policy metadata database.

    Returns the shared POLICY_METADATA_DATABASE.
    """
    database = DATA_SOURCES[POLICY_METADATA_DATABASE_KEY]
    assert isinstance(database, PolicyMetadataDatabase)
    return database
```

### Step 5: Add Makefile Target

**File**: `Makefile`

```makefile
PDPCLI_SEED_POLICY_METADATA_DOCKER_CMD = ./pdpcli cerbos seed_policy_metadata_cache

seed_policy_metadata_cache:
	${PDPCLI_SEED_POLICY_METADATA_DOCKER_CMD}
```

### Step 6: Add Jenkinsfile Stages

**File**: `Jenkinsfile`

Add two new stages — one for QA and one for Prod — each guarded by `isCerbosPolicyUpdated()`.
They should run **after** the corresponding deploy stage (so the new image with updated policies is
already deployed) and **alongside or before** the existing `Bludgeon *Cache` stages.

```groovy
stage('Seed QA Policy Metadata Cache') {
    when {
        allOf {
            branch 'master'
            expression { isCerbosPolicyUpdated() == true }
        }
    }
    steps {
        withEcr {
            withEnv([
                "REDIS_URL=${QA_REDIS_URL}",
                "CACHE_USE_REDIS=${CACHE_USE_REDIS}"
            ]) {
                withAWS(role: QA_DEPLOYMENT_ROLE, roleAccount: QA_ACCOUNT_ID, roleSessionName: SESSION_NAME, useNode: true) {
                    sh 'make seed_policy_metadata_cache'
                }
            }
        }
    }
}
```

```groovy
stage('Seed Prod Policy Metadata Cache') {
    when {
        allOf {
            branch 'master'
            expression { isCerbosPolicyUpdated() == true }
            expression { params.DEPLOY_TO_PROD == 'Yes' }
        }
    }
    steps {
        withEcr {
            withEnv([
                "REDIS_URL=${PROD_REDIS_URL}",
                "CACHE_USE_REDIS=${CACHE_USE_REDIS}"
            ]) {
                withAWS(role: PROD_DEPLOYMENT_ROLE, roleAccount: PROD_ACCOUNT_ID, roleSessionName: SESSION_NAME, useNode: true) {
                    sh 'make seed_policy_metadata_cache'
                }
            }
        }
    }
}
```

Place these stages immediately after the corresponding `Bludgeon *Cache` stages so the ordering
matches the existing bludgeon pattern.

### Step 7: Implement Resource Class Methods

**File**: `pdp/fastapi/schemas/identity.py`

The `Resource` class needs access to the database. Since Pydantic models can't easily have dependencies injected, we'll
pass the database as a parameter:

```python
class Resource(BaseModel):
    """Representation of a resource."""

    resource_id: Union[Annotated[str, StringConstraints(min_length=1)], int]
    resource_type: Annotated[str, StringConstraints(min_length=1)]
    attributes: Dict[str, Any] = {}

    model_config = ConfigDict(
        str_strip_whitespace=True,
    )

    def needs_workstation_role_lookup(
            self,
            policy_db: 'PolicyMetadataDatabase'
    ) -> bool:
        """Check if resource policy imports workstation_roles.

        Args:
            policy_db: The policy metadata database to query

        Returns:
            True if the resource type's policy imports workstation_roles
        """
        return policy_db.requires_workstation_roles(self.resource_type)

    def needs_account_feature_controls_lookup(
            self,
            policy_db: 'PolicyMetadataDatabase'
    ) -> bool:
        """Check if resource policy imports account_feature_controls.

        Args:
            policy_db: The policy metadata database to query

        Returns:
            True if the resource type's policy imports account_feature_controls
        """
        return policy_db.requires_account_feature_controls(self.resource_type)

    # ... rest of existing methods ...
```

**Alternative approach**: If we want to avoid passing the database explicitly, we could make these class methods:

```python
@classmethod
def resource_needs_workstation_role_lookup(
        cls,
        resource_type: str,
        policy_db: 'PolicyMetadataDatabase'
) -> bool:
    """Check if resource type requires workstation roles lookup."""
    return policy_db.requires_workstation_roles(resource_type)
```

### Step 8: Update _build_principal to Accept Workstation Roles

**File**: `pdp/logic/cerbos.py`

Update `_build_principal` to accept an optional `workstation_tenant_roles` parameter, following the same pattern as
`ows_permissions_tenant_roles`:

```python
@tracer.wrap()
def _build_principal(
        identity_uuid: str,
        pdp_tenant_roles: Dict[uuid.UUID, TenantRoles],
        ows_permissions_tenant_roles: Optional[Dict[uuid.UUID, TenantRoles]] = None,
        workstation_tenant_roles: Optional[Dict[uuid.UUID, TenantRoles]] = None,  # NEW
) -> CerbosPrincipal:
    """Get all data to represent an identity_uuid as a principal.

    Args:
        identity_uuid: UUID to match
        pdp_tenant_roles: Principal's tenant roles from PDP
        ows_permissions_tenant_roles: Principal's tenant roles from ows-permissions (optional)
        workstation_tenant_roles: Principal's workstation tenant roles (optional)
    """
    tenants = pdp_tenant_roles

    cerbos_principal_tenants: Dict[str, Any] = {}
    for tenant in tenants.values():
        cerbos_principal_tenants[str(tenant.tenant_uuid)] = (
            tenant.as_cerbos_principal_tenants_attribute()
        )

    # Merge ows_permissions_tenant_roles if provided (existing code)
    if ows_permissions_tenant_roles:
        for tenant_role in ows_permissions_tenant_roles.values():
            tenant_uuid = str(tenant_role.tenant_uuid)
            if tenant_uuid in cerbos_principal_tenants:
                cerbos_principal_tenants[tenant_uuid]["roles"].update(
                    tenant_role.as_cerbos_principal_tenants_attribute()["roles"]
                )
            else:
                cerbos_principal_tenants[tenant_uuid] = (
                    tenant_role.as_cerbos_principal_tenants_attribute()
                )

    # Merge workstation_tenant_roles if provided (NEW)
    if workstation_tenant_roles:
        for tenant_role in workstation_tenant_roles.values():
            tenant_uuid = str(tenant_role.tenant_uuid)
            if tenant_uuid in cerbos_principal_tenants:
                cerbos_principal_tenants[tenant_uuid]["roles"].update(
                    tenant_role.as_cerbos_principal_tenants_attribute()["roles"]
                )
            else:
                cerbos_principal_tenants[tenant_uuid] = (
                    tenant_role.as_cerbos_principal_tenants_attribute()
                )

    return CerbosPrincipal(
        id=str(identity_uuid),
        roles={"user"},
        attr={
            "type": "human",
            "tenants": cerbos_principal_tenants,
        },
    )
```

### Step 9: Create Function to Fetch Workstation Roles

**File**: `pdp/logic/cerbos.py`

Add a helper function to check if workstation roles are needed and fetch them:

```python
@tracer.wrap()
async def _fetch_workstation_roles_if_needed(
        check_resources_request: CheckResourcesRequest,
        policy_metadata_db: PolicyMetadataDatabase,
        identity_uuid: str,
        redis_connector: RedisConnector,
        # TODO: Add appropriate client for fetching workstation roles
        # ows_permissions_client: OwsPermissionsClient,
) -> Optional[Dict[uuid.UUID, TenantRoles]]:
    """Fetch workstation roles if any resource requires them.

    Args:
        check_resources_request: The resources being checked
        policy_metadata_db: Database of policy metadata
        identity_uuid: The identity to fetch roles for
        redis_connector: Redis client for caching

    Returns:
        Dict of workstation tenant roles, or None if not needed
    """
    # Check if any resource requires workstation roles
    needs_workstation_roles = False
    for check_resource_action in check_resources_request.resources:
        if check_resource_action.resource.needs_workstation_role_lookup(policy_metadata_db):
            needs_workstation_roles = True
            break

    if not needs_workstation_roles:
        return None

    # TODO: Implement fetching workstation roles from external service
    # This would call ows-permissions or another service
    # Example:
    # workstation_roles = await ows_permissions_client.get_workstation_roles(identity_uuid)
    # return workstation_roles
    logger.info(f"Fetching workstation roles for identity {identity_uuid}")

    # For now, return None until the external service is identified
    return None
```

### Step 10: Create VendorFeatureControlsProxy

**File**: `pdp/proxies/vendor_feature_controls_proxy.py`

Create a proxy class modelled on `UuidToIdExchangeTenantProxy`. It holds account and subaccount
UUID lists, calls the two dedicated ows-account feature endpoints concurrently, and exposes a
`get_feature_controls(uuid)` accessor. Keeping the fetch logic here isolates it from the
hydration function and makes it independently testable.

```python
"""VendorFeatureControlsProxy."""

import asyncio
import logging
from typing import Dict, List, Optional
from uuid import UUID

from ddtrace import tracer

from pdp.connectors.ows_account import (
    LookupSubaccountFeaturesResponse,
    LookupVendorFeaturesResponse,
    OwsAccountClient,
)
from pdp.proxies.helpers import lookup_with_error_handling

logger = logging.getLogger(__name__)


class FeatureControlsLookupError(Exception):
    """Feature controls lookup error."""

    pass


class VendorFeatureControlsProxy:
    """Fetches vendor feature controls for account and subaccount tenants.

    Calls two dedicated ows-account endpoints concurrently:
      - POST /lookup/vendors/features/uuids/     (account tenants)
      - POST /lookup/subaccount/features/uuids/  (subaccount tenants)

    ows-account handles UUID→vendor_id resolution internally for both endpoints,
    so no cross-type resolution is needed in PDP.
    """

    def __init__(
        self,
        account_uuids: List[UUID],
        subaccount_uuids: List[UUID],
        ows_account_client: OwsAccountClient,
    ):
        """VendorFeatureControlsProxy constructor."""
        self._account_uuids = list(set(account_uuids))
        self._subaccount_uuids = list(set(subaccount_uuids))
        self._ows_account_client = ows_account_client
        self._gathered_feature_controls: Optional[Dict[UUID, List[int]]] = None

    @property
    def gathered_feature_controls(self) -> Dict[UUID, List[int]]:
        """Access gathered feature controls, keyed by tenant UUID.

        Raises FeatureControlsLookupError if accessed before ows-account lookups.
        """
        if self._gathered_feature_controls is None:
            raise FeatureControlsLookupError(
                "Attempted to access feature controls before ows-account lookups."
            )
        return self._gathered_feature_controls

    @tracer.wrap()
    def get_feature_controls(self, uuid: UUID) -> List[int]:
        """Return feature control IDs for a tenant UUID. Returns [] if not found."""
        return self.gathered_feature_controls.get(uuid, [])

    @tracer.wrap()
    async def _fetch_vendor_feature_controls(self) -> Dict[UUID, List[int]]:
        """Fetch feature controls for account (vendor) tenants."""
        if not self._account_uuids:
            return {}

        response: LookupVendorFeaturesResponse = await lookup_with_error_handling(
            lookup_cb=self._ows_account_client.get_features_for_vendor_uuids(
                uuids=self._account_uuids,
            ),
            http_error_message_400_string="Bad request to ows-account for vendor feature controls lookup.",
            http_error_message_500_string="Vendor feature controls lookup by vendor_uuids failed.",
            unhandled_exception_string="Server error: vendor feature controls lookup by vendor_uuids failed.",
            lookup_error_cls=FeatureControlsLookupError,
        )

        return {UUID(k): v for k, v in response.vendors.items()}

    @tracer.wrap()
    async def _fetch_subaccount_feature_controls(self) -> Dict[UUID, List[int]]:
        """Fetch feature controls for subaccount tenants.

        ows-account resolves subaccount UUID → parent vendor internally.
        Results are keyed by the original subaccount UUID.
        """
        if not self._subaccount_uuids:
            return {}

        response: LookupSubaccountFeaturesResponse = await lookup_with_error_handling(
            lookup_cb=self._ows_account_client.get_features_for_subaccount_uuids(
                uuids=self._subaccount_uuids,
            ),
            http_error_message_400_string="Bad request to ows-account for subaccount feature controls lookup.",
            http_error_message_500_string="Subaccount feature controls lookup by subaccount_uuids failed.",
            unhandled_exception_string="Server error: subaccount feature controls lookup by subaccount_uuids failed.",
            lookup_error_cls=FeatureControlsLookupError,
        )

        return {UUID(k): v for k, v in response.subaccounts.items()}

    @tracer.wrap()
    async def gather_feature_controls(self) -> Dict[UUID, List[int]]:
        """Fetch feature controls for all account and subaccount tenants concurrently."""
        self._gathered_feature_controls = {}

        vendor_controls, subaccount_controls = await asyncio.gather(
            self._fetch_vendor_feature_controls(),
            self._fetch_subaccount_feature_controls(),
        )

        self._gathered_feature_controls.update(vendor_controls)
        self._gathered_feature_controls.update(subaccount_controls)

        return self._gathered_feature_controls
```

### Step 11: Create Function to Hydrate Resources with Feature Controls

**File**: `pdp/logic/cerbos.py`

Add a function to hydrate resources with account feature controls, similar to
`_hydrate_resources_with_hierarchy_as_needed`. The fetch logic is delegated to
`VendorFeatureControlsProxy`, keeping this function focused on collecting UUIDs and
writing results back into resource attributes.

```python
@tracer.wrap()
async def _hydrate_resources_with_feature_controls_as_needed(
        check_resources_request: CheckResourcesRequest,
        policy_metadata_db: PolicyMetadataDatabase,
        ows_account_client: OwsAccountClient,
) -> CheckResourcesRequest:
    """Hydrate resources with account feature controls as needed.

    This function checks if any resources require account_feature_controls data,
    fetches that data via VendorFeatureControlsProxy, and injects it into the
    resource attributes before calling Cerbos.

    Args:
        check_resources_request: The resources to check
        policy_metadata_db: Database of policy metadata
        ows_account_client: Client for fetching account data

    Returns:
        Updated check_resources_request with hydrated resource attributes
    """
    # Pass 1: Collect tenant UUIDs by type for resources that need feature controls.
    account_uuids: set[UUID] = set()
    subaccount_uuids: set[UUID] = set()
    for check_resource_action in check_resources_request.resources:
        if check_resource_action.resource.needs_account_feature_controls_lookup(policy_metadata_db):
            tenant = check_resource_action.get_tenant()
            if tenant:
                if tenant.tenant_type == TenantType.ACCOUNT:
                    account_uuids.add(tenant.tenant_uuid)
                elif tenant.tenant_type == TenantType.SUBACCOUNT:
                    subaccount_uuids.add(tenant.tenant_uuid)

    if not account_uuids and not subaccount_uuids:
        return check_resources_request

    # Pass 2: Fetch feature controls for all tenant UUIDs concurrently via proxy.
    proxy = VendorFeatureControlsProxy(
        account_uuids=list(account_uuids),
        subaccount_uuids=list(subaccount_uuids),
        ows_account_client=ows_account_client,
    )
    await proxy.gather_feature_controls()

    # Pass 3: Hydrate each resource.
    for check_resource_action in check_resources_request.resources:
        resource = check_resource_action.resource

        if resource.needs_account_feature_controls_lookup(policy_metadata_db):
            tenant = check_resource_action.get_tenant()
            if not tenant:
                continue

            if tenant.tenant_type not in (TenantType.ACCOUNT, TenantType.SUBACCOUNT):
                continue

            feature_ids = proxy.get_feature_controls(tenant.tenant_uuid)
            # Inject into resource.attributes.tenant.account_feature_controls
            if 'tenant' not in resource.attributes:
                resource.attributes['tenant'] = {}
            resource.attributes['tenant']['account_feature_controls'] = feature_ids

    return check_resources_request
```

**Key Architecture Decision**:

The hydration happens in two separate places because they affect different parts of the authorization check:

- **workstation_roles** affects the **Principal** → Fetched early and passed to `_build_principal`
- **account_feature_controls** affects the **Resource** → Hydrated into resource attributes similar to hierarchy

### Step 12: Integrate Hydration in check_resources

**File**: `pdp/logic/cerbos.py`

Update the `check_resources` function to call both hydration functions:

```python
async def check_resources(
        identity_uuid: str,
        check_resources_request: CheckResourcesRequest,
        cerbos_client: AsyncCerbosClient,
        pdp_tenant_roles: Dict[uuid.UUID, TenantRoles],
        ows_account_client: OwsAccountClient,
        ows_participant_client: OwsParticipantClient,
        redis_connector: RedisConnector,
        policy_metadata_db: PolicyMetadataDatabase,  # NEW PARAMETER
        ows_permissions_tenant_roles: Optional[Dict[uuid.UUID, TenantRoles]] = None,
        include_resource_attributes: bool = False,
        authenticated_identity_uuid: Optional[UUID4] = None,
        principal: Optional[Principal] = None,
) -> CheckResourcesResponse:
    """Make cerbos request to check identity's authorization to resources."""

    # STEP 1: Fetch workstation roles if needed (NEW)
    workstation_tenant_roles = await _fetch_workstation_roles_if_needed(
        check_resources_request=check_resources_request,
        policy_metadata_db=policy_metadata_db,
        identity_uuid=identity_uuid,
        redis_connector=redis_connector,
    )

    # STEP 2: Build principal with all role sources (UPDATED)
    if principal:
        cerbos_principal = principal.get_cerbos_principal()
    else:
        cerbos_principal = _build_principal(
            identity_uuid,
            pdp_tenant_roles=pdp_tenant_roles,
            ows_permissions_tenant_roles=ows_permissions_tenant_roles,
            workstation_tenant_roles=workstation_tenant_roles,  # NEW
        )

    # STEP 3: Hydrate resources with hierarchy (existing)
    check_resources_request = await _hydrate_resources_with_hierarchy_as_needed(
        check_resources_request=check_resources_request,
        redis_connector=redis_connector,
        ows_account_client=ows_account_client,
        ows_participant_client=ows_participant_client,
    )

    # STEP 4: Hydrate resources with feature controls (NEW)
    check_resources_request = await _hydrate_resources_with_feature_controls_as_needed(
        check_resources_request=check_resources_request,
        policy_metadata_db=policy_metadata_db,
        ows_account_client=ows_account_client,
    )

    # STEP 5: Call Cerbos (existing)
    result = await _paginated_resource_check(
        check_resources_request=check_resources_request,
        cerbos_client=cerbos_client,
        principal=cerbos_principal,
        include_resource_attributes=include_resource_attributes,
    )
    return result
```

**Flow Summary**:

1. Check if workstation roles are needed, fetch if yes
2. Build principal with pdp_tenant_roles + ows_permissions_tenant_roles + workstation_tenant_roles
3. Hydrate resources with hierarchy (existing functionality)
4. Hydrate resources with account feature controls
5. Call Cerbos with hydrated principal and resources

### Step 13: Update Endpoint to Inject Database

**File**: `pdp/fastapi/routers/identity.py`

Update `check_my_resources` to inject the policy metadata database:

```python
from pdp.fastapi.datasources import (
    get_async_cerbos_client,
    get_boto_connector,
    get_ows_account_client,
    get_ows_participant_client,
    get_redis_connector,
    get_splitio_client,
    get_policy_metadata_database,  # NEW
)


@router.post(
    "/self/check/resources/", response_model=identity_schema.CheckResourcesResponse
)
async def check_my_resources(
        check_resources_request: check_resources_schema.CheckResourcesRequest,
        identity_uuid: UUID4 = Depends(identity_uuid_from_scope),
        pdp_tenant_roles: Dict[uuid.UUID, identity_schema.TenantRoles] = Depends(
            get_principal_pdp_tenant_roles_from_scope
        ),
        impersonated_by_identity_uuid: UUID4 | None = Depends(
            impersonated_by_identity_uuid_from_scope
        ),
        user_type: str = Depends(user_type_from_scope),
        cerbos_client: AsyncCerbosClient = Depends(get_async_cerbos_client),
        ows_account_client: OwsAccountClient = Depends(get_ows_account_client),
        ows_participant_client: OwsParticipantClient = Depends(get_ows_participant_client),
        redis_connector: RedisConnector = Depends(get_redis_connector),
        splitio_client: SplitioClient = Depends(get_splitio_client),
        authenticated_identity_uuid: UUID = Depends(identity_uuid_from_scope),
        policy_metadata_db: PolicyMetadataDatabase = Depends(get_policy_metadata_database),  # NEW
) -> identity_schema.CheckResourcesResponse:
    """Check resources for the authenticated principal."""

    # ... existing code ...

    check_response = await cerbos.check_resources(
        identity_uuid=str(identity_uuid),
        check_resources_request=check_resources_request,
        pdp_tenant_roles=pdp_tenant_roles,
        cerbos_client=cerbos_client,
        policy_metadata_db=policy_metadata_db,  # NEW
        # ... other parameters ...
    )

    return check_response
```

Similarly update `check_identity_resources` and any other endpoints that call `cerbos.check_resources`.

## Implementation Considerations

### Workstation Roles vs Feature Controls

**Key difference**:

- **workstation_roles**: Affects the **principal** (user's derived roles)
    - Defined in `derived_roles/workstation_roles.yml`
    - Referenced in `importDerivedRoles`
    - Should be added to principal's tenant roles before calling Cerbos

- **account_feature_controls**: Affects the **resource** (tenant's feature flags)
    - Defined in `variables/account_feature_controls.yml`
    - Referenced in `variables.import`
    - Should be added to `resource.attributes.tenant.account_feature_controls` array

### Two-Part Hydration Architecture

The hydration follows a clear two-part pattern:

#### Part 1: Principal Hydration (Workstation Roles)

1. Check if any resource in the request requires `workstation_roles`
2. If yes, fetch workstation roles for the identity from external service
3. Pass `workstation_tenant_roles` to `_build_principal` alongside `pdp_tenant_roles` and `ows_permissions_tenant_roles`
4. `_build_principal` merges all three role sources into the Cerbos principal's tenant attributes

This follows the existing pattern where `_build_principal` already merges `ows_permissions_tenant_roles` into the
principal.

#### Part 2: Resource Hydration (Feature Controls)

1. Check if any resource requires `account_feature_controls`
2. If yes, collect all unique tenant UUIDs from those resources
3. Fetch feature controls for each tenant from ows-account
4. Inject feature controls into `resource.attributes.tenant.account_feature_controls` array

This follows the existing pattern of `_hydrate_resources_with_hierarchy_as_needed`.

### Caching Strategy

**Do NOT cache workstation roles or feature controls** - at least not initially.

Rationale:

- We don't have a cache invalidation mechanism when external services update these values
- This follows the existing pattern in `get_principal_ows_permissions_from_scope` (pdp/fastapi/auth.py:236)
  which fetches ows-permissions roles but does NOT cache them
- Caching without invalidation could lead to stale authorization decisions

What we DO cache:

- **Policy metadata database** - Built at startup, lives in memory for the application lifetime
- **pdp_tenant_roles** - These are cached because we control the data and can invalidate on updates

Future consideration:

- If performance becomes an issue, consider request-scoped caching (store in request.scope)
- This would cache for the duration of a single request but not across requests
- Similar to how `get_principal_ows_permissions_from_scope` uses `request.scope` for the request duration

### Error Handling

- If policy file parsing fails, log warning but continue startup
- If workstation roles fetch fails, should we deny access or proceed without?
- If feature controls fetch fails, should we proceed with empty array or fail?

### Testing Strategy

1. **Unit tests**:
    - Test policy parser with sample YAML files
    - Test database lookups
    - Test Resource methods

2. **Integration tests**:
    - Test full flow with mocked external service calls
    - Verify correct hydration of principals and resources

3. **Policy validation**:
    - Ensure all existing resource types are correctly identified
    - Verify workstation_roles and account_feature_controls are detected

## Open Questions

1. **Which service provides workstation roles?**
    - Need to identify the correct service and API endpoint
    - What format are the roles returned in? Should they be `Dict[uuid.UUID, TenantRoles]`?
    - How should they be cached?
    - Should we use ows-permissions client?
    - **Note from Notion doc**: The Workstation feature control is explicitly described as
      *"(Legacy)"* — the `ui_restriction` table "has no impact on permissions available" and
      "Accounts and Identities is not actively supporting this area." This may mean
      `workstation_roles` hydration is out of scope for this ticket, or is lower priority.
      Confirm with team before implementing.

2. ~~**Which service provides account feature controls?**~~ **RESOLVED**
    - Source: **ows-account**, via a new endpoint `POST /lookup/vendors/features/uuids/`
    - The endpoint accepts `{"uuids": ["uuid1", "uuid2"]}` and returns `{"vendors": {"uuid1": [39, 43], "uuid2": []}}`
    - Format: list of integers (feature_ids) per vendor UUID
    - Fetches all enabled features for a vendor (features not in `vendor_restricted_features`)
    - Feature controls are vendor-level only; **subaccounts are supported** by resolving to the
      parent vendor UUID via `OwsAccountClient.lookup_subaccounts_by_uuids` (which already returns
      `vendor_uuid`). The feature lookup then uses the resolved vendor UUID.
    - See [ows-account changes](#ows-account-changes) below

3. **Scope of implementation**:
    - Should this apply to all check_resources endpoints or just check_my_resources?
    - What about check_resource_type_actions endpoint?
    - Should the Principal class also support workstation roles in its `get_cerbos_principal()` method?

## ows-account Changes

Two new PIP (Policy Information Point) endpoints are needed in ows-account. The existing
`/vendor/<int:vendor_id>/features` endpoint only accepts integer `vendor_id`, but PDP works with
UUIDs. Rather than doing UUID→ID exchange in PDP, ows-account handles the join internally.
A dedicated subaccount endpoint also resolves to the parent vendor internally, so PDP can call
each endpoint directly without any cross-type resolution.

### New endpoint: `POST /lookup/vendors/features/uuids/`

For **account** tenants.

**Request**: `{"uuids": ["uuid1", "uuid2", ...]}`

**Response**: `{"vendors": {"uuid1": [39, 43], "uuid2": [], ...}}`

- Only UUIDs found in the vendor table are included in the response
- An empty list means the vendor exists but has no restricted features (all features enabled)
- Uses the same `NO access rule checks` pattern as other `/lookup/` endpoints (PIP)

### New endpoint: `POST /lookup/subaccount/features/uuids/`

For **subaccount** tenants. Resolves subaccount UUID → parent vendor internally before fetching
features, so the caller never needs to perform a separate subaccount lookup.

**Request**: `{"uuids": ["subaccount_uuid1", "subaccount_uuid2", ...]}`

**Response**: `{"subaccounts": {"subaccount_uuid1": [39, 43], "subaccount_uuid2": [], ...}}`

- Only UUIDs found in the subaccount table are included in the response
- Feature IDs are keyed by the original subaccount UUID (not the resolved vendor UUID)
- An empty list means the parent vendor has no restricted features
- Uses the same `NO access rule checks` pattern as other `/lookup/` endpoints (PIP)

### Files to modify in ows-account

1. **`account/models/feature.py`** — Add SQL queries:
    - `get_enabled_feature_ids_for_vendor_uuids(vendor_uuids, session)`:
        - Joins `vendor` (UUID → vendor_id) with `features` (CROSS JOIN)
        - Excludes rows present in `vendor_restricted_features`
        - Returns `dict[str, list[int]]`: vendor_uuid → [feature_id, ...]

        ```sql
        SELECT v.vendor_uuid, f.feature_id
        FROM vendor v
        CROSS JOIN features f
        WHERE v.vendor_uuid IN ({placeholders})
        AND NOT EXISTS (
            SELECT 1
            FROM vendor_restricted_features vrf
            WHERE vrf.vendor_id = v.vendor_id
            AND vrf.feature_id = f.feature_id
        )
        ORDER BY v.vendor_uuid, f.feature_id
        ```

    - `get_enabled_feature_ids_for_subaccount_uuids(subaccount_uuids, session)`:
        - Joins `subaccount` (UUID → vendor_id) with `features` (CROSS JOIN)
        - Excludes rows present in `vendor_restricted_features`
        - Returns `dict[str, list[int]]`: subaccount_uuid → [feature_id, ...]

        ```sql
        SELECT sa.subaccount_uuid, f.feature_id
        FROM subaccount sa
        CROSS JOIN features f
        WHERE sa.subaccount_uuid IN ({placeholders})
        AND NOT EXISTS (
            SELECT 1
            FROM vendor_restricted_features vrf
            WHERE vrf.vendor_id = sa.vendor_id
            AND vrf.feature_id = f.feature_id
        )
        ORDER BY sa.subaccount_uuid, f.feature_id
        ```

2. **`account/logic/feature.py`** — Add:
    - `lookup_features_by_vendor_uuids(uuids)` — wraps model call, returns `response.Response({"vendors": {...}})`
    - `lookup_features_by_subaccount_uuids(uuids)` — wraps model call, returns `response.Response({"subaccounts": {...}})`

3. **`account/validation/schemas/lookup.py`** — Add schemas:
    ```python
    class LookupVendorFeaturesByUuids(Schema):
        uuids = fields.List(fields.UUID(required=True))

    class LookupSubaccountFeaturesByUuids(Schema):
        uuids = fields.List(fields.UUID(required=True))
    ```

4. **`account/handlers/lookups.py`** — Add handlers:
    ```python
    @app.route('/lookup/vendors/features/uuids/', methods=['POST'])
    @validate_request_data(LookupVendorFeaturesByUuids())
    def lookup_vendors_features_by_uuids(deserialize_schema):
        """NOTE: No access rule checks — PIP endpoint for Permission Platform."""
        uuids = deserialize_schema['uuids']
        return flaskify(feature.lookup_features_by_vendor_uuids(uuids))

    @app.route('/lookup/subaccount/features/uuids/', methods=['POST'])
    @validate_request_data(LookupSubaccountFeaturesByUuids())
    def lookup_subaccount_features_by_uuids(deserialize_schema):
        """NOTE: No access rule checks — PIP endpoint for Permission Platform."""
        uuids = deserialize_schema['uuids']
        return flaskify(feature.lookup_features_by_subaccount_uuids(uuids))
    ```

## Files to Create/Modify

### New Files (ows-pdp)

1. `pdp/connectors/cerbos_policy_parser.py` - Parser, `PolicyMetadataDatabase` (incl. `to_json`/`from_json`)
2. `pdp/proxies/vendor_feature_controls_proxy.py` - `VendorFeatureControlsProxy` and `FeatureControlsLookupError`
3. `pdp/cli/commands/cerbos.py` - `seed_policy_metadata_cache` CLI command (invoked by Jenkins)

### Modified Files (ows-pdp)

1. `pdp/constants/constants.py` - Add `CACHE_ENTRY_CERBOS_POLICY_METADATA` constant and `CacheEntryType` member
2. `pdp/fastapi/datasources.py` - Read database from Redis at startup (not from filesystem)
3. `pdp/fastapi/schemas/identity.py` - Implement Resource methods
4. `pdp/logic/cerbos.py` - Add hydration function and update check_resources
5. `pdp/fastapi/routers/identity.py` - Inject database dependency
6. `pdp/connectors/ows_account.py` - Add `LookupVendorFeaturesResponse`, `LookupSubaccountFeaturesResponse`
   models and `get_features_for_vendor_uuids()`, `get_features_for_subaccount_uuids()` methods
7. `Makefile` - Add `seed_policy_metadata_cache` target
8. `Jenkinsfile` - Add `Seed QA Policy Metadata Cache` and `Seed Prod Policy Metadata Cache` stages
9. `pdp/cli/main.py` (or equivalent) - Register the new `cerbos` CLI group

### Modified Files (ows-account)

1. `account/models/feature.py` - Add SQL queries + `get_enabled_feature_ids_for_vendor_uuids()` and `get_enabled_feature_ids_for_subaccount_uuids()`
2. `account/logic/feature.py` - Add `lookup_features_by_vendor_uuids()` and `lookup_features_by_subaccount_uuids()`
3. `account/validation/schemas/lookup.py` - Add `LookupVendorFeaturesByUuids` and `LookupSubaccountFeaturesByUuids` schemas
4. `account/handlers/lookups.py` - Add handlers for `POST /lookup/vendors/features/uuids/` and `POST /lookup/subaccount/features/uuids/`

### Dependencies to Add

- `pyyaml` (if not already present in ows-pdp) - for parsing YAML policy files

## Next Steps

1. ~~Review this plan with the team~~ Done — Option 2 selected
2. Answer open questions about external services (workstation roles service still TBD)
3. Implement `cerbos_policy_parser.py` with `to_json`/`from_json` support
4. Add `CACHE_ENTRY_CERBOS_POLICY_METADATA` constant
5. Implement `seed_policy_metadata_cache` CLI command and Makefile target
6. Update `datasources.py` to read from Redis at startup
7. Add Jenkinsfile stages for QA and Prod
8. Implement `VendorFeatureControlsProxy` and hydration logic
9. Update endpoints and inject dependency
10. Add tests (unit tests for parser, integration tests for full flow)
11. Test CLI command against a local/QA Redis instance with real policy files