sql-jarvis 1.8.0

data/lib/blazer/adapters/athena_adapter.rb

@@ -0,0 +1,128 @@
+module Blazer
+  module Adapters
+    class AthenaAdapter < BaseAdapter
+      def run_statement(statement, comment)
+        require "digest/md5"
+
+        columns = []
+        rows = []
+        error = nil
+
+        begin
+          resp =
+            client.start_query_execution(
+              query_string: statement,
+              # use token so we fetch cached results after query is run
+              client_request_token: Digest::MD5.hexdigest(statement),
+              query_execution_context: {
+                database: database,
+              },
+              result_configuration: {
+                output_location: settings["output_location"]
+              }
+            )
+          query_execution_id = resp.query_execution_id
+
+          timeout = data_source.timeout || 300
+          stop_at = Time.now + timeout
+          resp = nil
+
+          begin
+            resp = client.get_query_results(
+              query_execution_id: query_execution_id
+            )
+          rescue Aws::Athena::Errors::InvalidRequestException => e
+            if e.message != "Query has not yet finished. Current state: RUNNING"
+              raise e
+            end
+            if Time.now < stop_at
+              sleep(3)
+              retry
+            end
+          end
+
+          if resp && resp.result_set
+            column_info = resp.result_set.result_set_metadata.column_info
+            columns = column_info.map(&:name)
+            column_types = column_info.map(&:type)
+
+            untyped_rows = []
+
+            # paginated
+            resp.each do |page|
+              untyped_rows.concat page.result_set.rows.map { |r| r.data.map(&:var_char_value) }
+            end
+
+            utc = ActiveSupport::TimeZone['Etc/UTC']
+
+            rows = untyped_rows[1..-1] || []
+            column_types.each_with_index do |ct, i|
+              # TODO more column_types
+              case ct
+              when "timestamp"
+                rows.each do |row|
+                  row[i] = utc.parse(row[i])
+                end
+              when "date"
+                rows.each do |row|
+                  row[i] = Date.parse(row[i])
+                end
+              when "bigint"
+                rows.each do |row|
+                  row[i] = row[i].to_i
+                end
+              when "double"
+                rows.each do |row|
+                  row[i] = row[i].to_f
+                end
+              end
+            end
+          elsif resp
+            error = fetch_error(query_execution_id)
+          else
+            error = Blazer::TIMEOUT_MESSAGE
+          end
+        rescue Aws::Athena::Errors::InvalidRequestException => e
+          error = e.message
+          if error == "Query did not finish successfully. Final query state: FAILED"
+            error = fetch_error(query_execution_id)
+          end
+        end
+
+        [columns, rows, error]
+      end
+
+      def tables
+        glue.get_tables(database_name: database).table_list.map(&:name).sort
+      end
+
+      def schema
+        glue.get_tables(database_name: database).table_list.map { |t| {table: t.name, columns: t.storage_descriptor.columns.map { |c| {name: c.name, data_type: c.type} }} }
+      end
+
+      def preview_statement
+        "SELECT * FROM {table} LIMIT 10"
+      end
+
+      private
+
+      def database
+        @database ||= settings["database"] || "default"
+      end
+
+      def fetch_error(query_execution_id)
+        client.get_query_execution(
+          query_execution_id: query_execution_id
+        ).query_execution.status.state_change_reason
+      end
+
+      def client
+        @client ||= Aws::Athena::Client.new
+      end
+
+      def glue
+        @glue ||= Aws::Glue::Client.new
+      end
+    end
+  end
+end

data/lib/blazer/adapters/base_adapter.rb

@@ -0,0 +1,53 @@
+module Blazer
+  module Adapters
+    class BaseAdapter
+      attr_reader :data_source
+
+      def initialize(data_source)
+        @data_source = data_source
+      end
+
+      def run_statement(statement, comment)
+        # the one required method
+      end
+
+      def tables
+        [] # optional, but nice to have
+      end
+
+      def schema
+        [] # optional, but nice to have
+      end
+
+      def preview_statement
+        "" # also optional, but nice to have
+      end
+
+      def reconnect
+        # optional
+      end
+
+      def cost(statement)
+        # optional
+      end
+
+      def explain(statement)
+        # optional
+      end
+
+      def cancel(run_id)
+        # optional
+      end
+
+      def cachable?(statement)
+        true # optional
+      end
+
+      protected
+
+      def settings
+        @data_source.settings
+      end
+    end
+  end
+end

