Without gcp

osprey_ui/src/components/bulk_actions/BulkActionStartModal.tsx

@@ -189,7 +189,7 @@ const BulkActionStartModalContainer = () => {
       job_name: values.jobName,
       job_description: values.jobDescription,
       workflow_name: workflow,
-      file_name: values.file.name,
+      file_name: values.file?.name,
       file: values.file,
       entity_type: values.entityType,
     });

osprey_worker/src/osprey/worker/cli/sinks.py

@@ -37,7 +37,7 @@
 from osprey.worker.lib.config import Config
 from osprey.worker.lib.osprey_engine import bootstrap_engine, bootstrap_engine_with_helpers, get_sources_provider
 from osprey.worker.lib.osprey_shared.logging import get_logger
-from osprey.worker.lib.publisher import PubSubPublisher
+from osprey.worker.lib.publisher import KafkaPublisher, PubSubPublisher
 from osprey.worker.lib.singletons import CONFIG, LABELS_PROVIDER
 from osprey.worker.lib.storage import postgres
 from osprey.worker.lib.storage.bigtable import osprey_bigtable
@@ -71,6 +71,40 @@ def gevent_liveliness_watcher() -> None:
             gevent.sleep(1 / 60.0)
 
 
+def get_analytics_publisher() -> PubSubPublisher | KafkaPublisher:
+    config = init_config()
+    with open('/sys/class/dmi/id/product_name', 'r') as pn:
+        if 'Google' in pn:
+            analytics_pubsub_project_id = config.get_str('PUBSUB_DATA_PROJECT_ID', 'osprey-dev')
+            analytics_pubsub_topic_id = config.get_str('PUBSUB_ANALYTICS_EVENT_TOPIC_ID', 'osprey-analytics')
+            analytics_publisher = PubSubPublisher(analytics_pubsub_project_id, analytics_pubsub_topic_id)
+            return analytics_publisher
+
+    bootstrap_servers = config.get_str_list('OSPREY_KAFKA_BOOTSTRAP_SERVERS', ['localhost'])
+    client_id = config.get_str('OSPREY_KAFKA_INPUT_STREAM_CLIENT_ID', 'localhost')
+    topic = config.get_str('OSPREY_KAFKA_ANALYTICS_TOPIC', 'osprey-analytics')
+    publisher = KafkaPublisher(bootstrap_servers, client_id, topic)
+
+    return publisher
+
+
+def get_webhooks_publisher() -> PubSubPublisher | KafkaPublisher:
+    config = init_config()
+    with open('/sys/class/dmi/id/product_name', 'r') as pn:
+        if 'Google' in pn:
+            osprey_webhook_pubsub_project = config.get_str('PUBSUB_OSPREY_WEBHOOKS_PROJECT_ID', 'osprey-dev')
+            osprey_webhook_pubsub_topic = config.get_str('PUBSUB_OSPREY_WEBHOOKS_TOPIC_ID', 'osprey-webhooks')
+            webhooks_publisher = PubSubPublisher(osprey_webhook_pubsub_project, osprey_webhook_pubsub_topic)
+            return webhooks_publisher
+
+    bootstrap_servers = config.get_str_list('OSPREY_KAFKA_BOOTSTRAP_SERVERS', ['localhost'])
+    client_id = config.get_str('OSPREY_KAFKA_INPUT_STREAM_CLIENT_ID', 'localhost')
+    topic = config.get_str('OSPREY_KAFKA_ANALYTICS_TOPIC', 'osprey-webhooks')
+    publisher = KafkaPublisher(bootstrap_servers, client_id, topic)
+
+    return publisher
+
+
 @click.group()
 def cli() -> None:
     pass
@@ -261,9 +295,7 @@ def run_bulk_label_sink(pooled: bool) -> None:
 
     engine = bootstrap_engine()
 
-    analytics_pubsub_project_id = config.get_str('PUBSUB_DATA_PROJECT_ID', 'osprey-dev')
-    analytics_pubsub_topic_id = config.get_str('PUBSUB_ANALYTICS_EVENT_TOPIC_ID', 'osprey-analytics')
-    analytics_publisher = PubSubPublisher(analytics_pubsub_project_id, analytics_pubsub_topic_id)
+    analytics_publisher = get_analytics_publisher()
 
     def factory() -> BulkLabelSink:
         # NOTE: It's very important the input stream is created per-webhook sink
