scrooge2.rb

require 'set'

module ActiveRecord
  class Base

    attr_accessor :scrooge_callsite,
                  :scrooge_callsite_set,
                  :scrooge_own_callsite_set

    @@scrooge_mutex = Mutex.new
    @@scrooge_callsites = {}
    @@scrooge_select_regexes = {}

    ScroogeBlankString = "".freeze
    ScroogeComma = ",".freeze 
    ScroogeRegexWhere = /WHERE.*/
    ScroogeCallsiteSample = 0..10

    class << self
 
      # Determine if a given SQL string is a candidate for callsite <=> columns
      # optimization.
      #     
      alias :find_by_sql_without_scrooge :find_by_sql
      def find_by_sql(sql)
        if scope_with_scrooge?( sql )
          find_by_sql_with_scrooge( sql )
        else
          find_by_sql_without_scrooge(sql)
        end
      end

      # Only scope n-1 rows by default.
      # Stephen: Temp. relaxed the LIMIT constraint - please advise.
      def scope_with_scrooge?( sql )
        sql =~ scrooge_select_regex #&& sql !~ /LIMIT 1$/
      end

      # Populate the storage for a given callsite signature
      #
      def scrooge_callsite_set!(callsite_signature, set)
        @@scrooge_callsites[table_name][callsite_signature.hash] = set
      end  

      # Reference storage for a given callsite signature
      #
      def scrooge_callsite_set(callsite_signature)
        @@scrooge_callsites[table_name] ||= {}
        @@scrooge_callsites[table_name][callsite_signature.hash] || false
      end
      
      # Augment a given callsite signature with a column / attribute.
      #
      def augment_scrooge_callsite!( callsite_signature, attr_name )
        @@scrooge_callsites[table_name][callsite_signature] << attr_name
      end

      # Generates a SELECT snippet for this Model from a given Set of columns
      #
      def scrooge_sql( set )
        set.map{|a| attribute_with_table( a ) }.join( ScroogeComma )
      end

      private
        
        # Find through callsites.
        #
        def find_by_sql_with_scrooge( sql )
          callsite_signature = caller[ScroogeCallsiteSample] << sql.gsub(ScroogeRegexWhere, ScroogeBlankString)
          callsite_set = set_for_callsite( callsite_signature )
          sql = sql.gsub(scrooge_select_regex, "SELECT #{scrooge_sql( callsite_set )}")
          connection.select_all(sanitize_sql(sql), "#{name} Load").collect! { |record| instantiate_with_callsite(record, callsite_signature, callsite_set) }          
        end

        # Return an attribute Set for a given callsite signature.
        # Respects already tracked columns and ensures at least the primary key
        # if this is a fresh callsite.
        #
        def set_for_callsite( callsite_signature )
          callsite_set = scrooge_callsite_set(callsite_signature)
          unless callsite_set
            callsite_set = Set.new([self.primary_key.to_s])
            @@scrooge_mutex.synchronize do
              scrooge_callsite_set!(callsite_signature, callsite_set) 
            end
          end
          callsite_set
        end
        
        # Generate a regex that respects the table name as well to catch
        # verbose SQL from JOINS etc.
        # 
        def scrooge_select_regex
          @@scrooge_select_regexes[table_name] ||= Regexp.compile( "SELECT (`?(?:#{table_name})?`?.?\\*)" )
        end

        # Link the column to it's table.
        #
        def attribute_with_table( attr_name )
          "#{quoted_table_name}.#{attr_name.to_s}"
        end

        # Specialized instantiation that enhances the record with a callsite signature
        # and any columns already known to the callsite.
        # 
        def instantiate_with_callsite(record, callsite_signature, callsite_set)
          record = instantiate(record)
          record.scrooge_callsite_set = callsite_set
          record.scrooge_own_callsite_set = callsite_set.dup  # remember what cols were loaded for this record
          record.scrooge_callsite = callsite_signature.hash
          record
        end

        # Shamelessly borrowed from AR. 
        #
        def define_read_method(symbol, attr_name, column)
          cast_code = column.type_cast_code('v') if column
          access_code = cast_code ? "(v=@attributes['#{attr_name}']) && #{cast_code}" : "@attributes['#{attr_name}']"

          unless attr_name.to_s == self.primary_key.to_s
            access_code = access_code.insert(0, "missing_attribute('#{attr_name}', caller) unless @attributes.has_key?('#{attr_name}'); ")
          end

          if cache_attribute?(attr_name)
            access_code = "@attributes_cache['#{attr_name}'] ||= (#{access_code})"
          end
          define_with_scrooge(symbol, attr_name, access_code, "")
        end

        # Shamelessly borrowed from AR. 
        #
        def define_read_method_for_serialized_attribute(attr_name)
          define_with_scrooge(attr_name, attr_name, "unserialize_attribute('#{attr_name}')", "")
        end

        # Shamelessly borrowed from AR. 
        #
        def define_read_method_for_time_zone_conversion(attr_name)
          method_body = <<-EOV
          cached = @attributes_cache['#{attr_name}']
          return cached if cached && !reload
          time = read_attribute('#{attr_name}')
          @attributes_cache['#{attr_name}'] = time.acts_like?(:time) ? time.in_time_zone : time
          EOV
          define_with_scrooge(attr_name, attr_name, method_body, "(reload = false)")
        end
        
        # Graceful missing attribute wrapper.
        #
        def define_with_scrooge(symbol, attr_name, access_code, args)
          method_def = <<-EOV
          def #{symbol}#{args}
            begin
              #{access_code}
            rescue ActiveRecord::MissingAttributeError
              scrooge_missing_attribute('#{attr_name}')
              #{access_code}
            end
          end
          EOV
          evaluate_attribute_method attr_name, method_def
        end
      end
      
      # Augment the callsite with a fresh column reference.
      #
      def augment_scrooge_attribute!(attr_name)
        @@scrooge_mutex.synchronize do
          @scrooge_callsite_set << attr_name
          self.class.augment_scrooge_callsite!( @scrooge_callsite, attr_name )
        end
        @scrooge_own_callsite_set << attr_name
      end

      # Handle a missing attribute - reload with columns already known + the given
      # missing column. 
      #
      def scrooge_missing_attribute(attr_name)
        augment_scrooge_attribute!(attr_name)
        @scrooge_fully_loaded = true
        reload_with =  self.class.scrooge_sql(scrooge_callsite_set << attr_name)
        reload(:select => reload_with ) # find a better way to signal the reload than a space!
      end
    
      # Wrap #read_attribute to gracefully handle missing attributes
      #
      def read_attribute(attr_name)
        if scrooge_attr_present?( attr_name )
          super(attr_name)
        else
          scrooge_missing_attribute(attr_name)
          super(attr_name)
        end
      end

      # Wrap #read_attribute_before_type_cast to gracefully handle missing attributes
      #
      def read_attribute_before_type_cast(attr_name)
        if scrooge_attr_present?( attr_name )
          super(attr_name)
        else
          scrooge_missing_attribute(attr_name)
          super(attr_name)
        end
      end
      
      # Is the given column known to Scrooge ?
      #
      def scrooge_attr_present?( attr_name )
        !@scrooge_callsite_set || @scrooge_fully_loaded || @scrooge_own_callsite_set.include?( attr_name )
      end      
      
    end
end