ros-apartment 2.4.0 → 2.7.1

data/lib/apartment/adapters/mysql2_adapter.rb

@@ -39,6 +39,8 @@ module Apartment
       #   Reset current tenant to the default_tenant
       #
       def reset
+        return unless default_tenant
+
         Apartment.connection.execute "use `#{default_tenant}`"
       end
 

data/lib/apartment/adapters/postgresql_adapter.rb

@@ -30,6 +30,10 @@ module Apartment
         reset
       end
 
+      def default_tenant
+        @default_tenant = Apartment.default_tenant || 'public'
+      end
+
       #   Reset schema search path to the default schema_search_path
       #
       #   @return {String} default schema search path
@@ -37,6 +41,12 @@ module Apartment
       def reset
         @current = default_tenant
         Apartment.connection.schema_search_path = full_search_path
+        reset_sequence_names
+      end
+
+      def init
+        super
+        Apartment.connection.schema_search_path = full_search_path
       end
 
       def current
@@ -62,24 +72,30 @@ module Apartment
       #
       def connect_to_new(tenant = nil)
         return reset if tenant.nil?
-        # rubocop:disable Style/RaiseArgs
-        raise ActiveRecord::StatementInvalid.new("Could not find schema #{tenant}") unless Apartment.connection.schema_exists?(tenant.to_s)
 
-        # rubocop:enable Style/RaiseArgs
+        tenant = tenant.to_s
+        raise ActiveRecord::StatementInvalid, "Could not find schema #{tenant}" unless tenant_exists?(tenant)
 
-        @current = tenant.to_s
+        @current = tenant
         Apartment.connection.schema_search_path = full_search_path
 
         # When the PostgreSQL version is < 9.3,
         # there is a issue for prepared statement with changing search_path.
         # https://www.postgresql.org/docs/9.3/static/sql-prepare.html
         Apartment.connection.clear_cache! if postgresql_version < 90_300
+        reset_sequence_names
       rescue *rescuable_exceptions
         raise TenantNotFound, "One of the following schema(s) is invalid: \"#{tenant}\" #{full_search_path}"
       end
 
       private
 
+      def tenant_exists?(tenant)
+        return true unless Apartment.tenant_presence_check
+
+        Apartment.connection.schema_exists?(tenant)
+      end
+
       def create_tenant_command(conn, tenant)
         conn.execute(%(CREATE SCHEMA "#{tenant}"))
       end
@@ -99,6 +115,19 @@ module Apartment
         # public from Rails 5.0.
         Apartment.connection.send(:postgresql_version)
       end
+
+      def reset_sequence_names
+        # sequence_name contains the schema, so it must be reset after switch
+        # There is `reset_sequence_name`, but that method actually goes to the database
+        # to find out the new name. Therefore, we do this hack to only unset the name,
+        # and it will be dynamically found the next time it is needed
+        ActiveRecord::Base.descendants
+                          .select { |c| c.instance_variable_defined?(:@sequence_name) }
+                          .reject { |c| c.instance_variable_defined?(:@explicit_sequence_name) && c.instance_variable_get(:@explicit_sequence_name) }
+                          .each do |c|
+                            c.remove_instance_variable :@sequence_name
+                          end
+      end
     end
 
     # Another Adapter for Postgresql when using schemas and SQL

data/lib/apartment/adapters/sqlite3_adapter.rb

@@ -33,6 +33,8 @@ module Apartment
       protected
 
       def connect_to_new(tenant)
+        return reset if tenant.nil?
+
         unless File.exist?(database_file(tenant))
           raise TenantNotFound,
                 "The tenant #{environmentify(tenant)} cannot be found."

data/lib/apartment/console.rb

@@ -35,3 +35,10 @@ def tenant_list
   tenant_list += Apartment.tenant_names
   tenant_list.uniq
 end
+
+def tenant_info_msg
+  # rubocop:disable Rails/Output
+  puts "Available Tenants: #{tenant_list}\n"
+  puts "Use `st 'tenant'` to switch tenants & `tenant_list` to see list\n"
+  # rubocop:enable Rails/Output
+end

