query_console 0.2.0 → 0.2.6

data/app/services/query_console/explain_runner.rb

@@ -96,6 +96,49 @@ module QueryConsole
     end
 
     def execute_with_timeout(sql)
+      case @config.timeout_strategy
+      when :database
+        execute_with_database_timeout(sql)
+      when :ruby
+        execute_with_ruby_timeout(sql)
+      else
+        # Auto-detect: use database timeout for PostgreSQL, Ruby timeout otherwise
+        if postgresql_connection?
+          execute_with_database_timeout(sql)
+        else
+          execute_with_ruby_timeout(sql)
+        end
+      end
+    end
+
+    # Database-level timeout (PostgreSQL only)
+    # Safer: database cancels the query cleanly, no orphan processes
+    def execute_with_database_timeout(sql)
+      conn = ActiveRecord::Base.connection
+      
+      unless postgresql_connection?
+        Rails.logger.warn("[QueryConsole] Database timeout strategy requires PostgreSQL, falling back to Ruby timeout")
+        return execute_with_ruby_timeout(sql)
+      end
+
+      conn.transaction do
+        # SET LOCAL scopes the timeout to this transaction only
+        conn.execute("SET LOCAL statement_timeout = '#{@config.timeout_ms}'")
+        conn.exec_query(sql)
+      end
+    rescue ActiveRecord::StatementInvalid => e
+      if e.message.include?("canceling statement due to statement timeout") ||
+         e.message.include?("query_canceled")
+        raise Timeout::Error, "Query timeout: exceeded #{@config.timeout_ms}ms limit"
+      else
+        raise
+      end
+    end
+
+    # Ruby-level timeout (fallback for non-PostgreSQL databases)
+    # Warning: The database query continues running as an orphan process
+    # even after Ruby times out. Can cause resource exhaustion.
+    def execute_with_ruby_timeout(sql)
       timeout_seconds = @config.timeout_ms / 1000.0
       
       Timeout.timeout(timeout_seconds) do
@@ -103,6 +146,10 @@ module QueryConsole
       end
     end
 
+    def postgresql_connection?
+      ActiveRecord::Base.connection.adapter_name.downcase.include?('postgresql')
+    end
+
     def format_explain_output(result)
       # For SQLite, the result has columns like: id, parent, notused, detail
       # For Postgres, it's usually a single "QUERY PLAN" column

data/app/services/query_console/runner.rb

@@ -3,15 +3,17 @@ require 'timeout'
 module QueryConsole
   class Runner
     class QueryResult
-      attr_reader :columns, :rows, :execution_time_ms, :row_count_shown, :truncated, :error
+      attr_reader :columns, :rows, :execution_time_ms, :row_count_shown, :truncated, :error, :is_dml, :rows_affected
 
-      def initialize(columns: [], rows: [], execution_time_ms: 0, row_count_shown: 0, truncated: false, error: nil)
+      def initialize(columns: [], rows: [], execution_time_ms: 0, row_count_shown: 0, truncated: false, error: nil, is_dml: false, rows_affected: nil)
         @columns = columns
         @rows = rows
         @execution_time_ms = execution_time_ms
         @row_count_shown = row_count_shown
         @truncated = truncated
         @error = error
+        @is_dml = is_dml
+        @rows_affected = rows_affected
       end
 
       def success?
@@ -25,6 +27,10 @@ module QueryConsole
       def truncated?
         @truncated
       end
+
+      def dml?
+        @is_dml
+      end
     end
 
     def initialize(sql, config = QueryConsole.configuration)
@@ -44,6 +50,7 @@ module QueryConsole
       end
 
       sanitized_sql = validation_result.sanitized_sql
+      is_dml = validation_result.dml?
 
       # Step 2: Apply row limit
       limiter = SqlLimiter.new(sanitized_sql, @config.max_rows, @config)
@@ -56,12 +63,20 @@ module QueryConsole
         result = execute_with_timeout(final_sql)
         execution_time = ((Time.now - start_time) * 1000).round(2)
 
