Skip to content

Commit

Permalink
add typing to aiokafka/cluster.py
Browse files Browse the repository at this point in the history
  • Loading branch information
@dimastbk
dimastbk committed Jul 2, 2024
1 parent b861a26 commit a219a68
Show file tree
Hide file tree
Showing 6 changed files with 185 additions and 67 deletions.
2 changes: 2 additions & 0 deletions Makefile
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ KAFKA_VERSION?=2.8.1
DOCKER_IMAGE=aiolibs/kafka:$(SCALA_VERSION)_$(KAFKA_VERSION)
DIFF_BRANCH=origin/master
FORMATTED_AREAS=\
aiokafka/cluster.py \
aiokafka/codec.py \
aiokafka/coordinator/ \
aiokafka/errors.py \
Expand All @@ -14,6 +15,7 @@ FORMATTED_AREAS=\
aiokafka/util.py \
aiokafka/protocol/ \
aiokafka/record/ \
tests/test_cluster.py \
tests/test_codec.py \
tests/test_helpers.py \
tests/test_protocol.py \
Expand Down
3 changes: 2 additions & 1 deletion aiokafka/client.py
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import asyncio
import contextlib
import enum
import logging
import random
import time
Expand Down Expand Up @@ -45,7 +46,7 @@ class ConnectionGroup:
COORDINATION = 1


class CoordinationType:
class CoordinationType(enum.IntEnum):
GROUP = 0
TRANSACTION = 1

Expand Down
160 changes: 109 additions & 51 deletions aiokafka/cluster.py
View file
Original file line number Diff line number Diff line change
@@ -1,18 +1,44 @@
from __future__ import annotations

import collections
import copy
import logging
import threading
import time
from concurrent.futures import Future
from typing import Optional, Set
from typing import Any, Callable, Dict, Iterable, Optional, Set, Tuple, Union

from typing_extensions import TypeAlias, TypedDict

from aiokafka import errors as Errors
from aiokafka.client import CoordinationType

Check failure

Code scanning / CodeQL

Module-level cyclic import Error

'CoordinationType' may not be defined if module
aiokafka.client
Loading
is imported before module
aiokafka.cluster
Loading
, as the
definition
Loading
of CoordinationType occurs after the cyclic
import
Loading
of aiokafka.cluster.
from aiokafka.conn import collect_hosts
from aiokafka.protocol.commit import (
GroupCoordinatorResponse_v0,
GroupCoordinatorResponse_v1,
)
from aiokafka.protocol.metadata import (
MetadataResponse_v0,
MetadataResponse_v1,
MetadataResponse_v2,
MetadataResponse_v3,
MetadataResponse_v4,
MetadataResponse_v5,
)
from aiokafka.structs import BrokerMetadata, PartitionMetadata, TopicPartition

ListenerCallable: TypeAlias = Callable[["ClusterMetadata"], None]
NodeId: TypeAlias = Union[str, int]

log = logging.getLogger(__name__)


class Config(TypedDict):
retry_backoff_ms: float
metadata_max_age_ms: float
bootstrap_servers: Union[str, Iterable[str]]


class ClusterMetadata:
"""
A class to manage kafka cluster metadata.
Expand All @@ -35,61 +61,67 @@ class ClusterMetadata:
specified, will default to localhost:9092.
"""

