activerecord-tenanted 0.1.0 → 0.2.0

data/lib/active_record/tenanted/tenant.rb

@@ -0,0 +1,257 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  module Tenanted
+    # instance methods common to both Tenant and Subtenant
+    module TenantCommon # :nodoc:
+      extend ActiveSupport::Concern
+
+      prepended do
+        attr_reader :tenant
+
+        before_save :ensure_tenant_context_safety
+      end
+
+      def cache_key
+        tenant ? "#{super}?tenant=#{tenant}" : super
+      end
+
+      def to_global_id(options = {})
+        super(options.merge(tenant: tenant))
+      end
+
+      def to_signed_global_id(options = {})
+        super(options.merge(tenant: tenant))
+      end
+
+      def association(name)
+        super.tap do |assoc|
+          if assoc.reflection.polymorphic? || assoc.reflection.klass.tenanted?
+            ensure_tenant_context_safety
+          end
+        end
+      end
+
+      alias to_gid to_global_id
+      alias to_sgid to_signed_global_id
+
+      private
+        # I would prefer to do this in an `after_initialize` callback, but some associations are
+        # created before those callbacks are invoked (for example, a `belongs_to` association) and
+        # we need to be able to ensure tenant context safety on all associations.
+        def init_internals
+          @tenant = self.class.current_tenant
+          super
+        end
+
+        def ensure_tenant_context_safety
+          self_tenant = self.tenant
+          current_tenant = self.class.current_tenant
+
+          if current_tenant.nil?
+            raise NoTenantError, "Cannot connect to a tenanted database while untenanted (#{self.class})"
+          elsif self_tenant != current_tenant
+            raise WrongTenantError,
+                  "#{self.class} model belongs to tenant #{self_tenant.inspect}, " \
+                  "but current tenant is #{current_tenant.inspect}"
+          end
+        end
+    end
+
+    module Tenant
+      extend ActiveSupport::Concern
+
+      # This is a sentinel value used to indicate that the class is not currently tenanted.
+      #
+      # It's the default value returned by `current_shard` when the class is not tenanted. The
+      # `current_tenant` method's job is to recognizes that sentinel value and return `nil`, because
+      # Active Record itself does not recognize `nil` as a valid shard value.
+      UNTENANTED_SENTINEL = Class.new do # :nodoc:
+        def inspect
+          "ActiveRecord::Tenanted::Tenant::UNTENANTED_SENTINEL"
+        end
+
+        def to_s
+          "(untenanted)"
+        end
+      end.new.freeze
+
+      CONNECTION_POOL_CREATION_LOCK = Thread::Mutex.new # :nodoc:
+
+      class_methods do
+        def tenanted?
+          true
+        end
+
+        def current_tenant
+          shard = current_shard
+          shard != UNTENANTED_SENTINEL ? shard : nil
+        end
+
+        def current_tenant=(tenant_name)
+          tenant_name = tenant_name.to_s unless tenant_name == UNTENANTED_SENTINEL
+
+          connection_class_for_self.connecting_to(shard: tenant_name, role: ActiveRecord.writing_role)
+        end
+
+        def tenant_exist?(tenant_name)
+          # this will have to be an adapter-specific implementation if we support other than sqlite
+          database_path = tenanted_root_config.database_path_for(tenant_name)
+
+          File.exist?(database_path) && !ActiveRecord::Tenanted::Mutex::Ready.locked?(database_path)
+        end
+
+        def with_tenant(tenant_name, prohibit_shard_swapping: true, &block)
+          tenant_name = tenant_name.to_s unless tenant_name == UNTENANTED_SENTINEL
+
+          if tenant_name == current_tenant
+            yield
+          else
+            connection_class_for_self.connected_to(shard: tenant_name, role: ActiveRecord.writing_role) do
+              prohibit_shard_swapping(prohibit_shard_swapping) do
+                log_tenant_tag(tenant_name, &block)
+              end
+            end
+          end
+        end
+
+        def create_tenant(tenant_name, if_not_exists: false, &block)
+          created_db = false
+          database_path = tenanted_root_config.database_path_for(tenant_name)
+
+          ActiveRecord::Tenanted::Mutex::Ready.lock(database_path) do
+            unless File.exist?(database_path)
+              # NOTE: This is obviously a sqlite-specific implementation.
+              # TODO: Add a `create_database` method upstream in the sqlite3 adapter, and call it.
+              #       Then this would delegate to the adapter and become adapter-agnostic.
+              FileUtils.touch(database_path)
+
+              with_tenant(tenant_name) do
+                connection_pool(schema_version_check: false)
+                ActiveRecord::Tenanted::DatabaseTasks.migrate_tenant(tenant_name)
+              end
+
+              created_db = true
+            end
+          rescue
+            FileUtils.rm_f(database_path)
+            raise
+          end
+
+          raise TenantExistsError unless created_db || if_not_exists
+
+          with_tenant(tenant_name) do
+            yield if block_given?
+          end
+        end
+
+        def destroy_tenant(tenant_name)
+          ActiveRecord::Base.logger.info "  DESTROY [tenant=#{tenant_name}] Destroying tenant database"
+
+          with_tenant(tenant_name, prohibit_shard_swapping: false) do
+            if retrieve_connection_pool(strict: false)
+              remove_connection
+            end
+          end
+
+          # NOTE: This is obviously a sqlite-specific implementation.
+          # TODO: Create a `drop_database` method upstream in the sqlite3 adapter, and call it.
+          #       Then this would delegate to the adapter and become adapter-agnostic.
+          FileUtils.rm_f(tenanted_root_config.database_path_for(tenant_name))
+        end
+
+        def tenants
+          # DatabaseConfigurations::RootConfig#tenants returns all tenants whose database files
+          # exist, but some of those may be getting initially migrated, so we perform an additional
+          # filter on readiness with `tenant_exist?`.
+          tenanted_root_config.tenants.select { |t| tenant_exist?(t) }
+        end
+
+        def with_each_tenant(**options, &block)
+          tenants.each { |tenant| with_tenant(tenant, **options) { yield tenant } }
+        end
+
+        # This method is really only intended to be used for testing.
+        def without_tenant(&block) # :nodoc:
+          with_tenant(ActiveRecord::Tenanted::Tenant::UNTENANTED_SENTINEL, prohibit_shard_swapping: false, &block)
+        end
+
+        def connection_pool(schema_version_check: true) # :nodoc:
+          if current_tenant
+            pool = retrieve_connection_pool(strict: false)
+
+            if pool.nil?
+              CONNECTION_POOL_CREATION_LOCK.synchronize do
+                # re-check now that we have the lock
+                pool = retrieve_connection_pool(strict: false)
+
+                if pool.nil?
+                  _create_tenanted_pool(schema_version_check: schema_version_check)
+                  pool = retrieve_connection_pool(strict: true)
+                end
+              end
+            end
+
+            pool
+          else
+            Tenanted::UntenantedConnectionPool.new(tenanted_root_config, self)
+          end
+        end
+
+        def tenanted_root_config # :nodoc:
+          ActiveRecord::Base.configurations.resolve(tenanted_config_name.to_sym)
+        end
+
+        def tenanted_config_name # :nodoc:
+          @tenanted_config_name ||= (superclass.respond_to?(:tenanted_config_name) ? superclass.tenanted_config_name : nil)
+        end
+
+        def _create_tenanted_pool(schema_version_check: true) # :nodoc:
+          # ensure all classes use the same connection pool
+          return superclass._create_tenanted_pool unless connection_class?
+
+          tenant = current_tenant
+          unless File.exist?(tenanted_root_config.database_path_for(tenant))
+            raise TenantDoesNotExistError, "The database file for tenant #{tenant.inspect} does not exist."
+          end
+
+          config = tenanted_root_config.new_tenant_config(tenant)
+          pool = establish_connection(config)
+
+          if schema_version_check
+            pending_migrations = pool.migration_context.open.pending_migrations
+            raise ActiveRecord::PendingMigrationError.new(pending_migrations: pending_migrations) if pending_migrations.any?
+          end
+
+          pool
+        end
+
+        private
+          def retrieve_connection_pool(strict:)
+            connection_handler.retrieve_connection_pool(connection_specification_name,
+                                                        role: current_role,
+                                                        shard: current_tenant,
+                                                        strict: strict)
+          end
+
+          def log_tenant_tag(tenant_name, &block)
+            if Rails.application.config.active_record_tenanted.log_tenant_tag
+              Rails.logger.tagged("tenant=#{tenant_name}", &block)
+            else
+              yield
+            end
+          end
+      end
+
+      prepended do
+        self.default_shard = ActiveRecord::Tenanted::Tenant::UNTENANTED_SENTINEL
+
+        prepend TenantCommon
+      end
+
+      def tenanted?
+        true
+      end
+    end
+  end
+end

