Database and Dataset Imports

lib/superset/dashboard/import.rb

@@ -1,29 +1,32 @@
-# Import the provided Dashboard zip file
+# frozen_string_literal: true
+
+# Import the provided Dashboard zip file or directory (aka source)
 # In the context of this API import process, assumption is that the database.yaml file details will match
 # an existing database in the Target Superset Environment.
 
 # Scenario 1: Export from Env1 -- Import to Env1 into the SAME Environment
-#             Will result in updating/over writing the dashboard with the contents of the zip file 
+#             Will result in updating/over writing the dashboard with the contents of the source
 
 # Scenario 2: Export from Env1 -- Import to Env2 into a DIFFERENT Environment
-#             Assumption is that the database.yaml will match a database configuration in the target env. 
-#             Initial import will result in creating a new dashboard with the contents of the zip file. 
-#             Subsequent imports will result in updating/over writing the previous imported dashboard with the contents of the zip file.
+#             Assumption is that the database.yaml will match a database configuration in the target env.
+#             Initial import will result in creating a new dashboard with the contents of the source.
+#             Subsequent imports will result in updating/over writing the previous imported dashboard with the contents of the source.
 
 # the overwrite flag will determine if the dashboard will be updated or created new
 # overwrite: false .. will result in an error if a dashboard with the same UUID already exists
 
 # Usage
-# Superset::Dashboard::Import.new(source_zip_file: '/tmp/dashboard.zip').perform
+# Superset::Dashboard::Import.new(source: '/tmp/dashboard.zip').perform
+# Superset::Dashboard::Import.new(source: '/tmp/dashboard').perform
 #
 
 module Superset
   module Dashboard
     class Import < Request
-      attr_reader :source_zip_file, :overwrite
+      attr_reader :source, :overwrite
 
-      def initialize(source_zip_file: , overwrite: true)
-        @source_zip_file = source_zip_file
+      def initialize(source:, overwrite: true)
+        @source = source
         @overwrite = overwrite
       end
 
@@ -42,16 +45,20 @@ def response
       private
 
       def validate_params
-        raise ArgumentError, 'source_zip_file is required' if source_zip_file.nil?
-        raise ArgumentError, 'source_zip_file does not exist' unless File.exist?(source_zip_file)
-        raise ArgumentError, 'source_zip_file is not a zip file' unless File.extname(source_zip_file) == '.zip'
-        raise ArgumentError, 'overwrite must be a boolean' unless [true, false].include?(overwrite)
-        raise ArgumentError, "zip target database does not exist: #{zip_database_config_not_found_in_superset}" if zip_database_config_not_found_in_superset.present?
+        raise ArgumentError, "source is required" if source.nil?
+        raise ArgumentError, "source does not exist" unless File.exist?(source)
+        raise ArgumentError, "source is not a zip file or directory" unless zip? || directory?
+        raise ArgumentError, "overwrite must be a boolean" unless [true, false].include?(overwrite)
+
+        return unless database_config_not_found_in_superset.present?
+
+        raise ArgumentError,
+              "target database does not exist: #{database_config_not_found_in_superset}"
       end
 
       def payload
         {
-          formData: Faraday::UploadIO.new(source_zip_file, 'application/zip'),
+          formData: Faraday::UploadIO.new(source_zip_file, "application/zip"),
           overwrite: overwrite.to_s
         }
       end
@@ -60,24 +67,56 @@ def route
         "dashboard/import/"
       end
 
-      def zip_database_config_not_found_in_superset
-        zip_databases_details.select {|s| !superset_database_uuids_found.include?(s[:uuid]) }
+      def zip?
+        File.extname(source) == ".zip"
       end
 
