puzzle-apartment 2.12.0

data/lib/apartment/elevators/host.rb

@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+require 'apartment/elevators/generic'
+
+module Apartment
+  module Elevators
+    #   Provides a rack based tenant switching solution based on the host
+    #   Assumes that tenant name should match host
+    #   Strips/ignores first subdomains in ignored_first_subdomains
+    #   eg. example.com       => example.com
+    #       www.example.bc.ca => www.example.bc.ca
+    #   if ignored_first_subdomains = ['www']
+    #       www.example.bc.ca => example.bc.ca
+    #       www.a.b.c.d.com   => a.b.c.d.com
+    #
+    class Host < Generic
+      def self.ignored_first_subdomains
+        @ignored_first_subdomains ||= []
+      end
+
+      # rubocop:disable Style/TrivialAccessors
+      def self.ignored_first_subdomains=(arg)
+        @ignored_first_subdomains = arg
+      end
+      # rubocop:enable Style/TrivialAccessors
+
+      def parse_tenant_name(request)
+        return nil if request.host.blank?
+
+        parts = request.host.split('.')
+        self.class.ignored_first_subdomains.include?(parts[0]) ? parts.drop(1).join('.') : request.host
+      end
+    end
+  end
+end

data/lib/apartment/elevators/host_hash.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+require 'apartment/elevators/generic'
+
+module Apartment
+  module Elevators
+    #   Provides a rack based tenant switching solution based on hosts
+    #   Uses a hash to find the corresponding tenant name for the host
+    #
+    class HostHash < Generic
+      def initialize(app, hash = {}, processor = nil)
+        super app, processor
+        @hash = hash
+      end
+
+      def parse_tenant_name(request)
+        unless @hash.key?(request.host)
+          raise TenantNotFound,
+                "Cannot find tenant for host #{request.host}"
+        end
+
+        @hash[request.host]
+      end
+    end
+  end
+end

data/lib/apartment/elevators/subdomain.rb

@@ -0,0 +1,66 @@
+# frozen_string_literal: true
+
+require 'apartment/elevators/generic'
+require 'public_suffix'
+
+module Apartment
+  module Elevators
+    #   Provides a rack based tenant switching solution based on subdomains
+    #   Assumes that tenant name should match subdomain
+    #
+    class Subdomain < Generic
+      def self.excluded_subdomains
+        @excluded_subdomains ||= []
+      end
+
+      # rubocop:disable Style/TrivialAccessors
+      def self.excluded_subdomains=(arg)
+        @excluded_subdomains = arg
+      end
+      # rubocop:enable Style/TrivialAccessors
+
+      def parse_tenant_name(request)
+        request_subdomain = subdomain(request.host)
+
+        # If the domain acquired is set to be excluded, set the tenant to whatever is currently
+        # next in line in the schema search path.
+        tenant = if self.class.excluded_subdomains.include?(request_subdomain)
+                   nil
+                 else
+                   request_subdomain
+                 end
+
+        tenant.presence
+      end
+
+      protected
+
+      # *Almost* a direct ripoff of ActionDispatch::Request subdomain methods
+
+      # Only care about the first subdomain for the database name
+      def subdomain(host)
+        subdomains(host).first
+      end
+
+      def subdomains(host)
+        host_valid?(host) ? parse_host(host) : []
+      end
+
+      def host_valid?(host)
+        !ip_host?(host) && domain_valid?(host)
+      end
+
+      def ip_host?(host)
+        !/^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}$/.match(host).nil?
+      end
+
+      def domain_valid?(host)
+        PublicSuffix.valid?(host, ignore_private: true)
+      end
+
+      def parse_host(host)
+        (PublicSuffix.parse(host, ignore_private: true).trd || '').split('.')
+      end
+    end
+  end
+end

data/lib/apartment/log_subscriber.rb

