Add support for LIMIT BY clause (#169)

lib/active_record/connection_adapters/clickhouse_adapter.rb

@@ -5,6 +5,7 @@
 require 'arel/nodes/grouping_sets'
 require 'arel/nodes/settings'
 require 'arel/nodes/using'
+require 'arel/nodes/limit_by'
 require 'active_record/connection_adapters/clickhouse/oid/array'
 require 'active_record/connection_adapters/clickhouse/oid/date'
 require 'active_record/connection_adapters/clickhouse/oid/date_time'
@@ -52,6 +53,7 @@ module ClassMethods
                :group_by_grouping_sets, :group_by_grouping_sets!,
                :settings, :settings!,
                :window, :window!,
+               :limit_by, :limit_by!,
                to: :all
 
       def is_view

lib/arel/nodes/limit_by.rb

@@ -0,0 +1,17 @@
+module Arel # :nodoc: all
+  module Nodes
+    class LimitBy < Arel::Nodes::Unary
+      attr_reader :column
+
+      def initialize(limit, column)
+        raise ArgumentError, 'Limit should be an integer' unless limit.is_a?(Integer)
+        raise ArgumentError, 'Limit should be a positive integer' unless limit >= 0
+        raise ArgumentError, 'Column should be a Symbol or String' unless column.is_a?(String) || column.is_a?(Symbol)
+
+        @column = column
+
+        super(limit)
+      end
+    end
+  end
+end

lib/arel/visitors/clickhouse.rb

@@ -25,6 +25,7 @@ def visit_Arel_Attributes_Attribute(o, collector)
       end
 
       def visit_Arel_Nodes_SelectOptions(o, collector)
+        maybe_visit o.limit_by, collector
         maybe_visit o.settings, super
       end
 
@@ -69,6 +70,11 @@ def visit_Arel_Nodes_Using o, collector
         collector
       end
 
+      def visit_Arel_Nodes_LimitBy(o, collector)
+        collector << "LIMIT #{o.expr} BY #{o.column}"
+        collector
+      end
+
       def visit_Arel_Nodes_Matches(o, collector)
         op = o.case_sensitive ? " LIKE " : " ILIKE "
         infix_value o, collector, op

lib/core_extensions/active_record/relation.rb

@@ -108,6 +108,24 @@ def window!(name, **opts)
         self
       end
 
+      # The LIMIT BY clause permit to improve deduplication based on a unique key, it has better performances than
+      # the GROUP BY clause
+      #
+      #   users = User.limit_by(1, id)
+      #   # SELECT users.* FROM users LIMIT 1 BY id
+      #
+      # An <tt>ActiveRecord::ActiveRecordError</tt> will be reaised if database is not Clickhouse.
+      # @param [Array] opts
+      def limit_by(*opts)
+        spawn.limit_by!(*opts)
+      end
+
+      # @param [Array] opts
+      def limit_by!(*opts)
+        @values[:limit_by] = *opts
+        self
+      end
+
       private
 
       def check_command(cmd)
@@ -122,6 +140,7 @@ def build_arel(connection_or_aliases = nil, aliases = nil)
         end
 
         arel.final! if @values[:final].present?
+        arel.limit_by(*@values[:limit_by]) if @values[:limit_by].present?
         arel.settings(@values[:settings]) if @values[:settings].present?
         arel.using(@values[:using]) if @values[:using].present?
         arel.windows(@values[:windows]) if @values[:windows].present?

lib/core_extensions/arel/nodes/select_statement.rb

@@ -2,15 +2,18 @@ module CoreExtensions
   module Arel # :nodoc: all
     module Nodes
       module SelectStatement
-        attr_accessor :settings
+        attr_accessor :limit_by, :settings
 
         def initialize(relation = nil)
           super
+          @limit_by = nil
           @settings = nil
         end
 
         def eql?(other)
-          super && settings == other.settings
+          super && 
+            limit_by == other.limit_by &&
+            settings == other.settings
         end
       end
     end

lib/core_extensions/arel/select_manager.rb

@@ -29,6 +29,11 @@ def using(*exprs)
         @ctx.source.right.last.right = ::Arel::Nodes::Using.new(::Arel.sql(exprs.join(',')))
         self
       end
+
+      def limit_by(*exprs)
+        @ast.limit_by = ::Arel::Nodes::LimitBy.new(*exprs)
+        self
+      end
     end
   end
 end

spec/single/model_spec.rb

@@ -294,6 +294,17 @@ class ModelPk < ActiveRecord::Base
       end
     end
 
+    describe '#limit_by' do
+      it 'works' do
+        sql = Model.limit_by(1, :event_name).to_sql
+        expect(sql).to eq('SELECT sample.* FROM sample LIMIT 1 BY event_name')
+      end
+
+      it 'works with limit' do
+        sql = Model.limit(1).limit_by(1, :event_name).to_sql
+        expect(sql).to eq('SELECT sample.* FROM sample LIMIT 1 BY event_name LIMIT 1')
+      end
+    end
 
     describe '#group_by_grouping_sets' do
       it 'raises an error with no arguments' do

spec/spec_helper.rb

@@ -40,8 +40,8 @@
     host: 'localhost',
     port: ENV['CLICKHOUSE_PORT'] || 8123,
     database: ENV['CLICKHOUSE_DATABASE'] || 'test',
-    username: nil,
-    password: nil,
+    username: ENV['CLICKHOUSE_USER'],
+    password: ENV['CLICKHOUSE_PASSWORD'],
     cluster_name: ENV['CLICKHOUSE_CLUSTER'],
   }
 )