@@ -297,18 +329,12 @@ def factory() -> BulkLabelSink:
 @click.option('--include-ids-from-file', type=click.File('r'))
 def rollback_bulk_label_effects(ctx: click.Context, task_id: int, include_ids_from_file: TextIO | None = None) -> None:
     # TODO: Clean up this copy pasta.
-    config = init_config()
     postgres.init_from_config('osprey_db')
 
     engine = bootstrap_engine()
 
-    analytics_pubsub_project_id = config.get_str('PUBSUB_DATA_PROJECT_ID', 'osprey-dev')
-    analytics_pubsub_topic_id = config.get_str('PUBSUB_ANALYTICS_EVENT_TOPIC_ID', 'osprey-analytics')
-    analytics_publisher = PubSubPublisher(analytics_pubsub_project_id, analytics_pubsub_topic_id)
-
-    osprey_webhook_pubsub_project = config.get_str('PUBSUB_OSPREY_WEBHOOKS_PROJECT_ID', 'osprey-dev')
-    osprey_webhook_pubsub_topic = config.get_str('PUBSUB_OSPREY_WEBHOOKS_TOPIC_ID', 'osprey-webhooks')
-    webhooks_publisher = PubSubPublisher(osprey_webhook_pubsub_project, osprey_webhook_pubsub_topic)
+    analytics_publisher = get_analytics_publisher()
+    webhooks_publisher = get_webhooks_publisher()
 
     task = BulkLabelTask.get_one(task_id)
     if task is None:

osprey_worker/src/osprey/worker/lib/acls/definitions/super_user.json

@@ -29,6 +29,10 @@
       "name": "CAN_BULK_LABEL",
       "allow_all": true
     },