data/lib/active_record/tenanted/tenant_selector.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+require "rack/contrib"
+
+module ActiveRecord
+  module Tenanted
+    #
+    # If config.active_record_tenanted.connection_class is set, this middleware will be loaded
+    # automatically, and will use config.active_record_tenanted.tenant_resolver to determine the
+    # appropriate tenant for the request.
+    #
+    # If no tenant is resolved, the request will be executed without wrapping it in a tenanted
+    # context. Application code will be free to set the tenant as needed.
+    #
+    # If a tenant is resolved and the tenant exists, the application will be locked to that
+    # tenant's database for the duration of the request.
+    #
+    # If a tenant is resolved, but the tenant does not exist, a 404 response will be returned.
+    #
+    class TenantSelector
+      attr_reader :app
+
+      def initialize(app)
+        @app = app
+      end
+
+      def call(env)
+        request = ActionDispatch::Request.new(env)
+        tenant_name = tenant_resolver.call(request)
+
+        if tenant_name.blank?
+          # run the request without wrapping it in a tenanted context
+          @app.call(env)
+        elsif tenanted_class.tenant_exist?(tenant_name)
+          tenanted_class.with_tenant(tenant_name) { @app.call(env) }
+        else
+          Rails.logger.info("ActiveRecord::Tenanted::TenantSelector: Tenant not found: #{tenant_name.inspect}")
+          Rack::NotFound.new(Rails.root.join("public/404.html")).call(env)
+        end
+      end
+
+      def tenanted_class
+        # Note: we'll probably want to cache this when we look at performance, but don't cache it
+        # when class reloading is enabled.
+        tenanted_class_name.constantize
+      end
+
+      def tenanted_class_name
+        @tenanted_class_name ||= Rails.application.config.active_record_tenanted.connection_class
+      end
+
+      def tenant_resolver
+        @tenanted_resolver ||= Rails.application.config.active_record_tenanted.tenant_resolver
+      end
+    end
+  end
+end

