KaneAI - World’s First E2E Software Testing Agent.Book a Demo Join TestMu Conference

How to use reset_partitions method in lisa

Best Python code snippet using lisa_python

bytes.py

Source:bytes.py

1"""This module contains implementation of extended confluent-kafka Consumer's API."""2import atexit3import datetime4from typing import Dict, List, Union, Optional5from wunderkafka.time import now6from wunderkafka.types import HowToSubscribe7from wunderkafka.config import ConsumerConfig8from wunderkafka.errors import ConsumerException9from wunderkafka.logger import logger10from wunderkafka.callbacks import reset_partitions11from wunderkafka.consumers.abc import Message, AbstractConsumer12from wunderkafka.consumers.subscription import TopicSubscription13from wunderkafka.hotfixes.watchdog.types import Watchdog14class BytesConsumer(AbstractConsumer):15    """Consumer implementation of extended interface for raw messages."""16    # FixMe (tribunsky.kir): add watchdog page reference17    def __init__(self, config: ConsumerConfig, sasl_watchdog: Optional[Watchdog] = None) -> None:18        """19        Init consumer.20        :param config:          Pydantic model with librdkafka consumer's configuration.21        :param sasl_watchdog:   Callable to handle global state of kerberos auth (see Watchdog).22        """23        super().__init__(config.dict())24        self.subscription_offsets: Optional[Dict[str, HowToSubscribe]] = None25        self._config = config26        self._last_poll_ts: int = now()27        self._sasl_watchdog = sasl_watchdog28        # ToDo (tribunsky-kir): make it configurable29        atexit.register(self.close)30    def __str__(self) -> str:31        """32        Get human-readable representation of consumer.33        :return:    string with consumer gid.34        """35        return '{0}:{1}'.format(self.__class__.__name__, self._config.group_id)36    def batch_poll(  # noqa: D102  # inherited from superclass.37        self,38        timeout: float = 1.0,39        num_messages: int = 1000000,40        *,41        raise_on_lost: bool = False,42    ) -> List[Message]:43        if self._sasl_watchdog is not None:44            self._sasl_watchdog()45        # ToDo (tribunsky.kir): naybe it better to use on_lost callback within subscribe()46        dt = now() - self._last_poll_ts47        if dt > self._config.max_poll_interval_ms:48            msg = 'Exceeded max.poll.interval.ms ({0}): {1}'.format(self._config.max_poll_interval_ms, dt)49            if raise_on_lost:50                # ToDo (tribunsky.kir): resubscribe by ts?51                raise ConsumerException(msg)52            logger.warning(msg)53        msgs = self.consume(num_messages=num_messages, timeout=timeout)54        self._last_poll_ts = now()55        return msgs56    # ToDo (tribunsky.kir): do not override original API and wrap it in superclass57    def subscribe(  # noqa: D102,WPS211  # inherited from superclass.58        self,59        topics: List[Union[str, TopicSubscription]],60        *,61        from_beginning: Optional[bool] = None,62        offset: Optional[int] = None,63        ts: Optional[int] = None,64        with_timedelta: Optional[datetime.timedelta] = None,65    ) -> None:66        subscriptions = {}67        for tpc in topics:68            if isinstance(tpc, str):69                tpc = TopicSubscription(70                    topic=tpc, from_beginning=from_beginning, offset=offset, ts=ts, with_timedelta=with_timedelta,71                )72            subscriptions[tpc.topic] = tpc.how73        # We have specific subscription at least once74        if any(subscriptions.values()):75            self.subscription_offsets = subscriptions76            # ToDo (tribunsky.kir): avoid mutation of self.subscription_offset and remove it as a field77            super().subscribe(topics=list(self.subscription_offsets), on_assign=reset_partitions)78        else:...

callbacks.py

Source:callbacks.py

...4from wunderkafka.logger import logger5from wunderkafka.structures import Timestamp6from wunderkafka.consumers.abc import AbstractConsumer7# FixMe (tribunsky.kir): do not mutate consumer from here, try to closure it in consumer itself.8def reset_partitions(consumer: AbstractConsumer, partitions: List[TopicPartition]) -> None:9    """10    Set specific offset for assignment after subscription.11    Depending on type of subscription, will set offset or timestamp.12    :param consumer:            Consumer, which is subscribes to topics.13    :param partitions:          List of TopicPartitions, which is returned from the underlying library.14    """15    new_offsets = consumer.subscription_offsets16    if new_offsets is None:17        logger.warning(18            '{0}: re-assigned (using auto.offset.reset={1})'.format(consumer, consumer.config.auto_offset_reset),19        )20        return21    by_offset = []22    by_ts = []...

partition_dataset.py

Source:partition_dataset.py

...31    print('Connecting to database...')32    dataset = mongo_dataset.MongoDataset()33    # Clear all34    print('Resetting partitions...')35    dataset.reset_partitions()36    print('Assigning documents to partitions...')37    query = mongo_dataset.year_range_query(start_year, end_year)38    filtered_ids = dataset.get_mongo_ids(query)39    split_ids = split_to_proportions(filtered_ids, list(proportions.values()))40    print('Updating partitions in database...')41    for name, ids in zip(proportions.keys(), split_ids):42        print(f'{name}: {len(ids)} posts')43        dataset.batch_update(ids, {'$set': {'partition': name}})44    print('Indexing partitions...')45    dataset.collection.create_index('partition', sparse=True)46if __name__=='__main__':47    # Partition names and sizes besides the defaults can be specified as a command line argument in JSON format.48    proportions = json.loads(sys.argv[1]) if len(sys.argv) > 1 else default_proportions49    main(proportions)

Automation Testing Tutorials

Learn to execute automation testing from scratch with LambdaTest Learning Hub. Right from setting up the prerequisites to run your first automation test, to following best practices and diving deeper into advanced test scenarios. LambdaTest Learning Hubs compile a list of step-by-step guides to help you be proficient with different test automation frameworks i.e. Selenium, Cypress, TestNG etc.