data/lib/blazer/adapters/bigquery_adapter.rb

@@ -0,0 +1,67 @@
+module Blazer
+  module Adapters
+    class BigQueryAdapter < BaseAdapter
+      def run_statement(statement, comment)
+        columns = []
+        rows = []
+        error = nil
+
+        begin
+          options = {}
+          options[:timeout] = data_source.timeout.to_i * 1000 if data_source.timeout
+          results = bigquery.query(statement, options) # ms
+          if results.complete?
+            columns = results.first.keys.map(&:to_s) if results.size > 0
+            rows = results.map(&:values)
+          else
+            error = Blazer::TIMEOUT_MESSAGE
+          end
+        rescue => e
+          error = e.message
+        end
+
+        [columns, rows, error]
+      end
+
+      def tables
+        table_refs.map { |t| "#{t.project_id}.#{t.dataset_id}.#{t.table_id}" }
+      end
+
+      def schema
+        table_refs.map do |table_ref|
+          {
+            schema: table_ref.dataset_id,
+            table: table_ref.table_id,
+            columns: table_columns(table_ref)
+          }
+        end
+      end
+
+      def preview_statement
+        "SELECT * FROM `{table}` LIMIT 10"
+      end
+
+      private
+
+      def bigquery
+        @bigquery ||= begin
+          require "google/cloud/bigquery"
+          Google::Cloud::Bigquery.new(
+            project: settings["project"],
+            keyfile: settings["keyfile"]
+          )
+        end
+      end
+
+      def table_refs
+        bigquery.datasets.map(&:tables).flat_map { |table_list| table_list.map(&:table_ref) }
+      end
+
+      def table_columns(table_ref)
+        schema = bigquery.service.get_table(table_ref.dataset_id, table_ref.table_id).schema
+        return [] if schema.nil?
+        schema.fields.map { |field| {name: field.name, data_type: field.type} }
+      end
+    end
+  end
+end

data/lib/blazer/adapters/drill_adapter.rb

@@ -0,0 +1,28 @@
+module Blazer
+  module Adapters
+    class DrillAdapter < BaseAdapter
+      def run_statement(statement, comment)
+        columns = []
+        rows = []
+        error = nil
+
+        begin
+          # remove trailing semicolon
+          response = drill.query(statement.sub(/;\s*\z/, ""))
+          rows = response.map { |r| r.values }
+          columns = rows.any? ? response.first.keys : []
+        rescue => e
+          error = e.message
+        end
+
+        [columns, rows, error]
+      end
+
+      private
+
+      def drill
+        @drill ||= ::Drill.new(url: settings["url"])
+      end
+    end
+  end
+end

data/lib/blazer/adapters/elasticsearch_adapter.rb

@@ -0,0 +1,49 @@
+module Blazer
+  module Adapters
+    class ElasticsearchAdapter < BaseAdapter
+      def run_statement(statement, comment)
+        columns = []
+        rows = []
+        error = nil
+
+        begin
+          header, body = statement.gsub(/\/\/.+/, "").strip.split("\n", 2)
+          body = JSON.parse(body)
+          body["timeout"] ||= data_source.timeout if data_source.timeout
+          response = client.msearch(body: [JSON.parse(header), body])["responses"].first
+          if response["error"]
+            error = response["error"]
+          else
+            hits = response["hits"]["hits"]
+            source_keys = hits.flat_map { |r| r["_source"].keys }.uniq
+            hit_keys = (hits.first.try(:keys) || []) - ["_source"]
+            columns = source_keys + hit_keys
+            rows =
+              hits.map do |r|
+                source = r["_source"]
+                source_keys.map { |k| source[k] } + hit_keys.map { |k| r[k] }
+              end
+          end
+        rescue => e
+          error = e.message
+        end
+
+        [columns, rows, error]
+      end
+
+      def tables
+        client.indices.get_aliases(name: "*").map { |k, v| [k, v["aliases"].keys] }.flatten.uniq.sort
+      end
+
+      def preview_statement
+        %!// header\n{"index": "{table}"}\n\n// body\n{"query": {"match_all": {}}, "size": 10}!
+      end
+
+      protected
+
+      def client
+        @client ||= Elasticsearch::Client.new(url: settings["url"])
+      end
+    end
+  end
+end