data/lib/active_record/tenanted/testing.rb

@@ -0,0 +1,121 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  module Tenanted
+    module Testing
+      module ActiveSupportTestCase
+        extend ActiveSupport::Concern
+
+        prepended do
+          if klass = ActiveRecord::Tenanted.connection_class
+            klass.current_tenant = Rails.application.config.active_record_tenanted.default_tenant
+            tenant_genesis(klass)
+
+            parallelize_setup do |worker|
+              # free up the connection pool for the tenant name
+              klass.destroy_tenant(klass.current_tenant)
+
+              # destroy and create tenant databases for this worker's unique id
+              klass.tenanted_root_config.test_worker_id = worker
+              tenant_genesis(klass)
+            end
+          end
+        end
+
+        class_methods do
+          private
+            # Destroy any existing tenants from the last test run, and create a fresh tenant
+            # database for the current tenant.
+            #
+            # Yes, this is a Star Trek reference.
+            def tenant_genesis(klass)
+              klass.tenants.each { |tenant| klass.destroy_tenant(tenant) }
+              klass.create_tenant(klass.current_tenant)
+            end
+        end
+      end
+
+      module ActionDispatchIntegrationTest
+        extend ActiveSupport::Concern
+
+        prepended do
+          setup do
+            if klass = ActiveRecord::Tenanted.connection_class
+              integration_session.host = "#{klass.current_tenant}.example.com"
+            end
+          end
+        end
+      end
+
+      module ActionDispatchIntegrationSession
+        def process(...)
+          if klass = ActiveRecord::Tenanted.connection_class
+            klass.without_tenant { super }
+          else
+            super
+          end
+        end
+      end
+
+      module ActionDispatchSystemTestCase
+        extend ActiveSupport::Concern
+
+        prepended do
+          setup do
+            if klass = ActiveRecord::Tenanted.connection_class
+              self.default_url_options = { host: "#{klass.current_tenant}.example.localhost" }
+            end
+          end
+        end
+      end
+
+      module ActiveRecordFixtures
+        def transactional_tests_for_pool?(pool)
+          config = pool.db_config
+
+          # Prevent the tenanted RootConfig from creating transactional fixtures on an unnecessary
+          # database, which would result in sporadic locking errors.
+          is_root_config = config.instance_of?(Tenanted::DatabaseConfigurations::RootConfig)
+
+          # Any tenanted database that isn't the default test fixture database should not be wrapped
+          # in a transaction, for a couple of reasons:
+          #
+          # 1. we migrate the database using a temporary pool, which will wrap the schema load in a
+          #    transaction that will not be visible to any connection used by the code under test to
+          #    insert data.
+          # 2. having an open transaction will prevent the test from being able to destroy the tenant.
+          is_non_default_tenant = (
+            config.instance_of?(Tenanted::DatabaseConfigurations::TenantConfig) &&
+            config.tenant != Rails.application.config.active_record_tenanted.default_tenant.to_s
+          )
+
+          return false if is_root_config || is_non_default_tenant
+
+          super
+        end
+      end
+
+      module ActiveJobTestCase
+        def perform_enqueued_jobs(...)
+          if klass = ActiveRecord::Tenanted.connection_class
+            klass.without_tenant { super }
+          else
+            super
+          end
+        end
+      end
+
+      module ActionCableTestCase
+        def connect(path = ActionCable.server.config.mount_path, **request_params)
+          if (klass = ActiveRecord::Tenanted.connection_class) && klass.current_tenant
+            env = request_params.fetch(:env, {})
+            env["HTTP_HOST"] ||= "#{klass.current_tenant}.example.com"
+            request_params[:env] = env
+          end
+
+          super
+        end
+      end
+    end
+  end
+end