data/lib/apartment/custom_console.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require_relative 'console'
+
 module Apartment
   module CustomConsole
     begin
@@ -10,17 +12,31 @@ module Apartment
       # rubocop:enable Rails/Output
     end
 
-    if Pry::Prompt.respond_to?(:add)
-      desc = "Includes the current Rails environment and project folder name.\n" \
-            '[1] [project_name][Rails.env][Apartment::Tenant.current] pry(main)>'
+    desc = "Includes the current Rails environment and project folder name.\n" \
+          '[1] [project_name][Rails.env][Apartment::Tenant.current] pry(main)>'
+
+    prompt_procs = [
+      proc { |target_self, nest_level, pry| prompt_contents(pry, target_self, nest_level, '>') },
+      proc { |target_self, nest_level, pry| prompt_contents(pry, target_self, nest_level, '*') }
+    ]
 
+    if Gem::Version.new(Pry::VERSION) >= Gem::Version.new('0.13')
+      Pry.config.prompt = Pry::Prompt.new 'ros', desc, prompt_procs
+    else
       Pry::Prompt.add 'ros', desc, %w[> *] do |target_self, nest_level, pry, sep|
-        "[#{pry.input_ring.size}] [#{PryRails::Prompt.formatted_env}][#{Apartment::Tenant.current}] " \
-        "#{pry.config.prompt_name}(#{Pry.view_clip(target_self)})" \
-        "#{":#{nest_level}" unless nest_level.zero?}#{sep} "
+        prompt_contents(pry, target_self, nest_level, sep)
       end
-
       Pry.config.prompt = Pry::Prompt[:ros][:value]
     end
+
+    Pry.config.hooks.add_hook(:when_started, 'startup message') do
+      tenant_info_msg
+    end
+
+    def self.prompt_contents(pry, target_self, nest_level, sep)
+      "[#{pry.input_ring.size}] [#{PryRails::Prompt.formatted_env}][#{Apartment::Tenant.current}] " \
+      "#{pry.config.prompt_name}(#{Pry.view_clip(target_self)})" \
+      "#{":#{nest_level}" unless nest_level.zero?}#{sep} "
+    end
   end
 end

data/lib/apartment/model.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module Apartment
+  module Model
+    extend ActiveSupport::Concern
+
+    module ClassMethods
+      # NOTE: key can either be an array of symbols or a single value.
+      # E.g. If we run the following query:
+      # `Setting.find_by(key: 'something', value: 'amazing')` key will have an array of symbols: `[:key, :something]`
+      # while if we run:
+      # `Setting.find(10)` key will have the value 'id'
+      def cached_find_by_statement(key, &block)
+        # Modifying the cache key to have a reference to the current tenant,
+        # so the cached statement is referring only to the tenant in which we've
+        # executed this
+        cache_key = if key.is_a? String
+                      "#{Apartment::Tenant.current}_#{key}"
+                    else
+                      [Apartment::Tenant.current] + key
+                    end
+        cache = @find_by_statement_cache[connection.prepared_statements]
+        cache.compute_if_absent(cache_key) { ActiveRecord::StatementCache.create(connection, &block) }
+      end
+    end
+  end
+end

data/lib/apartment/railtie.rb

@@ -18,31 +18,35 @@ module Apartment
         config.seed_after_create = false
         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
     end
 
-    #   Hook into ActionDispatch::Reloader to ensure Apartment is properly initialized
-    #   Note that this doens't entirely work as expected in Development, because this is called before classes are reloaded
-    #   See the middleware/console declarations below to help with this. Hope to fix that soon.
-    #
+    # Make sure Apartment is reconfigured when the code is reloaded
     config.to_prepare do
-      next if ARGV.any? { |arg| arg =~ /\Aassets:(?:precompile|clean)\z/ }
-      next if ARGV.any? { |arg| arg == 'webpacker:compile' }
+      Apartment::Tenant.reinitialize
+    end
 
