Mirror of https://github.com/roostorg/osprey
github.com/roostorg/osprey
pass config as an arg to register_input_stream (#14)
+6
-5
+1
-1
osprey_worker/src/osprey/worker/adaptor/hookspecs/osprey_hooks.py
+1
-1
osprey_worker/src/osprey/worker/adaptor/hookspecs/osprey_hooks.py
···
43
43
44
44
45
45
@hookspec(firstresult=True)
46
-
def register_input_stream() -> BaseInputStream[BaseAckingContext[Action]]:
46
+
def register_input_stream(config: Config) -> BaseInputStream[BaseAckingContext[Action]]:
47
47
raise NotImplementedError('register_input_stream must be implemented by the plugin')
+2
-2
osprey_worker/src/osprey/worker/adaptor/plugin_manager.py
+2
-2
osprey_worker/src/osprey/worker/adaptor/plugin_manager.py
···
82
82
return None
83
83
84
84
85
-
def bootstrap_input_stream() -> BaseInputStream[BaseAckingContext[Action]] | None:
85
+
def bootstrap_input_stream(config: Config) -> BaseInputStream[BaseAckingContext[Action]] | None:
86
86
load_all_osprey_plugins()
87
87
88
88
# spec has firstresult=True set, so it will return the first registered stream if one is registered
89
-
stream = plugin_manager.hook.register_input_stream()
89
+
stream = plugin_manager.hook.register_input_stream(config=config)
90
90
if stream:
91
91
return stream
92
92
else:
+3
-2
osprey_worker/src/osprey/worker/sinks/input_stream_chooser.py
+3
-2
osprey_worker/src/osprey/worker/sinks/input_stream_chooser.py
···
24
24
"""Based on the `input_stream_source` constructs a configured input stream that can be used to source events to
25
25
classify. For more details, see `InputStreamSource`."""
26
26
27
+
config = CONFIG.instance()
28
+
27
29
if input_stream_source == InputStreamSource.PUBSUB:
28
-
config = CONFIG.instance()
29
30
gcloud_project = config.get_str('PUBSUB_OSPREY_PROJECT_ID', 'osprey-dev')
30
31
pubsub_subscription = config.get_str('PUBSUB_OSPREY_RULES_SINK_SUBSCRIPTION', 'rules-sink')
31
32
subscriber = pubsub_v1.SubscriberClient()
···
121
122
kafka_consumer=consumer,
122
123
)
123
124
elif input_stream_source == InputStreamSource.PLUGIN:
124
-
stream = bootstrap_input_stream()
125
+
stream = bootstrap_input_stream(config=config)
125
126
if stream is None:
126
127
raise AssertionError('No input stream plugin registered')
127
128
return stream