data/lib/active_record/tenanted/untenanted_connection_pool.rb

@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  module Tenanted
+    # In an untenanted context, instances of this class are returned by `Tenant.connection_pool`.
+    #
+    # Many places in Rails assume that `.connection_pool` can be called and will return an object,
+    # and so we can't just raise an exception if it's called while untenanted.
+    #
+    # Instead, this class exists to provide a minimal set of features that don't need a database
+    # connection, and that will raise if a connection is attempted.
+    class UntenantedConnectionPool < ActiveRecord::ConnectionAdapters::NullPool # :nodoc:
+      attr_reader :db_config
+
+      def initialize(db_config, model)
+        super()
+
+        @db_config = db_config
+        @model = model
+      end
+
+      def schema_reflection
+        schema_cache_path = ActiveRecord::Tasks::DatabaseTasks.cache_dump_filename(db_config)
+        ActiveRecord::ConnectionAdapters::SchemaReflection.new(schema_cache_path)
+      end
+
+      def schema_cache
+        ActiveRecord::ConnectionAdapters::BoundSchemaReflection.new(schema_reflection, self)
+      end
+
+      def lease_connection(...)
+        raise Tenanted::NoTenantError, "Cannot connect to a tenanted database while untenanted (#{@model})."
+      end
+
+      def checkout(...)
+        raise Tenanted::NoTenantError, "Cannot connect to a tenanted database while untenanted (#{@model})."
+      end
+
+      def with_connection(...)
+        raise Tenanted::NoTenantError, "Cannot connect to a tenanted database while untenanted (#{@model})."
+      end
+
+      def new_connection(...)
+        raise Tenanted::NoTenantError, "Cannot connect to a tenanted database while untenanted (#{@model})."
+      end
+    end
+  end
+end

data/lib/{activerecord → active_record}/tenanted/version.rb

@@ -1,7 +1,7 @@
 # frozen_string_literal: true
 
-module Activerecord
+module ActiveRecord
   module Tenanted
-    VERSION = "0.1.0"
+    VERSION = "0.2.0"
   end
 end

data/lib/active_record/tenanted.rb