+        # For DML queries, capture the number of affected rows
+        rows_affected = nil
+        if is_dml
+          rows_affected = get_affected_rows_count(result)
+        end
+
         QueryResult.new(
           columns: result.columns,
           rows: result.rows,
           execution_time_ms: execution_time,
           row_count_shown: result.rows.length,
-          truncated: truncated
+          truncated: truncated,
+          is_dml: is_dml,
+          rows_affected: rows_affected
         )
       rescue Timeout::Error
         QueryResult.new(
@@ -79,11 +94,96 @@ module QueryConsole
     attr_reader :sql, :config
 
     def execute_with_timeout(sql)
+      case @config.timeout_strategy
+      when :database
+        execute_with_database_timeout(sql)
+      when :ruby
+        execute_with_ruby_timeout(sql)
+      else
+        # Auto-detect: use database timeout for PostgreSQL, Ruby timeout otherwise
+        if postgresql_connection?
+          execute_with_database_timeout(sql)
+        else
+          execute_with_ruby_timeout(sql)
+        end
+      end
+    end
+
+    # Database-level timeout (PostgreSQL only)
+    # Safer: database cancels the query cleanly, no orphan processes
+    def execute_with_database_timeout(sql)
+      conn = ActiveRecord::Base.connection
+      
+      unless postgresql_connection?
+        Rails.logger.warn("[QueryConsole] Database timeout strategy requires PostgreSQL, falling back to Ruby timeout")
+        return execute_with_ruby_timeout(sql)
+      end
+
+      conn.transaction do
+        # SET LOCAL scopes the timeout to this transaction only
+        conn.execute("SET LOCAL statement_timeout = '#{@config.timeout_ms}'")
+        conn.exec_query(sql)
+      end
+    rescue ActiveRecord::StatementInvalid => e
+      if e.message.include?("canceling statement due to statement timeout") ||
+         e.message.include?("query_canceled")
+        raise Timeout::Error, "Query timeout: exceeded #{@config.timeout_ms}ms limit"
+      else
+        raise
+      end
+    end
+
+    # Ruby-level timeout (fallback for non-PostgreSQL databases)
+    # Warning: The database query continues running as an orphan process
+    # even after Ruby times out. Can cause resource exhaustion.
+    def execute_with_ruby_timeout(sql)
       timeout_seconds = @config.timeout_ms / 1000.0
       
       Timeout.timeout(timeout_seconds) do
         ActiveRecord::Base.connection.exec_query(sql)
       end
     end
+
+    def postgresql_connection?
+      ActiveRecord::Base.connection.adapter_name.downcase.include?('postgresql')
+    end
+
+    # Get the number of rows affected by a DML query
+    # This is database-specific, so we try different approaches
+    def get_affected_rows_count(result)
+      conn = ActiveRecord::Base.connection
+      
+      # For SQLite, use the raw connection's changes method
+      if conn.adapter_name.downcase.include?('sqlite')
+        return conn.raw_connection.changes
+      end
+      
+      # For PostgreSQL, MySQL, and others, check if result has rows_affected
+      # Note: exec_query doesn't always provide this, but we can try
+      if result.respond_to?(:rows_affected)
+        return result.rows_affected
+      end
+      
+      # Fallback: try to get it from the connection's last result
+      begin
+        if conn.respond_to?(:raw_connection)
+          raw_conn = conn.raw_connection
+          
+          # PostgreSQL
+          if raw_conn.respond_to?(:cmd_tuples)
+            return raw_conn.cmd_tuples
+          end
+          
+          # MySQL
+          if raw_conn.respond_to?(:affected_rows)
+            return raw_conn.affected_rows
+          end
+        end
+      rescue
+        # If we can't determine affected rows, return nil
+      end
+      
+      nil
+    end
   end
 end

data/app/services/query_console/sql_limiter.rb

@@ -20,6 +20,11 @@ module QueryConsole
     end
 
     def apply_limit
+      # Skip limiting for DML queries (INSERT, UPDATE, DELETE, MERGE)
+      if is_dml_query?
+        return LimitResult.new(sql: @sql, truncated: false)
+      end
+      
       # Check if query already has a LIMIT clause
       if sql_has_limit?
         LimitResult.new(sql: @sql, truncated: false)
@@ -33,6 +38,11 @@ module QueryConsole
 
     attr_reader :sql, :max_rows, :config
 
+    def is_dml_query?
+      # Check if query is a DML operation (INSERT, UPDATE, DELETE, MERGE)
+      @sql.strip.downcase.match?(/\A(insert|update|delete|merge)\b/)
+    end
+
     def sql_has_limit?
       # Check for LIMIT clause (case-insensitive)
       # Match: "LIMIT", " LIMIT ", etc.

data/app/services/query_console/sql_validator.rb

@@ -1,12 +1,13 @@
 module QueryConsole
   class SqlValidator
     class ValidationResult
-      attr_reader :valid, :error, :sanitized_sql
+      attr_reader :valid, :error, :sanitized_sql, :is_dml
 
-      def initialize(valid:, sanitized_sql: nil, error: nil)
+      def initialize(valid:, sanitized_sql: nil, error: nil, is_dml: false)
         @valid = valid
         @sanitized_sql = sanitized_sql
         @error = error
+        @is_dml = is_dml
       end
 
       def valid?
@@ -16,6 +17,10 @@ module QueryConsole
       def invalid?
         !@valid
       end
+
+      def dml?
+        @is_dml
+      end
     end
 
     def initialize(sql, config = QueryConsole.configuration)
@@ -41,16 +46,35 @@ module QueryConsole
 
       # Check if query starts with allowed keywords
       normalized_start = sanitized.downcase
-      unless @config.allowed_starts_with.any? { |keyword| normalized_start.start_with?(keyword) }
+      
+      # Define DML-specific keywords that are conditionally allowed
+      dml_keywords = %w[insert update delete merge]
+      
+      # Expand allowed_starts_with if DML is enabled
+      effective_allowed = if @config.enable_dml
+        @config.allowed_starts_with + dml_keywords
+      else
+        @config.allowed_starts_with
+      end
+      
+      unless effective_allowed.any? { |keyword| normalized_start.start_with?(keyword) }
         return ValidationResult.new(
           valid: false,
-          error: "Query must start with one of: #{@config.allowed_starts_with.join(', ').upcase}"
+          error: "Query must start with one of: #{effective_allowed.join(', ').upcase}"
         )
       end
 
       # Check for forbidden keywords
       normalized_query = sanitized.downcase
-      forbidden = @config.forbidden_keywords.find do |keyword|
+      
+      # Filter forbidden keywords based on DML enablement
+      effective_forbidden = if @config.enable_dml
+        @config.forbidden_keywords.reject { |kw| dml_keywords.include?(kw) || kw == 'replace' || kw == 'into' }
+      else
+        @config.forbidden_keywords
+      end
+      
+      forbidden = effective_forbidden.find do |keyword|
         # Match whole words to avoid false positives (e.g., "updates" table name)
         normalized_query.match?(/\b#{Regexp.escape(keyword.downcase)}\b/)
       end
@@ -62,7 +86,10 @@ module QueryConsole
         )
       end
 
-      ValidationResult.new(valid: true, sanitized_sql: sanitized)
+      # Detect if this is a DML query
+      is_dml_query = sanitized.downcase.match?(/\A(insert|update|delete|merge)\b/)
+
+      ValidationResult.new(valid: true, sanitized_sql: sanitized, is_dml: is_dml_query)
     end
 
     private

data/app/views/query_console/queries/_results.html.erb

@@ -116,10 +116,33 @@
       color: #999;
       font-style: italic;
     }
