add a helper class for publishing to shelter

Author: nora-shap

codecov_auth/signals.py

@@ -1,5 +1,6 @@
 import json
 import logging
+from typing import Any, Dict, Optional, Type
 
 from django.conf import settings
 from django.db.models.signals import post_save
@@ -13,73 +14,81 @@
 
 @receiver(post_save, sender=Owner)
 def create_owner_profile_when_owner_is_created(
-    sender, instance: Owner, created, **kwargs
-):
+    sender: Type[Owner], instance: Owner, created: bool, **kwargs: Dict[str, Any]
+) -> Optional[OwnerProfile]:
     if created:
         return OwnerProfile.objects.create(owner_id=instance.ownerid)
 
 
-_pubsub_publisher = None
+class ShelterPubsub:
+    pubsub_publisher = None
+    _instance = None
 
+    @classmethod
+    def get_instance(cls) -> "ShelterPubsub":
+        """
+        This class needs the Django settings to be fully loaded before it can be instantiated,
+        therefore use this method to get an instance rather than instantiating directly.
+        """
+        if cls._instance is None:
+            cls._instance = cls()
+        return cls._instance
 
-def _get_pubsub_publisher():
-    global _pubsub_publisher
-    if not _pubsub_publisher:
-        _pubsub_publisher = pubsub_v1.PublisherClient()
-    return _pubsub_publisher
+    def __init__(self) -> None:
+        if not self.pubsub_publisher:
+            self.pubsub_publisher = pubsub_v1.PublisherClient()
+        pubsub_project_id: str = settings.SHELTER_PUBSUB_PROJECT_ID
+
+        # topic_id has REPO in the name but it is used for all types of objects
+        topic_id: str = settings.SHELTER_PUBSUB_SYNC_REPO_TOPIC_ID
+        self.topic_path = self.pubsub_publisher.topic_path(pubsub_project_id, topic_id)
+
+    def publish(self, data: Dict[str, Any]) -> None:
+        try:
+            self.pubsub_publisher.publish(
+                self.topic_path,
+                json.dumps(data).encode("utf-8"),
+            )
+        except Exception as e:
+            log.warning(
+                "Failed to publish a message",
+                extra=dict(data_to_publish=data, error=e),
+            )
 
 
 @receiver(
     post_save, sender=OrganizationLevelToken, dispatch_uid="shelter_sync_org_token"
 )
-def update_org_token(sender, instance: OrganizationLevelToken, **kwargs):
-    pubsub_project_id = settings.SHELTER_PUBSUB_PROJECT_ID
-    topic_id = settings.SHELTER_PUBSUB_SYNC_REPO_TOPIC_ID
-    if pubsub_project_id and topic_id:
-        publisher = _get_pubsub_publisher()
-        topic_path = publisher.topic_path(pubsub_project_id, topic_id)
-        publisher.publish(
-            topic_path,
-            json.dumps(
-                {
-                    "type": "org_token",
-                    "sync": "one",
-                    "id": instance.id,
-                }
-            ).encode("utf-8"),
-        )
+def update_org_token(
+    sender: Type[OrganizationLevelToken],
+    instance: OrganizationLevelToken,
+    **kwargs: Dict[str, Any],
+) -> None:
+    data = {
+        "type": "org_token",
+        "sync": "one",
+        "id": instance.id,
+    }
+    ShelterPubsub.get_instance().publish(data)
 
 
 @receiver(post_save, sender=Owner, dispatch_uid="shelter_sync_owner")
-def update_owner(sender, instance: Owner, **kwargs):
+def update_owner(
+    sender: Type[Owner], instance: Owner, **kwargs: Dict[str, Any]
+) -> None:
     """
     Shelter tracks a limited set of Owner fields - only update if those fields have changed.
     """
-    created = kwargs["created"]
+    created: bool = kwargs["created"]
     tracked_fields = [
         "upload_token_required_for_public_repos",
         "username",
         "service",
     ]
     if created or any(instance.tracker.has_changed(field) for field in tracked_fields):
