Merge pull request #473 from Multiwoven/cherry-pick-ce-commit-4bc5750…

…b05cc7f08c8a72750e210e69f07f8c680

feat(CE): Fetch json_schema from model

server/app/models/catalog.rb

@@ -50,4 +50,9 @@ def stream_to_protocol(stream)
       request_rate_concurrency:
     )
   end
+
+  def json_schema(stream_name)
+    stream = find_stream_by_name(stream_name)
+    stream&.[]("json_schema") || {}
+  end
 end

server/app/models/schema_validations/models/configuration_aiml.json

@@ -3,7 +3,7 @@
   "$schema": "http://json-schema.org/draft-04/schema#",
   "properties": {
     "harvesters": { "type": "array" },
-    "json_schema": { "type": "array" }
+    "json_schema": { "type": "object" }
   },
   "required": ["harvesters"]
 }

server/app/models/schema_validations/models/configuration_dynamic_sql.json

@@ -3,7 +3,7 @@
   "$schema": "http://json-schema.org/draft-04/schema#",
   "properties": {
     "harvesters": { "type": "array" },
-    "json_schema": { "type": "array" }
+    "json_schema": { "type": "object" }
   },
   "required": ["harvesters", "json_schema"]
 }

server/app/serializers/model_serializer.rb

@@ -6,4 +6,16 @@ class ModelSerializer < ActiveModel::Serializer
   attribute :connector do
     ConnectorSerializer.new(object.connector).attributes
   end
+
+  def configuration
+    if object.ai_ml?
+      connector = object.connector
+      json_schema = connector.catalog.json_schema(connector.connector_name)
+      existing_config = object.configuration
+
+      existing_config.merge({ json_schema: })
+    else
+      object.configuration
+    end
+  end
 end

server/spec/contracts/model_contracts_spec.rb

@@ -64,7 +64,7 @@
             query: "SELECT * FROM table;",
             query_type: "dynamic_sql",
             primary_key: "id",
-            configuration: { "json_schema" => [], "harvesters" => [] }
+            configuration: { "json_schema" => {}, "harvesters" => [] }
           }
         }
       end

server/spec/factories/catalogs.rb

@@ -20,7 +20,14 @@
       {
         "streams" => [
           { "name" => "profile", "batch_support" => false, "batch_size" => 1, "json_schema" => {} },
-          { "name" => "customer", "batch_support" => false, "batch_size" => 1, "json_schema" => {} }
+          { "name" => "customer", "batch_support" => false, "batch_size" => 1, "json_schema" => {} },
+          { "name" => "DatabricksModel", "batch_support" => false, "batch_size" => 1, "json_schema" =>
+            {
+              "input" => [{ "name" => "inputs.0", "type" => "string", "value" => "dynamic", "value_type" => "dynamic" },
+                          { "name" => "inputs.0", "type" => "number", "value" => "9522", "value_type" => "static" }],
+              "output" => [{ "name" => "predictions.col1.0", "type" => "string" },
+                           { "name" => "predictions.col1.1", "type" => "number" }]
+            } }
         ],
         "request_rate_limit" => 60,
         "request_rate_limit_unit" => "minute",

server/spec/models/catalog_spec.rb

@@ -101,4 +101,29 @@
       end
     end
   end
+
+  describe "#json_schema" do
+    let(:json_schema) do
+      { "input" => [{ "name" => "inputs.0", "type" => "string", "value" => "dynamic", "value_type" => "dynamic" },
+                    { "name" => "inputs.0", "type" => "number", "value" => "9522", "value_type" => "static" }],
+        "output" => [{ "name" => "predictions.col1.0", "type" => "string" },
+                     { "name" => "predictions.col1.1", "type" => "number" }] }
+    end
+    let(:catalog) do
+      create(:catalog, catalog: { "streams" => [{
+               "name" => "DatabricksModel",
+               "json_schema" => json_schema
+             }] })
+    end
+
+    it "returns the correct json_schema when it exists" do
+      json_schema_from_catalog = catalog.json_schema("DatabricksModel")
+      expect(json_schema_from_catalog).to eq(json_schema)
+    end
+
+    it "returns an empty json when no valid schema available" do
+      json_schema_from_catalog = catalog.json_schema("NotExist")
+      expect(json_schema_from_catalog).to eq({})
+    end
+  end
 end

server/spec/models/model_spec.rb

@@ -116,7 +116,7 @@
               name: "test_model",
               query_type: :dynamic_sql, connector_id: source.id,
               workspace_id: source.workspace_id,
-              configuration: { "harvesters": [], "json_schema": [] }
+              configuration: { "harvesters": [], "json_schema": {} }
             )
             expect(dynamic_sql_model).to be_valid
           end
