def delete_stream(name: str): async def wrapper(): async with LiftBridgeClient() as client: await client.delete_stream(client.get_offset_stream(name)) await client.delete_stream(name) run_sync(wrapper)
def handle_benchmark_subscriber(self, name: str, cursor: Optional[str] = None, *args, **kwargs): async def subscriber(): async with LiftBridgeClient() as client: report_interval = 1.0 t0 = perf_counter() total_msg = last_msg = 0 total_size = last_size = 0 async for msg in client.subscribe(name): total_msg += 1 total_size += len(msg.value) t = perf_counter() dt = t - t0 if dt >= report_interval: self.print("%d msg/sec, %d bytes/sec" % ((total_msg - last_msg) / dt, (total_size - last_size) / dt)) t0 = t last_msg = total_msg last_size = total_size if cursor: await client.set_cursor(stream=name, partition=0, cursor_id=cursor, offset=msg.offset) run_sync(subscriber)
def handle_subscribe( self, name: str, partition: int = 0, cursor: str = "", start_offset: int = 0, start_ts: int = None, *args, **kwargs, ): async def subscribe(): async with LiftBridgeClient() as client: async for msg in client.subscribe( stream=name, partition=partition, start_offset=start_offset, cursor_id=cursor or None, start_timestamp=start_ts, ): print( "# Subject: %s Partition: %s Offset: %s Timestamp: %s Key: %s Headers: %s" % ( msg.subject, msg.partition, msg.offset, msg.timestamp, msg.key, msg.headers, )) print(msg.value) if start_ts: start_offset = None start_ts *= 1000000000 run_sync(subscribe)
def handle(self, input, addresses, jobs, *args, **options): async def runner(): nonlocal lock lock = asyncio.Lock() tasks = [ asyncio.create_task(ping_worker(), name=f"ping-{i}") for i in range(min(jobs, len(addr_list))) ] await asyncio.gather(*tasks) async def ping_worker(): while True: async with lock: if not addr_list: break # Done addr = addr_list.pop(0) rtt, attempts = await ping.ping_check_rtt(addr, count=1, timeout=1000) if rtt: self.stdout.write(f"{addr} {rtt * 1000:.2f}ms\n") else: self.stdout.write(f"{addr} FAIL\n") # Run ping addr_list = self.get_addresses(addresses, input) lock: Optional[asyncio.Lock] = None ping = Ping() setup_asyncio() run_sync(runner)
def handle_load(self, fields, input, chunk, rm, *args, **kwargs): async def upload(table: str, data: List[bytes]): CHUNK = 1000 n_parts = len(config.clickhouse.cluster_topology.split(",")) async with LiftBridgeClient() as client: while data: chunk, data = data[:CHUNK], data[CHUNK:] await client.publish( b"\n".join(chunk), stream=f"ch.{table}", partition=random.randint(0, n_parts - 1), ) for fn in input: # Read data self.print("Reading file %s" % fn) if fn.endswith(".gz"): with gzip.GzipFile(fn) as f: records = f.read().replace("\r", "").splitlines() else: with open(fn) as f: records = f.read().replace("\r", "").splitlines() table = fn.split("-", 1)[0] run_sync(partial(upload, table, records)) if rm: os.unlink(fn)
def handle_show_metadata(self, *args, **options): async def get_meta() -> Metadata: async with LiftBridgeClient() as client: return await client.fetch_metadata() async def get_partition_meta(stream, partition) -> PartitionMetadata: async with LiftBridgeClient() as client: return await client.fetch_partition_metadata(stream, partition) meta: Metadata = run_sync(get_meta) self.print("# Brokers (%d)" % len(meta.brokers)) self.print("%-20s | %s" % ("ID", "HOST:PORT")) for broker in meta.brokers: self.print("%-20s | %s:%s" % (broker.id, broker.host, broker.port)) self.print("# Streams") for stream in meta.metadata: print(" ## Name: %s Subject: %s" % (stream.name, stream.subject)) for p in sorted(stream.partitions): print(" ### Partition: %d" % p) try: p_meta: PartitionMetadata = run_sync( functools.partial(get_partition_meta, stream.name, p)) except Exception as e: print("[%s|%s] Failed getting data for partition: %s" % (stream.name, p, e)) continue print(" Leader : %s" % p_meta.leader) print(" Replicas : %s" % ", ".join(sorted(p_meta.replicas, key=alnum_key))) print(" ISR : %s" % ", ".join(sorted(p_meta.isr, key=alnum_key))) print(" HighWatermark : %s" % p_meta.high_watermark) print(" NewestOffset : %s" % p_meta.newest_offset)
def handle_benchmark_publisher( self, name: str, num_messages: int, payload_size: int = 64, batch=1, wait_for_stream=False, *args, **kwargs, ): async def publisher(): async with LiftBridgeClient() as client: payload = b" " * payload_size t0 = perf_counter() for _ in self.progress(range(num_messages), num_messages): await client.publish(payload, stream=name, wait_for_stream=wait_for_stream) dt = perf_counter() - t0 self.print("%d messages sent in %.2fms" % (num_messages, dt * 1000)) self.print("%d msg/sec, %d bytes/sec" % (num_messages / dt, num_messages * payload_size / dt)) async def batch_publisher(): async with LiftBridgeClient() as client: payload = b" " * payload_size t0 = perf_counter() out = [] n_acks = 0 for _ in self.progress(range(num_messages), num_messages): out += [client.get_publish_request(payload, stream=name)] if len(out) == batch: async for ack in client.publish_async(out): n_acks += 1 out = [] if out: async for _ in client.publish_async(out): n_acks += 1 out = [] dt = perf_counter() - t0 self.print("%d messages sent in %.2fms (%d acks)" % (num_messages, dt * 1000, n_acks)) self.print("%d msg/sec, %d bytes/sec" % (num_messages / dt, num_messages * payload_size / dt)) if batch == 1: run_sync(publisher) else: print("batch") run_sync(batch_publisher)
def handle_fetch_cursor(self, name: str, stream: str, partition: int = 0, *args, **kwargs): async def fetch_cursor(): async with LiftBridgeClient() as client: cursor = await client.fetch_cursor(stream=stream, partition=partition, cursor_id=name) print(cursor) run_sync(fetch_cursor)
def publish( value: bytes, stream: str, partition: Optional[int] = None, key: Optional[bytes] = None, headers: Optional[Dict[str, bytes]] = None, ): async def wrap(): async with LiftBridgeClient() as client: await client.publish( value=value, stream=stream, partition=partition, key=key, headers=headers ) run_sync(wrap)
def handle_set_cursor(self, name: str, stream: str, partition: int = 0, offset: int = 0, *args, **kwargs): async def set_cursor(): async with LiftBridgeClient() as client: await client.set_cursor(stream=stream, partition=partition, cursor_id=name, offset=offset) run_sync(set_cursor)
def count(self, oid, filter=None, version=None): """ Iterate MIB subtree and count matching instances :param oid: OID :param filter: Callable accepting oid and value and returning boolean """ async def run(): try: r = await snmp_count( address=self.script.credentials["address"], oid=oid, community=str(self.script.credentials["snmp_ro"]), bulk=self.script.has_snmp_bulk(), filter=filter, tos=self.script.tos, udp_socket=self.get_socket(), version=version, rate_limit=self.rate_limit, ) return r except SNMPError as e: if e.code == TIMED_OUT: raise self.TimeOutError() else: raise if "snmp_ro" not in self.script.credentials: raise SNMPError(code=ERR_SNMP_BAD_COMMUNITY) version = self._get_snmp_version(version) return run_sync(run, close_all=False)
def set(self, *args): """ Perform SNMP GET request :param oid: string or list of oids :returns: eigther result scalar or dict of name -> value """ async def run(): try: r = await snmp_set( address=self.script.credentials["address"], varbinds=varbinds, community=str(self.script.credentials["snmp_rw"]), tos=self.script.tos, udp_socket=self.get_socket(), rate_limit=self.rate_limit, ) return r except SNMPError as e: if e.code == TIMED_OUT: raise self.TimeOutError() else: raise if len(args) == 1: varbinds = args elif len(args) == 2: varbinds = [(args[0], args[1])] else: raise ValueError("Invalid varbinds") if "snmp_ro" not in self.script.credentials: raise SNMPError(code=ERR_SNMP_BAD_COMMUNITY) return run_sync(run, close_all=False)
def publish( self, value: bytes, stream: str, partition: Optional[int] = None, key: Optional[bytes] = None, headers: Optional[Dict[str, bytes]] = None, ): async def wrap(): async with LiftBridgeClient() as client: await client.publish( value=value, stream=stream, partition=partition, key=key, headers=headers, auto_compress=bool(config.liftbridge.compression_method), ) run_sync(wrap)
def handle_create_stream( self, name: str, subject: Optional[str] = None, partitions: int = 1, rf: int = 1, *args, **kwargs, ): async def create(): async with LiftBridgeClient() as client: await client.create_stream( name=name, subject=subject, partitions=partitions, replication_factor=rf, ) subject = subject or name run_sync(create)
def create_stream(name: str, n_partitions: int, replication_factor: int): base_name = name.split(".")[0] minisr = 0 if base_name == "ch": replication_factor = min( config.liftbridge.stream_ch_replication_factor, replication_factor ) minisr = min(2, replication_factor) async def wrapper(): async with LiftBridgeClient() as client: await client.create_stream( subject=name, name=name, partitions=n_partitions, minisr=minisr, replication_factor=replication_factor, retention_max_bytes=getattr( config.liftbridge, f"stream_{base_name}_retention_max_age", 0 ), retention_max_age=getattr( config.liftbridge, f"stream_{base_name}_retention_max_bytes", 0 ), segment_max_bytes=getattr( config.liftbridge, f"stream_{base_name}_segment_max_bytes", 0 ), segment_max_age=getattr( config.liftbridge, f"stream_{base_name}_segment_max_age", 0 ), auto_pause_time=getattr( config.liftbridge, f"stream_{base_name}_auto_pause_time", 0 ), auto_pause_disable_if_subscribers=getattr( config.liftbridge, f"stream_{base_name}_auto_pause_disable_if_subscribers", False, ), ) run_sync(wrapper)
def iter_ch_cursors(self, stream, partition): # Parse cluster = config.clickhouse.cluster_topology.split(",") for replica in range(0, int(cluster[partition])): cursor = run_sync( functools.partial(self.fetch_cursor, stream, partition, f"chwriter-{replica}")) yield ( "stream_cursor_offset", ("name", stream), ("partition", partition), ("cursor_id", f"chwriter-{replica}"), ), cursor
def getnext( self, oid, community_suffix=None, filter=None, cached=False, only_first=False, bulk=None, max_repetitions=None, version=None, max_retries=0, timeout=10, raw_varbinds=False, display_hints=None, ): async def run(): try: r = await snmp_getnext( address=self.script.credentials["address"], oid=oid, community=str(self.script.credentials["snmp_ro"]), bulk=self.script.has_snmp_bulk() if bulk is None else bulk, max_repetitions=max_repetitions, filter=filter, only_first=only_first, tos=self.script.tos, udp_socket=self.get_socket(), version=version, max_retries=max_retries, timeout=timeout, raw_varbinds=raw_varbinds, display_hints=display_hints, response_parser=self.script.profile. get_snmp_response_parser(self.script), rate_limit=self.rate_limit, ) return r except SNMPError as e: if e.code == TIMED_OUT: raise self.TimeOutError() else: raise if "snmp_ro" not in self.script.credentials: raise SNMPError(code=ERR_SNMP_BAD_COMMUNITY) if display_hints is None: display_hints = self._get_display_hints() version = self._get_snmp_version(version) return run_sync(run, close_all=False)
def get(self, oids, cached=False, version=None, raw_varbinds=False, display_hints=None): """ Perform SNMP GET request :param oid: string or list of oids :param cached: True if get results can be cached during session :param raw_varbinds: Return value in BER encoding :param display_hints: Dict of oid -> render_function. See BaseProfile.snmp_display_hints for details :returns: eigther result scalar or dict of name -> value """ async def run(): try: r = await snmp_get( address=self.script.credentials["address"], oids=oids, community=str(self.script.credentials["snmp_ro"]), tos=self.script.tos, udp_socket=self.get_socket(), version=version, raw_varbinds=raw_varbinds, display_hints=display_hints, response_parser=self.script.profile. get_snmp_response_parser(self.script), rate_limit=self.rate_limit, ) self.timeouts = self.timeouts_limit return r except SNMPError as e: if e.code == TIMED_OUT: if self.timeouts_limit: self.timeouts -= 1 if not self.timeouts: raise self.FatalTimeoutError() raise self.TimeOutError() else: raise if "snmp_ro" not in self.script.credentials: raise SNMPError(code=ERR_SNMP_BAD_COMMUNITY) if display_hints is None: display_hints = self._get_display_hints() version = self._get_snmp_version(version) return run_sync(run, close_all=False)
def resolve_sync(self, name, hint=None, wait=True, timeout=None, full_result=False): """ Returns *hint* when service is active or new service instance, :param name: :param hint: :param full_result: :return: """ async def _resolve(): r = await self.resolve( name, hint=hint, wait=wait, timeout=timeout, full_result=full_result ) return r return run_sync(_resolve)
def resolve(name, hint=None, wait=True, timeout=None, full_result=False, near=False, critical=False): """ Returns *hint* when service is active or new service instance, :param name: :param hint: :param wait: :param timeout: :param full_result: :param near: :return: """ async def _resolve(): dcs = get_dcs() try: if near: r = await dcs.resolve_near( name, hint=hint, wait=wait, timeout=timeout, full_result=full_result, critical=critical, ) else: r = await dcs.resolve( name, hint=hint, wait=wait, timeout=timeout, full_result=full_result, critical=critical, track=False, ) finally: dcs.stop() return r return run_sync(_resolve)
def fetch_sync( url: str, method: str = "GET", headers=None, body: Optional[bytes] = None, connect_timeout=DEFAULT_CONNECT_TIMEOUT, request_timeout=DEFAULT_REQUEST_TIMEOUT, resolver=resolve, max_buffer_size=DEFAULT_BUFFER_SIZE, follow_redirects=False, max_redirects=DEFAULT_MAX_REDIRECTS, validate_cert=config.http_client.validate_certs, allow_proxy: bool = False, proxies=None, user: Optional[str] = None, password: Optional[str] = None, content_encoding: Optional[str] = None, eof_mark: Optional[bytes] = None, ): async def _fetch(): return await fetch( url, method=method, headers=headers, body=body, connect_timeout=connect_timeout, request_timeout=request_timeout, resolver=resolver, max_buffer_size=max_buffer_size, follow_redirects=follow_redirects, max_redirects=max_redirects, validate_cert=validate_cert, allow_proxy=allow_proxy, proxies=proxies, user=user, password=password, content_encoding=content_encoding, eof_mark=eof_mark, ) return run_sync(_fetch)
def get_mx_partitions() -> int: """ Get number of MX stream partitions :return: """ async def wrap(): async with LiftBridgeClient() as client: r = await client.fetch_metadata(MX_STREAM, wait_for_stream=True) for m in r.metadata: if m.name == MX_STREAM: return len(m.partitions) global _mx_partitions if _mx_partitions: return _mx_partitions with _mx_lock: if _mx_partitions: return _mx_partitions # Set by concurrent thread _mx_partitions = run_sync(wrap) return _mx_partitions
def save_avatar(user: User = Depends(get_current_user), image: UploadFile = File(...)): async def read_file() -> bytes: return smart_bytes(await image.read(config.ui.max_avatar_size + 1)) data = run_sync(read_file) if len(data) > config.ui.max_avatar_size: raise HTTPException(status_code=413) content_type = ContentType.from_content_type(image.content_type) if content_type is None: raise HTTPException(status_code=421) avatar = Avatar.objects.filter(user_id=str(user.id)).first() if avatar: # Update avatar.data = data avatar.content_type = content_type else: # Create avatar = Avatar(user_id=str(user.id), data=data, content_type=content_type) avatar.save() return StatusResponse(status=True)
def iter_limits(self) -> Tuple[str, int]: async def get_slot_limits(): nonlocal slot_name return await dcs.get_slot_limit(slot_name) dcs = get_dcs() # Plain streams for stream_name in self.STREAMS: yield stream_name, 1 # Slot-based streams for slot_name, stream_name in self.iter_slot_streams(): n_partitions = run_sync(get_slot_limits) if n_partitions: yield stream_name, n_partitions # Metric scopes n_ch_shards = len(config.clickhouse.cluster_topology.split(",")) for scope in MetricScope.objects.all(): yield f"ch.{scope.table_name}", n_ch_shards # BI models for name in bi_loader: bi_model = bi_loader[name] if not bi_model: continue yield f"ch.{bi_model._meta.db_table}", n_ch_shards
def get_meta(self) -> Metadata: async def get_meta() -> Metadata: async with LiftBridgeClient() as client: return await client.fetch_metadata() return run_sync(get_meta)
async def wrapper(): self.print("Altering stream %s" % name) async with LiftBridgeClient() as client: # Create temporary stream with same structure, as original one tmp_stream = "__tmp-%s" % name self.print("Creating temporary stream %s" % tmp_stream) await client.create_stream( subject=tmp_stream, name=tmp_stream, partitions=old_partitions, replication_factor=replication_factor, ) # Copy all unread data to temporary stream as is for partition in range(old_partitions): self.print( "Copying partition %s:%s to %s:%s" % (name, partition, tmp_stream, partition) ) n_msg[partition] = 0 # Get current offset p_meta = run_sync(functools.partial(get_partition_meta, stream, partition)) newest_offset = p_meta.newest_offset or 0 # Fetch cursor current_offset = await client.fetch_cursor( stream=stream, partition=partition, cursor_id=self.CURSOR_STREAM[name.split(".")[0]], ) if current_offset > newest_offset: # Fix if cursor not set properly current_offset = newest_offset self.print( "Start copying from current_offset: %s to newest offset: %s" % (current_offset, newest_offset) ) if current_offset < newest_offset: async for msg in client.subscribe( stream=name, partition=partition, start_offset=current_offset ): await client.publish( msg.value, stream=tmp_stream, partition=partition, ) n_msg[partition] += 1 if msg.offset == newest_offset: break if n_msg[partition]: self.print(" %d messages has been copied" % n_msg[partition]) else: self.print(" nothing to copy") # Drop original stream self.print("Dropping original stream %s" % name) await client.delete_stream(name) # Create new stream with required structure self.print("Creating stream %s" % name) await client.create_stream( subject=name, name=name, partitions=new_partitions, replication_factor=replication_factor, ) # Copy data from temporary stream to a new one for partition in range(old_partitions): self.print( "Restoring partition %s:%s to %s" % (tmp_stream, partition, new_partitions) ) # Re-route dropped partitions to partition 0 dest_partition = partition if partition < new_partitions else 0 n = n_msg[partition] if n > 0: async for msg in client.subscribe( stream=tmp_stream, partition=partition, start_position=StartPosition.EARLIEST, ): await client.publish( msg.value, stream=name, partition=dest_partition ) n -= 1 if not n: break self.print(" %d messages restored" % n_msg[partition]) else: self.print(" nothing to restore") # Drop temporary stream self.print("Dropping temporary stream %s" % tmp_stream) await client.delete_stream(tmp_stream) # Uh-oh self.print("Stream %s has been altered" % name)
def whois(query, fields=None): async def _whois(): return await whois_async(query, fields) return run_sync(_whois)
def load(self): run_sync(self.load_async)
def iter_metrics(self): meta: Metadata = run_sync(self.get_meta) for stream in meta.metadata: for p in sorted(stream.partitions): if stream.name.startswith("_"): continue name, pool, cursor = stream.name, None, -1 if name.startswith("ch"): # Chwriter streams for r in self.iter_ch_cursors(name, p): yield r elif "." in name: name, pool = name.split(".", 1) if name in self.CURSOR_STREAM: cursor = run_sync( functools.partial(self.fetch_cursor, stream.name, p, self.CURSOR_STREAM[name])) if pool: yield ( "stream_cursor_offset", ("name", stream.name), ("partition", p), ("cursor_id", self.CURSOR_STREAM[name]), ("pool", pool), ), cursor else: yield ( "stream_cursor_offset", ("name", stream.name), ("partition", p), ("cursor_id", self.CURSOR_STREAM[name]), ), cursor try: p_meta: PartitionMetadata = run_sync( functools.partial( self.get_partition_meta, stream.name, p, )) except Exception as e: self.logger.error( "[%s|%s] Failed getting data for partition: %s", stream.name, p, e) continue if pool: yield ( "stream_newest_offset", ("name", stream.name), ("partition", p), ("pool", pool), ), p_meta.newest_offset yield ( "stream_high_watermark", ("name", stream.name), ("partition", p), ("pool", pool), ), p_meta.high_watermark else: yield ( "stream_newest_offset", ("name", stream.name), ("partition", p), ), p_meta.newest_offset yield ( "stream_high_watermark", ("name", stream.name), ("partition", p), ), p_meta.high_watermark
def handle_delete_stream(self, name: str, *args, **kwargs): async def delete(): async with LiftBridgeClient() as client: await client.delete_stream(name) run_sync(delete)