DEFAULT_CONFIG = {
DEFAULT_CONFIG: Config = {
"retry_backoff_ms": 100,
"metadata_max_age_ms": 300000,
"bootstrap_servers": [],
}

def __init__(self, **configs):
self._brokers = {} # node_id -> BrokerMetadata
self._partitions = {} # topic -> partition -> PartitionMetadata
def __init__(self, **configs: Any) -> None:
self._brokers: Dict[NodeId, BrokerMetadata] = {} # node_id -> BrokerMetadata
self._partitions: Dict[
str, Dict[int, PartitionMetadata]
] = {} # topic -> partition -> PartitionMetadata
# node_id -> {TopicPartition...}
self._broker_partitions = collections.defaultdict(set)
self._groups = {} # group_name -> node_id
self._last_refresh_ms = 0
self._last_successful_refresh_ms = 0
self._broker_partitions: Dict[NodeId, Set[TopicPartition]] = (
collections.defaultdict(set)
)
self._groups: Dict[str, NodeId] = {} # group_name -> node_id
self._last_refresh_ms: float = 0
self._last_successful_refresh_ms: float = 0
self._need_update = True
self._future = None
self._listeners = set()
self._future: Optional[Future[ClusterMetadata]] = None
self._listeners: Set[ListenerCallable] = set()
self._lock = threading.Lock()
self.need_all_topic_metadata = False
self.unauthorized_topics = set()
self.internal_topics = set()
self.controller = None
self.unauthorized_topics: Set[str] = set()
self.internal_topics: Set[str] = set()
self.controller: Optional[BrokerMetadata] = None

self.config = copy.copy(self.DEFAULT_CONFIG)
self.config: Config = copy.copy(self.DEFAULT_CONFIG)
for key in self.config:
if key in configs:
self.config[key] = configs[key]
self.config[key] = configs[key] # type: ignore[literal-required]

self._bootstrap_brokers = self._generate_bootstrap_brokers()
self._coordinator_brokers = {}
self._coordinators = {}
self._coordinator_by_key = {}
self._bootstrap_brokers: Dict[NodeId, BrokerMetadata] = (
self._generate_bootstrap_brokers()
)
self._coordinator_brokers: Dict[NodeId, BrokerMetadata] = {}
self._coordinators: Dict[NodeId, BrokerMetadata] = {}
self._coordinator_by_key: Dict[Tuple[CoordinationType, str], NodeId] = {}

def _generate_bootstrap_brokers(self):
def _generate_bootstrap_brokers(self) -> Dict[NodeId, BrokerMetadata]:
# collect_hosts does not perform DNS, so we should be fine to re-use
bootstrap_hosts = collect_hosts(self.config["bootstrap_servers"])

brokers = {}
brokers: Dict[NodeId, BrokerMetadata] = {}
for i, (host, port, _) in enumerate(bootstrap_hosts):
node_id = "bootstrap-%s" % i
brokers[node_id] = BrokerMetadata(node_id, host, port, None)
return brokers

def is_bootstrap(self, node_id):
def is_bootstrap(self, node_id: NodeId) -> bool:
return node_id in self._bootstrap_brokers

def brokers(self):
def brokers(self) -> Set[BrokerMetadata]:
"""Get all BrokerMetadata
Returns:
set: {BrokerMetadata, ...}
"""
return set(self._brokers.values()) or set(self._bootstrap_brokers.values())

def broker_metadata(self, broker_id):
def broker_metadata(self, broker_id: NodeId) -> Optional[BrokerMetadata]:
"""Get BrokerMetadata
Arguments:
Expand Down Expand Up @@ -117,7 +149,7 @@ def partitions_for_topic(self, topic: str) -> Optional[Set[int]]:
return None
return set(self._partitions[topic].keys())

def available_partitions_for_topic(self, topic):
def available_partitions_for_topic(self, topic: str) -> Optional[Set[int]]:
"""Return set of partitions with known leaders
Arguments:
Expand All @@ -135,7 +167,7 @@ def available_partitions_for_topic(self, topic):
if metadata.leader != -1
}

def leader_for_partition(self, partition):
def leader_for_partition(self, partition: PartitionMetadata) -> Optional[int]:
"""Return node_id of leader, -1 unavailable, None if unknown."""
if partition.topic not in self._partitions:
return None
Expand All @@ -144,7 +176,7 @@ def leader_for_partition(self, partition):
return None
return partitions[partition.partition].leader

def partitions_for_broker(self, broker_id):
def partitions_for_broker(self, broker_id: NodeId) -> Optional[Set[TopicPartition]]:
"""Return TopicPartitions for which the broker is a leader.
Arguments:
Expand All @@ -156,7 +188,7 @@ def partitions_for_broker(self, broker_id):
"""
return self._broker_partitions.get(broker_id)

def coordinator_for_group(self, group):
def coordinator_for_group(self, group: str) -> Optional[NodeId]:
"""Return node_id of group coordinator.
Arguments:
Expand All @@ -168,7 +200,7 @@ def coordinator_for_group(self, group):
"""
return self._groups.get(group)

def request_update(self):
def request_update(self) -> Future[ClusterMetadata]:
"""Flags metadata for update, return Future()
Actual update must be handled separately. This method will only
Expand All @@ -179,11 +211,11 @@ def request_update(self):
"""
with self._lock:
self._need_update = True
if not self._future or self._future.is_done:
if not self._future or self._future.done():
self._future = Future()
return self._future