@@ -208,7 +208,7 @@
                             "preprocess": ""
                           }
                         ],
-                        "json_schema": [
+                        "json_schema":
                           {
                             "input": [
                               {
@@ -225,7 +225,6 @@
                               }
                             ]
                           }
-                        ]
                       })
     end
 

server/spec/requests/api/v1/models_controller_spec.rb

@@ -10,6 +10,10 @@
     create(:connector, workspace:, connector_type: "destination", name: "klavio1", connector_name: "Klaviyo")
   end
 
+  let(:source_connector) do
+    create(:connector, workspace:, connector_type: "source", name: "DatabricksModel", connector_name: "DatabricksModel")
+  end
+
   let(:connector_without_catalog) do
     create(:connector, workspace:, connector_type: "destination", name: "klavio1", connector_name: "Klaviyo")
   end
@@ -26,11 +30,15 @@
   let!(:ai_ml_model) do
     create(:model, query_type: :ai_ml, connector:, configuration: { harvesters: [] }, workspace:)
   end
+  let!(:ai_ml_source_model) do
+    create(:model, query_type: :ai_ml, connector: source_connector, configuration: { harvesters: [] }, workspace:)
+  end
   let(:viewer_role) { create(:role, :viewer) }
   let(:member_role) { create(:role, :member) }
 
   before do
     create(:catalog, connector:)
+    create(:catalog, connector: source_connector)
     user.confirm
   end
 
@@ -47,7 +55,7 @@
         get "/api/v1/models", headers: auth_headers(user, workspace_id)
         expect(response).to have_http_status(:ok)
         response_hash = JSON.parse(response.body).with_indifferent_access
-        expect(response_hash[:data].count).to eql(6)
+        expect(response_hash[:data].count).to eql(7)
         expect(response_hash.dig(:data, 0, :type)).to eq("models")
         expect(response_hash.dig(:links, :first)).to include("http://www.example.com/api/v1/models?page=1")
 
@@ -68,7 +76,7 @@
         get "/api/v1/models", headers: auth_headers(user, workspace_id)
         expect(response).to have_http_status(:ok)
         response_hash = JSON.parse(response.body).with_indifferent_access
-        expect(response_hash[:data].count).to eql(6)
+        expect(response_hash[:data].count).to eql(7)
         expect(response_hash.dig(:data, 0, :type)).to eq("models")
         expect(response_hash.dig(:links, :first)).to include("http://www.example.com/api/v1/models?page=1")
       end
@@ -78,7 +86,7 @@
         get "/api/v1/models", headers: auth_headers(user, workspace_id)
         expect(response).to have_http_status(:ok)
         response_hash = JSON.parse(response.body).with_indifferent_access
-        expect(response_hash[:data].count).to eql(6)
+        expect(response_hash[:data].count).to eql(7)
         expect(response_hash.dig(:data, 0, :type)).to eq("models")
         expect(response_hash.dig(:links, :first)).to include("http://www.example.com/api/v1/models?page=1")
 
@@ -134,7 +142,7 @@
         get "/api/v1/models", headers: auth_headers(user, workspace_id)
         expect(response).to have_http_status(:ok)
         model_ids = JSON.parse(response.body)["data"].map { |m| m["id"] }
-        expect(model_ids.count).to eql(6)
+        expect(model_ids.count).to eql(7)
 
         audit_log = AuditLog.last
         expect(audit_log).not_to be_nil
@@ -183,6 +191,34 @@
         expect(audit_log.updated_at).not_to be_nil
       end
 
