Heshamaly/add column selector extension

README.md

@@ -2,15 +2,25 @@
 
 # blueprinter-activerecord
 
-*blueprinter-activerecord* is a [blueprinter](https://github.com/procore-oss/blueprinter) extension to help you easily preload the associations from your Blueprints, N levels deep. It also provides logging extensions so you can measure how effective the primary extension is/will be.
+*blueprinter-activerecord* is a [blueprinter](https://github.com/procore-oss/blueprinter) extension to help you easily preload associations from your Blueprints (N levels deep) and optimize your queries by selecting only the necessary columns. It also provides logging extensions so you can measure how effective these extensions are.
+
+The two primary extensions are:
+- **`Preloader`**: Automatically preloads ActiveRecord associations defined in your Blueprints.
+- **`ColumnSelector`**: Intelligently selects only the database columns required by your Blueprint's fields, reducing data transfer and memory usage.
 
 ## Installation
 
 Add `blueprinter-activerecord` to your Gemfile and enable the extension using one of the configurations below.
 
 ## Configurations
 
-### Automatic mode
+This section covers configurations for both `Preloader` and `ColumnSelector`.
+
+### Preloader Extension
+
+The `Preloader` extension automatically loads associated records to prevent N+1 query problems.
+
+#### Automatic mode (Preloader)
 
 In automatic mode, every query (`ActiveRecord::Relation`) passed to a Blueprint will automatically be preloaded during render.
 
@@ -26,7 +36,7 @@ json = WidgetBlueprint.render(widgets)
 
 If you'd prefer to use `includes` rather than `preload`, pass `use: :includes` to the initializer.
 
-### Dynamic mode
+#### Dynamic mode (Preloader)
 
 In dynamic mode, each query passed to a Blueprint is evaluated by the block. If it returns `true`, the query will be preloaded during render.
 
@@ -45,7 +55,7 @@ json = WidgetBlueprint.render(widgets)
 
 If you'd prefer to use `includes` rather than `preload`, pass `use: :includes` to the initializer.
 
-### Manual mode
+#### Manual mode (Preloader)
 
 In manual mode, nothing happens automatically; you'll need to opt individual queries into preloading.
 
@@ -73,9 +83,90 @@ If you'd prefer to use `includes` or `eager_load` rather than `preload`, pass th
   preload_blueprint(use: :includes)
 ```
 
+### ColumnSelector Extension
+
+The `ColumnSelector` extension optimizes your database queries by automatically selecting only the columns that are actually used by your Blueprint for a given view. This can significantly reduce data transfer from the database and memory footprint in your Ruby application.
+
+**Benefits:**
+- **Reduced Data Transfer**: Fetches only necessary data, speeding up database interaction.
+- **Lower Memory Usage**: ActiveRecord instantiates objects with fewer attributes.
+- **Improved Performance**: Can lead to faster query execution and serialization.
+
+#### Automatic mode (ColumnSelector)
+
+In automatic mode, `ColumnSelector` will optimize column selection for every `ActiveRecord::Relation` passed to a Blueprint during render.
+
+```ruby
+Blueprinter.configure do |config|
+  config.extensions << BlueprinterActiveRecord::ColumnSelector.new(auto: true)
+end
+
+# Column selection will always happen during render
+widgets = Widget.where(...).order(...)
+# SQL will be like: SELECT "widgets"."id", "widgets"."name", ... FROM "widgets"
+json = WidgetBlueprint.render(widgets, view: :compact)
+```
+
+#### Dynamic mode (ColumnSelector)
+
+In dynamic mode, `ColumnSelector` evaluates each query using the provided block. If the block returns `true`, column selection will be applied.
+
+```ruby
+Blueprinter.configure do |config|
+  config.extensions << BlueprinterActiveRecord::ColumnSelector.new do |q, blueprint, view, options|
+    # examine q, q.model, blueprint, view, or options and return true or false
+  end
+end
+
+# Column selection will happen if the above block returns true
+widgets = Widget.where(...).order(...)
+json = WidgetBlueprint.render(widgets)
+```
+
+#### Manual mode (ColumnSelector)
+
+In manual mode, `ColumnSelector` does nothing automatically. You need to explicitly opt-in queries for column selection using the `select_blueprint_columns` method on an `ActiveRecord::Relation`.
+
+```ruby
+Blueprinter.configure do |config|
+  config.extensions << BlueprinterActiveRecord::ColumnSelector.new
+end
+```
+
+Opt-in specific queries:
+```ruby
+q = Widget.
+  where(...).
+  order(...).
+  select_blueprint_columns(WidgetBlueprint, :compact) # Specify Blueprint and view
+
+# Column selection will be applied before rendering
+json = WidgetBlueprint.render(q, view: :compact)
+```
+
+### Combining Preloader and ColumnSelector
+
+`Preloader` and `ColumnSelector` are designed to work together seamlessly. For optimal performance and correctness, it's generally recommended to **configure `Preloader` before `ColumnSelector`**.
+
+```ruby
+Blueprinter.configure do |config|
+  # 1. Preloader first: to identify associations and their foreign keys
+  config.extensions << BlueprinterActiveRecord::Preloader.new(auto: true)
+  # 2. ColumnSelector second: to select only needed columns, including foreign keys for preloads
+  config.extensions << BlueprinterActiveRecord::ColumnSelector.new(auto: true)
+end
+```
+This order ensures that:
+1. `Preloader` determines which associations need to be loaded.
+2. `ColumnSelector` then analyzes the blueprint to select the minimum set of columns.
+
 ## Annotations
 
-Sometimes a field in your blueprint is a method or block. This extension can't "see" into methods or blocks, meaning it can't preload any associations inside. In these cases, annotate your blueprint so the extension knows what to preload.
+Annotations help both extensions understand fields or associations that aren't straightforward.
+
+### Preloader Annotations (`preload: ...`)
+
+Sometimes a field in your blueprint is a method or block. `Preloader` can't "see" into methods or blocks, meaning it can't preload any associations inside. In these cases, annotate your blueprint so the extension knows what to preload.
 
 ```ruby
 # Here is a model with some instance methods
@@ -111,6 +202,32 @@ class WidgetBlueprint < Blueprinter::Base
 end
 ```
 
+### ColumnSelector Annotations (`select: ...`)
+
+Similarly, if a field in your blueprint accesses specific database columns that are not directly mapped as attributes or standard associations (e.g., through a method on the model that reads specific columns), you can guide `ColumnSelector` using the `select:` annotation.
+
+This is **distinct** from `Preloader`'s `preload:` annotation.
+- `preload: :association_name` tells `Preloader` to load an association.
+- `select: [:column1, :column2]` tells `ColumnSelector` to ensure these specific columns are included in the `SELECT` statement for the primary model.
+
+```ruby
+class WidgetBlueprint < Blueprinter::Base
+  identifier :id
+  field :name
+  field :description
+
+  # Assume 'Widget' model has 'custom_field_a' and 'custom_field_b' in its table,
+  # which are used by the 'computed_details' method but not defined as separate fields.
+  field :computed_widget_details, select: [:custom_field_a, :custom_field_b] do |widget|
+    widget.compute_details # method on Widget model that uses self.custom_field_a
+  end
+end
+
+# When rendering WidgetBlueprint, ColumnSelector will ensure 'custom_field_a', 'custom_field_b'
+# are included in the SELECT statement for the main query, in addition to 'id', 'name', etc.
+```
+**Note**: The `select:` annotation is for columns on the **primary model** (or its direct table) being blueprinted. Future enhancements may allow more granular control over association column selection.
+
 ## Recursive Blueprints
 
 Sometimes a model, and its blueprint, will have recursive associations. Think of a nested Category model:
@@ -137,7 +254,7 @@ association :children, blueprint: CategoryBlueprint, max_recursion: 20
 
 ### Pass the *query* to render, not query *results*
 
-If the query runs _before_ being passed to `render`, it's too late for preloading to happen.
+This is critical for both `Preloader` and `ColumnSelector`. If the query runs _before_ being passed to `render`, it's too late for either extension to modify it.
 
 ```ruby
 widgets = Widget.where(...)
@@ -175,9 +292,9 @@ WidgetBlueprint.render(widgets)
 
 ## Logging
 
-There are two different logging extensions. You can use them together or separately to measure how much the Preloder extension is, or can, help your application.
+There are two different logging extensions available for `Preloader` You can use them together or separately to measure how much the Preloder extension is, or can, help your application. More specific logging for `ColumnSelector` may be added in the future.
 
-### Missing Preloads Logger
+### Missing Preloads Logger (For Preloader)
 
 This extension is useful for measuring how helpful `BlueprinterActiveRecord::Preloader` will be for your application. It can be used with or without `Preloader`. Any Blueprint-rendered queries *not* caught by the `Preloader` extension will be caught by this logger.
 
@@ -202,7 +319,7 @@ Blueprinter.configure do |config|
 end
 ```
 
-### Added Preloads Logger
+### Added Preloads Logger (For Preloader)
 
 This extension measures how many missing preloads are being found & fixed by the preloader. Any query caught by this extension *won't* end up in `MissingPreloadsLogger`.
 

lib/blueprinter-activerecord.rb

@@ -6,6 +6,7 @@
 module BlueprinterActiveRecord
   autoload :QueryMethods, 'blueprinter-activerecord/query_methods'
   autoload :Preloader, 'blueprinter-activerecord/preloader'
+  autoload :ColumnSelector, 'blueprinter-activerecord/column_selector'
   autoload :AddedPreloadsLogger, 'blueprinter-activerecord/added_preloads_logger'
   autoload :MissingPreloadsLogger, 'blueprinter-activerecord/missing_preloads_logger'
   autoload :PreloadInfo, 'blueprinter-activerecord/preload_info'

lib/blueprinter-activerecord/column_selector.rb

@@ -0,0 +1,117 @@
+# frozen_string_literal: true
+
+module BlueprinterActiveRecord
+  # A Blueprinter extension to automatically select only the columns needed by a Blueprint view
+  class ColumnSelector < Blueprinter::Extension
+    attr_reader :auto, :auto_proc
+
+    #
+    # Initialize and configure the extension.
+    #
+    # @param auto [true|false] When true, select columns for EVERY ActiveRecord::Relation passed to a Blueprint
+    # @yield [Object, Class, Symbol, Hash] Instead of passing `auto` as a boolean, you may define a block that accepts the object to render, the blueprint class, the view, and options. If the block returns true, auto column selection will take place.
+    #
+    def initialize(auto: false, &auto_proc)
+      @auto = auto
+      @auto_proc = auto_proc
+    end
+
+    #
+    # Implements the "pre_render" Blueprinter Extension to select only needed columns from a view.
+    # If auto is true, all ActiveRecord::Relation objects
+    # will have their columns optimized. If auto is false, only queries that have called `.select_blueprint_columns`
+    # will be optimized.
+    #
+    def pre_render(object, blueprint, view, options)
+      if object.is_a?(ActiveRecord::Relation) && !object.loaded?
+        if object.select_blueprint_columns_method || auto || auto_proc&.call(object, blueprint, view, options) == true
+          # Store original select values for logging
+          object.before_select_blueprint_columns = extract_selects(object)
+
+          # Get the columns needed by this blueprint view
+          blueprint_columns = self.class.columns(blueprint, view, model: object.model)
+
+          # Apply the column selection
+          object.select(*blueprint_columns) if blueprint_columns.any?
+        else
+          object
+        end
+      else
+        object
+      end
+    end
+
+    private
+
+    #
+    # Returns an array of column names needed by the Blueprint and view.
+    #
+    # Columns are found from:
+    # 1. Blueprint field names that match model column names
+    # 2. Foreign key columns for associations (always included when associations are present)
+    # 3. Explicit :select options on fields
+    #
+    # Example:
+    #
+    #   columns = BlueprinterActiveRecord::ColumnSelector.columns(WidgetBlueprint, :extended, model: Widget)
+    #   q = Widget.where(...).select(*columns)
+    #
+    # @param blueprint [Class] The Blueprint class
+    # @param view_name [Symbol] Name of the view in blueprint
+    # @param model [Class|:polymorphic] The ActiveRecord model class that blueprint represents
+    # @return [Array<String>] Array of column names to select
+    #
+    def self.columns(blueprint, view_name, model:)
+      view = blueprint.reflections.fetch(view_name)
+      columns = Set.new
+
+      # Process fields
+      view.fields.each do |field_name, field|
+        columns.add(field.name.to_s) if model.column_names.include?(field.name.to_s)
+
+        # Check if field has explicit select option
+        if (custom_select = field.options[:select])
+          case custom_select
+          when String, Symbol
+            raise ArgumentError, "Blueprint field #{field_name} has explicit select option #{custom_select} that does not match a column on the model #{model.name}" unless model.column_names.include?(custom_select.to_s)
+            columns.add(custom_select.to_s)
+          when Array
+            columns.merge(custom_select.map(&:to_s))
+          end
+        end
+      end
+
+      # Process associations to get foreign keys
+      view.associations.each do |assoc_name, assoc|
+        if model.respond_to?(:reflections) && (ref = model.reflections[assoc_name.to_s])
+          # Add foreign key columns for belongs_to associations
+          if ref.belongs_to?
+            columns.add(ref.foreign_key) if ref.foreign_key
+            columns.add(ref.foreign_type) if ref.foreign_type # in case it was polymorphic
+          elsif ref.through_reflection
+            # For has_many/one through belongs_to associations, we need the foreign key of the intermediate association
+            # Example: has_many :customer_projects, through: :customer
+            # We need the customer_id foreign key
+            through_ref = ref.through_reflection
+            if through_ref.belongs_to?
+              columns.add(through_ref.foreign_key) if through_ref.foreign_key
+              columns.add(through_ref.foreign_type) if through_ref.foreign_type # in case it was polymorphic
+            end
+          end
+        end
+      end
+      # Databases often cache query plans. Consistent column order means:
+      # - Same SQL string → same cached plan
+      # - Better database performance
+      # - Less chance of unexpected changes in query results
+      columns.to_a.flatten.sort
+    end
+
+    #
+    # Extract existing select values from a query
+    #
+    def extract_selects(query)
+      query.values[:select] || []
+    end
+  end
+end