fix: More appropriate redis stats, run black

platform_plugin_aspects/management/commands/monitor_load_test_tracking.py

@@ -62,7 +62,7 @@ def run(self) -> None:
             if collect_celery:
                 current_stats["celery"] = self.get_celery_stats()
             # if collect_kafka_bus:
-            # current_stats["kafka_bus"] = self.get_kafka_stats()
+            # current_stats["kafka_bus"] = self.get_kafka_bus_stats()
             # if collect_vector:
             # current_stats["vector"] = self.get_vector_stats()
 
@@ -160,15 +160,74 @@ def get_redis_bus_stats(self):
              'pending': []}],
         """
 
+        lag = []
+        for g in info["groups"]:
+            lag.append({g["name"]: g["lag"]})
+
         consumer_stats = {
             "total_events": info["length"],
-            "queue_length": info["lag"],
+            "queue_lengths": lag,
         }
 
         print(f"Redis bus queue length: {consumer_stats['queue_length']}")
 
         return consumer_stats
 
+    def get_kafka_bus_stats(self):
+        import sys
+        import confluent_kafka
+
+        if len(sys.argv) < 4:
+            sys.stderr.write("Usage: {} <brokers> <group.id> <topic> <topic2..>\n".format(sys.argv[0]))
+            sys.exit(1)
+
+        brokers, group = sys.argv[1:3]
+
+        # Create consumer.
+        # This consumer will not join the group, but the group.id is required by
+        # committed() to know which group to get offsets for.
+        consumer = confluent_kafka.Consumer({'bootstrap.servers': brokers,
+                                             'group.id': group})
+
+        print("%-50s  %9s  %9s" % ("Topic [Partition]", "Committed", "Lag"))
+        print("=" * 72)
+
+        for topic in sys.argv[3:]:
+            # Get the topic's partitions
+            metadata = consumer.list_topics(topic, timeout=10)
+            if metadata.topics[topic].error is not None:
+                raise confluent_kafka.KafkaException(metadata.topics[topic].error)
+
+            # Construct TopicPartition list of partitions to query
+            partitions = [confluent_kafka.TopicPartition(topic, p) for p in metadata.topics[topic].partitions]
+
+            # Query committed offsets for this group and the given partitions
+            committed = consumer.committed(partitions, timeout=10)
+
+            for partition in committed:
+                # Get the partitions low and high watermark offsets.
+                (lo, hi) = consumer.get_watermark_offsets(partition, timeout=10, cached=False)
+
+                if partition.offset == confluent_kafka.OFFSET_INVALID:
+                    offset = "-"
+                else:
+                    offset = "%d" % (partition.offset)
+
+                if hi < 0:
+                    lag = "no hwmark"  # Unlikely
+                elif partition.offset < 0:
+                    # No committed offset, show total message count as lag.
+                    # The actual message count may be lower due to compaction
+                    # and record deletions.
+                    lag = "%d" % (hi - lo)
+                else:
+                    lag = "%d" % (hi - partition.offset)
+
+                print("%-50s  %9s  %9s" % (
+                    "{} [{}]".format(partition.topic, partition.partition), offset, lag))
+
+        consumer.close()
+
 
 class Command(BaseCommand):
     """