+    {
+      "name": "CAN_BULK_ACTION",
+      "allow_all": true
+    },
     {
       "name": "CAN_VIEW_SAVED_QUERIES",
       "allow_all": true

osprey_worker/src/osprey/worker/lib/publisher.py

@@ -2,6 +2,7 @@
 from typing import TypeVar
 
 from google.cloud import pubsub_v1
+from kafka import KafkaProducer
 from osprey.worker.lib.instruments import metrics
 from osprey.worker.lib.pubsub.publisher_client import BatchPubsubPublisherClient
 from pydantic import BaseModel
@@ -103,3 +104,32 @@ def publish(self, data: str, attributes: dict[str, str] | None = None) -> None:
             attributes = {}
 
         super().publish(data, attributes)  # type: ignore[type-var]
+
+
+class KafkaPublisher(BasePublisher):
+    def __init__(self, bootstrap_servers: list[str], client_id: str, kafka_topic: str):
+        self._kafka_topic = kafka_topic
+        self._bootstrap_servers = bootstrap_servers
+        self._client_id = client_id
+
+        self._publisher = KafkaProducer(bootstrap_servers=self._bootstrap_servers, client_id=self._client_id)
+
+    def _prepare_data(self, data: _PydanticModelT) -> bytes:
+        """
+        Convert the data to bytes for publishing. Some topics will require packing to bytes in a specific way.
+        Override this method if needed.
+        """
+        return data.json(exclude_none=True).encode()
+
+    def publish(self, data: _PydanticModelT, attributes: dict[str, str] | None = None) -> None:
+        if attributes is None:
+            attributes = {}
+
+        try:
+            processed_data = self._prepare_data(data)
+            self._publisher.send(self._kafka_topic, value=processed_data, **attributes)
+        except Exception:
+            pass
+
+    def stop(self) -> None:
+        self._publisher.close()

osprey_worker/src/osprey/worker/ui_api/osprey/singletons.py

@@ -1,4 +1,6 @@
-from osprey.worker.lib.publisher import PubSubPublisher
+from typing import cast
+
+from osprey.worker.lib.publisher import BasePublisher, KafkaPublisher, PubSubPublisher
 from osprey.worker.lib.singleton import Singleton
 from osprey.worker.lib.singletons import CONFIG
 
@@ -14,4 +16,19 @@ def _init_analytics_publisher() -> PubSubPublisher:
     return PubSubPublisher(project, topic)
 
 
-ANALYTICS_PUBLISHER: Singleton[PubSubPublisher] = Singleton(_init_analytics_publisher)
+def _init_kafka_analytics_publisher() -> KafkaPublisher:
+    config = CONFIG.instance()
+    bootstrap_servers = config.get_str_list('OSPREY_KAFKA_BOOTSTRAP_SERVERS', ['localhost'])
+    client_id = config.get_str('OSPREY_KAFKA_INPUT_STREAM_CLIENT_ID', 'osprey-dev')
+    topic = config.get_str('OSPREY_KAFKA_ANALYTICS_TOPIC', 'osprey-analytics')
+    return KafkaPublisher(bootstrap_servers, client_id, topic)
+
+
+def get_analytics_publisher() -> Singleton[BasePublisher]:
+    with open('/sys/class/dmi/id/product_name', 'r') as pn:
+        if 'Google' in pn:
+            ANALYTICS_PUBLISHER: Singleton[PubSubPublisher] = Singleton(_init_analytics_publisher)
+        else:
+            ANALYTICS_PUBLISHER: Singleton[KafkaPublisher] = Singleton(_init_kafka_analytics_publisher)  # type: ignore[no-redef]
+
+    return cast(Singleton[BasePublisher], ANALYTICS_PUBLISHER)

osprey_worker/src/osprey/worker/ui_api/osprey/views/bulk_actions.py

@@ -1,12 +1,16 @@
 from typing import Any
 
-from flask import Blueprint, abort, jsonify
-from osprey.worker.lib.storage.bulk_action_task import BulkActionJob
+from flask import Blueprint, current_app, jsonify, request
+from osprey.worker.lib.snowflake import generate_snowflake
+from osprey.worker.lib.storage.bulk_action_files import BulkActionFileManager
+from osprey.worker.lib.storage.bulk_action_task import BulkActionJob, BulkActionJobStatus
 from osprey.worker.ui_api.osprey.lib.abilities import CanBulkAction, require_ability
 from osprey.worker.ui_api.osprey.lib.decorators import require_multipart_form_data
 from osprey.worker.ui_api.osprey.lib.marshal import JsonBodyMarshaller, marshal_with
 from pydantic import BaseModel
 
+from ..lib.auth import get_current_user_email
+
 blueprint = Blueprint('bulk_actions', __name__)
 
 
@@ -18,42 +22,35 @@ class StartBulkActionJobRequest(BaseModel, JsonBodyMarshaller):
     entity_type: str
 
 
-# NOTE(ayubun): Bulk action requires GCS to upload the files to process for the actions.
-#               I have modified the endpoints to return 501 for now.
-
-
 @blueprint.route('/bulk_action/start', methods=['POST'])
 @require_ability(CanBulkAction)
 @marshal_with(StartBulkActionJobRequest)
 def start_bulk_job(start_bulk_action_job_request: StartBulkActionJobRequest) -> Any:
-    # TODO(ayubun): Support bulk action service
-    return abort(501, 'Not Implemented')
-
-    # file_manager: BulkActionFileManager = current_app.bulk_action_file_manager
-
-    # job_id = generate_snowflake().to_int()
-    # BulkActionJob.create_job(
-    #     job_id=job_id,
-    #     user_id=get_current_user_email(),
-    #     gcs_path=f'{job_id}/{start_bulk_action_job_request.file_name}',
-    #     original_filename=start_bulk_action_job_request.file_name,
-    #     total_rows=0,
-    #     action_workflow_name=start_bulk_action_job_request.workflow_name,
-    #     entity_type=start_bulk_action_job_request.entity_type,
-    #     name=start_bulk_action_job_request.job_name,
-    #     description=start_bulk_action_job_request.job_description,
-    # )
-
-    # url = file_manager.generate_upload_url(f'{job_id}/{start_bulk_action_job_request.file_name}')
-    # return (
-    #     jsonify(
-    #         {
-    #             'id': str(job_id),
-    #             'url': url,
-    #         }
-    #     ),
-    #     200,
-    # )
+    file_manager: BulkActionFileManager = current_app.bulk_action_file_manager
+
+    job_id = generate_snowflake().to_int()
+    BulkActionJob.create_job(
+        job_id=job_id,
+        user_id=get_current_user_email(),
+        gcs_path=f'{job_id}/{start_bulk_action_job_request.file_name}',
+        original_filename=start_bulk_action_job_request.file_name,
+        total_rows=0,
+        action_workflow_name=start_bulk_action_job_request.workflow_name,
+        entity_type=start_bulk_action_job_request.entity_type,
+        name=start_bulk_action_job_request.job_name,
+        description=start_bulk_action_job_request.job_description,
+    )
+
+    url = file_manager.generate_upload_url(f'{job_id}/{start_bulk_action_job_request.file_name}')
+    return (
+        jsonify(
+            {
+                'id': str(job_id),
+                'url': url,
+            }
+        ),
+        200,
+    )
 
 
 class UploadCompletedRequest(BaseModel, JsonBodyMarshaller):
@@ -64,39 +61,33 @@ class UploadCompletedRequest(BaseModel, JsonBodyMarshaller):
 @require_ability(CanBulkAction)
 @marshal_with(UploadCompletedRequest)
 def upload_completed(upload_completed_request: UploadCompletedRequest) -> Any:
-    # TODO(ayubun): Support bulk action service
-    return abort(501, 'Not Implemented')
-
-    # job = BulkActionJob.get_one(int(upload_completed_request.job_id))
-    # if not job:
-    #     return jsonify({'error': 'Job not found'}), 404
+    job = BulkActionJob.get_one(int(upload_completed_request.job_id))
+    if not job:
+        return jsonify({'error': 'Job not found'}), 404
 
-    # job.update_job(status=BulkActionJobStatus.UPLOADED)
+    job.update_job(status=BulkActionJobStatus.UPLOADED)
 
-    # return (
-    #     jsonify(
-    #         {
-    #             'id': upload_completed_request.job_id,
-    #         }
-    #     ),
-    #     200,
-    # )
+    return (
+        jsonify(
+            {
+                'id': upload_completed_request.job_id,
+            }
+        ),
+        200,
+    )
 
 
 @blueprint.route('/bulk_action/upload/<job_id>/<file_name>', methods=['PUT', 'POST'])
 @require_multipart_form_data
 @require_ability(CanBulkAction)
 def upload_file(job_id, file_name) -> Any:
-    # TODO(ayubun): Support bulk action service
-    return abort(501, 'Not Implemented')
+    file = request.files.get('file')
+    if file is None or len(request.files) != 1:
+        return jsonify({'error': 'Invalid file'}), 400
 
-    # file = request.files.get('file')
-    # if file is None or len(request.files) != 1:
-    #     return jsonify({'error': 'Invalid file'}), 400
-
-    # file_manager: BulkActionFileManager = current_app.bulk_action_file_manager
-    # file_manager.upload_file(f'{job_id}/{file_name}', file)
-    # return jsonify({}), 200
+    file_manager: BulkActionFileManager = current_app.bulk_action_file_manager
+    file_manager.upload_file(f'{job_id}/{file_name}', file)
+    return jsonify({}), 200
 
 
 @blueprint.route('/bulk_action/jobs', methods=['GET'])

osprey_worker/src/osprey/worker/ui_api/osprey/views/rules_visualizer.py

@@ -8,9 +8,10 @@
 from osprey.worker.ui_api.osprey.lib.marshal import JsonBodyMarshaller, marshal_with
 from pydantic.main import BaseModel
 
-from ..singletons import ANALYTICS_PUBLISHER
+from ..singletons import get_analytics_publisher
 
 blueprint = Blueprint('rules_visualizer', __name__)
+ANALYTICS_PUBLISHER = get_analytics_publisher()
 
 
 class BaseActionsViewQuery(BaseModel, JsonBodyMarshaller):

osprey_worker/src/osprey/worker/ui_api/osprey/views/tests/test_bulk_actions.py

@@ -38,10 +38,6 @@ def test_bulk_action_start_job_with_bulk_action_ability(app: Flask, client: 'Fla
         },
     )
 
-    # TODO(caidanw): update this test when the bulk action feature is re-implemented
-    assert res.status_code == 501
-    return
-
     assert res.status_code == 200
     assert res.json['id'] is not None
     assert res.json['url'] is not None
@@ -75,10 +71,6 @@ def test_bulk_action_upload_completed_with_bulk_action_ability(app: Flask, clien
         },
     )
 
-    # TODO(caidanw): update this test when the bulk action feature is re-implemented
-    assert res.status_code == 501
-    return
-
     assert res.status_code == 200
     assert res.json['id'] is not None
     assert res.json['url'] is not None