+
+    .dml-warning {
+      background-color: #fff3cd;
+      border-left: 4px solid #ffc107;
+      padding: 12px 16px;
+      margin-bottom: 16px;
+      border-radius: 4px;
+    }
+
+    .dml-warning-icon {
+      color: #ff6b6b;
+      font-weight: bold;
+      margin-right: 8px;
+    }
   </style>
 
   <div class="results-container">
     <% result_to_render = local_assigns[:result] || @result %>
+    <% is_dml_result = local_assigns[:is_dml] || @is_dml %>
+    
+    <% if is_dml_result %>
+      <div class="dml-warning">
+        <span class="dml-warning-icon">ℹ️</span>
+        <strong>Data Modified:</strong> 
+        This query has modified the database. All changes are logged.
+      </div>
+    <% end %>
     
     <% if result_to_render.error %>
       <div class="error-message">
@@ -136,8 +159,13 @@
           <span class="metadata-value"><%= result_to_render.execution_time_ms %>ms</span>
         </div>
         <div class="metadata-item">
-          <span class="metadata-label">Rows:</span>
-          <span class="metadata-value"><%= result_to_render.row_count_shown %></span>
+          <% if is_dml_result && result_to_render.rows_affected %>
+            <span class="metadata-label">Rows Affected:</span>
+            <span class="metadata-value"><%= result_to_render.rows_affected %></span>
+          <% else %>
+            <span class="metadata-label">Rows:</span>
+            <span class="metadata-value"><%= result_to_render.row_count_shown %></span>
+          <% end %>
         </div>
         <% if result_to_render.truncated %>
           <div class="metadata-item">
@@ -148,7 +176,11 @@
 
       <% if result_to_render.rows.empty? %>
         <div class="empty-results">
-          No rows returned
+          <% if is_dml_result && result_to_render.rows_affected %>
+            <%= result_to_render.rows_affected %> row(s) affected
+          <% else %>
+            No rows returned
+          <% end %>
         </div>
       <% else %>
         <div class="table-wrapper">