[NEP-18090]: Splitting the class into 2. Dashboard::Warmup class warms up all datasets of the dashboard and Dataset::Warmup class warms_up individual datasets

Author: soundarya-mv

lib/superset/dashboard/warm_up_cache.rb

@@ -0,0 +1,42 @@
+module Superset
+  module Dashboard
+    class WarmUpCache < Superset::Request
+
+      attr_reader :dashboard_id
+      
+      def initialize(dashboard_id:)
+        @dashboard_id = dashboard_id
+      end
+
+      def perform
+        validate_dashboard_id
+        response
+      end
+
+      def response
+        dataset_details = fetch_dataset_details(dashboard_id)
+        dataset_details.each do |dataset|
+          begin
+            warm_up_dataset(dataset["datasource_name"], dataset["name"])
+          rescue => e
+            Rollbar.error("Warm up cache failed for the dashboard #{dashboard_id.to_s} and for the dataset #{dataset["datasource_name"]} - #{e}")
+          end 
+        end
+      end
+
+      def warm_up_dataset(dataset_name, db_name)
+        Superset::Dataset::WarmUpCache.new(dashboard_id: dashboard_id, table_name: dataset_name, db_name: db_name).perform
+      end
+
+      private
+      
+      def validate_dashboard_id
+        raise InvalidParameterError, "dashboard_id must be present and must be an integer" unless dashboard_id.present? && dashboard_id.is_a?(Integer)
+      end
+
+      def fetch_dataset_details(dashboard_id)
+        Superset::Dashboard::Datasets::List.new(dashboard_id).datasets_details.map { |dataset| dataset['database'].slice('name').merge(dataset.slice('datasource_name'))}
+      end
+    end
+  end
+end

lib/superset/dataset/warm_up_cache.rb

@@ -2,50 +2,32 @@ module Superset
   module Dataset
     class WarmUpCache < Superset::Request
 
-      attr_reader :dashboard_id
+      attr_reader :dashboard_id, :table_name, :db_name
 
-      def initialize(dashboard_id:)
+      def initialize(dashboard_id:, table_name:, db_name:)
         @dashboard_id = dashboard_id
+        @table_name = table_name
+        @db_name = db_name
       end
 
       def perform
-        validate_dashboard_id
         response
       end
 
       def response
-        dataset_details = fetch_dataset_details(dashboard_id)
-        dataset_details.each do |dataset|
-          begin
-            logger.info("Hitting #{route} for warming up the cache for the dashboard #{dashboard_id.to_s} and for the dataset #{dataset["datasource_name"]}")
-            api_response(dataset["datasource_name"], dataset["name"])
-          rescue => e
-            Rollbar.error("Warm up cache failed for the dashboard #{dashboard_id.to_s} and for the dataset #{dataset["datasource_name"]} - #{e}")
-          end 
-        end
+        logger.info("Hitting #{route} for warming up the cache for the dashboard #{dashboard_id.to_s} and for the dataset #{table_name}")
+        client.put(route, params(dashboard_id, table_name, db_name))
       end
 
-      def api_response(dataset_name, db_name)
-        client.put(route, params(dashboard_id, dataset_name, db_name))
-      end
-
-      def params(dashboard_id, dataset_name, db_name)
+      def params(dashboard_id, table_name, db_name)
         {
           "dashboard_id" => dashboard_id,
-          "table_name" => dataset_name,
+          "table_name" => table_name,
           "db_name" => db_name
         }
       end
 
       private
-      
-      def validate_dashboard_id
-        raise InvalidParameterError, "dashboard_id must be present and must be an integer" unless dashboard_id.present? && dashboard_id.is_a?(Integer)
-      end
-
-      def fetch_dataset_details(dashboard_id)
-        Superset::Dashboard::Datasets::List.new(dashboard_id).datasets_details.map { |dataset| dataset['database'].slice('name').merge(dataset.slice('datasource_name'))}
-      end
 
       def route
         "dataset/warm_up_cache"

spec/superset/dashboard/warm_up_cache_spec.rb