def topics(self, exclude_internal_topics=True):
def topics(self, exclude_internal_topics: bool = True) -> Set[str]:
"""Get set of known topics.
Arguments:
Expand All @@ -201,18 +233,28 @@ def topics(self, exclude_internal_topics=True):
else:
return topics

def failed_update(self, exception):
def failed_update(self, exception: Exception) -> None:
"""Update cluster state given a failed MetadataRequest."""
f = None
with self._lock:
if self._future:
f = self._future
self._future = None
if f:
f.failure(exception)
f.set_exception(exception)
self._last_refresh_ms = time.time() * 1000

def update_metadata(self, metadata):
def update_metadata(
self,
metadata: Union[
MetadataResponse_v0,
MetadataResponse_v1,
MetadataResponse_v2,
MetadataResponse_v3,
MetadataResponse_v4,
MetadataResponse_v5,
],
) -> None:
"""Update cluster state given a MetadataResponse.
Arguments:
Expand All @@ -225,27 +267,29 @@ def update_metadata(self, metadata):
self.failed_update(Errors.MetadataEmptyBrokerList(metadata))
return

_new_brokers = {}
_new_brokers: Dict[NodeId, BrokerMetadata] = {}
for broker in metadata.brokers:
if metadata.API_VERSION == 0:
if isinstance(metadata, MetadataResponse_v0):
node_id, host, port = broker
rack = None
else:
node_id, host, port, rack = broker
_new_brokers.update({node_id: BrokerMetadata(node_id, host, port, rack)})

if metadata.API_VERSION == 0:
if isinstance(metadata, MetadataResponse_v0):
_new_controller = None
else:
_new_controller = _new_brokers.get(metadata.controller_id)

_new_partitions = {}
_new_broker_partitions = collections.defaultdict(set)
_new_unauthorized_topics = set()
_new_internal_topics = set()
_new_partitions: Dict[str, Dict[int, PartitionMetadata]] = {}
_new_broker_partitions: Dict[NodeId, Set[TopicPartition]] = (
collections.defaultdict(set)
)
_new_unauthorized_topics: Set[str] = set()
_new_internal_topics: Set[str] = set()

for topic_data in metadata.topics:
if metadata.API_VERSION == 0:
if isinstance(metadata, MetadataResponse_v0):
error_code, topic, partitions = topic_data
is_internal = False
else:
Expand Down Expand Up @@ -307,7 +351,7 @@ def update_metadata(self, metadata):
self._last_successful_refresh_ms = now

if f:
f.success(self)
f.set_result(self)
log.debug("Updated cluster metadata to %s", self)

for listener in self._listeners:
Expand All @@ -320,15 +364,19 @@ def update_metadata(self, metadata):
# another fetch should be unnecessary.
self._need_update = False

def add_listener(self, listener):
def add_listener(self, listener: ListenerCallable) -> None:
"""Add a callback function to be called on each metadata update"""
self._listeners.add(listener)

def remove_listener(self, listener):
def remove_listener(self, listener: ListenerCallable) -> None:
"""Remove a previously added listener callback"""
self._listeners.remove(listener)

def add_group_coordinator(self, group, response):
def add_group_coordinator(
self,
group: str,
response: Union[GroupCoordinatorResponse_v0, GroupCoordinatorResponse_v1],
) -> Optional[str]:
"""Update with metadata for a group coordinator
Arguments:
Expand All @@ -355,7 +403,9 @@ def add_group_coordinator(self, group, response):
self._groups[group] = node_id
return node_id

def with_partitions(self, partitions_to_add):
def with_partitions(
self, partitions_to_add: Iterable[PartitionMetadata]
) -> ClusterMetadata:
"""Returns a copy of cluster metadata with partitions added"""
new_metadata = ClusterMetadata(**self.config)
new_metadata._brokers = copy.deepcopy(self._brokers)
Expand All @@ -375,10 +425,18 @@ def with_partitions(self, partitions_to_add):

return new_metadata

def coordinator_metadata(self, node_id):
def coordinator_metadata(self, node_id: NodeId) -> Optional[BrokerMetadata]:
return self._coordinators.get(node_id)

def add_coordinator(self, node_id, host, port, rack=None, *, purpose):
def add_coordinator(
self,
node_id: int,
host: str,
port: int,
rack: Optional[str] = None,
*,
purpose: Tuple[CoordinationType, str],
) -> None:
"""Keep track of all coordinator nodes separately and remove them if
a new one was elected for the same purpose (For example group
coordinator for group X).
Expand All @@ -390,7 +448,7 @@ def add_coordinator(self, node_id, host, port, rack=None, *, purpose):
self._coordinators[node_id] = BrokerMetadata(node_id, host, port, rack)
self._coordinator_by_key[purpose] = node_id

def __str__(self):
def __str__(self) -> str:
return "ClusterMetadata(brokers: %d, topics: %d, groups: %d)" % (
len(self._brokers),
len(self._partitions),
Expand Down
Loading

0 comments on commit a219a68

Please sign in to comment.