Merge pull request #13 from getlago/support-limit-by

feat(limit_by): Add Limit By clause support

lib/active_record/connection_adapters/clickhouse_adapter.rb

@@ -4,6 +4,7 @@
 require 'arel/nodes/final'
 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'
@@ -48,7 +49,7 @@ def is_view
 
   module ModelSchema
     module ClassMethods
-      delegate :final, :final!, :settings, :settings!, :window, :window!, to: :all
+      delegate :final, :final!, :settings, :settings!, :window, :window!, :limit_by, :limit_by!, to: :all
 
       def is_view
         @is_view || false

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
 
@@ -64,6 +65,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

@@ -81,6 +81,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)
@@ -95,6 +113,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

@@ -293,6 +293,18 @@ class ModelPk < ActiveRecord::Base
         expect(Model.final.where(date: '2023-07-21').to_sql).to eq('SELECT sample.* FROM sample FINAL WHERE sample.date = \'2023-07-21\'')
       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
   end
 
   context 'sample with id column' 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'],
   }
 )