@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+
+require 'active_record/log_subscriber'
+
+module Apartment
+  # Custom Log subscriber to include database name and schema name in sql logs
+  class LogSubscriber < ActiveRecord::LogSubscriber
+    # NOTE: for some reason, if the method definition is not here, then the custom debug method is not called
+    # rubocop:disable Lint/UselessMethodDefinition
+    def sql(event)
+      super(event)
+    end
+    # rubocop:enable Lint/UselessMethodDefinition
+
+    private
+
+    def debug(progname = nil, &block)
+      progname = "  #{apartment_log}#{progname}" unless progname.nil?
+
+      super(progname, &block)
+    end
+
+    def apartment_log
+      database = color("[#{database_name}] ", ActiveSupport::LogSubscriber::MAGENTA, true)
+      schema = current_search_path
+      schema = color("[#{schema.tr('"', '')}] ", ActiveSupport::LogSubscriber::YELLOW, true) unless schema.nil?
+      "#{database}#{schema}"
+    end
+
+    def current_search_path
+      if Apartment.connection.respond_to?(:schema_search_path)
+        Apartment.connection.schema_search_path
+      else
+        Apartment::Tenant.current # all others
+      end
+    end
+
+    def database_name
+      db_name = Apartment.connection.raw_connection.try(:db) # PostgreSQL, PostGIS
+      db_name ||= Apartment.connection.raw_connection.try(:query_options)&.dig(:database) # Mysql
+      db_name ||= Apartment.connection.current_database # Failover
+      db_name
+    end
+  end
+end

data/lib/apartment/migrator.rb

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+require 'apartment/tenant'
+
+module Apartment
+  module Migrator
+    extend self
+
+    # Migrate to latest
+    def migrate(database)
+      Tenant.switch(database) do
+        version = ENV['VERSION'] ? ENV['VERSION'].to_i : nil
+
+        migration_scope_block = ->(migration) { ENV['SCOPE'].blank? || (ENV['SCOPE'] == migration.scope) }
+
+        if activerecord_below_5_2?
+          ActiveRecord::Migrator.migrate(ActiveRecord::Migrator.migrations_paths, version, &migration_scope_block)
+        else
+          ActiveRecord::Base.connection.migration_context.migrate(version, &migration_scope_block)
+        end
+      end
+    end
+
+    # Migrate up/down to a specific version
+    def run(direction, database, version)
+      Tenant.switch(database) do
+        if activerecord_below_5_2?
+          ActiveRecord::Migrator.run(direction, ActiveRecord::Migrator.migrations_paths, version)
+        else
+          ActiveRecord::Base.connection.migration_context.run(direction, version)
+        end
+      end
+    end
+
+    # rollback latest migration `step` number of times
+    def rollback(database, step = 1)
+      Tenant.switch(database) do
+        if activerecord_below_5_2?
+          ActiveRecord::Migrator.rollback(ActiveRecord::Migrator.migrations_paths, step)
+        else
+          ActiveRecord::Base.connection.migration_context.rollback(step)
+        end
+      end
+    end
+
+    private
+
+    def activerecord_below_5_2?
+      ActiveRecord.version.release < Gem::Version.new('5.2.0')
+    end
+  end
+end

data/lib/apartment/model.rb

@@ -0,0 +1,29 @@
+# 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
+                      # NOTE: In Rails 6.0.4 we start receiving an ActiveRecord::Reflection::BelongsToReflection
+                      # as the key, which wouldn't work well with an array.
+                      [Apartment::Tenant.current] + Array.wrap(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

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+require 'rails'
+require 'apartment/tenant'
+
+module Apartment
+  class Railtie < Rails::Railtie
+    #
+    #   Set up our default config options
+    #   Do this before the app initializers run so we don't override custom settings
+    #
+    config.before_initialize do
+      Apartment.configure do |config|
+        config.excluded_models = []
+        config.use_schemas = true
+        config.tenant_names = []
+        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 doesn'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.
+    #
+    config.to_prepare do
+      next if ARGV.any? { |arg| arg =~ /\Aassets:(?:precompile|clean)\z/ }
+      next if ARGV.any?('webpacker:compile')
+      next if ENV['APARTMENT_DISABLE_INIT']
+
+      begin
+        Apartment.connection_class.connection_pool.with_connection do
+          Apartment::Tenant.init
+        end
+      rescue ::ActiveRecord::NoDatabaseError
+        # 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.
+      end
+    end
+
+    config.after_initialize do
+      # NOTE: Load the custom log subscriber if enabled
+      if Apartment.active_record_log
+        ActiveSupport::Notifications.notifier.listeners_for('sql.active_record').each do |listener|
+          next unless listener.instance_variable_get('@delegate').is_a?(ActiveRecord::LogSubscriber)
+
+          ActiveSupport::Notifications.unsubscribe listener
+        end
+
+        Apartment::LogSubscriber.attach_to :active_record
+      end
+    end
+
+    #
+    #   Ensure rake tasks are loaded
+    #
+    rake_tasks do
+      load 'tasks/apartment.rake'
+      require 'apartment/tasks/enhancements' if Apartment.db_migrate_tenants
+    end
+  end
+end

data/lib/apartment/tasks/enhancements.rb

@@ -0,0 +1,55 @@
+# frozen_string_literal: true
+
+# Require this file to append Apartment rake tasks to ActiveRecord db rake tasks
+# Enabled by default in the initializer
+
+module Apartment
+  class RakeTaskEnhancer
+    module TASKS
+      ENHANCE_BEFORE = %w[db:drop].freeze
+      ENHANCE_AFTER  = %w[db:migrate db:rollback db:migrate:up db:migrate:down db:migrate:redo db:seed].freeze
+      freeze
+    end
+
+    # This is a bit convoluted, but helps solve problems when using Apartment within an engine
+    # See spec/integration/use_within_an_engine.rb
+
+    class << self
+      def enhance!
+        return unless should_enhance?
+
+        # insert task before
+        TASKS::ENHANCE_BEFORE.each do |name|
+          task = Rake::Task[name]
+          enhance_before_task(task)
+        end
+
+        # insert task after
+        TASKS::ENHANCE_AFTER.each do |name|
+          task = Rake::Task[name]
+          enhance_after_task(task)
+        end
+      end
+
+      def should_enhance?
+        Apartment.db_migrate_tenants
+      end
+
+      def enhance_before_task(task)
+        task.enhance([inserted_task_name(task)])
+      end
+
+      def enhance_after_task(task)
+        task.enhance do
+          Rake::Task[inserted_task_name(task)].invoke
+        end
+      end
+
+      def inserted_task_name(task)
+        task.name.sub(/db:/, 'apartment:')
+      end
+    end
+  end
+end
+
+Apartment::RakeTaskEnhancer.enhance!

data/lib/apartment/tasks/task_helper.rb

@@ -0,0 +1,52 @@
+# 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'
+
+      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
+
+    def self.create_tenant(tenant_name)
+      puts("Creating #{tenant_name} tenant")
+      Apartment::Tenant.create(tenant_name)
+    rescue Apartment::TenantExists => e
+      puts "Tried to create already existing tenant: #{e}"
+    end
+
+    def self.migrate_tenant(tenant_name)
+      strategy = Apartment.db_migrate_tenant_missing_strategy
+      create_tenant(tenant_name) if strategy == :create_tenant
+
+      puts("Migrating #{tenant_name} tenant")
+      Apartment::Migrator.migrate tenant_name
+    rescue Apartment::TenantNotFound => e
+      raise e if strategy == :raise_exception
+
+      puts e.message
+    end
+  end
+end

data/lib/apartment/tenant.rb

@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+require 'forwardable'
+
+module Apartment
+  #   The main entry point to Apartment functions
+  #
+  module Tenant
+    extend self
+    extend Forwardable
+
+    def_delegators :adapter, :create, :drop, :switch, :switch!, :current, :each,
+                   :reset, :init, :set_callback, :seed, :default_tenant, :environmentify
+
+    attr_writer :config
+
+    #   Fetch the proper multi-tenant adapter based on Rails config
+    #
+    #   @return {subclass of Apartment::AbstractAdapter}
+    #
+    def adapter
+      Thread.current[:apartment_adapter] ||= begin
+        adapter_method = "#{config[:adapter]}_adapter"
+
+        if defined?(JRUBY_VERSION)
+          case config[:adapter]
+          when /mysql/
+            adapter_method = 'jdbc_mysql_adapter'
+          when /postgresql/
+            adapter_method = 'jdbc_postgresql_adapter'
+          end
+        end
+
+        begin
+          require "apartment/adapters/#{adapter_method}"
+        rescue LoadError
+          raise "The adapter `#{adapter_method}` is not yet supported"
+        end
+
+        unless respond_to?(adapter_method)
+          raise AdapterNotFound, "database configuration specifies nonexistent #{config[:adapter]} adapter"
+        end
+
+        send(adapter_method, config)
+      end
+    end
+
+    #   Reset config and adapter so they are regenerated
+    #
+    def reload!(config = nil)
+      Thread.current[:apartment_adapter] = nil
+      @config = config
+    end
+
+    private
+
+    #   Fetch the rails database configuration
+    #
+    def config
+      @config ||= Apartment.connection_config
+    end
+  end
+end

data/lib/apartment/version.rb

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

data/lib/apartment.rb

@@ -0,0 +1,159 @@
+# frozen_string_literal: true
+
+require 'apartment/railtie' if defined?(Rails)
+require 'active_support/core_ext/object/blank'
+require 'forwardable'
+require 'active_record'
+require 'apartment/tenant'
+
+require_relative 'apartment/log_subscriber'
+
+if ActiveRecord.version.release >= Gem::Version.new('6.0')
+  require_relative 'apartment/active_record/connection_handling'
+end
+
+if ActiveRecord.version.release >= Gem::Version.new('6.1')
+  require_relative 'apartment/active_record/schema_migration'
+  require_relative 'apartment/active_record/internal_metadata'
+end
+
+# Apartment main definitions
+module Apartment
+  class << self
+    extend Forwardable
+
+    ACCESSOR_METHODS = %i[use_schemas use_sql seed_after_create prepend_environment default_tenant
+                          append_environment with_multi_server_setup tenant_presence_check active_record_log].freeze
+
+    WRITER_METHODS = %i[tenant_names database_schema_file excluded_models
+                        persistent_schemas connection_class
+                        db_migrate_tenants db_migrate_tenant_missing_strategy seed_data_file
+                        parallel_migration_threads pg_excluded_names].freeze
+
+    attr_accessor(*ACCESSOR_METHODS)
+    attr_writer(*WRITER_METHODS)
+
+    if ActiveRecord.version.release >= Gem::Version.new('6.1')
+      def_delegators :connection_class, :connection, :connection_db_config, :establish_connection
+
+      def connection_config
+        connection_db_config.configuration_hash
+      end
+    else
+      def_delegators :connection_class, :connection, :connection_config, :establish_connection
+    end
+
+    # configure apartment with available options
+    def configure
+      yield self if block_given?
+    end
+
+    def tenant_names
+      extract_tenant_config.keys.map(&:to_s)
+    end
+
+    def tenants_with_config
+      extract_tenant_config
+    end
+
+    def tld_length=(_)
+      Apartment::Deprecation.warn('`config.tld_length` have no effect because it was removed in https://github.com/influitive/apartment/pull/309')
+    end
+
+    def db_config_for(tenant)
+      (tenants_with_config[tenant] || connection_config)
+    end
+
+    # Whether or not db:migrate should also migrate tenants
+    # defaults to true
+    def db_migrate_tenants
+      return @db_migrate_tenants if defined?(@db_migrate_tenants)
+
+      @db_migrate_tenants = true
+    end
+
+    # How to handle tenant missing on db:migrate
+    # defaults to :rescue_exception
+    # available options: rescue_exception, raise_exception, create_tenant
+    def db_migrate_tenant_missing_strategy
+      valid = %i[rescue_exception raise_exception create_tenant]
+      value = @db_migrate_tenant_missing_strategy || :rescue_exception
+
+      return value if valid.include?(value)
+
+      key_name  = 'config.db_migrate_tenant_missing_strategy'
+      opt_names = valid.join(', ')
+
+      raise ApartmentError, "Option #{value} not valid for `#{key_name}`. Use one of #{opt_names}"
+    end
+
+    # Default to empty array
+    def excluded_models
+      @excluded_models || []
+    end
+
+    def parallel_migration_threads
+      @parallel_migration_threads || 0
+    end
+
+    def persistent_schemas
+      @persistent_schemas || []
+    end
+
+    def connection_class
+      @connection_class || ActiveRecord::Base
+    end
+
+    def database_schema_file
+      return @database_schema_file if defined?(@database_schema_file)
+
+      @database_schema_file = Rails.root.join('db/schema.rb')
+    end
+
+    def seed_data_file
+      return @seed_data_file if defined?(@seed_data_file)
+
+      @seed_data_file = Rails.root.join('db/seeds.rb')
+    end
+
+    def pg_excluded_names
+      @pg_excluded_names || []
+    end
+
+    # Reset all the config for Apartment
+    def reset
+      (ACCESSOR_METHODS + WRITER_METHODS).each do |method|
+        remove_instance_variable(:"@#{method}") if instance_variable_defined?(:"@#{method}")
+      end
+    end
+
+    def extract_tenant_config
+      return {} unless @tenant_names
+
+      values = @tenant_names.respond_to?(:call) ? @tenant_names.call : @tenant_names
+      unless values.is_a? Hash
+        values = values.each_with_object({}) do |tenant, hash|
+          hash[tenant] = connection_config
+        end
+      end
+      values.with_indifferent_access
+    rescue ActiveRecord::StatementInvalid
+      {}
+    end
+  end
+
+  # Exceptions
+  ApartmentError = Class.new(StandardError)
+
+  # Raised when apartment cannot find the adapter specified in <tt>config/database.yml</tt>
+  AdapterNotFound = Class.new(ApartmentError)
+
+  # Raised when apartment cannot find the file to be loaded
+  FileNotFound = Class.new(ApartmentError)
+
+  # Tenant specified is unknown
+  TenantNotFound = Class.new(ApartmentError)
+
+  # The Tenant attempting to be created already exists
+  TenantExists = Class.new(ApartmentError)
+end

data/lib/generators/apartment/install/USAGE

@@ -0,0 +1,5 @@
+Description:
+  Creates an initializer for apartment.
+
+Example:
+  `rails generate apartment:install`

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

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+module Apartment
+  class InstallGenerator < Rails::Generators::Base
+    source_root File.expand_path('templates', __dir__)
+
+    def copy_files
+      template 'apartment.rb', File.join('config', 'initializers', 'apartment.rb')
+    end
+  end
+end