-      def superset_database_uuids_found
-        @superset_database_uuids_found ||= begin
-          zip_databases_details.map {|i| i[:uuid]}.map do |uuid|
-            uuid if Superset::Database::List.new(uuid_equals: uuid).result.present?
-          end.compact
+      def directory?
+        File.directory?(source)
+      end
+
+      def source_zip_file
+        return source if zip?
+
+        Zip::File.open(new_zip_file, Zip::File::CREATE) do |zipfile|
+          Dir[File.join(source, "**", "**")].each do |file|
+            zipfile.add(file.sub("#{source}/", "#{File.basename(source)}/"), file) if File.file?(file)
+          end
         end
+        new_zip_file
+      end
+
+      def new_zip_file
+        new_database_name = dashboard_config[:databases].first[:content][:database_name]
+        File.join(source, "dashboard_import.zip")
       end
 
-      def zip_databases_details
-        zip_dashboard_config[:databases].map{|d| {uuid: d[:content][:uuid], name: d[:content][:database_name]} }
+      def database_config_not_found_in_superset
+        databases_details.reject { |s| superset_database_uuids_found.include?(s[:uuid]) }
+      end
+
+      def superset_database_uuids_found
+        @superset_database_uuids_found ||= databases_details.map { |i| i[:uuid] }.map do |uuid|
+          uuid if Superset::Database::List.new(uuid_equals: uuid).result.present?
+        end.compact
+      end
+
+      def databases_details
+        dashboard_config[:databases].map { |d| { uuid: d[:content][:uuid], name: d[:content][:database_name] } }
+      end
+
+      def dashboard_config
+        @dashboard_config ||= zip? ? zip_dashboard_config : directory_dashboard_config
       end
 
       def zip_dashboard_config
-        @zip_dashboard_config ||= Superset::Services::DashboardLoader.new(dashboard_export_zip: source_zip_file).perform
+        Superset::Services::DashboardLoader.new(dashboard_export_zip: source).perform
+      end
+
+      def directory_dashboard_config
+        Superset::Services::DashboardLoader::DashboardConfig.new(
+          dashboard_export_zip: "", tmp_uniq_dashboard_path: source
+        ).config
       end
     end
   end

lib/superset/database/import.rb

