lib/pg_search/features/tsearch.rb in pg_search-1.0.5 vs lib/pg_search/features/tsearch.rb in pg_search-1.0.6

@@ -1,23 +1,19 @@
 require "pg_search/compatibility"
 require "active_support/core_ext/module/delegation"
 
 module PgSearch
   module Features
-    class TSearch < Feature
+    class TSearch < Feature # rubocop:disable Metrics/ClassLength
       def self.valid_options
-        super + [:dictionary, :prefix, :negation, :any_word, :normalization, :tsvector_column]
+        super + [:dictionary, :prefix, :negation, :any_word, :normalization, :tsvector_column, :highlight]
       end
 
       def initialize(*args)
         super
-
-        if options[:prefix] && model.connection.send(:postgresql_version) < 80400 # rubocop:disable Style/GuardClause
-          raise PgSearch::NotSupportedForPostgresqlVersion.new(<<-MESSAGE.strip_heredoc)
-            Sorry, {:using => {:tsearch => {:prefix => true}}} only works in PostgreSQL 8.4 and above.")
-          MESSAGE
-        end
+        checks_for_highlight
+        checks_for_prefix
       end
 
       def conditions
         Arel::Nodes::Grouping.new(
           Arel::Nodes::InfixOperation.new("@@", arel_wrap(tsdocument), arel_wrap(tsquery))
@@ -26,10 +22,49 @@
 
       def rank
         arel_wrap(tsearch_rank)
       end
 
+      def highlight
+        arel_wrap(ts_headline)
+      end
+
       private
+
+      def checks_for_prefix
+        if options[:prefix] && model.connection.send(:postgresql_version) < 80400
+          raise PgSearch::NotSupportedForPostgresqlVersion.new(<<-MESSAGE.strip_heredoc)
+            Sorry, {:using => {:tsearch => {:prefix => true}}} only works in PostgreSQL 8.4 and above.")
+          MESSAGE
+        end
+      end
+
+      def checks_for_highlight
+        if options[:highlight] && model.connection.send(:postgresql_version) < 90000
+          raise PgSearch::NotSupportedForPostgresqlVersion.new(<<-MESSAGE.strip_heredoc)
+            Sorry, {:using => {:tsearch => {:highlight => true}}} only works in PostgreSQL 9.0 and above.")
+          MESSAGE
+        end
+      end
+
+      def ts_headline
+        "ts_headline((#{document}), (#{tsquery}), '#{ts_headline_options}')"
+      end
+
+      def ts_headline_options
+        return nil unless options[:highlight].is_a?(Hash)
+
+        headline_options = map_headline_options
+        headline_options.map{|key, value| "#{key} = #{value}" if value }.compact.join(", ")
+      end
+
+      def map_headline_options
+        {
+          "StartSel" => options[:highlight][:start_sel],
+          "StopSel" => options[:highlight][:stop_sel],
+          "MaxFragments" => options[:highlight][:max_fragments]
+        }
+      end
 
       DISALLOWED_TSQUERY_CHARACTERS = /['?\\:]/
 
       def tsquery_for_term(unsanitized_term) # rubocop:disable Metrics/AbcSize
         if options[:negation] && unsanitized_term.start_with?("!")