@@ -0,0 +1,56 @@
+require 'spec_helper'
+
+RSpec.describe Superset::Dashboard::WarmUpCache do
+  subject { described_class.new(dashboard_id: dashboard_id) }
+  let(:dashboard_id) { 1 }
+
+  describe '.perform' do
+    context "Dataset count is not considered" do
+      let(:response) { nil }
+      before do
+        allow(subject).to receive(:response).and_return(response)
+      end
+      context 'when dashboard_id is not present' do
+        let(:dashboard_id) { nil }
+
+        it 'raises an error' do
+          expect { subject.perform }.to raise_error(Superset::Request::InvalidParameterError, "dashboard_id must be present and must be an integer")
+        end
+      end
+
+      context 'when dashboard_id is not an integer' do
+        let(:dashboard_id) { 'string' }
+
+        it 'raises an error' do
+          expect { subject.perform }.to raise_error(Superset::Request::InvalidParameterError, "dashboard_id must be present and must be an integer")
+        end
+      end
+
+      context 'when dashboard_id is an integer' do
+        let(:response) { 'Dashboard warmed up' }
+
+        it 'warms up the dashboard' do
+          expect(subject.perform).to eq response
+        end
+      end
+    end
+
+    context 'when dashboard has multiple datasets' do
+      let(:dataset_details) do
+        [
+            {"name"=>"client database 1", "datasource_name"=>"datasource 101"},
+            {"name"=>"client database 2", "datasource_name"=>"datasource 102"},
+        ]
+      end
+      let(:api_response) { "Datasets warmed up" }
+      before do
+        allow(subject).to receive(:fetch_dataset_details).with(dashboard_id) { dataset_details } 
+        allow(subject).to receive(:warm_up_dataset).and_return(api_response)
+      end
+      it 'warms up all the datasets' do
+        subject.response
+        expect(subject).to have_received(:warm_up_dataset).exactly(dataset_details.count).times
+      end
+    end
+  end
+end

spec/superset/dataset/warm_up_cache_spec.rb

@@ -1,56 +1,19 @@
 require 'spec_helper'
 
 RSpec.describe Superset::Dataset::WarmUpCache do
-  subject { described_class.new(dashboard_id: dashboard_id) }
+  subject { described_class.new(dashboard_id: dashboard_id, table_name: table_name, db_name: db_name) }
   let(:dashboard_id) { 1 }
+  let(:table_name) { "Dataset 101"}
+  let(:db_name) { "Client Database 1" }
+  before do
+    allow(subject).to receive(:response).and_return(response)
+  end
 
   describe '.perform' do
-    context "Dataset count is not considered" do
-      let(:response) { nil }
-      before do
-        allow(subject).to receive(:response).and_return(response)
-      end
-      context 'when dashboard_id is not present' do
-        let(:dashboard_id) { nil }
-
-        it 'raises an error' do
-          expect { subject.perform }.to raise_error(Superset::Request::InvalidParameterError, "dashboard_id must be present and must be an integer")
-        end
-      end
-
-      context 'when dashboard_id is not an integer' do
-        let(:dashboard_id) { 'string' }
-
-        it 'raises an error' do
-          expect { subject.perform }.to raise_error(Superset::Request::InvalidParameterError, "dashboard_id must be present and must be an integer")
-        end
-      end
-
-      context 'when dashboard_id is an integer' do
-        let(:response) { 'Dashboard warmed up' }
-
-        it 'warms up the dashboard' do
-          expect(subject.perform).to eq response
-        end
-      end
-    end
+    let(:response) { 'Dataset warmed up' }
 
-    context 'when dashboard has multiple datasets' do
-      let(:dataset_details) do
-        [
-            {"name"=>"client database 1", "datasource_name"=>"datasource 101"},
-            {"name"=>"client database 2", "datasource_name"=>"datasource 102"},
-        ]
-      end
-      let(:api_response) { "Dataset warmed up" }
-      before do
-        allow(subject).to receive(:fetch_dataset_details).with(dashboard_id) { dataset_details } 
-        allow(subject).to receive(:api_response).and_return(api_response)
-      end
-      it 'warms up all the datasets' do
-        subject.response
-        expect(subject).to have_received(:api_response).exactly(dataset_details.count).times
-      end
+    it 'warms up the dataset' do
+      expect(subject.perform).to eq response
     end
   end
 end