@@ -0,0 +1,119 @@
+# frozen_string_literal: true
+
+# Import the provided Database zip file or directory (aka source)
+
+# Scenario 1: Export from Env1 -- Import to Env1 into the SAME Environment
+#             Will result in updating/over writing the database with the contents of the source
+
+# Scenario 2: Export from Env1 -- Import to Env2 into a DIFFERENT Environment
+#             Initial import will result in creating a new database with the contents of the source.
+#             Subsequent imports will result in updating/over writing the previous imported database with the contents of the source.
+
+# the overwrite flag will determine if the database will be updated or created new
+# overwrite: false .. will result in an error if a database with the same UUID already exists
+
+# passwords can be set by passing in an hash in the form {"databases/MyDatabase.yaml": "my_password", "databases/db2.yaml": "other_pass"}
+# Usage
+# Superset::Database::Import.new(source: '/tmp/database.zip').perform
+# Superset::Database::Import.new(source: '/tmp/database').perform
+#
+
+require "json"
+require "zip"
+require "superset/file_utilities"
+
+module Superset
+  module Database
+    class Import < Request
+      include FileUtilities
+
+      attr_reader :source, :overwrite, :passwords, :ssh_tunnel_passwords,
+                :ssh_tunnel_private_key_passwords, :ssh_tunnel_private_keys
+
+      def initialize(source:, overwrite: true, passwords: {},
+                     ssh_tunnel_passwords: {}, ssh_tunnel_private_key_passwords: {},
+                     ssh_tunnel_private_keys: {})
+        @source = source
+        @overwrite = overwrite
+        @passwords = passwords
+        @ssh_tunnel_passwords = ssh_tunnel_passwords
+        @ssh_tunnel_private_key_passwords = ssh_tunnel_private_key_passwords
+        @ssh_tunnel_private_keys = ssh_tunnel_private_keys
+      end
+
+      def perform
+        validate_params
+        response
+      end
+
+      def response
+        @response ||= client(use_json: false).post(
+          route,
+          payload
+        )
+      end
+
+      private
+
+      def validate_params
+        raise ArgumentError, "source is required" if source.nil?
+        raise ArgumentError, "source does not exist" unless File.exist?(source)
+        raise ArgumentError, "source is not a zip file or directory" unless zip? || directory?
+        raise ArgumentError, "overwrite must be a boolean" unless [true, false].include?(overwrite)
+      end
+
+      def payload
+        {
+          formData: Faraday::UploadIO.new(source_zip_file, "application/zip"),
+          overwrite: overwrite.to_s,
+          passwords: passwords.to_json,
+          ssh_tunnel_passwords: ssh_tunnel_passwords.to_json,
+          ssh_tunnel_private_key_passwords: ssh_tunnel_private_key_passwords.to_json,
+          ssh_tunnel_private_keys: ssh_tunnel_private_keys.to_json
+        }
+      end
+
+      def route
+        "database/import/"
+      end
+
+      def zip?
+        File.extname(source) == ".zip"
+      end
+
+      def directory?
+        File.directory?(source)
+      end
+
+      def source_zip_file
+        return source if zip?
+
+        Zip::File.open(new_zip_file, Zip::File::CREATE) do |zipfile|
+          Dir[File.join(source, "**", "**")].each do |file|
+            zipfile.add(file.sub("#{source}/", "#{File.basename(source)}/"), file) if File.file?(file)
+          end
+        end
+        new_zip_file
+      end
+
+      def new_zip_file
+        database_name = database_config[:databases].first[:content][:database_name]
+        File.join(source, "database_import.zip")
+      end
+
+      def database_config
+        @database_config ||= zip? ? zip_database_config : directory_database_config
+      end
+
+      def zip_database_config
+        Superset::Services::Loader.new(export_zip: source).perform
+      end
+
+      def directory_database_config
+        Superset::Services::Loader::Config.new(
+          export_zip: "", tmp_uniq_path: source
+        ).config
+      end
+    end
+  end
+end

lib/superset/services/loader.rb

@@ -0,0 +1,70 @@
+# Given a path, load all yaml files
+
+require "superset/file_utilities"
+require "yaml"
+
+module Superset
+  module Services
+    class Loader
+      include FileUtilities
+
+      TMP_PATH = "/tmp/superset_imports".freeze
+
+      attr_reader :export_zip
+
+      def initialize(export_zip:)
+        @export_zip = export_zip
+      end
+
+      def perform
+        unzip_source_file
+        config
+      end
+
+      def config
+        @config ||= Config.new(
+          export_zip: export_zip,
+          tmp_uniq_path: tmp_uniq_path
+        ).config
+      end
+
+      private
+
+      def unzip_source_file
+        @extracted_files = unzip_file(export_zip, tmp_uniq_path)
+      end
+
+      def tmp_uniq_path
+        @tmp_uniq_path ||= File.join(TMP_PATH, uuid)
+      end
+
+      def uuid
+        SecureRandom.uuid
+      end
+
+      class Config < ::OpenStruct
+        def config
+          {
+            tmp_uniq_path: tmp_uniq_path,
+            dashboards: load_yamls_for("dashboards"),
+            databases: load_yamls_for("databases"),
+            datasets: load_yamls_for("datasets"),
+            charts: load_yamls_for("charts"),
+            metadata: load_yamls_for("metadata.yaml", pattern_sufix: nil)
+          }
+        end
+
+        def load_yamls_for(object_path, pattern_sufix: "**/*.yaml")
+          pattern = File.join([tmp_uniq_path, "**", object_path, pattern_sufix].compact)
+          Dir.glob(pattern).map do |file|
+            { filename: file, content: load_yaml_and_symbolize_keys(file) } if File.file?(file)
+          end.compact
+        end
+
+        def load_yaml_and_symbolize_keys(path)
+          YAML.load_file(path).deep_symbolize_keys
+        end
+      end
+    end
+  end
+end