railsblazer 2.0.0

data/lib/blazer/adapters/bigquery_adapter.rb

@@ -0,0 +1,68 @@
+module Blazer
+  module Adapters
+    class BigQueryAdapter < BaseAdapter
+      def run_statement(statement, comment)
+        columns = []
+        rows = []
+        error = nil
+
+        begin
+          results = bigquery.query(statement)
+
+          # complete? was removed in google-cloud-bigquery 0.29.0
+          # code is for backward compatibility
+          if !results.respond_to?(:complete?) || 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/cassandra_adapter.rb

@@ -0,0 +1,59 @@
+module Blazer
+  module Adapters
+    class CassandraAdapter < BaseAdapter
+      def run_statement(statement, comment)
+        columns = []
+        rows = []
+        error = nil
+
+        begin
+          response = session.execute("#{statement} /*#{comment}*/")
+          rows = response.map { |r| r.values }
+          columns = rows.any? ? response.first.keys : []
+        rescue => e
+          error = e.message
+        end
+
+        [columns, rows, error]
+      end
+
+      def tables
+        session.execute("SELECT table_name FROM system_schema.tables WHERE keyspace_name = '#{keyspace}'").map { |r| r["table_name"] }
+      end
+
+      def schema
+        result = session.execute("SELECT keyspace_name, table_name, column_name, type, position FROM system_schema.columns WHERE keyspace_name = '#{keyspace}'")
+        result.map(&:values).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
+        "SELECT * FROM {table} LIMIT 10"
+      end
+
+      private
+
+      def cluster
+        @cluster ||= begin
+          require "cassandra"
+          options = {hosts: [uri.host]}
+          options[:port] = uri.port if uri.port
+          options[:username] = uri.user if uri.user
+          options[:password] = uri.password if uri.password
+          ::Cassandra.cluster(options)
+        end
+      end
+
+      def session
+        @session ||= cluster.connect(keyspace)
+      end
+
+      def uri
+        @uri ||= URI.parse(data_source.settings["url"])
+      end
+
+      def keyspace
+        @keyspace ||= uri.path[1..-1]
+      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/druid_adapter.rb

@@ -0,0 +1,67 @@
+module Blazer
+  module Adapters
+    class DruidAdapter < BaseAdapter
+      TIMESTAMP_REGEX = /\A\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.\d{3}Z\z/
+
+      def run_statement(statement, comment)
+        columns = []
+        rows = []
+        error = nil
+
+        header = {"Content-Type" => "application/json", "Accept" => "application/json"}
+        timeout = data_source.timeout ? data_source.timeout.to_i : 300
+        data = {
+          query: statement,
+          context: {
+            timeout: timeout * 1000
+          }
+        }
+
+        uri = URI.parse("#{settings["url"]}/druid/v2/sql/")
+        http = Net::HTTP.new(uri.host, uri.port)
+        http.read_timeout = timeout
+
+        begin
+          response = JSON.parse(http.post(uri.request_uri, data.to_json, header).body)
+          if response.is_a?(Hash)
+            error = response["errorMessage"] || "Unknown error: #{response.inspect}"
+            if error.include?("timed out")
+              error = Blazer::TIMEOUT_MESSAGE
+            end
+          else
+            columns = (response.first || {}).keys
+            rows = response.map { |r| r.values }
+
+            # Druid doesn't return column types
+            # and no timestamp type in JSON
+            rows.each do |row|
+              row.each_with_index do |v, i|
+                if v.is_a?(String) && TIMESTAMP_REGEX.match(v)
+                  row[i] = Time.parse(v)
+                end
+              end
+            end
+          end
+         rescue => e
+           error = e.message
+         end
+
+        [columns, rows, error]
+      end
+
+      def tables
+        result = data_source.run_statement("SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA NOT IN ('INFORMATION_SCHEMA') ORDER BY TABLE_NAME")
+        result.rows.map(&:first)
+      end
+
+      def schema
+        result = data_source.run_statement("SELECT TABLE_SCHEMA, TABLE_NAME, COLUMN_NAME, DATA_TYPE, ORDINAL_POSITION FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_SCHEMA NOT IN ('INFORMATION_SCHEMA') ORDER BY 1, 2")
+        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
+        "SELECT * FROM {table} LIMIT 10"
+      end
+    end
+  end
+end

