feat: pass modular input script as parameter to input helper stream_e…

…vents
splunk · Sep 16, 2024 · 3411ff2 · 3411ff2
1 parent bc6a36b
commit 3411ff2
Show file tree

Hide file tree

Showing 9 changed files with 15 additions and 10 deletions.
diff --git a/docs/inputs/helper.md b/docs/inputs/helper.md
@@ -37,7 +37,7 @@ The input script `example_input_one` will import this module and call its method
 The file must contain the following two functions:
 
 - `validate_input(definition: smi.ValidationDefinition)`
-- `stream_events(inputs: smi.InputDefinition, event_writer: smi.EventWriter)`
+- `stream_events(inputs: smi.InputDefinition, event_writer: smi.EventWriter, modular_input: smi.Script)`
 
 ```python
 from splunklib import modularinput as smi
@@ -47,7 +47,7 @@ def validate_input(definition: smi.ValidationDefinition):
     ...
 
 
-def stream_events(inputs: smi.InputDefinition, event_writer: smi.EventWriter):
+def stream_events(inputs: smi.InputDefinition, event_writer: smi.EventWriter, modular_input: smi.Script):
     ...
 ```
 

diff --git a/splunk_add_on_ucc_framework/templates/input.template b/splunk_add_on_ucc_framework/templates/input.template
@@ -1,6 +1,8 @@
 import import_declare_test
 
-{% if not input_helper_module -%}
+{% if input_helper_module -%}
+import inspect
+{% else -%}
 import json
 {% endif -%}
 import sys
@@ -49,6 +51,9 @@ class {{class_name}}(smi.Script):
 
     def stream_events(self, inputs: smi.InputDefinition, ew: smi.EventWriter):
         {%- if input_helper_module %}
+        stream_events_signature = inspect.signature(stream_events)
+        if len(stream_events_signature.parameters) == 3:
+            return stream_events(inputs, ew, self)
         return stream_events(inputs, ew)
         {%- else %}
         input_items = [{'count': len(inputs.inputs)}]

diff --git a/...expected_addons/expected_addon_after_init/demo_addon_for_splunk/package/bin/demo_input.py b/...expected_addons/expected_addon_after_init/demo_addon_for_splunk/package/bin/demo_input.py
@@ -31,7 +31,7 @@ def validate_input(self, definition: smi.ValidationDefinition):
         return validate_input(definition)
 
     def stream_events(self, inputs: smi.InputDefinition, ew: smi.EventWriter):
-        return stream_events(inputs, ew)
+        return stream_events(inputs, ew, self)
 
 
 if __name__ == '__main__':

diff --git a/...d_addons/expected_addon_after_init/demo_addon_for_splunk/package/bin/demo_input_helper.py b/...d_addons/expected_addon_after_init/demo_addon_for_splunk/package/bin/demo_input_helper.py
@@ -40,7 +40,7 @@ def validate_input(definition: smi.ValidationDefinition):
     return
 
 
-def stream_events(inputs: smi.InputDefinition, event_writer: smi.EventWriter):
+def stream_events(inputs: smi.InputDefinition, event_writer: smi.EventWriter, modular_input: smi.Script):
     # inputs.inputs is a Python dictionary object like:
     # {
     #   "demo_input://<input_name>": {

diff --git a/...ns/expected_output_global_config_everything/Splunk_TA_UCCExample/bin/example_input_one.py b/...ns/expected_output_global_config_everything/Splunk_TA_UCCExample/bin/example_input_one.py
@@ -121,7 +121,7 @@ def validate_input(self, definition: smi.ValidationDefinition):
         return validate_input(definition)
 
     def stream_events(self, inputs: smi.InputDefinition, ew: smi.EventWriter):
-        return stream_events(inputs, ew)
+        return stream_events(inputs, ew, self)
 
 
 if __name__ == '__main__':

diff --git a/...ns/expected_output_global_config_everything/Splunk_TA_UCCExample/bin/example_input_two.py b/...ns/expected_output_global_config_everything/Splunk_TA_UCCExample/bin/example_input_two.py
@@ -91,7 +91,7 @@ def validate_input(self, definition: smi.ValidationDefinition):
         return validate_input(definition)
 
     def stream_events(self, inputs: smi.InputDefinition, ew: smi.EventWriter):
-        return stream_events(inputs, ew)
+        return stream_events(inputs, ew, self)
 
 
 if __name__ == '__main__':

diff --git a/...ta/test_addons/package_global_config_everything/package/bin/example_input_three_helper.py b/...ta/test_addons/package_global_config_everything/package/bin/example_input_three_helper.py
@@ -7,7 +7,7 @@ def validate_input(definition: smi.ValidationDefinition):
     return
 
 
-def stream_events(inputs: smi.InputDefinition, ew: smi.EventWriter):
+def stream_events(inputs: smi.InputDefinition, ew: smi.EventWriter, modular_input: smi.Script):
     # INPUT 3
     input_items = [{'count': len(inputs.inputs)}]
     for input_name, input_item in inputs.inputs.items():

diff --git a/tests/testdata/test_addons/package_global_config_everything/package/bin/helper_one.py b/tests/testdata/test_addons/package_global_config_everything/package/bin/helper_one.py
@@ -7,7 +7,7 @@ def validate_input(definition: smi.ValidationDefinition):
     return
 
 
-def stream_events(inputs: smi.InputDefinition, ew: smi.EventWriter):
+def stream_events(inputs: smi.InputDefinition, ew: smi.EventWriter, modular_input: smi.Script):
     # INPUT 1
     input_items = [{'count': len(inputs.inputs)}]
     for input_name, input_item in inputs.inputs.items():

diff --git a/tests/testdata/test_addons/package_global_config_everything/package/bin/helper_two.py b/tests/testdata/test_addons/package_global_config_everything/package/bin/helper_two.py
@@ -7,7 +7,7 @@ def validate_input(definition: smi.ValidationDefinition):
     return
 
 
-def stream_events(inputs: smi.InputDefinition, ew: smi.EventWriter):
+def stream_events(inputs: smi.InputDefinition, ew: smi.EventWriter, modular_input: smi.Script):
     # INPUT 2
     input_items = [{'count': len(inputs.inputs)}]
     for input_name, input_item in inputs.inputs.items():