Merge pull request #28 from rdytech/NEP-18090-new_endpoint_for_warm_up_cache

[NEP-18090]: Adding a class for warming up the dataset cache

Author: soundarya-mv

CHANGELOG.md

@@ -1,5 +1,9 @@
 ## Change Log
 
+## 0.1.6 - 2024-07-10
+
+* added a class **WarmUpCache** to hit the 'api/v1/dataset/warm_up_cache' endpoint to warm up the cache of all the datasets for a particular dashaboard being passed to the class - https://github.com/rdytech/superset-client/pull/28
+
 ## 0.1.5 - 2024-05-10
 
 * add multi config for multi env creds https://github.com/rdytech/superset-client/pull/22

Gemfile.lock

@@ -14,7 +14,7 @@ GIT
 PATH
   remote: .
   specs:
-    superset (0.1.5)
+    superset (0.1.6)
       dotenv (~> 2.7)
       enumerate_it (~> 1.7.0)
       faraday (~> 1.0)

bin/console

@@ -11,6 +11,7 @@ require "terminal-table"
 require "yaml"
 require "enumerate_it"
 require "superset"
+require "rollbar"
 
 Dir["./lib/**/*.rb"].each { |f| require f }
 # You can add fixtures and/or initialization code here to make experimenting

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

@@ -0,0 +1,41 @@
+module Superset
+  module Dataset
+    class WarmUpCache < Superset::Request
+
+      attr_reader :dashboard_id, :table_name, :db_name
+      
+      def initialize(dashboard_id:, table_name:, db_name:)
+        @dashboard_id = dashboard_id
+        @table_name = table_name
+        @db_name = db_name
+      end
+
+      def perform
+        response
+      end
+
+      def response
+        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 params(dashboard_id, table_name, db_name)
+        {
+          "dashboard_id" => dashboard_id,
+          "table_name" => table_name,
+          "db_name" => db_name
+        }
+      end
+
+      private
+
+      def route
+        "dataset/warm_up_cache"
+      end
+
+      def logger
+        @logger ||= Superset::Logger.new
+      end
+    end
+  end
+end

lib/superset/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Superset
-  VERSION = "0.1.5"
+  VERSION = "0.1.6"
 end

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

@@ -0,0 +1,19 @@
+require 'spec_helper'
+
+RSpec.describe Superset::Dataset::WarmUpCache do
+  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
+    let(:response) { 'Dataset warmed up' }
+
+    it 'warms up the dataset' do
+      expect(subject.perform).to eq response
+    end
+  end
+end