data/lib/blazer/adapters/elasticsearch_adapter.rb

@@ -0,0 +1,46 @@
+module Blazer
+  module Adapters
+    class ElasticsearchAdapter < BaseAdapter
+      def run_statement(statement, comment)
+        columns = []
+        rows = []
+        error = nil
+
+        begin
+          response = client.xpack.sql.query(body: {query: "#{statement} /*#{comment}*/"})
+          columns = response["columns"].map { |v| v["name"] }
+          # Elasticsearch does not differentiate between dates and times
+          date_indexes = response["columns"].each_index.select { |i| response["columns"][i]["type"] == "date" }
+          if columns.any?
+            rows = response["rows"]
+            date_indexes.each do |i|
+              rows.each do |row|
+                row[i] = Time.parse(row[i])
+              end
+            end
+          end
+        rescue => e
+          error = e.message
+        end
+
+        [columns, rows, error]
+      end
+
+      def tables
+        indices = client.cat.indices(format: "json").map { |v| v["index"] }
+        aliases = client.cat.aliases(format: "json").map { |v| v["alias"] }
+        (indices + aliases).uniq.sort
+      end
+
+      def preview_statement
+        "SELECT * FROM \"{table}\" LIMIT 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()", nolock: true}).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/snowflake_adapter.rb

@@ -0,0 +1,73 @@
+module Blazer
+  module Adapters
+    class SnowflakeAdapter < SqlAdapter
+      def initialize(data_source)
+        @data_source = data_source
+
+        @@registered ||= begin
+          require "active_record/connection_adapters/odbc_adapter"
+          require "odbc_adapter/adapters/postgresql_odbc_adapter"
+
+          ODBCAdapter.register(/snowflake/, ODBCAdapter::Adapters::PostgreSQLODBCAdapter) do
+            # Explicitly turning off prepared statements as they are not yet working with
+            # snowflake + the ODBC ActiveRecord adapter
+            def prepared_statements
+              false
+            end
+
+            # Quoting needs to be changed for snowflake
+            def quote_column_name(name)
+              name.to_s
+            end
+
+            private
+
+            # Override dbms_type_cast to get the values encoded in UTF-8
+            def dbms_type_cast(columns, values)
+              int_column = {}
+              columns.each_with_index do |c, i|
+                int_column[i] = c.type == 3 && c.scale == 0
+              end
+
+              float_column = {}
+              columns.each_with_index do |c, i|
+                float_column[i] = c.type == 3 && c.scale != 0
+              end
+
+              values.each do |row|
+                row.each_index do |idx|
+                  val = row[idx]
+                  if val
+                    if int_column[idx]
+                      row[idx] = val.to_i
+                    elsif float_column[idx]
+                      row[idx] = val.to_f
+                    elsif val.is_a?(String)
+                      row[idx] = val.force_encoding('UTF-8')
+                    end
+                  end
+                end
+              end
+            end
+          end
+        end
+
+        @connection_model =
+          Class.new(Blazer::Connection) do
+            def self.name
+              "Blazer::Connection::SnowflakeAdapter#{object_id}"
+            end
+            if data_source.settings["conn_str"]
+              establish_connection(adapter: "odbc", conn_str: data_source.settings["conn_str"])
+            elsif data_source.settings["dsn"]
+              establish_connection(adapter: "odbc", dsn: data_source.settings["dsn"])
+            end
+        end
+      end
+
+      def cancel(run_id)
+        # todo
+      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