diff --git a/osprey_ui/src/components/bulk_actions/BulkActionStartModal.tsx b/osprey_ui/src/components/bulk_actions/BulkActionStartModal.tsx index 57951c27..f9957bb7 100644 --- a/osprey_ui/src/components/bulk_actions/BulkActionStartModal.tsx +++ b/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, }); diff --git a/osprey_worker/src/osprey/worker/cli/sinks.py b/osprey_worker/src/osprey/worker/cli/sinks.py index 25b6d31b..15a12671 100644 --- a/osprey_worker/src/osprey/worker/cli/sinks.py +++ b/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: diff --git a/osprey_worker/src/osprey/worker/lib/acls/definitions/super_user.json b/osprey_worker/src/osprey/worker/lib/acls/definitions/super_user.json index c824783c..527f5b7f 100644 --- a/osprey_worker/src/osprey/worker/lib/acls/definitions/super_user.json +++ b/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 diff --git a/osprey_worker/src/osprey/worker/lib/publisher.py b/osprey_worker/src/osprey/worker/lib/publisher.py index cb7a3ed0..3f09887b 100644 --- a/osprey_worker/src/osprey/worker/lib/publisher.py +++ b/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() diff --git a/osprey_worker/src/osprey/worker/ui_api/osprey/singletons.py b/osprey_worker/src/osprey/worker/ui_api/osprey/singletons.py index 10fecf34..38b14dc8 100644 --- a/osprey_worker/src/osprey/worker/ui_api/osprey/singletons.py +++ b/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) diff --git a/osprey_worker/src/osprey/worker/ui_api/osprey/views/bulk_actions.py b/osprey_worker/src/osprey/worker/ui_api/osprey/views/bulk_actions.py index 4b0284d1..455afc55 100644 --- a/osprey_worker/src/osprey/worker/ui_api/osprey/views/bulk_actions.py +++ b/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//', 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']) diff --git a/osprey_worker/src/osprey/worker/ui_api/osprey/views/rules_visualizer.py b/osprey_worker/src/osprey/worker/ui_api/osprey/views/rules_visualizer.py index d6ecde27..60a6c0bc 100644 --- a/osprey_worker/src/osprey/worker/ui_api/osprey/views/rules_visualizer.py +++ b/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): diff --git a/osprey_worker/src/osprey/worker/ui_api/osprey/views/tests/test_bulk_actions.py b/osprey_worker/src/osprey/worker/ui_api/osprey/views/tests/test_bulk_actions.py index 1a32286a..54609a24 100644 --- a/osprey_worker/src/osprey/worker/ui_api/osprey/views/tests/test_bulk_actions.py +++ b/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