Add ability to trace columns usage in custom code

CHANGELOG.md

@@ -1,5 +1,6 @@
 ## master (unreleased)
 
+  - Add ability to trace columns usage in custom code
   - Trace unused columns in mailers
 
 ## 0.2.0 (2023-10-13)

README.md

@@ -25,7 +25,7 @@ $ bundle install
 
 ## Usage
 
-Hit a controller action or run `ActiveJob` (or `Sidekiq`) job, open `log/columns_trace.log`,
+Hit a controller or email action or run `ActiveJob` (or `Sidekiq`) job, open `log/columns_trace.log`,
 and see the output:
 
 ```
@@ -54,6 +54,18 @@ ImportProjectJob
   ↳ app/jobs/import_project_job.rb:24:in `perform'
 ```
 
+### Tracing custom code
+
+To get columns usage in the custom code, you can manually wrap it by `ColumnsTrace.report`:
+
+```ruby
+task my_rake_task: :environment do
+  ColumnsTrace.report("my_rake_task") do
+    # do stuff
+  end
+end
+```
+
 ## Configuration
 
 You can override the following default options:

lib/columns_trace.rb

@@ -11,6 +11,23 @@
 
 module ColumnsTrace
   class << self
+    # Manually trace columns usage in an arbitrary code.
+    #
+    # @param title [String] title of the reporting, e.g. controller action etc
+    #
+    # @example
+    #   task my_rake_task: :environment do
+    #     ColumnsTrace.report("my_rake_task") do
+    #       # do stuff
+    #     end
+    #   end
+    #
+    def report(title)
+      Registry.clear
+      yield
+      reporter.report(title, Registry.created_records)
+    end
+
     # @private
     attr_reader :ignored_models
 

lib/columns_trace/rails_integration.rb

@@ -14,34 +14,20 @@ module ColumnsTrace
   end
 
   ActiveSupport.on_load(:action_controller) do
-    before_action { Registry.clear }
-
-    after_action do
-      ColumnsTrace.reporter.report("#{self.class.name}##{action_name}", Registry.created_records)
+    around_action do |controller, action|
+      ColumnsTrace.report("#{controller.class.name}##{action_name}", &action)
     end
   end
 
   ActiveSupport.on_load(:action_mailer) do
-    before_action { Registry.clear }
-
-    after_action do
-      ColumnsTrace.reporter.report("#{self.class.name}##{action_name}", Registry.created_records)
+    around_action do |mailer, action|
+      ColumnsTrace.report("#{mailer.class.name}##{action_name}", &action)
     end
   end
 
   ActiveSupport.on_load(:active_job) do
-    before_perform { Registry.clear }
-
-    after_perform do
-      ColumnsTrace.reporter.report("#{self.class.name}#perform", Registry.created_records)
-    end
-  end
-
-  ActiveSupport.on_load(:action_mailer) do
-    before_action { Registry.clear }
-
-    after_action do
-      ColumnsTrace.reporter.report("#{self.class.name}##{action_name}", Registry.created_records)
+    around_perform do |job, perform|
+      ColumnsTrace.report("#{job.class.name}#perform", &perform)
     end
   end
 end

lib/columns_trace/sidekiq_integration.rb

@@ -3,10 +3,8 @@
 module ColumnsTrace
   # @private
   class SidekiqMiddleware
-    def call(worker, _job, _queue)
-      Registry.clear
-      yield
-      ColumnsTrace.reporter.report("#{worker.class.name}#perform", Registry.created_records)
+    def call(worker, _job, _queue, &block)
+      ColumnsTrace.report("#{worker.class.name}#perform", &block)
     end
   end
 end