-        pubsub_project_id = settings.SHELTER_PUBSUB_PROJECT_ID
-        topic_id = settings.SHELTER_PUBSUB_SYNC_REPO_TOPIC_ID
-        if pubsub_project_id and topic_id:
-            try:
-                publisher = _get_pubsub_publisher()
-                topic_path = publisher.topic_path(pubsub_project_id, topic_id)
-                publisher.publish(
-                    topic_path,
-                    json.dumps(
-                        {
-                            "type": "owner",
-                            "sync": "one",
-                            "id": instance.ownerid,
-                        }
-                    ).encode("utf-8"),
-                )
-            except Exception as e:
-                log.warning(
-                    "Failed to publish message for owner",
-                    extra=dict(owner_id=instance.pk, error=e),
-                )
+        data = {
+            "type": "owner",
+            "sync": "one",
+            "id": instance.ownerid,
+        }
+        ShelterPubsub.get_instance().publish(data)

codecov_auth/tests/test_signals.py

@@ -128,6 +128,9 @@ def test_sync_error(self, mock_log, mock_publish):
         )
 
         mock_log.assert_called_once_with(
-            "Failed to publish message for owner",
-            extra=dict(owner_id=12345, error=mock_publish.side_effect),
+            "Failed to publish a message",
+            extra=dict(
+                data_to_publish={"type": "owner", "sync": "one", "id": 12345},
+                error=mock_publish.side_effect,
+            ),
         )

core/signals.py

@@ -1,68 +1,43 @@
-import json
 import logging
+from typing import Any, Dict, List, Type
 
-from django.conf import settings
 from django.db.models.signals import post_save
 from django.dispatch import receiver
-from google.cloud import pubsub_v1
 from shared.django_apps.core.models import Commit
 
+from codecov_auth.signals import ShelterPubsub
 from core.models import Repository
 
-_pubsub_publisher = None
 log = logging.getLogger(__name__)
 
 
-def _get_pubsub_publisher():
-    global _pubsub_publisher
-    if not _pubsub_publisher:
-        _pubsub_publisher = pubsub_v1.PublisherClient()
-    return _pubsub_publisher
-
-
 @receiver(post_save, sender=Repository, dispatch_uid="shelter_sync_repo")
-def update_repository(sender, instance: Repository, **kwargs):
+def update_repository(
+    sender: Type[Repository], instance: Repository, **kwargs: Dict[str, Any]
+) -> None:
     log.info(f"Signal triggered for repository {instance.repoid}")
-    created = kwargs["created"]
-    changes = instance.tracker.changed()
-    if created or any([field in changes for field in ["name", "upload_token"]]):
-        try:
-            pubsub_project_id = settings.SHELTER_PUBSUB_PROJECT_ID
-            topic_id = settings.SHELTER_PUBSUB_SYNC_REPO_TOPIC_ID
-            if pubsub_project_id and topic_id:
-                publisher = _get_pubsub_publisher()
-                topic_path = publisher.topic_path(pubsub_project_id, topic_id)
-                publisher.publish(
-                    topic_path,
-                    json.dumps(
-                        {
-                            "type": "repo",
-                            "sync": "one",
-                            "id": instance.repoid,
-                        }
-                    ).encode("utf-8"),
-                )
-            log.info(f"Message published for repository {instance.repoid}")
-        except Exception as e:
-            log.warning(f"Failed to publish message for repo {instance.repoid}: {e}")
+    created: bool = kwargs["created"]
+    changes: Dict[str, Any] = instance.tracker.changed()
+    tracked_fields: List[str] = ["name", "upload_token"]
+
+    if created or any([field in changes for field in tracked_fields]):
+        data = {
+            "type": "repo",
+            "sync": "one",
+            "id": instance.repoid,
+        }
+        ShelterPubsub.get_instance().publish(data)
 
 
 @receiver(post_save, sender=Commit, dispatch_uid="shelter_sync_commit")
-def update_commit(sender, instance: Commit, **kwargs):
-    branch = instance.branch
+def update_commit(
+    sender: Type[Commit], instance: Commit, **kwargs: Dict[str, Any]
+) -> None:
+    branch: str = instance.branch
     if branch and ":" in branch:
-        pubsub_project_id = settings.SHELTER_PUBSUB_PROJECT_ID
-        topic_id = settings.SHELTER_PUBSUB_SYNC_REPO_TOPIC_ID
-        if pubsub_project_id and topic_id:
-            publisher = _get_pubsub_publisher()
-            topic_path = publisher.topic_path(pubsub_project_id, topic_id)
-            publisher.publish(
-                topic_path,
-                json.dumps(
-                    {
-                        "type": "commit",
-                        "sync": "one",
-                        "id": instance.id,
-                    }
-                ).encode("utf-8"),
-            )
+        data = {
+            "type": "commit",
+            "sync": "one",
+            "id": instance.id,
+        }
+        ShelterPubsub.get_instance().publish(data)