From e41c682d37a0572898ad4cec7902b4c158a1734e Mon Sep 17 00:00:00 2001
From: Heber Uriegas <woohoou@users.noreply.github.com>
Date: Thu, 28 May 2020 10:08:26 -0500
Subject: [PATCH] Add database and schema to active record log

---
 lib/apartment.rb                              |  3 +-
 lib/apartment/active_record/log_subscriber.rb | 41 +++++++++++++++++++
 lib/apartment/railtie.rb                      |  1 +
 .../apartment/install/templates/apartment.rb  |  5 +++
 4 files changed, 49 insertions(+), 1 deletion(-)
 create mode 100644 lib/apartment/active_record/log_subscriber.rb

diff --git a/lib/apartment.rb b/lib/apartment.rb
index 4b051e6b..854a0250 100644
--- a/lib/apartment.rb
+++ b/lib/apartment.rb
@@ -8,6 +8,7 @@
 
 # require_relative 'apartment/arel/visitors/postgresql'
 
+require_relative 'apartment/active_record/log_subscriber'
 require_relative 'apartment/active_record/connection_handling' if ActiveRecord.version.release >= Gem::Version.new('6.0')
 
 if ActiveRecord.version.release >= Gem::Version.new('6.1')
@@ -20,7 +21,7 @@ class << self
     extend Forwardable
 
     ACCESSOR_METHODS = %i[use_schemas use_sql seed_after_create prepend_environment
-                          append_environment with_multi_server_setup tenant_presence_check].freeze
+                          append_environment with_multi_server_setup tenant_presence_check active_record_log].freeze
 
     WRITER_METHODS = %i[tenant_names database_schema_file excluded_models
                         default_schema persistent_schemas connection_class
diff --git a/lib/apartment/active_record/log_subscriber.rb b/lib/apartment/active_record/log_subscriber.rb
new file mode 100644
index 00000000..ca090264
--- /dev/null
+++ b/lib/apartment/active_record/log_subscriber.rb
@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  class LogSubscriber
+    def apartment_log
+      return unless Apartment.active_record_log
+
+      database = color("[#{Apartment.connection.current_database}] ", ActiveSupport::LogSubscriber::MAGENTA, true)
+      schema = nil
+      schema = color("[#{Apartment.connection.schema_search_path.tr('"', '')}] ", ActiveSupport::LogSubscriber::YELLOW, true) unless Apartment.connection.schema_search_path.nil?
+      "#{database}#{schema}"
+    end
+
+    def payload_binds(binds, type_casted_binds)
+      return unless (binds || []).empty?
+
+      casted_params = type_casted_binds(type_casted_binds)
+      binds = '  ' + binds.zip(casted_params).map { |attr, value| render_bind(attr, value) }.inspect
+      binds
+    end
+
+    def sql(event)
+      self.class.runtime += event.duration
+      return unless logger.debug?
+
+      payload = event.payload
+
+      return if IGNORE_PAYLOAD_NAMES.include?(payload[:name])
+
+      name  = "#{payload[:name]} (#{event.duration.round(1)}ms)"
+      name  = "CACHE #{name}" if payload[:cached]
+      sql = payload[:sql]
+      binds = payload_binds(payload[:binds], payload[:type_casted_binds])
+
+      name = colorize_payload_name(name, payload[:name])
+      sql  = color(sql, sql_color(sql), true) if colorize_logging
+
+      debug "    #{apartment_log}#{name} #{sql}#{binds}"
+    end
+  end
+end
diff --git a/lib/apartment/railtie.rb b/lib/apartment/railtie.rb
index ddc3e18e..101c8b43 100644
--- a/lib/apartment/railtie.rb
+++ b/lib/apartment/railtie.rb
@@ -19,6 +19,7 @@ class Railtie < Rails::Railtie
         config.prepend_environment = false
         config.append_environment = false
         config.tenant_presence_check = true
+        config.active_record_log = false
       end
 
       ActiveRecord::Migrator.migrations_paths = Rails.application.paths['db/migrate'].to_a
diff --git a/lib/generators/apartment/install/templates/apartment.rb b/lib/generators/apartment/install/templates/apartment.rb
index 2a774737..ece13534 100644
--- a/lib/generators/apartment/install/templates/apartment.rb
+++ b/lib/generators/apartment/install/templates/apartment.rb
@@ -96,6 +96,11 @@
   # the new tenant
   #
   # config.pg_excluded_names = ["uuid_generate_v4"]
+
+  # Specifies whether the database and schema (when using PostgreSQL schemas) will prepend in ActiveRecord log.
+  # Uncomment the line below if you want to enable this behavior.
+  #
+  # config.active_record_log = true
 end
 
 # Setup a custom Tenant switching middleware. The Proc should return the name of the Tenant that