lib/builder_apm/models/instrumenter.rb in builder_apm-0.2.5 vs lib/builder_apm/models/instrumenter.rb in builder_apm-0.3.0

@@ -52,18 +52,29 @@
           class_name: ''
         }
       end
 
       def store_sql_query_data(sql_query_data)
-        Thread.current[:sql_event_id] = sql_query_data[:sql_id]
-
-        if Thread.current[:stack]&.any?
-          Thread.current[:stack].last[:sql_events].push(sql_query_data)
-        else
-          (Thread.current[:stack] ||= []).push({sql_events: [sql_query_data], children: []})
-        end
+      Thread.current[:sql_event_id] = sql_query_data[:sql_id]
+      
+      # Create the stack if it doesn't exist yet
+      stack = (Thread.current[:stack] ||= [])
+      
+      if stack&.any?
+        stack.last[:sql_events].push(sql_query_data)
+      else
+        stack.push({sql_events: [sql_query_data], children: []})
       end
+      # Do the N+1 check if it wasn't done yet
+      if BuilderApm.configuration.enable_n_plus_one_profiler && Thread.current[:has_n_plus_one] == false
+        start_time = Time.now.to_f * 1000
+        perform_n_plus_one_check(stack)
+        duration = (Time.now.to_f * 1000) - start_time
+        Thread.current[:n_plus_one_duration] ||= 0
+        Thread.current[:n_plus_one_duration] += duration
+      end
+    end
 
       def update_last_sql_query_data_with_instantiation_info(event)
         stack = Thread.current[:stack]
         request_id = Thread.current[:request_id]
       
@@ -79,9 +90,33 @@
           end
           Thread.current[:stack].last[:sql_events].push(last_sql)
         ensure
           Thread.current[:sql_event_id] = nil
         end
-      end      
+      end 
+      
+      # Add this method to perform the N+1 check
+      def perform_n_plus_one_check(stack)
+        sql_queries = stack.map { |frame| frame[:sql_events] }.flatten
+
+        # Group queries count by table and triggering_line
+        queries_count = Hash.new { |h, k| h[k] = {count: 0, indices: []} }
+        sql_queries.each_with_index do |query, index|
+          match = query[:sql].match(/FROM ['`"]*([^ '`"]+)['`"]*/i)
+          if match
+            table_and_line = "#{match[1]}|#{query[:triggering_line]}"
+            queries_count[table_and_line][:count] += 1
+            queries_count[table_and_line][:indices].push(index)
+          end
+        end
+
+        # If any N+1 issue is found, set 'has_n_plus_one' to true and return
+        queries_count.each do |_, value|
+          if value[:count] > 1
+            Thread.current[:has_n_plus_one] = true
+            return
+          end
+        end
+      end
     end
   end
 end