[NEP-18090]: Adding en endpoint for warming up the dataset cache

Adding new loggers and spec cases

Author: soundarya-mv

lib/superset/dataset/warm_up_cache.rb

@@ -17,15 +17,16 @@ def response
         dataset_details = fetch_dataset_details(dashboard_id)
         dataset_details.each do |dataset|
           begin
-            client.put(route, params(dashboard_id, dataset["datasource_name"], dataset["name"]))
+            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(e.message)
+            Rollbar.error("Warm up cache failed for the dashboard #{dashboard_id.to_s} and for the dataset #{dataset["datasource_name"]} - #{e.message}")
           end 
         end
       end
 
-      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)
+      def api_response(dataset_name, database_name)
+        client.put(route, params(dashboard_id, dataset_name, database_name))
       end
 
       def params(dashboard_id, dataset_name, db_name)
@@ -36,15 +37,23 @@ def params(dashboard_id, dataset_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
 
-      private
-
       def route
         "dataset/warm_up_cache"
       end
+
+      def logger
+        @logger ||= Superset::Logger.new
+      end
     end
   end
 end

spec/superset/dataset/warm_up_cache_spec.rb

@@ -5,33 +5,52 @@
   let(:dashboard_id) { 1 }
   let(:response) { nil }
 
-  before do
-    allow(subject).to receive(:response).and_return(response)
-  end
-
   describe '.perform' do
-    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 dataset' do
-        expect(subject.perform).to eq response
-      end
+		context "Dataset count is not considered" do
+			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
+				[
+					{"database"=>"client database 1", "datasource_name"=>"datasource 101"},
+					{"database"=>"client database 2", "datasource_name"=>"datasource 102"},
+			 	]
+			end
+			let(:api_response) { "Dataset warmed up" }
+			before do
+				allow(subject).to receive(:fetch_dataset_details).with(1) { dataset_details } 
+				allow(subject).to receive(:api_response).and_return(api_response)
+			end
+			it 'warms up both the dataset' do
+				subject.response
+				expect(subject).to have_received(:api_response).twice
+			end
     end
   end
 end