Add support for header passing to sinks (#60)

* Add support for header passing to sinks.

* Use json file for custom file.

* Fix call to producer helper.

* Update comment.

Author: doru1004

examples/cloud_object_storage/cos_sink.py

@@ -18,6 +18,7 @@
 import rayvens
 import sys
 import time
+import json
 
 # This example demonstrates how to send objects to the AWS S3 or
 # IBM Cloud Object Storage.
@@ -37,6 +38,11 @@
 
 # Initialize Ray and Rayvens
 ray.init()
+
+# TODO: make header setting work for Kafka transport. Currently the
+# Camel-K component for Kafka does not propagate message headers. This
+# will be fixed by Camel-K 1.8.0 release.
+# rayvens.init(transport="kafka")
 rayvens.init()
 
 # Create an object stream
@@ -47,8 +53,7 @@
                    bucket_name=bucket,
                    access_key_id=access_key_id,
                    secret_access_key=secret_access_key,
-                   endpoint=endpoint,
-                   file_name="test-file-name.txt")
+                   endpoint=endpoint)
 
 if region is not None:
     sink_config['region'] = region
@@ -57,7 +62,10 @@
 sink = stream.add_sink(sink_config)
 
 # Send file contents to Cloud Object Storage:
-stream << "File contents sample!"
+json_content = ['test', {'json': ('content', None, 1.0, 2)}]
+event = rayvens.OutputEvent(json.dumps(json_content),
+                            {"CamelAwsS3Key": "custom_file.json"})
+stream << event
 
 # Run for a while
 time.sleep(10)

rayvens/__init__.py

@@ -15,5 +15,6 @@
 #
 
 from rayvens.api import init, Stream
+from rayvens.core.common import OutputEvent
 
-__all__ = ['init', 'Stream']
+__all__ = ['init', 'Stream', 'OutputEvent']

rayvens/core/catalog.py

@@ -455,7 +455,7 @@ def cos_sink(config):
     secret_access_key = config['secret_access_key']
     endpoint = config['endpoint']
 
-    file_name = None
+    file_name = "default.txt"
     if 'file_name' in config:
         file_name = config['file_name']
 
@@ -578,18 +578,19 @@ def cos_sink(config):
 
     # This is the default behavior when user application data is
     # written into a COS file directly.
-    if 'file_name' not in config:
-        raise TypeError('Created cloud object name is required.')
-    regular_spec = {
-        'steps': [{
-            'set-header': {
-                'name': 'CamelAwsS3Key',
-                'simple': f"{file_name}"
-            }
-        }, {
-            'to': uri
-        }]
-    }
+    if 'file_name' in config:
+        regular_spec = {
+            'steps': [{
+                'set-header': {
+                    'name': 'CamelAwsS3Key',
+                    'simple': f"{file_name}"
+                }
+            }, {
+                'to': uri
+            }]
+        }
+    else:
+        regular_spec = {'steps': [{'to': uri}]}
     spec_list.append((regular_spec, None))
     return spec_list
 

rayvens/core/common.py

@@ -28,6 +28,12 @@
 from ray.remote_function import RemoteFunction
 
 
+class OutputEvent:
+    def __init__(self, data, headers={}):
+        self.data = data
+        self.headers = headers
+
+
 def get_run_mode(camel_mode, check_port):
     if camel_mode == 'auto' or camel_mode == 'local':
         mode.run_mode = RayvensMode.LOCAL
@@ -116,11 +122,22 @@ def __init__(self, url):
         self.url = url
 
     def append(self, data):
+        event_data = data
+        event_headers = {}
+        if isinstance(data, OutputEvent):
+            event_data = data.data
+            event_headers = data.headers
         try:
             if isinstance(data, Path):
-                requests.post(self.url, str(data), timeout=(5, None))
+                requests.post(self.url,
+                              str(event_data),
+                              timeout=(5, None),
+                              headers=event_headers)
             else:
-                requests.post(self.url, data, timeout=(5, None))
+                requests.post(self.url,
+                              event_data,
+                              timeout=(5, None),
+                              headers=event_headers)
         except requests.exceptions.ConnectionError:
             pass
 
@@ -170,20 +187,27 @@ def __init__(self, name):
         self.producer = ProducerHelper(name)
 
     def append(self, data):
+        event_data = data
+        event_headers = {}
+        if isinstance(data, OutputEvent):
+            event_data = data.data
+            event_headers = data.headers
         if isinstance(data, Path):
-            self.producer.produce(str(data))
+            self.producer.produce(str(event_data), headers=event_headers)
         else:
-            self.producer.produce(data)
+            self.producer.produce(event_data, headers=event_headers)
 
 
 class ProducerHelper:
     def __init__(self, name):
         self.name = name
         self.producer = Producer({'bootstrap.servers': brokers()})
 
-    def produce(self, data):
+    def produce(self, data, headers={}):
         if data is not None:
-            self.producer.produce(self.name, data.encode('utf-8'))
+            self.producer.produce(self.name,
+                                  data.encode('utf-8'),
+                                  headers=headers)
 
 
 @ray.remote(num_cpus=0.05)