+      it "returns success and fetch ai_ml model with configuration from catalog" do
+        get "/api/v1/models/#{ai_ml_source_model.id}", headers: auth_headers(user, workspace_id)
+        expect(response).to have_http_status(:ok)
+        response_hash = JSON.parse(response.body).with_indifferent_access
+        expect(response_hash.dig(:data, :id)).to be_present
+        expect(response_hash.dig(:data, :id)).to eq(ai_ml_source_model.id.to_s)
+        expect(response_hash.dig(:data, :type)).to eq("models")
+        expect(response_hash.dig(:data, :attributes, :name)).to eq(ai_ml_source_model.name)
+        expect(response_hash.dig(:data, :attributes, :query)).to eq(ai_ml_source_model.query)
+        expect(response_hash.dig(:data, :attributes, :query_type)).to eq(ai_ml_source_model.query_type)
+        expect(response_hash.dig(:data, :attributes, :primary_key)).to eq(ai_ml_source_model.primary_key)
+        expected_configuration = {
+          "harvesters" => [],
+          "json_schema" =>
+          {
+            "input" => [
+              { "name" => "inputs.0", "type" => "string", "value" => "dynamic", "value_type" => "dynamic" },
+              { "name" => "inputs.0", "type" => "number", "value" => "9522", "value_type" => "static" }
+            ],
+            "output" => [
+              { "name" => "predictions.col1.0", "type" => "string" },
+              { "name" => "predictions.col1.1", "type" => "number" }
+            ]
+          }
+        }
+        expect(response_hash.dig(:data, :attributes, :configuration)).to eq(expected_configuration)
+      end
+
       it "returns success and fetch model for viewer role" do
         workspace.workspace_users.first.update(role: viewer_role)
         get "/api/v1/models/#{models.first.id}", headers: auth_headers(user, workspace_id)
@@ -364,7 +400,8 @@
           expect(response_hash.dig(:data, :id)).to be_present
           expect(response_hash.dig(:data, :attributes, :name)).to eq(request_body.dig(:model, :name))
           expect(response_hash.dig(:data, :attributes, :query_type)).to eq("ai_ml")
-          expect(response_hash.dig(:data, :attributes, :configuration)).to eq(request_body.dig(:model, :configuration))
+          expected_configuration = { "harvesters" => [], "json_schema" => {} }
+          expect(response_hash.dig(:data, :attributes, :configuration)).to eq(expected_configuration)
 
           audit_log = AuditLog.last
           expect(audit_log).not_to be_nil
@@ -397,7 +434,7 @@
                     "preprocess" => ""
                   }
                 ],
-                "json_schema" => [
+                "json_schema" =>
                   {
                     "input" => [
                       {
@@ -414,7 +451,6 @@
                       }
                     ]
                   }
-                ]
               }
             }
           }
@@ -428,7 +464,34 @@
           expect(response_hash.dig(:data, :id)).to be_present
           expect(response_hash.dig(:data, :attributes, :name)).to eq(request_body.dig(:model, :name))
           expect(response_hash.dig(:data, :attributes, :query_type)).to eq("dynamic_sql")
-          expect(response_hash.dig(:data, :attributes, :configuration)).to eq(request_body.dig(:model, :configuration))
+          expected_configuration = {
+            "harvesters" => [
+              {
+                "value" => "dynamic test",
+                "method" => "dom",
+                "selector" => "dom_id",
+                "preprocess" => ""
+              }
+            ],
+            "json_schema" =>
+              {
+                "input" => [
+                  {
+                    "name" => "risk_level",
+                    "type" => "string",
+                    "value" => "",
+                    "value_type" => "dynamic"
+                  }
+                ],
+                "output" => [
+                  {
+                    "name" => "data.col0.calculated_risk",
+                    "type" => "string"
+                  }
+                ]
+              }
+          }
+          expect(response_hash.dig(:data, :attributes, :configuration)).to eq(expected_configuration)
 
           audit_log = AuditLog.last
           expect(audit_log).not_to be_nil
@@ -576,7 +639,8 @@
           expect(response_hash.dig(:data, :id)).to eq(models.second.id.to_s)
           expect(response_hash.dig(:data, :attributes, :name)).to eq(request_body.dig(:model, :name))
           expect(response_hash.dig(:data, :attributes, :query_type)).to eq("ai_ml")
-          expect(response_hash.dig(:data, :attributes, :configuration)).to eq(request_body.dig(:model, :configuration))
+          expected_configuration = { "harvesters" => [], "json_schema" => {} }
+          expect(response_hash.dig(:data, :attributes, :configuration)).to eq(expected_configuration)
 
           audit_log = AuditLog.last
           expect(audit_log).not_to be_nil