Clean up controller

app/controllers/blazer/queries_controller.rb

@@ -74,87 +74,10 @@ module Blazer
 
         data_source = @query.data_source if @query && @query.data_source
         @data_source = Blazer.data_sources[data_source]
-        Blazer.transform_statement.call(@data_source, @statement) if Blazer.transform_statement
-
-        # audit
-        if Blazer.audit
-          audit = Blazer::Audit.new(statement: @statement)
-          audit.query = @query
-          audit.data_source = data_source
-          audit.user = blazer_user
-          audit.save!
-        end
-
-        start_time = Time.now
-        @columns, @rows, @error, @cached_at, @just_cached = @data_source.run_statement(@statement, user: blazer_user, query: @query, refresh_cache: params[:check], with_just_cached: true)
-        duration = Time.now - start_time
-
-        if Blazer.audit
-          audit.duration = duration if audit.respond_to?(:duration=)
-          audit.error = @error if audit.respond_to?(:error=)
-          audit.timed_out = @error == Blazer::TIMEOUT_MESSAGE if audit.respond_to?(:timed_out=)
-          audit.cached = @cached_at.present? if audit.respond_to?(:cached=)
-          if !@cached_at && duration >= 10
-            audit.cost = @data_source.cost(@statement) if audit.respond_to?(:cost=)
-          end
-          audit.save! if audit.changed?
-        end
 
-        if @query && @error != Blazer::TIMEOUT_MESSAGE
-          @query.checks.each do |check|
-            check.update_state(@rows, @error)
-          end
-        end
-
-        @first_row = @rows.first || []
-        @column_types = []
-        if @rows.any?
-          @columns.each_with_index do |column, i|
-            @column_types << (
-              case @first_row[i]
-              when Integer
-                "int"
-              when Float
-                "float"
-              else
-                "string-ins"
-              end
-            )
-          end
-        end
+        @columns, @rows, @error, @cached_at, @just_cached = @data_source.run_main_statement(@statement, user: blazer_user, query: @query, refresh_cache: params[:check])
 
-        @filename = @query.name.parameterize if @query
-        @min_width_types = @columns.each_with_index.select { |c, i| @first_row[i].is_a?(Time) || @first_row[i].is_a?(String) || @data_source.smart_columns[c] }
-
-        @boom = {}
-        @columns.each_with_index do |key, i|
-          query = @data_source.smart_columns[key]
-          if query
-            values = @rows.map { |r| r[i] }.compact.uniq
-            columns, rows, error, cached_at = @data_source.run_statement(ActiveRecord::Base.send(:sanitize_sql_array, [query.sub("{value}", "(?)"), values]))
-            @boom[key] = Hash[rows.map { |k, v| [k.to_s, v] }]
-          end
-        end
-
-        @linked_columns = @data_source.linked_columns
-
-        @markers = []
-        [["latitude", "longitude"], ["lat", "lon"]].each do |keys|
-          lat_index = @columns.index(keys.first)
-          lon_index = @columns.index(keys.last)
-          if lat_index && lon_index
-            @markers =
-              @rows.select do |r|
-                r[lat_index] && r[lon_index]
-              end.map do |r|
-                {
-                  title: r.each_with_index.map{ |v, i| i == lat_index || i == lon_index ? nil : "<strong>#{@columns[i]}:</strong> #{v}" }.compact.join("<br />").truncate(140),
-                  latitude: r[lat_index],
-                  longitude: r[lon_index]
-                }
-              end
-          end
-        end
+        render_run
       end
 
       respond_to do |format|
@@ -203,6 +126,58 @@ module Blazer
 
     private
 
+    def render_run
+      @first_row = @rows.first || []
+      @column_types = []
+      if @rows.any?
+        @columns.each_with_index do |column, i|
+          @column_types << (
+            case @first_row[i]
+            when Integer
+              "int"
+            when Float
+              "float"
+            else
+              "string-ins"
+            end
+          )
+        end
+      end
+
+      @filename = @query.name.parameterize if @query
+      @min_width_types = @columns.each_with_index.select { |c, i| @first_row[i].is_a?(Time) || @first_row[i].is_a?(String) || @data_source.smart_columns[c] }
+
+      @boom = {}
+      @columns.each_with_index do |key, i|
+        query = @data_source.smart_columns[key]
+        if query
+          values = @rows.map { |r| r[i] }.compact.uniq
+          columns, rows, error, cached_at = @data_source.run_statement(ActiveRecord::Base.send(:sanitize_sql_array, [query.sub("{value}", "(?)"), values]))
+          @boom[key] = Hash[rows.map { |k, v| [k.to_s, v] }]
+        end
+      end
+
+      @linked_columns = @data_source.linked_columns
+
+      @markers = []
+      [["latitude", "longitude"], ["lat", "lon"]].each do |keys|
+        lat_index = @columns.index(keys.first)
+        lon_index = @columns.index(keys.last)
+        if lat_index && lon_index
+          @markers =
+            @rows.select do |r|
+              r[lat_index] && r[lon_index]
+            end.map do |r|
+              {
+                title: r.each_with_index.map{ |v, i| i == lat_index || i == lon_index ? nil : "<strong>#{@columns[i]}:</strong> #{v}" }.compact.join("<br />").truncate(140),
+                latitude: r[lat_index],
+                longitude: r[lon_index]
+              }
+            end
+        end
+      end
+    end
+
     def set_queries(limit = nil)
       @my_queries =
         if limit && blazer_user

lib/blazer/data_source.rb

@@ -94,6 +94,43 @@ module Blazer
       end
     end
 
+    def run_main_statement(statement, options = {})
+      query = options[:query]
+      Blazer.transform_statement.call(self, statement) if Blazer.transform_statement
+
+      # audit
+      if Blazer.audit
+        audit = Blazer::Audit.new(statement: statement)
+        audit.query = query
+        audit.data_source = id
+        audit.user = options[:user]
+        audit.save!
+      end
+
+      start_time = Time.now
+      columns, rows, error, cached_at, just_cached = run_statement(statement, options.merge(with_just_cached: true))
+      duration = Time.now - start_time
+
+      if Blazer.audit
+        audit.duration = duration if audit.respond_to?(:duration=)
+        audit.error = error if audit.respond_to?(:error=)
+        audit.timed_out = error == Blazer::TIMEOUT_MESSAGE if audit.respond_to?(:timed_out=)
+        audit.cached = cached_at.present? if audit.respond_to?(:cached=)
+        if !cached_at && duration >= 10
+          audit.cost = cost(statement) if audit.respond_to?(:cost=)
+        end
+        audit.save! if audit.changed?
+      end
+
+      if query && error != Blazer::TIMEOUT_MESSAGE
+        query.checks.each do |check|
+          check.update_state(rows, error)
+        end
+      end
+
+      [columns, rows, error, cached_at, just_cached]
+    end
+
     def run_statement(statement, options = {})
       columns = nil
       rows = nil