-      begin
-        Apartment.connection_class.connection_pool.with_connection do
-          Apartment::Tenant.init
-        end
-      rescue ::ActiveRecord::NoDatabaseError, PG::ConnectionBad
-        # Since `db:create` and other tasks invoke this block from Rails 5.2.0,
-        # we need to swallow the error to execute `db:create` properly.
-        Rails.logger.warn do
-          'Failed to initialize Apartment because a database connection could not be established.'
+    #
+    # Ensure that Apartment::Tenant.init is called when
+    # a new connection is requested.
+    #
+    module ApartmentInitializer
+      def connection
+        super.tap do
+          Apartment::Tenant.init_once
         end
       end
+
+      def arel_table
+        Apartment::Tenant.init_once
+        super
+      end
     end
+    ActiveRecord::Base.singleton_class.prepend ApartmentInitializer
 
     #
     #   Ensure rake tasks are loaded

data/lib/apartment/tasks/task_helper.rb

@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+module Apartment
+  module TaskHelper
+    def self.each_tenant(&block)
+      Parallel.each(tenants_without_default, in_threads: Apartment.parallel_migration_threads) do |tenant|
+        block.call(tenant)
+      end
+    end
+
+    def self.tenants_without_default
+      tenants - [Apartment.default_tenant]
+    end
+
+    def self.tenants
+      ENV['DB'] ? ENV['DB'].split(',').map(&:strip) : Apartment.tenant_names || []
+    end
+
+    def self.warn_if_tenants_empty
+      return unless tenants.empty? && ENV['IGNORE_EMPTY_TENANTS'] != 'true'
+
+      # rubocop:disable Rails/Output
+      puts <<-WARNING
+        [WARNING] - The list of tenants to migrate appears to be empty. This could mean a few things:
+
+          1. You may not have created any, in which case you can ignore this message
+          2. You've run `apartment:migrate` directly without loading the Rails environment
+            * `apartment:migrate` is now deprecated. Tenants will automatically be migrated with `db:migrate`
+
+        Note that your tenants currently haven't been migrated. You'll need to run `db:migrate` to rectify this.
+      WARNING
+      # rubocop:enable Rails/Output
+    end
+  end
+end

data/lib/apartment/tenant.rb

@@ -9,14 +9,24 @@ module Apartment
     extend self
     extend Forwardable
 
-    def_delegators :adapter, :create, :drop, :switch, :switch!, :current, :each, :reset, :set_callback, :seed, :current_tenant, :default_tenant, :environmentify
+    def_delegators :adapter, :create, :drop, :switch, :switch!, :current, :each,
+                   :reset, :init, :set_callback, :seed, :current_tenant,
+                   :default_tenant, :environmentify
 
     attr_writer :config
 
-    #   Initialize Apartment config options such as excluded_models
-    #
-    def init
-      adapter.process_excluded_models
+    def init_once
+      return if @already_initialized
+
+      # To avoid infinite loops in work init is doing,
+      # we need to set @already_initialized to true
+      # before init is called
+      @already_initialized = true
+      init
+    end
+
+    def reinitialize
+      @already_initialized = false
     end
 
     #   Fetch the proper multi-tenant adapter based on Rails config
@@ -51,6 +61,7 @@ module Apartment
     #
     def reload!(config = nil)
       Thread.current[:apartment_adapter] = nil
+      reinitialize
       @config = config
     end
 

data/lib/apartment/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Apartment
-  VERSION = '2.4.0'
+  VERSION = '2.7.1'
 end

data/lib/generators/apartment/install/templates/apartment.rb

@@ -96,6 +96,11 @@ Apartment.configure do |config|
   # 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

data/lib/tasks/apartment.rake

@@ -1,12 +1,15 @@
 # frozen_string_literal: true
 
 require 'apartment/migrator'
+require 'apartment/tasks/task_helper'
 require 'parallel'
 
 apartment_namespace = namespace :apartment do
   desc 'Create all tenants'
   task :create do
-    tenants.each do |tenant|
+    Apartment::TaskHelper.warn_if_tenants_empty
+
+    Apartment::TaskHelper.tenants.each do |tenant|
       begin
         puts("Creating #{tenant} tenant")
         Apartment::Tenant.create(tenant)
@@ -18,7 +21,7 @@ apartment_namespace = namespace :apartment do
 
   desc 'Drop all tenants'
   task :drop do
