From 095776ddf089ecd9ad07372087ca7cbade93a673 Mon Sep 17 00:00:00 2001 From: Chihurumnaya Ibiam Date: Thu, 2 Jul 2026 00:43:51 +0100 Subject: [PATCH 1/5] [WIP] Make Osprey runnable without GCP credentials. I've created a KafkaPublisher that's used in place of PubSubPublisher when osprey isn't being run on GCP. The published data doesn't have a subscriber/consumer yet, it's just published, the rules visualizer now shows nodes like it should. I need to better handle when an exception is raised in KafkaPublisher.publish. Signed-off-by: Chihurumnaya Ibiam --- .../src/osprey/worker/lib/publisher.py | 30 +++++++++++++++++++ .../osprey/worker/ui_api/osprey/singletons.py | 19 ++++++++++-- .../ui_api/osprey/views/rules_visualizer.py | 3 +- 3 files changed, 49 insertions(+), 3 deletions(-) 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..56dc1681 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,4 @@ -from osprey.worker.lib.publisher import PubSubPublisher +from osprey.worker.lib.publisher import KafkaPublisher, PubSubPublisher from osprey.worker.lib.singleton import Singleton from osprey.worker.lib.singletons import CONFIG @@ -14,4 +14,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[PubSubPublisher] | Singleton[KafkaPublisher]: + 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 ANALYTICS_PUBLISHER 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): From cf194974c924816ce9158e9a108abfef60dc457d Mon Sep 17 00:00:00 2001 From: Chihurumnaya Ibiam Date: Thu, 2 Jul 2026 02:38:41 +0100 Subject: [PATCH 2/5] Remove dependence on GCP in sinks. This should complete Tier1 of issue #343 Signed-off-by: Chihurumnaya Ibiam --- osprey_worker/src/osprey/worker/cli/sinks.py | 50 +++++++++++++++----- 1 file changed, 38 insertions(+), 12 deletions(-) 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: From 1dd7402b496a25fae729c5059f470ebc0577307a Mon Sep 17 00:00:00 2001 From: Chihurumnaya Ibiam Date: Tue, 7 Jul 2026 16:06:40 +0100 Subject: [PATCH 3/5] Cast return value to Singleton[BasePublisher] This makes it easier for any future subclasses of BasePublisher to work as a return value, so there's no need to explicitly state the return types like I did. Signed-off-by: Chihurumnaya Ibiam --- docs/book.toml | 2 +- .../src/osprey/worker/ui_api/osprey/singletons.py | 8 +++++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/docs/book.toml b/docs/book.toml index c663872b..8c47edbb 100644 --- a/docs/book.toml +++ b/docs/book.toml @@ -8,7 +8,7 @@ src = "." build-dir = "./book" [output.html] -site-url = "/osprey/" +site-url = "/osprey/latest/" edit-url-template = "https://github.com/roostorg/osprey/edit/main/docs/{path}" git-repository-url = "https://github.com/roostorg/osprey" 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 56dc1681..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 KafkaPublisher, 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 @@ -22,11 +24,11 @@ def _init_kafka_analytics_publisher() -> KafkaPublisher: return KafkaPublisher(bootstrap_servers, client_id, topic) -def get_analytics_publisher() -> Singleton[PubSubPublisher] | Singleton[KafkaPublisher]: +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 ANALYTICS_PUBLISHER + return cast(Singleton[BasePublisher], ANALYTICS_PUBLISHER) From 5ae3e71c7fd87da9fd5e0048fbde5d18f30fced5 Mon Sep 17 00:00:00 2001 From: Chihurumnaya Ibiam Date: Wed, 8 Jul 2026 00:37:12 +0100 Subject: [PATCH 4/5] Use optional chaining to avoid error. Signed-off-by: Chihurumnaya Ibiam --- osprey_ui/src/components/bulk_actions/BulkActionStartModal.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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, }); From 2ae506f8c5101aff2a5ec69ed5885d32df03e350 Mon Sep 17 00:00:00 2001 From: Chihurumnaya Ibiam Date: Fri, 10 Jul 2026 17:56:30 +0100 Subject: [PATCH 5/5] Enable Bulk action. It doesn't require GCS to upload the files to process for the actions as there's a LocalBulkActionFileManager that it defaults to for dev environments. Signed-off-by: Chihurumnaya Ibiam --- docs/book.toml | 2 +- .../lib/acls/definitions/super_user.json | 4 + .../ui_api/osprey/views/bulk_actions.py | 107 ++++++++---------- .../osprey/views/tests/test_bulk_actions.py | 8 -- 4 files changed, 54 insertions(+), 67 deletions(-) diff --git a/docs/book.toml b/docs/book.toml index 8c47edbb..c663872b 100644 --- a/docs/book.toml +++ b/docs/book.toml @@ -8,7 +8,7 @@ src = "." build-dir = "./book" [output.html] -site-url = "/osprey/latest/" +site-url = "/osprey/" edit-url-template = "https://github.com/roostorg/osprey/edit/main/docs/{path}" git-repository-url = "https://github.com/roostorg/osprey" 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/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/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