Add Kafka download metrics

.env.template

@@ -62,3 +62,9 @@ IS_CYBERSTORM_ENABLED=
 SHOW_CYBERSTORM_API_DOCS=
 
 DECOMPILER_API_HOST=http://host.docker.internal:8000
+
+KAFKA_ENABLED=False
+KAFKA_BOOTSTRAP_SERVERS=
+KAFKA_USERNAME=
+KAFKA_PASSWORD=
+KAFKA_CA_CERT=

django/pyproject.toml

@@ -35,6 +35,7 @@ markdown-it-py = {version = "2.1.0", extras = ["linkify"]}
 abyss = {git = "https://github.com/akx/abyss.git", rev = "4352e557f25a303f718ec1bf82ea0ca545ebc077"}
 lxml = "^4.9.1"
 pygments = "^2.16.1"
+confluent-kafka = "^2.0"
 
 [tool.poetry.group.plugins]
 optional = true

django/thunderstore/core/settings.py

@@ -136,6 +136,11 @@
     CACHALOT_TIMEOUT_SECONDS=(int, 60 * 15),  # 15 minutes by default
     CACHALOT_ENABLED=(bool, True),
     DOWNLOAD_METRICS_TTL_SECONDS=(int, 60 * 10),
+    KAFKA_BOOTSTRAP_SERVERS=(str, ""),
+    KAFKA_USERNAME=(str, ""),
+    KAFKA_PASSWORD=(str, ""),
+    KAFKA_CA_CERT=(str, ""),
+    KAFKA_ENABLED=(bool, False),
     # FEATURE FLAGS UNDER HERE
     IS_CYBERSTORM_ENABLED=(bool, False),
     SHOW_CYBERSTORM_API_DOCS=(bool, False),
@@ -283,6 +288,7 @@ def load_db_certs():
         "thunderstore.moderation",
         "thunderstore.permissions",
         "thunderstore.ts_reports",
+        "thunderstore.ts_analytics",
     ]
 )
 
@@ -391,6 +397,7 @@ class CeleryQueues:
     BackgroundCache = "background.cache"
     BackgroundTask = "background.task"
     BackgroundLongRunning = "background.long_running"
+    Analytics = "analytics"
 
 
 CELERY_BROKER_URL = env.str("CELERY_BROKER_URL")
@@ -584,6 +591,21 @@ def get_redis_cache(env_key: str, fallback_key: Optional[str] = None):
 
 THUMBNAIL_QUALITY = 95
 
+# Kafka configuration
+
+# Whether Kafka is enabled
+KAFKA_ENABLED = env.bool("KAFKA_ENABLED")
+
+KAFKA_CONFIG = {
+    "bootstrap.servers": env.str("KAFKA_BOOTSTRAP_SERVERS"),
+    "security.protocol": "SASL_SSL",
+    "sasl.mechanism": "SCRAM-SHA-256",
+    "sasl.username": env.str("KAFKA_USERNAME"),
+    "sasl.password": env.str("KAFKA_PASSWORD"),
+    "ssl.ca.pem": env.str("KAFKA_CA_CERT"),
+    "client.id": "thunderstore-analytics",
+}
+
 #######################################
 #               STORAGE               #
 #######################################

django/thunderstore/core/tests/test_celery.py

@@ -41,6 +41,7 @@ def test_task():
     "thunderstore.repository.tasks.cleanup_package_submissions",
     "thunderstore.repository.tasks.log_version_download",
     "thunderstore.webhooks.tasks.process_audit_event",
+    "thunderstore.analytics.send_kafka_message_async",
 )
 
 

django/thunderstore/ts_analytics/__init__.py

@@ -0,0 +1 @@
+default_app_config = "thunderstore.ts_analytics.apps.AnalyticsAppConfig"

django/thunderstore/ts_analytics/apps.py