@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+require "active_record"
+
+require "zeitwerk"
+loader = Zeitwerk::Loader.for_gem_extension(ActiveRecord)
+loader.setup
+
+module ActiveRecord
+  module Tenanted
+    # Base exception class for the library.
+    class Error < StandardError; end
+
+    # Raised when database access is attempted without a current tenant having been set.
+    class NoTenantError < Error; end
+
+    # Raised when database access is attempted on a record whose tenant does not match the current tenant.
+    class WrongTenantError < Error; end
+
+    # Raised when attempting to locate a GlobalID without a tenant.
+    class MissingTenantError < Error; end
+
+    # Raised when attempting to create a tenant that already exists.
+    class TenantExistsError < Error; end
+
+    # Raised when attempting to create a tenant with illegal characters in it.
+    class BadTenantNameError < Error; end
+
+    # Raised when the application's tenant configuration is invalid.
+    class TenantConfigurationError < Error; end
+
+    # Raised when implicit creation is disabled and a tenant is referenced that does not exist
+    class TenantDoesNotExistError < Error; end
+
+    # Raised when the Rails integration is being invoked but has not been configured.
+    class IntegrationNotConfiguredError < Error; end
+
+    def self.connection_class
+      # TODO: cache this / speed this up
+      Rails.application.config.active_record_tenanted.connection_class&.constantize
+    end
+  end
+end
+
+loader.eager_load
+
+ActiveSupport.run_load_hooks :active_record_tenanted, ActiveRecord::Tenanted

data/lib/activerecord-tenanted.rb

@@ -0,0 +1,3 @@
+# frozen_string_literal: true
+
+require_relative "active_record/tenanted"

data/lib/tasks/active_record/tenanted_tasks.rake

@@ -0,0 +1,78 @@
+# frozen_string_literal: true
+
+namespace :db do
+  desc "Migrate the database for tenant ARTENANT"
+  task "migrate:tenant" => "load_config" do
+    unless ActiveRecord::Tenanted.connection_class
+      warn "ActiveRecord::Tenanted integration is not configured via connection_class"
+      next
+    end
+
+    unless ActiveRecord::Tenanted::DatabaseTasks.root_database_config
+      warn "WARNING: No tenanted database found, skipping tenanted migration"
+    else
+      begin
+        verbose_was = ActiveRecord::Migration.verbose
+        ActiveRecord::Migration.verbose = ActiveRecord::Tenanted::DatabaseTasks.verbose?
+
+        ActiveRecord::Tenanted::DatabaseTasks.migrate_tenant
+      ensure
+        ActiveRecord::Migration.verbose = verbose_was
+      end
+    end
+  end
+
+  desc "Migrate the database for all existing tenants"
+  task "migrate:tenant:all" => "load_config" do
+    unless ActiveRecord::Tenanted.connection_class
+      warn "ActiveRecord::Tenanted integration is not configured via connection_class"
+      next
+    end
+
+    verbose_was = ActiveRecord::Migration.verbose
+    ActiveRecord::Migration.verbose = ActiveRecord::Tenanted::DatabaseTasks.verbose?
+
+    ActiveRecord::Tenanted::DatabaseTasks.migrate_all
+  ensure
+    ActiveRecord::Migration.verbose = verbose_was
+  end
+
+  desc "Drop and recreate all tenant databases from their schema for the current environment"
+  task "reset:tenant" => [ "db:drop:tenant", "db:migrate:tenant" ]
+
+  desc "Drop all tenanted databases for the current environment"
+  task "drop:tenant" => "load_config" do
+    unless ActiveRecord::Tenanted::DatabaseTasks.root_database_config
+      warn "WARNING: No tenanted database found, skipping tenanted reset"
+    else
+      begin
+        verbose_was = ActiveRecord::Migration.verbose
+        ActiveRecord::Migration.verbose = ActiveRecord::Tenanted::DatabaseTasks.verbose?
+
+        ActiveRecord::Tenanted::DatabaseTasks.drop_all
+      ensure
+        ActiveRecord::Migration.verbose = verbose_was
+      end
+    end
+  end
+
+  desc "Set the current tenant to ARTENANT if present, else the environment default"
+  task "tenant" => "load_config" do
+    unless ActiveRecord::Tenanted.connection_class
+      warn "ActiveRecord::Tenanted integration is not configured via connection_class"
+      next
+    end
+
+    ActiveRecord::Tenanted::DatabaseTasks.set_current_tenant
+  end
+end
+
+# Decorate database tasks with the tenanted version.
+task "db:migrate" => "db:migrate:tenant:all"
+task "db:prepare" => "db:migrate:tenant:all"
+task "db:reset"   => "db:reset:tenant"
+task "db:drop"    => "db:drop:tenant"
+
+# Ensure a default tenant is set for database tasks that may need it.
+task "db:fixtures:load" => "db:tenant"
+task "db:seed"          => "db:tenant"