-    tenants.each do |tenant|
+    Apartment::TaskHelper.tenants.each do |tenant|
       begin
         puts("Dropping #{tenant} tenant")
         Apartment::Tenant.drop(tenant)
@@ -30,8 +33,8 @@ apartment_namespace = namespace :apartment do
 
   desc 'Migrate all tenants'
   task :migrate do
-    warn_if_tenants_empty
-    each_tenant do |tenant|
+    Apartment::TaskHelper.warn_if_tenants_empty
+    Apartment::TaskHelper.each_tenant do |tenant|
       begin
         puts("Migrating #{tenant} tenant")
         Apartment::Migrator.migrate tenant
@@ -42,10 +45,10 @@ apartment_namespace = namespace :apartment do
   end
 
   desc 'Seed all tenants'
-  task :seed do
-    warn_if_tenants_empty
+  task seed: :create do
+    Apartment::TaskHelper.warn_if_tenants_empty
 
-    each_tenant do |tenant|
+    Apartment::TaskHelper.each_tenant do |tenant|
       begin
         puts("Seeding #{tenant} tenant")
         Apartment::Tenant.switch(tenant) do
@@ -59,11 +62,11 @@ apartment_namespace = namespace :apartment do
 
   desc 'Rolls the migration back to the previous version (specify steps w/ STEP=n) across all tenants.'
   task :rollback do
-    warn_if_tenants_empty
+    Apartment::TaskHelper.warn_if_tenants_empty
 
     step = ENV['STEP'] ? ENV['STEP'].to_i : 1
 
-    each_tenant do |tenant|
+    Apartment::TaskHelper.each_tenant do |tenant|
       begin
         puts("Rolling back #{tenant} tenant")
         Apartment::Migrator.rollback tenant, step
@@ -76,12 +79,12 @@ apartment_namespace = namespace :apartment do
   namespace :migrate do
     desc 'Runs the "up" for a given migration VERSION across all tenants.'
     task :up do
-      warn_if_tenants_empty
+      Apartment::TaskHelper.warn_if_tenants_empty
 
       version = ENV['VERSION'] ? ENV['VERSION'].to_i : nil
       raise 'VERSION is required' unless version
 
-      each_tenant do |tenant|
+      Apartment::TaskHelper.each_tenant do |tenant|
         begin
           puts("Migrating #{tenant} tenant up")
           Apartment::Migrator.run :up, tenant, version
@@ -93,12 +96,12 @@ apartment_namespace = namespace :apartment do
 
     desc 'Runs the "down" for a given migration VERSION across all tenants.'
     task :down do
-      warn_if_tenants_empty
+      Apartment::TaskHelper.warn_if_tenants_empty
 
       version = ENV['VERSION'] ? ENV['VERSION'].to_i : nil
       raise 'VERSION is required' unless version
 
-      each_tenant do |tenant|
+      Apartment::TaskHelper.each_tenant do |tenant|
         begin
           puts("Migrating #{tenant} tenant down")
           Apartment::Migrator.run :down, tenant, version
@@ -119,28 +122,4 @@ apartment_namespace = namespace :apartment do
       end
     end
   end
-
-  def each_tenant(&block)
-    Parallel.each(tenants, in_threads: Apartment.parallel_migration_threads) do |tenant|
-      block.call(tenant)
-    end
-  end
-
-  def tenants
-    ENV['DB'] ? ENV['DB'].split(',').map(&:strip) : Apartment.tenant_names || []
-  end
-
-  def warn_if_tenants_empty
-    return unless tenants.empty? && ENV['IGNORE_EMPTY_TENANTS'] != 'true'
-
-    puts <<-WARNING
-      [WARNING] - The list of tenants to migrate appears to be empty. This could mean a few things:
-
-        1. You may not have created any, in which case you can ignore this message
-        2. You've run `apartment:migrate` directly without loading the Rails environment
-          * `apartment:migrate` is now deprecated. Tenants will automatically be migrated with `db:migrate`
-
-      Note that your tenants currently haven't been migrated. You'll need to run `db:migrate` to rectify this.
-    WARNING
-  end
 end