Move notification helper code into the root module

buildman/jobutil/buildjob.py

@@ -2,7 +2,7 @@ import json
 import logging
 
 from cachetools import lru_cache
-from notifications.notificationhelper import spawn_notification
+from notifications import spawn_notification
 from data import model
 from util.imagetree import ImageTree
 from util.morecollections import AttrDict

endpoints/api/repositorynotification.py

@@ -7,10 +7,10 @@ from endpoints.api import (RepositoryParamResource, nickname, resource, require_
                            log_action, validate_json_request, request_error,
                            path_param, disallow_for_app_repositories)
 from endpoints.exception import NotFound
+from notifications import build_notification_data
 from notifications.notificationevent import NotificationEvent
 from notifications.notificationmethod import (NotificationMethod,
                                               CannotValidateNotificationMethodException)
-from notifications.notificationhelper import build_notification_data
 from workers.notificationworker.models_pre_oci import notification
 from endpoints.api.repositorynotification_models_pre_oci import pre_oci_model as model
 

endpoints/building.py

@@ -9,7 +9,7 @@ from app import app, dockerfile_build_queue, metric_queue
 from data import model
 from data.database import db
 from auth.auth_context import get_authenticated_user
-from notifications.notificationhelper import spawn_notification
+from notifications import spawn_notification
 from util.names import escape_tag
 from util.morecollections import AttrDict
 

endpoints/v1/index.py

@@ -17,7 +17,7 @@ from endpoints.decorators import anon_protect, anon_allowed, parse_repository_na
 from endpoints.notificationhelper import spawn_notification
 from endpoints.v1 import v1_bp
 from endpoints.v1.models_pre_oci import pre_oci_model as model
-from notifications.notificationhelper import spawn_notification
+from notifications import spawn_notification
 from util.audit import track_and_log
 from util.http import abort
 from util.names import REPOSITORY_NAME_REGEX

endpoints/v2/manifest.py

@@ -20,7 +20,7 @@ from endpoints.v2.labelhandlers import handle_label
 from image.docker import ManifestException
 from image.docker.schema1 import DockerSchema1Manifest, DockerSchema1ManifestBuilder
 from image.docker.schema2 import DOCKER_SCHEMA2_CONTENT_TYPES
-from notifications.notificationhelper import spawn_notification
+from notifications import spawn_notification
 from util.audit import track_and_log
 from util.names import VALID_TAG_PATTERN
 from util.registry.replication import queue_replication_batch

notifications/__init__.py

@@ -0,0 +1,84 @@
+import json
+
+from contextlib import contextmanager
+
+from app import app, notification_queue
+from data import model
+from auth.auth_context import get_authenticated_user, get_validated_oauth_token
+
+
+DEFAULT_BATCH_SIZE = 1000
+
+
+def build_event_data(repo, extra_data=None, subpage=None):
+  repo_string = '%s/%s' % (repo.namespace_name, repo.name)
+  homepage = '%s://%s/repository/%s' % (app.config['PREFERRED_URL_SCHEME'],
+                                        app.config['SERVER_HOSTNAME'],
+                                        repo_string)
+
+  if subpage:
+    if not subpage.startswith('/'):
+      subpage = '/' + subpage
+
+    homepage = homepage + subpage
+
+  event_data = {
+    'repository': repo_string,
+    'namespace': repo.namespace_name,
+    'name': repo.name,
+    'docker_url': '%s/%s' % (app.config['SERVER_HOSTNAME'], repo_string),
+    'homepage': homepage,
+  }
+
+  event_data.update(extra_data or {})
+  return event_data
+
+def build_notification_data(notification, event_data, performer_data=None):
+  if not performer_data:
+    performer_data = {}
+
+    oauth_token = get_validated_oauth_token()
+    if oauth_token:
+      performer_data['oauth_token_id'] = oauth_token.id
+      performer_data['oauth_token_application_id'] = oauth_token.application.client_id
+      performer_data['oauth_token_application'] = oauth_token.application.name
+
+    performer_user = get_authenticated_user()
+    if performer_user:
+      performer_data['entity_id'] = performer_user.id
+      performer_data['entity_name'] = performer_user.username
+
+  return {
+    'notification_uuid': notification.uuid,
+    'event_data': event_data,
+    'performer_data': performer_data,
+  }
+
+
+@contextmanager
+def notification_batch(batch_size=DEFAULT_BATCH_SIZE):
+  """
+  Context manager implementation which returns a target callable with the same signature
+  as spawn_notification. When the the context block exits the notifications generated by
+  the callable will be bulk inserted into the queue with the specified batch size.
+  """
+  with notification_queue.batch_insert(batch_size) as queue_put:
+    def spawn_notification_batch(repo, event_name, extra_data=None, subpage=None, pathargs=None,
+                                 performer_data=None):
+      event_data = build_event_data(repo, extra_data=extra_data, subpage=subpage)
+
+      notifications = model.notification.list_repo_notifications(repo.namespace_name,
+                                                                 repo.name,
+                                                                 event_name=event_name)
+      path = [repo.namespace_name, repo.name, event_name] + (pathargs or [])
+      for notification in list(notifications):
+        notification_data = build_notification_data(notification, event_data, performer_data)
+        queue_put(path, json.dumps(notification_data))
+
+    yield spawn_notification_batch
+
+
+def spawn_notification(repo, event_name, extra_data=None, subpage=None, pathargs=None,
+                       performer_data=None):
+  with notification_batch(1) as batch_spawn:
+    batch_spawn(repo, event_name, extra_data, subpage, pathargs, performer_data)

notifications/notificationevent.py

@@ -3,7 +3,7 @@ import time
 import re
 
 from datetime import datetime
-from notifications.notificationhelper import build_event_data
+from notifications import build_event_data
 from util.jinjautil import get_template_env
 from util.secscan import PRIORITY_LEVELS, get_priority_for_index
 

util/secscan/analyzer.py

@@ -8,7 +8,7 @@ import features
 from data.database import ExternalNotificationEvent, IMAGE_NOT_SCANNED_ENGINE_VERSION, Image
 from data.model.tag import filter_tags_have_repository_event, get_tags_for_image
 from data.model.image import set_secscan_status, get_image_with_storage_and_parent_base
-from notifications.notificationhelper import spawn_notification
+from notifications import spawn_notification
 from util.secscan import PRIORITY_LEVELS
 from util.secscan.api import (APIRequestFailure, AnalyzeLayerException, MissingParentLayerException,
                               InvalidLayerException, AnalyzeLayerRetryException)

util/secscan/notifier.py

@@ -10,7 +10,7 @@ from data.model.tag import (filter_has_repository_event, filter_tags_have_reposi
 
 from data.database import (Image, ImageStorage, ExternalNotificationEvent, Repository,
                            RepositoryTag)
-from notifications.notificationhelper import notification_batch
+from notifications import notification_batch
 from util.secscan import PRIORITY_LEVELS
 from util.secscan.api import APIRequestFailure
 from util.morecollections import AttrDict, StreamingDiffTracker, IndexedStreamingDiffTracker