From 8c395c7597948317fb4c1b884690298786029027 Mon Sep 17 00:00:00 2001
From: "github-actions[bot]"
 <41898282+github-actions[bot]@users.noreply.github.com>
Date: Fri, 27 Dec 2024 15:55:49 +0530
Subject: [PATCH] chore(CE): add payload limit to databricks model (#554)

Co-authored-by: TivonB-AI2 <124182151+TivonB-AI2@users.noreply.github.com>
---
 integrations/Gemfile.lock                     |  2 +-
 .../lib/multiwoven/integrations/rollout.rb    |  2 +-
 .../source/databrics_model/client.rb          |  8 +++++--
 .../source/databricks_model/client_spec.rb    | 21 +++++++++++++++++++
 4 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/integrations/Gemfile.lock b/integrations/Gemfile.lock
index cd85a166..37a81655 100644
--- a/integrations/Gemfile.lock
+++ b/integrations/Gemfile.lock
@@ -7,7 +7,7 @@ GIT
 PATH
   remote: .
   specs:
-    multiwoven-integrations (0.16.1)
+    multiwoven-integrations (0.16.2)
       MailchimpMarketing
       activesupport
       async-websocket
diff --git a/integrations/lib/multiwoven/integrations/rollout.rb b/integrations/lib/multiwoven/integrations/rollout.rb
index 6f6409fb..d833d1d1 100644
--- a/integrations/lib/multiwoven/integrations/rollout.rb
+++ b/integrations/lib/multiwoven/integrations/rollout.rb
@@ -2,7 +2,7 @@
 
 module Multiwoven
   module Integrations
-    VERSION = "0.16.1"
+    VERSION = "0.16.2"
 
     ENABLED_SOURCES = %w[
       Snowflake
diff --git a/integrations/lib/multiwoven/integrations/source/databrics_model/client.rb b/integrations/lib/multiwoven/integrations/source/databrics_model/client.rb
index 9132d001..33420e8b 100644
--- a/integrations/lib/multiwoven/integrations/source/databrics_model/client.rb
+++ b/integrations/lib/multiwoven/integrations/source/databrics_model/client.rb
@@ -62,8 +62,12 @@ def run_model(connection_config, payload)
 
       def process_response(response)
         if success?(response)
-          data = JSON.parse(response.body)
-          [RecordMessage.new(data: data, emitted_at: Time.now.to_i).to_multiwoven_message]
+          begin
+            data = JSON.parse(response.body)
+            [RecordMessage.new(data: data, emitted_at: Time.now.to_i).to_multiwoven_message]
+          rescue JSON::ParserError
+            create_log_message("DATABRICKS MODEL:RUN_MODEL", "error", "parsing failed: please send a valid payload")
+          end
         else
           create_log_message("DATABRICKS MODEL:RUN_MODEL", "error", "request failed: #{response.body}")
         end
diff --git a/integrations/spec/multiwoven/integrations/source/databricks_model/client_spec.rb b/integrations/spec/multiwoven/integrations/source/databricks_model/client_spec.rb
index de785f16..f17b6e9b 100644
--- a/integrations/spec/multiwoven/integrations/source/databricks_model/client_spec.rb
+++ b/integrations/spec/multiwoven/integrations/source/databricks_model/client_spec.rb
@@ -177,6 +177,27 @@
         expect(records.first.record.data).to eq(JSON.parse(response_body))
       end
     end
+
+    context "when the payload is invalid in read" do
+      let(:response_body) { "{\"key\": invalid_json}" }.to_json
+      before do
+        response = Net::HTTPSuccess.new("1.1", "200", "Unauthorized")
+        response.content_type = "application/json"
+        allow(response).to receive(:body).and_return(response_body)
+        allow(Multiwoven::Integrations::Core::HttpClient).to receive(:request)
+          .with("https://test-host.databricks.com/serving-endpoints/test/invocations",
+                "POST",
+                payload: JSON.parse(payload.to_json),
+                headers: headers)
+          .and_return(response)
+      end
+      it "handles exceptions during reading" do
+        records = client.read(sync_config)
+        expect(records.log).to be_a(Multiwoven::Integrations::Protocol::LogMessage)
+        expect(records.log.message).to eq("parsing failed: please send a valid payload")
+      end
+    end
+
     context "when the read is failed" do
       it "handles exceptions during reading" do
         error_instance = StandardError.new("test error")