py: allow specifying update format in SQLContext.input_json (feldera#2144)

abhizer · web-flow · commit 8e2323026640 · 2024-08-02T16:24:10.000+05:45
Fixes: feldera#2143 Signed-off-by: Abhinav Gyawali <22275402+abhizer@users.noreply.github.com>
diff --git a/python/feldera/sql_context.py b/python/feldera/sql_context.py
@@ -283,18 +283,31 @@ def input_pandas(self, table_name: str, df: pandas.DataFrame, force: bool = Fals
                 )
             return
 
-    def input_json(self, table_name: str, data: Dict | list, force: bool = False):
+    def input_json(self, table_name: str, data: Dict | list, update_format: str = "raw", force: bool = False):
         """
         Push this JSON data to the specified table of the pipeline.
 
         :param table_name: The name of the table to push data into.
         :param data: The JSON encoded data to be pushed to the pipeline. The data should be in the form:
             `{'col1': 'val1', 'col2': 'val2'}` or `[{'col1': 'val1', 'col2': 'val2'}, {'col1': 'val1', 'col2': 'val2'}]`
+        :param update_format: The update format of the JSON data to be pushed to the pipeline. Must be one of:
+            "raw", "insert_delete". <https://www.feldera.com/docs/api/json#the-insertdelete-format>
         :param force: `True` to push data even if the pipeline is paused. `False` by default.
         """
 
+        if update_format not in ["raw", "insert_delete"]:
+            ValueError("update_format must be one of raw or insert_delete")
+
         array = True if isinstance(data, list) else False
-        self.client.push_to_pipeline(self.pipeline_name, table_name, "json", data, array=array, force=force)
+        self.client.push_to_pipeline(
+            self.pipeline_name,
+            table_name,
+           "json",
+            data,
+            update_format=update_format,
+            array=array,
+            force=force
+        )
 
     def register_local_view(self, name: str, query: str):
         """
diff --git a/python/tests/test_wireframes.py b/python/tests/test_wireframes.py
@@ -430,18 +430,18 @@ def test_input_json0(self):
         sql.register_table(TBL_NAME, SQLSchema({"id": "INT", "name": "STRING"}))
         sql.register_materialized_view(VIEW_NAME, f"SELECT * FROM {TBL_NAME}")
 
-        data = {'id': 1, 'name': 'a'}
+        data = {"insert": {'id': 1, 'name': 'a'}}
 
         out = sql.listen(VIEW_NAME)
 
         sql.start()
-        sql.input_json(TBL_NAME, data)
+        sql.input_json(TBL_NAME, data, update_format="insert_delete")
         sql.wait_for_completion(True)
 
         out_data = out.to_dict()
 
-        data.update({"insert_delete": 1})
-        assert out_data == [data]
+        data["insert"].update({"insert_delete": 1})
+        assert out_data == [data["insert"]]
 
     def test_input_json1(self):
         sql = SQLContext("test_input_json", TEST_CLIENT).get_or_create()
