wehub-resource-sync
818 lines
35 KiB
Python
818 lines
35 KiB
Python
"""Kafka datasource for bounded data reads.
|
|
|
|
This module provides a Kafka datasource implementation for Ray Data that supports
|
|
bounded reads with offset-based range queries.
|
|
|
|
Message keys and values are returned as raw bytes to support any serialization format
|
|
(JSON, Avro, Protobuf, etc.). Users can decode them using map operations.
|
|
|
|
Requires:
|
|
- confluent-kafka: https://docs.confluent.io/platform/current/clients/confluent-kafka-python/html/
|
|
"""
|
|
|
|
import logging
|
|
import time
|
|
import warnings
|
|
from dataclasses import dataclass
|
|
from datetime import datetime, timezone
|
|
from typing import (
|
|
TYPE_CHECKING,
|
|
Any,
|
|
Dict,
|
|
Iterable,
|
|
List,
|
|
Literal,
|
|
Optional,
|
|
Set,
|
|
Tuple,
|
|
Union,
|
|
)
|
|
|
|
import pyarrow as pa
|
|
|
|
if TYPE_CHECKING:
|
|
from confluent_kafka import Consumer, TopicPartition
|
|
|
|
from ray.data._internal.output_buffer import BlockOutputBuffer, OutputBlockSizeOption
|
|
from ray.data._internal.util import _check_import
|
|
from ray.data.block import Block, BlockMetadata
|
|
from ray.data.context import DataContext
|
|
from ray.data.datasource import Datasource, ReadTask
|
|
|
|
PartitionOffsets = Dict[int, Union[int, str]]
|
|
PerPartitionOffsets = Dict[str, PartitionOffsets]
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Mapping from kafka-python style KafkaAuthConfig fields to Confluent/librdkafka config.
|
|
# TODO(youcheng): Remove this mapping and use consumer_config directly.
|
|
_KAFKA_AUTH_TO_CONFLUENT: Dict[str, str] = {
|
|
"security_protocol": "security.protocol",
|
|
"sasl_mechanism": "sasl.mechanism",
|
|
"sasl_plain_username": "sasl.username",
|
|
"sasl_plain_password": "sasl.password",
|
|
"sasl_kerberos_service_name": "sasl.kerberos.service.name",
|
|
"sasl_kerberos_name": "sasl.kerberos.principal",
|
|
"ssl_cafile": "ssl.ca.location",
|
|
"ssl_certfile": "ssl.certificate.location",
|
|
"ssl_keyfile": "ssl.key.location",
|
|
"ssl_password": "ssl.key.password",
|
|
"ssl_ciphers": "ssl.cipher.suites",
|
|
"ssl_crlfile": "ssl.crl.location",
|
|
# Note: ssl_check_hostname is intentionally NOT mapped due to semantics mismatch.
|
|
}
|
|
|
|
|
|
KAFKA_TOPIC_METADATA_TIMEOUT_S = 10
|
|
KAFKA_QUERY_OFFSET_TIMEOUT_S = 10
|
|
# Cap each consume timeout to keep responsiveness of timeout/position checks.
|
|
KAFKA_CONSUME_TIMEOUT_MAX_S = 10 # 10 seconds per consume call
|
|
|
|
KAFKA_MSG_SCHEMA = pa.schema(
|
|
[
|
|
("offset", pa.int64()),
|
|
("key", pa.binary()),
|
|
("value", pa.binary()),
|
|
("topic", pa.string()),
|
|
("partition", pa.int32()),
|
|
("timestamp", pa.int64()), # Kafka timestamp in milliseconds
|
|
(
|
|
"timestamp_type",
|
|
pa.int32(),
|
|
), # 0=TIMESTAMP_NOT_AVAILABLE, 1=TIMESTAMP_CREATE_TIME, 2=TIMESTAMP_LOG_APPEND_TIME
|
|
("headers", pa.map_(pa.string(), pa.binary())), # Message headers
|
|
]
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class KafkaAuthConfig:
|
|
"""Authentication configuration for Kafka connections.
|
|
|
|
Uses standard kafka-python parameter names. See kafka-python documentation
|
|
for full details: https://kafka-python.readthedocs.io/
|
|
|
|
Note: Ray Data maps these options to Confluent/librdkafka config under the hood.
|
|
Some options have different semantics or are unsupported by the Confluent client;
|
|
see notes below for those fields. Prefer passing Confluent options directly via
|
|
``consumer_config`` where possible.
|
|
|
|
security_protocol: Protocol used to communicate with brokers.
|
|
Valid values are: PLAINTEXT, SSL, SASL_PLAINTEXT, SASL_SSL.
|
|
Default: PLAINTEXT.
|
|
sasl_mechanism: Authentication mechanism when security_protocol
|
|
is configured for SASL_PLAINTEXT or SASL_SSL. Valid values are:
|
|
PLAIN, GSSAPI, OAUTHBEARER, SCRAM-SHA-256, SCRAM-SHA-512.
|
|
sasl_plain_username: username for sasl PLAIN and SCRAM authentication.
|
|
Required if sasl_mechanism is PLAIN or one of the SCRAM mechanisms.
|
|
sasl_plain_password: password for sasl PLAIN and SCRAM authentication.
|
|
Required if sasl_mechanism is PLAIN or one of the SCRAM mechanisms.
|
|
sasl_kerberos_name: Constructed gssapi.Name for use with
|
|
sasl mechanism handshake. If provided, sasl_kerberos_service_name and
|
|
sasl_kerberos_domain name are ignored. Default: None.
|
|
sasl_kerberos_service_name: Service name to include in GSSAPI
|
|
sasl mechanism handshake. Default: 'kafka'
|
|
sasl_kerberos_domain_name: kerberos domain name to use in GSSAPI
|
|
sasl mechanism handshake. Default: one of bootstrap servers.
|
|
Note (Confluent): This option is not supported by Confluent/librdkafka
|
|
and will be ignored when building the client configuration. Prefer specifying
|
|
an explicit principal via ``sasl_kerberos_name`` or rely on defaults.
|
|
sasl_oauth_token_provider: OAuthBearer token provider instance. Default: None.
|
|
Note (Confluent): Not supported directly; use ``consumer_config`` with
|
|
``sasl.oauthbearer.*`` options instead.
|
|
ssl_context: Pre-configured SSLContext for wrapping socket connections. If provided,
|
|
all other ssl_* configurations will be ignored. Default: None.
|
|
Note (Confluent): Passing an SSLContext object is not supported and will be
|
|
ignored. Use ``ssl_cafile``, ``ssl_certfile``, and ``ssl_keyfile`` instead.
|
|
ssl_check_hostname: Flag to configure whether ssl handshake should verify that the
|
|
certificate matches the broker's hostname. Default: True.
|
|
Note (Confluent): There is no 1:1 equivalent; disabling hostname verification
|
|
via ``enable.ssl.certificate.verification=False`` would also disable the entire
|
|
certificate chain verification. To avoid weakening security, this flag is not
|
|
mapped when False. If you need to disable only hostname verification, set
|
|
``ssl.endpoint.identification.algorithm=none`` via ``consumer_config`` (if supported
|
|
by your librdkafka version).
|
|
ssl_cafile: Optional filename of ca file to use in certificate verification. Default: None.
|
|
ssl_certfile: Optional filename of file in pem format containing the client certificate,
|
|
as well as any ca certificates needed to establish the certificate's authenticity. Default: None.
|
|
ssl_keyfile: Optional filename containing the client private key. Default: None.
|
|
ssl_password: Optional password to be used when loading the certificate chain. Default: None.
|
|
ssl_crlfile: Optional filename containing the CRL to check for certificate expiration. By default,
|
|
no CRL check is done. When providing a file, only the leaf certificate will be checked against
|
|
this CRL. The CRL can only be checked with Python 3.4+ or 2.7.9+. Default: None.
|
|
ssl_ciphers: optionally set the available ciphers for ssl connections. It should be a string in the
|
|
OpenSSL cipher list format. If no cipher can be selected (because compile-time options or other
|
|
configuration forbids use of all the specified ciphers), an ssl.SSLError will be raised.
|
|
See ssl.SSLContext.set_ciphers.
|
|
"""
|
|
|
|
# Security protocol
|
|
security_protocol: Optional[str] = None
|
|
|
|
# SASL configuration
|
|
sasl_mechanism: Optional[str] = None
|
|
sasl_plain_username: Optional[str] = None
|
|
sasl_plain_password: Optional[str] = None
|
|
sasl_kerberos_name: Optional[str] = None
|
|
sasl_kerberos_service_name: Optional[str] = None
|
|
sasl_kerberos_domain_name: Optional[str] = None
|
|
sasl_oauth_token_provider: Optional[Any] = None
|
|
|
|
# SSL configuration
|
|
ssl_context: Optional[Any] = None
|
|
ssl_check_hostname: Optional[bool] = None
|
|
ssl_cafile: Optional[str] = None
|
|
ssl_certfile: Optional[str] = None
|
|
ssl_keyfile: Optional[str] = None
|
|
ssl_password: Optional[str] = None
|
|
ssl_ciphers: Optional[str] = None
|
|
ssl_crlfile: Optional[str] = None
|
|
|
|
|
|
def _handle_deprecated_configs(kafka_auth_config: KafkaAuthConfig) -> None:
|
|
# Handle special fields with warnings
|
|
if kafka_auth_config.ssl_context is not None:
|
|
logger.warning(
|
|
"ssl_context is not supported by Confluent. Skipping. "
|
|
"Use KafkaAuthConfig fields ssl_cafile, ssl_certfile, ssl_keyfile instead."
|
|
)
|
|
if kafka_auth_config.sasl_oauth_token_provider is not None:
|
|
logger.warning(
|
|
"sasl_oauth_token_provider is not supported by Confluent. Skipping. "
|
|
"Use consumer_config with sasl.oauthbearer.* options instead."
|
|
)
|
|
if kafka_auth_config.sasl_kerberos_domain_name is not None:
|
|
logger.warning(
|
|
"sasl_kerberos_domain_name is not supported by Confluent and will be ignored. "
|
|
"Set sasl_kerberos_name (principal) or rely on defaults."
|
|
)
|
|
if kafka_auth_config.ssl_check_hostname is False:
|
|
logger.warning(
|
|
"ssl_check_hostname=False cannot be mapped safely to Confluent; "
|
|
"setting enable.ssl.certificate.verification=False would disable all certificate verification. "
|
|
"Ignoring ssl_check_hostname. If you need to disable only hostname verification, "
|
|
"configure the client directly via consumer_config (e.g., ssl.endpoint.identification.algorithm=none)."
|
|
)
|
|
|
|
|
|
def _add_authentication_to_config(
|
|
config: Dict[str, Any], kafka_auth_config: Optional[KafkaAuthConfig]
|
|
) -> None:
|
|
"""Map KafkaAuthConfig (kafka-python style) into Confluent/librdkafka config.
|
|
|
|
Special cases:
|
|
- ssl_context: unsupported; warn and ignore
|
|
- sasl_oauth_token_provider: unsupported; warn and ignore
|
|
- sasl_kerberos_domain_name: unsupported; warn and ignore
|
|
- ssl_check_hostname: not mapped due to semantics; if False, warn and ignore
|
|
"""
|
|
if not kafka_auth_config:
|
|
return
|
|
warnings.warn(
|
|
"kafka_auth_config (kafka-python style) is deprecated and will be removed in a future release. "
|
|
"Please provide Confluent/librdkafka options via consumer_config instead.",
|
|
DeprecationWarning,
|
|
stacklevel=2,
|
|
)
|
|
|
|
_handle_deprecated_configs(kafka_auth_config)
|
|
|
|
# Map directly compatible fields
|
|
for key, confluent_key in _KAFKA_AUTH_TO_CONFLUENT.items():
|
|
val = getattr(kafka_auth_config, key, None)
|
|
if val is not None:
|
|
config[confluent_key] = val
|
|
|
|
|
|
def _build_confluent_config(
|
|
bootstrap_servers: List[str],
|
|
kafka_auth_config: Optional[KafkaAuthConfig] = None,
|
|
extra: Optional[Dict[str, Any]] = None,
|
|
user_config: Optional[Dict[str, Any]] = None,
|
|
) -> Dict[str, Any]:
|
|
"""Build Confluent config with bootstrap servers and auth.
|
|
|
|
Args:
|
|
bootstrap_servers: List of Kafka broker addresses.
|
|
kafka_auth_config: Authentication configuration (kafka-python style). Deprecated; prefer consumer_config with Confluent keys. Mutually exclusive with consumer_config.
|
|
extra: Additional config options.
|
|
user_config: User-provided config options.
|
|
|
|
Returns:
|
|
Confluent configuration dict.
|
|
"""
|
|
config: Dict[str, Any] = {
|
|
"bootstrap.servers": ",".join(bootstrap_servers),
|
|
}
|
|
# Map kafka-python-style auth if provided
|
|
_add_authentication_to_config(config, kafka_auth_config)
|
|
if extra:
|
|
config.update(extra)
|
|
if user_config:
|
|
if (
|
|
"bootstrap.servers" in user_config
|
|
and user_config["bootstrap.servers"] != config["bootstrap.servers"]
|
|
):
|
|
logger.warning(
|
|
"Ignoring 'bootstrap.servers' from consumer_config; use bootstrap_servers parameter instead."
|
|
)
|
|
for k, v in user_config.items():
|
|
if k == "bootstrap.servers":
|
|
continue
|
|
config[k] = v
|
|
return config
|
|
|
|
|
|
def _build_consumer_config_for_read(
|
|
bootstrap_servers: List[str],
|
|
kafka_auth_config: Optional[KafkaAuthConfig] = None,
|
|
consumer_config: Optional[Dict[str, Any]] = None,
|
|
) -> Dict[str, Any]:
|
|
"""Build Consumer config for reading messages (Confluent)."""
|
|
return _build_confluent_config(
|
|
bootstrap_servers,
|
|
extra={
|
|
"enable.auto.commit": False,
|
|
# Confluent requires a group.id even when using manual assign.
|
|
"group.id": "ray-data-kafka-reader",
|
|
},
|
|
user_config=consumer_config,
|
|
kafka_auth_config=kafka_auth_config,
|
|
)
|
|
|
|
|
|
def _datetime_to_ms(dt: datetime) -> int:
|
|
"""Convert a datetime to milliseconds since epoch (UTC).
|
|
|
|
If the datetime has no timezone info (i.e., ``tzinfo is None``),
|
|
it is assumed to be UTC. Timezone-aware datetimes are converted to
|
|
UTC automatically via ``datetime.timestamp()``.
|
|
|
|
Args:
|
|
dt: A datetime object, with or without timezone info.
|
|
|
|
Returns:
|
|
Milliseconds since Unix epoch.
|
|
"""
|
|
if dt.tzinfo is None:
|
|
dt = dt.replace(tzinfo=timezone.utc)
|
|
return int(dt.timestamp() * 1000)
|
|
|
|
|
|
def _validate_offsets(
|
|
start_offset: Union[int, datetime, Literal["earliest"], PerPartitionOffsets],
|
|
end_offset: Union[int, datetime, Literal["latest"], PerPartitionOffsets],
|
|
) -> None:
|
|
if isinstance(start_offset, dict):
|
|
for topic, partition_map in start_offset.items():
|
|
if not isinstance(partition_map, dict):
|
|
raise ValueError(
|
|
f"start_offset[{topic!r}] must be a dict mapping "
|
|
f"partition_id (int) to offset (int or str)."
|
|
)
|
|
for partition_id, offset in partition_map.items():
|
|
if not isinstance(partition_id, int):
|
|
raise ValueError(
|
|
f"start_offset[{topic!r}] keys must be integers "
|
|
f"(partition IDs), got {type(partition_id).__name__!r}."
|
|
)
|
|
if isinstance(offset, str) and offset == "latest":
|
|
raise ValueError(
|
|
f"start_offset[{topic!r}][{partition_id}] cannot be 'latest'."
|
|
)
|
|
else:
|
|
if isinstance(start_offset, int) and isinstance(end_offset, int):
|
|
if start_offset > end_offset:
|
|
raise ValueError("start_offset must be less than end_offset")
|
|
if isinstance(start_offset, datetime) and isinstance(end_offset, datetime):
|
|
if _datetime_to_ms(start_offset) > _datetime_to_ms(end_offset):
|
|
raise ValueError("start_offset must be less than end_offset")
|
|
if isinstance(start_offset, str) and start_offset == "latest":
|
|
raise ValueError("start_offset cannot be 'latest'")
|
|
|
|
if isinstance(end_offset, dict):
|
|
for topic, partition_map in end_offset.items():
|
|
if not isinstance(partition_map, dict):
|
|
raise ValueError(
|
|
f"end_offset[{topic!r}] must be a dict mapping "
|
|
f"partition_id (int) to offset (int or str)."
|
|
)
|
|
for partition_id, offset in partition_map.items():
|
|
if not isinstance(partition_id, int):
|
|
raise ValueError(
|
|
f"end_offset[{topic!r}] keys must be integers "
|
|
f"(partition IDs), got {type(partition_id).__name__!r}."
|
|
)
|
|
if isinstance(offset, str) and offset == "earliest":
|
|
raise ValueError(
|
|
f"end_offset[{topic!r}][{partition_id}] cannot be 'earliest'."
|
|
)
|
|
else:
|
|
if isinstance(end_offset, str) and end_offset == "earliest":
|
|
raise ValueError("end_offset cannot be 'earliest'")
|
|
|
|
|
|
def _resolve_datetime_to_offset(
|
|
consumer: "Consumer",
|
|
topic_partition: "TopicPartition",
|
|
dt: datetime,
|
|
fallback_offset: int,
|
|
) -> int:
|
|
"""Resolve a datetime to an integer offset via offsets_for_times.
|
|
|
|
Returns fallback_offset when offsets_for_times returns empty or offset < 0
|
|
(e.g., datetime in future or no messages at that time).
|
|
"""
|
|
from confluent_kafka import TopicPartition
|
|
|
|
timestamp_ms = _datetime_to_ms(dt)
|
|
tp_with_ts = TopicPartition(
|
|
topic_partition.topic, topic_partition.partition, timestamp_ms
|
|
)
|
|
result = consumer.offsets_for_times(
|
|
[tp_with_ts], timeout=KAFKA_QUERY_OFFSET_TIMEOUT_S
|
|
)
|
|
if result and result[0].offset >= 0:
|
|
return result[0].offset
|
|
return fallback_offset
|
|
|
|
|
|
def _resolve_offsets(
|
|
consumer: "Consumer",
|
|
topic_partition: "TopicPartition",
|
|
start_offset: Union[int, datetime, Literal["earliest"]],
|
|
end_offset: Union[int, datetime, Literal["latest"]],
|
|
) -> Tuple[int, int]:
|
|
"""Resolve start and end offsets to actual integer offsets.
|
|
|
|
Handles int offsets, "earliest"/"latest" strings, and datetime objects.
|
|
For datetime objects, uses ``consumer.offsets_for_times()`` to find the
|
|
earliest offset whose timestamp is >= the given datetime.
|
|
|
|
Args:
|
|
consumer: Confluent Kafka consumer instance.
|
|
topic_partition: TopicPartition to resolve offsets for.
|
|
start_offset: Start offset (int, datetime, or "earliest").
|
|
end_offset: End offset (int, datetime, or "latest").
|
|
|
|
Returns:
|
|
Tuple of (resolved_start_offset, resolved_end_offset).
|
|
"""
|
|
# TODO(youcheng): add retry logic for this call.
|
|
low, high = consumer.get_watermark_offsets(
|
|
topic_partition, timeout=KAFKA_QUERY_OFFSET_TIMEOUT_S
|
|
)
|
|
earliest_offset = low
|
|
latest_offset = high
|
|
|
|
# Keep original values for error messages
|
|
original_start = start_offset
|
|
original_end = end_offset
|
|
|
|
if start_offset == "earliest" or start_offset is None:
|
|
start_offset = earliest_offset
|
|
elif isinstance(start_offset, datetime):
|
|
# fallback to latest_offset if the start_offset is in the future, so the read range is empty (start == end).
|
|
start_offset = _resolve_datetime_to_offset(
|
|
consumer, topic_partition, start_offset, latest_offset
|
|
)
|
|
|
|
if end_offset == "latest" or end_offset is None:
|
|
end_offset = latest_offset
|
|
elif isinstance(end_offset, datetime):
|
|
end_offset = _resolve_datetime_to_offset(
|
|
consumer, topic_partition, end_offset, latest_offset
|
|
)
|
|
|
|
# Clamp end_offset to the high watermark so we never try to read beyond
|
|
# what is currently available. This prevents the read loop from polling
|
|
# indefinitely when a user-supplied integer end_offset exceeds the number
|
|
# of messages in the partition.
|
|
end_offset = min(end_offset, latest_offset)
|
|
|
|
if isinstance(start_offset, int) and start_offset < earliest_offset:
|
|
logger.warning(
|
|
f"start_offset ({start_offset}) is below the earliest available offset "
|
|
f"({earliest_offset}) for partition {topic_partition.partition} in topic "
|
|
f"{topic_partition.topic} (data may have been deleted by Kafka retention). "
|
|
f"Falling back to earliest available offset ({earliest_offset})."
|
|
)
|
|
start_offset = earliest_offset
|
|
|
|
if start_offset > end_offset:
|
|
start_str = (
|
|
f"{original_start}"
|
|
if original_start == start_offset
|
|
else f"{original_start} (resolved to {start_offset})"
|
|
)
|
|
end_str = (
|
|
f"{original_end}"
|
|
if original_end == end_offset
|
|
else f"{original_end} (resolved to {end_offset})"
|
|
)
|
|
raise ValueError(
|
|
f"start_offset ({start_str}) > end_offset ({end_str}) "
|
|
f"for partition {topic_partition.partition} in topic {topic_partition.topic}"
|
|
)
|
|
|
|
return start_offset, end_offset
|
|
|
|
|
|
class KafkaDatasource(Datasource):
|
|
"""Kafka datasource for reading from Kafka topics with bounded reads."""
|
|
|
|
# Batch size for incremental block yielding
|
|
BATCH_SIZE_FOR_YIELD = 1000
|
|
|
|
def __init__(
|
|
self,
|
|
topics: Union[str, List[str]],
|
|
bootstrap_servers: Union[str, List[str]],
|
|
start_offset: Union[int, datetime, Literal["earliest"], PerPartitionOffsets],
|
|
end_offset: Union[int, datetime, Literal["latest"], PerPartitionOffsets],
|
|
kafka_auth_config: Optional[KafkaAuthConfig] = None,
|
|
consumer_config: Optional[Dict[str, Any]] = None,
|
|
timeout_ms: Optional[int] = None,
|
|
):
|
|
"""Initialize Kafka datasource.
|
|
|
|
Args:
|
|
topics: Kafka topic name(s) to read from.
|
|
bootstrap_servers: Kafka broker addresses (string or list of strings).
|
|
start_offset: Starting position. Can be:
|
|
- int: Offset number
|
|
- datetime: Read from the first message at or after this time.
|
|
datetimes with no timezone info are treated as UTC.
|
|
- str: "earliest"
|
|
end_offset: Ending position (exclusive). Can be:
|
|
- int: Offset number
|
|
- datetime: Read up to (but not including) the first message
|
|
at or after this time. datetimes with no timezone info are treated as UTC.
|
|
- str: "latest"
|
|
kafka_auth_config: Authentication configuration (kafka-python style). Deprecated; prefer consumer_config with Confluent keys. Mutually exclusive with consumer_config.
|
|
consumer_config: Confluent/librdkafka consumer configuration dict.
|
|
Keys and values are passed through to the underlying client. The
|
|
`bootstrap.servers` option is derived from `bootstrap_servers` and
|
|
cannot be overridden here.
|
|
timeout_ms: Optional timeout in milliseconds for every read task to poll until reaching end_offset.
|
|
If None (default), no task-level timeout is applied and each read task
|
|
will poll until it reaches end_offset. If set, the read task will stop
|
|
polling after the timeout and return the messages it has read so far.
|
|
|
|
Raises:
|
|
ValueError: If required configuration is missing.
|
|
ImportError: If confluent-kafka is not installed.
|
|
"""
|
|
_check_import(self, module="confluent_kafka", package="confluent-kafka")
|
|
|
|
if not topics:
|
|
raise ValueError("topics cannot be empty")
|
|
|
|
if not bootstrap_servers:
|
|
raise ValueError("bootstrap_servers cannot be empty")
|
|
|
|
if timeout_ms is not None and timeout_ms <= 0:
|
|
raise ValueError("timeout_ms must be positive")
|
|
|
|
_validate_offsets(start_offset, end_offset)
|
|
|
|
# Validate bootstrap_servers format
|
|
if isinstance(bootstrap_servers, str):
|
|
if not bootstrap_servers or ":" not in bootstrap_servers:
|
|
raise ValueError(
|
|
f"Invalid bootstrap_servers format: {bootstrap_servers}. "
|
|
"Expected 'host:port' or list of 'host:port' strings."
|
|
)
|
|
elif isinstance(bootstrap_servers, list):
|
|
if not bootstrap_servers:
|
|
raise ValueError("bootstrap_servers cannot be empty list")
|
|
for server in bootstrap_servers:
|
|
if not isinstance(server, str) or ":" not in server:
|
|
raise ValueError(
|
|
f"Invalid bootstrap_servers format: {server}. "
|
|
"Expected 'host:port' string."
|
|
)
|
|
|
|
# Disallow specifying both config styles at once to avoid ambiguity.
|
|
if kafka_auth_config is not None and consumer_config is not None:
|
|
raise ValueError(
|
|
"Provide only one of kafka_auth_config (deprecated) or consumer_config, not both."
|
|
)
|
|
|
|
self._topics = topics if isinstance(topics, list) else [topics]
|
|
self._bootstrap_servers = (
|
|
bootstrap_servers
|
|
if isinstance(bootstrap_servers, list)
|
|
else [bootstrap_servers]
|
|
)
|
|
self._start_offset = start_offset
|
|
self._end_offset = end_offset
|
|
self._kafka_auth_config = kafka_auth_config
|
|
self._consumer_config = consumer_config
|
|
self._timeout_ms = timeout_ms
|
|
self._target_max_block_size = DataContext.get_current().target_max_block_size
|
|
|
|
def estimate_inmemory_data_size(self) -> Optional[int]:
|
|
"""Return an estimate of the in-memory data size, or None if unknown."""
|
|
return None
|
|
|
|
def get_read_tasks(
|
|
self,
|
|
parallelism: int,
|
|
per_task_row_limit: Optional[int] = None,
|
|
data_context: Optional["DataContext"] = None,
|
|
) -> List[ReadTask]:
|
|
"""Create read tasks for Kafka partitions.
|
|
|
|
Creates one read task per partition.
|
|
Each task reads from a single partition of a single topic.
|
|
|
|
Args:
|
|
parallelism: This argument is deprecated.
|
|
per_task_row_limit: Maximum number of rows per read task.
|
|
data_context: The data context to use to get read tasks. This is not used by this datasource.
|
|
|
|
Returns:
|
|
List of ReadTask objects, one per partition.
|
|
"""
|
|
from confluent_kafka import Consumer
|
|
|
|
consumer_config = _build_consumer_config_for_read(
|
|
self._bootstrap_servers, self._kafka_auth_config, self._consumer_config
|
|
)
|
|
discovery_consumer = Consumer(consumer_config)
|
|
try:
|
|
metadata = discovery_consumer.list_topics(
|
|
timeout=KAFKA_TOPIC_METADATA_TIMEOUT_S
|
|
)
|
|
|
|
topic_partitions: List[Tuple[str, int]] = []
|
|
for topic in self._topics:
|
|
if topic not in metadata.topics:
|
|
raise ValueError(
|
|
f"Topic {topic} has no partitions or doesn't exist"
|
|
)
|
|
topic_meta = metadata.topics[topic]
|
|
if not topic_meta.partitions:
|
|
raise ValueError(
|
|
f"Topic {topic} has no partitions or doesn't exist"
|
|
)
|
|
for partition_id in topic_meta.partitions.keys():
|
|
topic_partitions.append((topic, partition_id))
|
|
finally:
|
|
discovery_consumer.close()
|
|
|
|
bootstrap_servers = self._bootstrap_servers
|
|
start_offset = self._start_offset
|
|
end_offset = self._end_offset
|
|
timeout_ms = self._timeout_ms
|
|
target_max_block_size = self._target_max_block_size
|
|
|
|
# Validate that any partitions referenced in a per-partition dict
|
|
# actually exist on the broker. Check once per topic before the loop.
|
|
actual_partition_ids: Dict[str, Set[int]] = {}
|
|
for topic_name, partition_id in topic_partitions:
|
|
actual_partition_ids.setdefault(topic_name, set()).add(partition_id)
|
|
|
|
for param_name, offset in (
|
|
("start_offset", start_offset),
|
|
("end_offset", end_offset),
|
|
):
|
|
if isinstance(offset, dict):
|
|
for topic, partition_map in offset.items():
|
|
existing_partitions = actual_partition_ids.get(topic, set())
|
|
for pid in partition_map:
|
|
if pid not in existing_partitions:
|
|
raise ValueError(
|
|
f"{param_name} references partition {pid} in topic "
|
|
f"{topic!r}, but that partition does not exist."
|
|
)
|
|
|
|
tasks = []
|
|
for topic_name, partition_id in topic_partitions:
|
|
|
|
def create_kafka_read_fn(
|
|
topic_name: str = topic_name,
|
|
partition_id: int = partition_id,
|
|
bootstrap_servers: List[str] = bootstrap_servers,
|
|
start_offset: Optional[
|
|
Union[int, datetime, Literal["earliest"]]
|
|
] = start_offset,
|
|
end_offset: Optional[
|
|
Union[int, datetime, Literal["latest"]]
|
|
] = end_offset,
|
|
kafka_auth_config: Optional[KafkaAuthConfig] = self._kafka_auth_config,
|
|
user_consumer_config: Optional[Dict[str, Any]] = self._consumer_config,
|
|
timeout_ms: Optional[int] = timeout_ms,
|
|
target_max_block_size: int = target_max_block_size,
|
|
):
|
|
"""Create a Kafka read function with captured variables."""
|
|
|
|
def kafka_read_fn() -> Iterable[Block]:
|
|
"""Read function for a single Kafka partition using confluent-kafka."""
|
|
from confluent_kafka import (
|
|
Consumer,
|
|
KafkaError,
|
|
KafkaException,
|
|
TopicPartition,
|
|
)
|
|
|
|
built_consumer_config = _build_consumer_config_for_read(
|
|
bootstrap_servers, kafka_auth_config, user_consumer_config
|
|
)
|
|
|
|
consumer = Consumer(built_consumer_config)
|
|
try:
|
|
topic_partition = TopicPartition(topic_name, partition_id)
|
|
resolved_start, resolved_end = _resolve_offsets(
|
|
consumer, topic_partition, start_offset, end_offset
|
|
)
|
|
|
|
records = []
|
|
|
|
output_buffer = BlockOutputBuffer(
|
|
OutputBlockSizeOption.of(
|
|
target_max_block_size=target_max_block_size
|
|
)
|
|
)
|
|
|
|
if resolved_start < resolved_end:
|
|
start_time = time.perf_counter()
|
|
timeout_seconds = (
|
|
timeout_ms / 1000.0 if timeout_ms is not None else None
|
|
)
|
|
|
|
tp_with_offset = TopicPartition(
|
|
topic_name, partition_id, resolved_start
|
|
)
|
|
consumer.assign([tp_with_offset])
|
|
|
|
next_offset = resolved_start
|
|
|
|
partition_done = False
|
|
while not partition_done:
|
|
if next_offset >= resolved_end:
|
|
break
|
|
|
|
if timeout_seconds is not None:
|
|
elapsed_time_s = time.perf_counter() - start_time
|
|
if elapsed_time_s >= timeout_seconds:
|
|
logger.warning(
|
|
f"Kafka read task timed out after {timeout_ms}ms while reading partition {partition_id} of topic {topic_name}; "
|
|
f"end_offset {resolved_end} was not reached. Returning {len(records)} messages collected in this read task so far."
|
|
)
|
|
break
|
|
remaining_timeout_s = (
|
|
timeout_seconds - elapsed_time_s
|
|
)
|
|
consume_timeout_s = min(
|
|
remaining_timeout_s, KAFKA_CONSUME_TIMEOUT_MAX_S
|
|
)
|
|
else:
|
|
consume_timeout_s = KAFKA_CONSUME_TIMEOUT_MAX_S
|
|
|
|
remaining = resolved_end - next_offset
|
|
batch_size = min(
|
|
remaining,
|
|
KafkaDatasource.BATCH_SIZE_FOR_YIELD,
|
|
)
|
|
msgs = consumer.consume(
|
|
num_messages=batch_size,
|
|
timeout=consume_timeout_s,
|
|
)
|
|
|
|
if not msgs:
|
|
continue
|
|
|
|
for msg in msgs:
|
|
if msg.error():
|
|
# In confluent-kafka, errors are delivered
|
|
# as messages. Only skip partition EOF
|
|
# events, raise others.
|
|
err = msg.error()
|
|
if err.code() == KafkaError._PARTITION_EOF:
|
|
continue
|
|
raise KafkaException(err)
|
|
|
|
# Stop once we reached the end offset
|
|
# (exclusive).
|
|
if msg.offset() >= resolved_end:
|
|
partition_done = True
|
|
break
|
|
|
|
ts_type, ts_ms = msg.timestamp()
|
|
headers_list = msg.headers() or []
|
|
headers_dict = dict(headers_list)
|
|
records.append(
|
|
{
|
|
"offset": msg.offset(),
|
|
"key": msg.key(),
|
|
"value": msg.value(),
|
|
"topic": msg.topic(),
|
|
"partition": msg.partition(),
|
|
"timestamp": ts_ms,
|
|
"timestamp_type": ts_type,
|
|
"headers": headers_dict,
|
|
}
|
|
)
|
|
|
|
next_offset = msg.offset() + 1
|
|
|
|
if (
|
|
len(records)
|
|
>= KafkaDatasource.BATCH_SIZE_FOR_YIELD
|
|
):
|
|
table = pa.Table.from_pylist(records)
|
|
output_buffer.add_block(table)
|
|
yield from output_buffer.iter_ready_blocks()
|
|
records = []
|
|
|
|
# Yield any remaining records
|
|
if records:
|
|
table = pa.Table.from_pylist(records)
|
|
output_buffer.add_block(table)
|
|
|
|
output_buffer.finalize()
|
|
yield from output_buffer.iter_ready_blocks()
|
|
|
|
finally:
|
|
consumer.close()
|
|
|
|
return kafka_read_fn
|
|
|
|
# TODO: We could output the offset range for every partition after the read is done.
|
|
metadata = BlockMetadata(
|
|
num_rows=None,
|
|
size_bytes=None,
|
|
input_files=[f"kafka://{topic_name}/{partition_id}"],
|
|
exec_stats=None,
|
|
)
|
|
|
|
effective_start = (
|
|
start_offset.get(topic_name, {}).get(partition_id, "earliest")
|
|
if isinstance(start_offset, dict)
|
|
else start_offset
|
|
)
|
|
effective_end = (
|
|
end_offset.get(topic_name, {}).get(partition_id, "latest")
|
|
if isinstance(end_offset, dict)
|
|
else end_offset
|
|
)
|
|
kafka_read_fn = create_kafka_read_fn(
|
|
topic_name,
|
|
partition_id,
|
|
start_offset=effective_start,
|
|
end_offset=effective_end,
|
|
)
|
|
|
|
# Create read task
|
|
task = ReadTask(
|
|
read_fn=kafka_read_fn,
|
|
metadata=metadata,
|
|
schema=KAFKA_MSG_SCHEMA,
|
|
per_task_row_limit=per_task_row_limit,
|
|
)
|
|
tasks.append(task)
|
|
|
|
return tasks
|