@@ -0,0 +1,28 @@
+from django.apps import AppConfig
+from django.db.models.signals import post_save
+
+
+class AnalyticsAppConfig(AppConfig):
+    name = "thunderstore.ts_analytics"
+    label = "ts_analytics"
+
+    def ready(self):
+        # Connect the signal handlers
+        from thunderstore.community.models import Community, PackageListing
+        from thunderstore.metrics.models import PackageVersionDownloadEvent
+        from thunderstore.repository.models import Package, PackageVersion
+        from thunderstore.ts_analytics.signals import (
+            community_post_save,
+            package_listing_post_save,
+            package_post_save,
+            package_version_download_event_post_save,
+            package_version_post_save,
+        )
+
+        post_save.connect(
+            package_version_download_event_post_save, sender=PackageVersionDownloadEvent
+        )
+        post_save.connect(package_post_save, sender=Package)
+        post_save.connect(package_version_post_save, sender=PackageVersion)
+        post_save.connect(package_listing_post_save, sender=PackageListing)
+        post_save.connect(community_post_save, sender=Community)

django/thunderstore/ts_analytics/kafka.py

@@ -0,0 +1,83 @@
+import json
+from enum import Enum
+from functools import lru_cache
+from typing import Any, Dict, Optional, Union
+
+from celery import shared_task
+from confluent_kafka import Producer
+from django.conf import settings
+
+from thunderstore.core.settings import CeleryQueues
+
+
+class KafkaTopic(str, Enum):
+    PACKAGE_DOWNLOADED = "ts.package.downloaded"
+    PACKAGE_UPDATED = "ts.package.updated"
+    PACKAGE_VERSION_UPDATED = "ts.package.version.updated"
+    PACKAGE_LISTING_UPDATED = "ts.package.listing.updated"
+    COMMUNITY_UPDATED = "ts.community.updated"
+
+
+def send_kafka_message(
+    topic: Union[KafkaTopic, str], payload: dict, key: Optional[str] = None
+):
+    client = get_kafka_client()
+    client.send(topic=topic, payload=payload, key=key)
+
+
+@shared_task(
+    queue=CeleryQueues.Analytics,
+    name="thunderstore.analytics.send_kafka_message_async",
+)
+def send_kafka_message_async(
+    topic: Union[KafkaTopic, str], payload: dict, key: Optional[str] = None
+):
+    send_kafka_message(topic=topic, payload=payload, key=key)
+
+
+class KafkaClient:
+    def __init__(self, config: Dict[str, Any]):
+        self._producer = Producer(config)
+
+    def send(
+        self,
+        topic: Union[KafkaTopic, str],
+        payload: dict,
+        key: Optional[str] = None,
+    ):
+        value_bytes = json.dumps(payload).encode("utf-8")
+        key_bytes = key.encode("utf-8") if key else None
+
+        topic_str = topic.value if isinstance(topic, KafkaTopic) else topic
+        self._producer.produce(
+            topic=topic_str,
+            value=value_bytes,
+            key=key_bytes,
+        )
+
+        self._producer.poll(0)
+
+
+class DummyKafkaClient:
+    """A dummy Kafka client that does nothing when Kafka is disabled."""
+
+    def send(
+        self, topic: Union[KafkaTopic, str], payload: dict, key: Optional[str] = None
+    ):
+        pass
+
+
+@lru_cache(maxsize=1)
+def get_kafka_client() -> Union[KafkaClient, DummyKafkaClient]:
+    # Return dummy client if Kafka is disabled
+    if not getattr(settings, "KAFKA_ENABLED", False):
+        return DummyKafkaClient()
+
+    config = getattr(settings, "KAFKA_CONFIG", None)
+    if not config:
+        raise RuntimeError("KAFKA_CONFIG is not configured.")
+
+    if not config.get("bootstrap.servers"):
+        raise RuntimeError("Kafka bootstrap servers are not configured.")
+
+    return KafkaClient(config)