Source code for faust.assignor.cluster_assignment

"""Cluster assignment."""
from typing import List, MutableMapping, Sequence, Set, cast
from faust.models import Record
from .client_assignment import (
    ClientAssignment,
    ClientMetadata,
    CopartitionedAssignment,
)

__all__ = ['CopartMapping', 'ClusterAssignment']

CopartMapping = MutableMapping[str, CopartitionedAssignment]


class ClusterAssignment(Record,
                        serializer='json',
                        include_metadata=False,
                        namespace='@ClusterAssignment'):
    """Cluster assignment state."""

    # These are optional, but should never be set to None
    subscriptions: MutableMapping[str, Sequence[str]] = cast(
        MutableMapping[str, Sequence[str]], None)
    assignments: MutableMapping[str, ClientAssignment] = cast(
        MutableMapping[str, ClientAssignment], None)

    def __post_init__(self) -> None:
        if self.subscriptions is None:
            self.subscriptions = {}
        if self.assignments is None:
            self.assignments = {}

    def topics(self) -> Set[str]:
        # All topics subscribed to in the cluster
        return {topic for sub in self.subscriptions.values() for topic in sub}

    def add_client(self, client: str, subscription: List[str],
                   metadata: ClientMetadata) -> None:
        self.subscriptions[client] = list(subscription)
        self.assignments[client] = metadata.assignment

    def copartitioned_assignments(
            self, copartitioned_topics: Set[str]) -> CopartMapping:
        # We only pick clients that subscribe to all copartitioned topics
        subscribed_clis = {
            cli for cli, sub in self.subscriptions.items()
            if copartitioned_topics.issubset(sub)
        }
        return {
            cli: assignment.copartitioned_assignment(copartitioned_topics)
            for cli, assignment in self.assignments.items()
            if cli in subscribed_clis
        }