data/lib/blazer/adapters/mongodb_adapter.rb

@@ -0,0 +1,39 @@
+module Blazer
+  module Adapters
+    class MongodbAdapter < BaseAdapter
+      def run_statement(statement, comment)
+        columns = []
+        rows = []
+        error = nil
+
+        begin
+          documents = db.command({:$eval => "#{statement.strip}.toArray()"}).documents.first["retval"]
+          columns = documents.flat_map { |r| r.keys }.uniq
+          rows = documents.map { |r| columns.map { |c| r[c] } }
+        rescue => e
+          error = e.message
+        end
+
+        [columns, rows, error]
+      end
+
+      def tables
+        db.collection_names
+      end
+
+      def preview_statement
+        "db.{table}.find().limit(10)"
+      end
+
+      protected
+
+      def client
+        @client ||= Mongo::Client.new(settings["url"], connect_timeout: 1, socket_timeout: 1, server_selection_timeout: 1)
+      end
+
+      def db
+        @db ||= client.database
+      end
+    end
+  end
+end

data/lib/blazer/adapters/presto_adapter.rb

@@ -0,0 +1,45 @@
+module Blazer
+  module Adapters
+    class PrestoAdapter < BaseAdapter
+      def run_statement(statement, comment)
+        columns = []
+        rows = []
+        error = nil
+
+        begin
+          columns, rows = client.run("#{statement} /*#{comment}*/")
+          columns = columns.map(&:name)
+        rescue => e
+          error = e.message
+        end
+
+        [columns, rows, error]
+      end
+
+      def tables
+        _, rows = client.run("SHOW TABLES")
+        rows.map(&:first)
+      end
+
+      def preview_statement
+        "SELECT * FROM {table} LIMIT 10"
+      end
+
+      protected
+
+      def client
+        @client ||= begin
+          uri = URI.parse(settings["url"])
+          query = uri.query ? CGI::parse(uri.query) : {}
+          Presto::Client.new(
+            server: "#{uri.host}:#{uri.port}",
+            catalog: uri.path.to_s.sub(/\A\//, ""),
+            schema: query["schema"] || "public",
+            user: uri.user,
+            http_debug: false
+          )
+        end
+      end
+    end
+  end
+end

data/lib/blazer/adapters/sql_adapter.rb

@@ -0,0 +1,182 @@
+module Blazer
+  module Adapters
+    class SqlAdapter < BaseAdapter
+      attr_reader :connection_model
+
+      def initialize(data_source)
+        super
+
+        @connection_model =
+          Class.new(Blazer::Connection) do
+            def self.name
+              "Blazer::Connection::Adapter#{object_id}"
+            end
+            establish_connection(data_source.settings["url"]) if data_source.settings["url"]
+          end
+      end
+
+      def run_statement(statement, comment)
+        columns = []
+        rows = []
+        error = nil
+
+        begin
+          in_transaction do
+            set_timeout(data_source.timeout) if data_source.timeout
+
+            result = select_all("#{statement} /*#{comment}*/")
+            columns = result.columns
+            cast_method = Rails::VERSION::MAJOR < 5 ? :type_cast : :cast_value
+            result.rows.each do |untyped_row|
+              rows << (result.column_types.empty? ? untyped_row : columns.each_with_index.map { |c, i| untyped_row[i] ? result.column_types[c].send(cast_method, untyped_row[i]) : untyped_row[i] })
+            end
+          end
+        rescue => e
+          error = e.message.sub(/.+ERROR: /, "")
+          error = Blazer::TIMEOUT_MESSAGE if Blazer::TIMEOUT_ERRORS.any? { |e| error.include?(e) }
+          reconnect if error.include?("PG::ConnectionBad")
+        end
+
+        [columns, rows, error]
+      end
+
+      def tables
+        result = data_source.run_statement(connection_model.send(:sanitize_sql_array, ["SELECT table_name FROM information_schema.tables WHERE table_schema IN (?) ORDER BY table_name", schemas]), refresh_cache: true)
+        result.rows.map(&:first)
+      end
+
+      def schema
+        result = data_source.run_statement(connection_model.send(:sanitize_sql_array, ["SELECT table_schema, table_name, column_name, data_type, ordinal_position FROM information_schema.columns WHERE table_schema IN (?) ORDER BY 1, 2", schemas]))
+        result.rows.group_by { |r| [r[0], r[1]] }.map { |k, vs| {schema: k[0], table: k[1], columns: vs.sort_by { |v| v[2] }.map { |v| {name: v[2], data_type: v[3]} }} }
+      end
+
+      def preview_statement
+        if postgresql?
+          "SELECT * FROM \"{table}\" LIMIT 10"
+        elsif sqlserver?
+          "SELECT TOP (10) * FROM {table}"
+        else
+          "SELECT * FROM {table} LIMIT 10"
+        end
+      end
+
+      def reconnect
+        connection_model.establish_connection(settings["url"])
+      end
+
+      def cost(statement)
+        result = explain(statement)
+        if sqlserver?
+          result["TotalSubtreeCost"]
+        else
+          match = /cost=\d+\.\d+..(\d+\.\d+) /.match(result)
+          match[1] if match
+        end
+      end
+
+      def explain(statement)
+        if postgresql? || redshift?
+          select_all("EXPLAIN #{statement}").rows.first.first
+        elsif sqlserver?
+          begin
+            execute("SET SHOWPLAN_ALL ON")
+            result = select_all(statement).each.first
+          ensure
+            execute("SET SHOWPLAN_ALL OFF")
+          end
+          result
+        end
+      rescue
+        nil
+      end
+
+      def cancel(run_id)
+        if postgresql?
+          select_all("SELECT pg_cancel_backend(pid) FROM pg_stat_activity WHERE pid <> pg_backend_pid() AND query LIKE '%,run_id:#{run_id}%'")
+        elsif redshift?
+          first_row = select_all("SELECT pid FROM stv_recents WHERE status = 'Running' AND query LIKE '%,run_id:#{run_id}%'").first
+          if first_row
+            select_all("CANCEL #{first_row["pid"].to_i}")
+          end
+        end
+      end
+
+      def cachable?(statement)
+        !%w[CREATE ALTER UPDATE INSERT DELETE].include?(statement.split.first.to_s.upcase)
+      end
+
+      protected
+
+      def select_all(statement)
+        connection_model.connection.select_all(statement)
+      end
+
+      # seperate from select_all to prevent mysql error
+      def execute(statement)
+        connection_model.connection.execute(statement)
+      end
+
+      def postgresql?
+        ["PostgreSQL", "PostGIS"].include?(adapter_name)
+      end
+
+      def redshift?
+        ["Redshift"].include?(adapter_name)
+      end
+
+      def mysql?
+        ["MySQL", "Mysql2", "Mysql2Spatial"].include?(adapter_name)
+      end
+
+      def sqlserver?
+        ["SQLServer", "tinytds", "mssql"].include?(adapter_name)
+      end
+
+      def adapter_name
+        # prevent bad data source from taking down queries/new
+        connection_model.connection.adapter_name rescue nil
+      end
+
+      def schemas
+        settings["schemas"] || [connection_model.connection_config[:schema] || default_schema]
+      end
+
+      def default_schema
+        if postgresql? || redshift?
+          "public"
+        elsif sqlserver?
+          "dbo"
+        else
+          connection_model.connection_config[:database]
+        end
+      end
+
+      def set_timeout(timeout)
+        if postgresql? || redshift?
+          execute("SET #{use_transaction? ? "LOCAL " : ""}statement_timeout = #{timeout.to_i * 1000}")
+        elsif mysql?
+          execute("SET max_execution_time = #{timeout.to_i * 1000}")
+        else
+          raise Blazer::TimeoutNotSupported, "Timeout not supported for #{adapter_name} adapter"
+        end
+      end
+
+      def use_transaction?
+        settings.key?("use_transaction") ? settings["use_transaction"] : true
+      end
+
+      def in_transaction
+        connection_model.connection_pool.with_connection do
+          if use_transaction?
+            connection_model.transaction do
+              yield
+              raise ActiveRecord::Rollback
+            end
+          else
+            yield
+          end
+        end
+      end
+    end
+  end
+end