ec-pg 0.1.0

data/lib/ec/pg/rls_mixin.rb

@@ -0,0 +1,80 @@
+# frozen_string_literal: true
+
+module Ec 
+  module Pg
+    # Mixin for ActiveRecord models that participate in multi-tenancy.
+    #
+    # Include via +acts_as_rls+ in your model or add it to ApplicationRecord.
+    #
+    # == Usage
+    #
+    #   class ApplicationRecord < ActiveRecord::Base
+    #     include Ec::Pg::RlsMixin
+    #     acts_as_rls
+    #   end
+    #
+    # == What it does
+    #
+    # * Adds an +around_action+-compatible callback that applies RLS before
+    #   every query on the model when a tenant is active in the thread context.
+    # * Provides a +with_rls+ class method for explicit scoping.
+    # * Optionally raises TenantNotSet when required: true and no tenant is set.
+    #
+    module RlsMixin
+      class RlsError < StandardError; end
+
+      extend ActiveSupport::Concern
+      RlsModes = %i(local session)
+
+      included do 
+        class_attribute :is_rls, default: false
+        class_attribute :registered_variables, default: {}
+        class_attribute :rls_mode, default: :local
+      end
+
+      class_methods do
+        def acts_as_rls(mode: nil, variables: {})
+          validate_variables!(variables)
+
+          self.is_rls = true
+          self.rls_mode = mode || Ec::Pg.configuration.rls_mode || :local
+          self.registered_variables = variables
+        end
+
+        # Executes +block+ with the AR connection switched to +shard_name+.
+        #
+        # @param shard_name [Symbol] the registered shard key
+        # @param role       [Symbol] :writing (default) or :reading
+        # @param klass      [Class]  the AR base class whose connection pool to switch
+        #                            (default: ActiveRecord::Base)
+        # @yield  block to run in shard context
+        # @return the return value of +block+
+        def with_rls(variables: {}, &block)
+          RlsManager.with_rls(
+            rls_mode: self.rls_mode, 
+            registered_variables: self.registered_variables,
+            variables: variables, 
+            connection: self.connection, 
+            &block
+          )
+        end
+
+        def clear!
+          self.registered_variables = {}
+        end
+
+        private
+        # Prevents variable name injection. Only allows dot-separated identifiers.
+        ValidVariablesRegexp = /\A[a-zA-Z_][a-zA-Z0-9_.]*\z/
+
+        def validate_variables!(variables)
+          variables.each do |key, variable|
+            unless ValidVariablesRegexp.match?(variable)
+              raise RlsError, "Invalid RLS variable name: #{key}: #{variable.inspect}" 
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/ec/pg/schema_manager.rb

@@ -0,0 +1,126 @@
+# frozen_string_literal: true
+
+module Ec
+  module Pg
+    # Manages PostgreSQL schema switching via the +search_path+ session variable.
+    #
+    # When you have one Postgres database with a schema-per-tenant layout:
+    #
+    #   CREATE SCHEMA tenant_abc;
+    #   CREATE TABLE tenant_abc.records (...);
+    #
+    # This manager updates the connection's search_path so that ActiveRecord's
+    # unqualified table references resolve to the correct tenant schema.
+    #
+    # == Usage
+    #
+    #   SchemaManager.with_schema("tenant_abc") { Record.all }
+    #
+    # The search_path is reset to the previous value (or the configured default)
+    # when the block exits, even if an exception is raised.
+    #
+    # == Thread safety
+    #
+    # +with_schema+ stores the schema name in the thread context AND updates the
+    # underlying database connection.  When using connection pooling, each thread
+    # checks out its own connection, so search_path changes are thread-safe.
+    #
+    module SchemaManager
+      class InvalidSchema < StandardError; end
+
+      module_function
+
+      # Executes +block+ with the Postgres search_path set to +schema_name+,
+      # followed by any +shared_schemas+ from the configuration.
+      #
+      # @param schema_name [String] tenant schema (e.g. "tenant_abc")
+      # @yield  block to execute within the schema context
+      # @return the return value of the block
+      def with_schema(schema_name, &block)
+        schema_name = schema_name.to_s.strip
+        validate_schema_name!(schema_name)
+
+        Context.with(schema: schema_name) do
+          apply_schema(schema_name)
+          block.call
+        end
+      ensure
+        # Restore previous schema on the same connection if we changed it.
+        # Guard against the case where the connection was already returned to
+        # the pool during the block.
+        restore_schema
+      end
+
+      # Returns the schema currently active in thread context.
+      def current_schema
+        Context.schema || Ec::Pg.configuration.default_schema
+      end
+
+      # Returns true if a non-default schema is active.
+      def active?
+        !Context.schema.nil?
+      end
+
+      # Applies the search_path on an existing connection without block scoping.
+      # Useful in middleware or Rack apps that manage connection lifecycle manually.
+      def apply!(schema_name)
+        schema_name = schema_name.to_s.strip
+        validate_schema_name!(schema_name)
+
+        Context.set(schema: schema_name)
+        apply_schema(schema_name)
+      end
+
+      # Resets search_path to the configured default schema on +connection+.
+      def reset!
+        Context.delete(:schema)
+        restore_schema
+      end
+
+      private
+
+      # Builds the search_path string and executes SET search_path on the connection.
+      def apply_schema(schema_name)
+        registered_connections.each do |connection|
+          shared = Ec::Pg.configuration.shared_schemas
+          path   = ([schema_name] + shared).uniq.join(", ")
+          connection.schema_search_path = path
+        end
+      end
+
+      # Restores the search_path to the configured default schema(s).
+      def restore_schema
+        default = Ec::Pg.configuration.default_schema
+        shared = Ec::Pg.configuration.shared_schemas
+        path   = ([default] + shared).uniq.join(", ")
+
+        registered_connections.each do |connection|
+          if connection && connection.pool.present?
+            connection.schema_search_path = path
+          end
+        end
+      rescue StandardError
+        # Swallow errors during cleanup (connection may have been returned to pool)
+      end
+
+      def registered_connections
+        Ec::Pg::SchemaMixin
+          .registered_models
+          .map(&:connection)
+      end
+
+      # Guards against SQL injection via schema names.
+      ValidSchemaRegex = /\A[a-zA-Z][a-zA-Z0-9_$]*\z/
+
+      def validate_schema_name!(name)
+        return if ValidSchemaRegex.match?(name)
+
+        raise InvalidSchema,
+              "Invalid PostgreSQL schema name: #{name.inspect}. " \
+              "Only alphanumeric characters and underscores are allowed."
+      end
+
+      module_function :apply_schema, :restore_schema, :validate_schema_name!, :registered_connections
+    end
+  end
+end

data/lib/ec/pg/schema_mixin.rb

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+module Ec 
+  module Pg
+    # Mixin for ActiveRecord models that participate in multi-tenancy.
+    #
+    # Include via +acts_as_sharded+ in your model or add it to ApplicationRecord.
+    #
+    # == Usage
+    #
+    #   class ApplicationRecord < ActiveRecord::Base
+    #     include Ec::Pg::SchemaMixin
+    #     acts_as_namespaced
+    #   end
+    #
+    # == What it does
+    #
+    # * Adds an +around_action+-compatible callback that applies RLS before
+    #   every query on the model when a tenant is active in the thread context.
+    # * Provides a +with_tenant+ class method for explicit scoping.
+    # * Optionally raises TenantNotSet when required: true and no tenant is set.
+    #
+    module SchemaMixin
+      extend ActiveSupport::Concern
+      mattr_accessor :registered_models, default: []
+      included do 
+        class_attribute :is_namespaced, default: false
+      end
+
+      class_methods do
+        def acts_as_namespaced
+          self.is_namespaced = true
+
+          register(self)
+        end
+
+        def registered_models 
+          Ec::Pg::SchemaMixin.registered_models
+        end
+
+        def clear!
+          Ec::Pg::SchemaMixin.registered_models = []
+        end
+
+        private
+        def register(klass)
+          Ec::Pg::SchemaMixin.registered_models += [klass]
+        end
+      end
+    end
+  end
+end

data/lib/ec/pg/shard_manager.rb

@@ -0,0 +1,94 @@
+# frozen_string_literal: true
+module Ec
+  module Pg
+    # Manages ActiveRecord database shard switching.
+    #
+    # Leverages ActiveRecord's native multi-database support (AR 6.1+) via
+    # +connected_to(shard:)+.  Falls back to manual connection swapping for
+    # connection configurations that are not registered as named shards.
+    #
+    # == Configuration in database.yml (Rails multi-db style)
+    #
+    #   production:
+    #     primary:
+    #       <<: *default
+    #       database: app_primary
+    #     db_sharded_shard_1:
+    #       <<: *default
+    #       database: app_shard_one
+    #       migrations_paths: db/migrate_shards
+    #
+    # == Usage
+    #
+    #   ShardManager.with_shard(:shard_one) { User.all }
+    #   ShardManager.with_shard(:shard_one, role: :reading) { User.all }
+    #
+    module ShardManager
+      class ShardNotFound < StandardError; end
+      class UnsupportedActiveRecordVersion < StandardError; end
+
+      MinimumARVersion = Gem::Version.new("7.1")
+
+      module_function
+
+      # Executes +block+ with the AR connection switched to +shard_name+.
+      #
+      # @param shard_name [Symbol] the registered shard key
+      # @param role       [Symbol] :writing (default) or :reading
+      # @param klass      [Class]  the AR base class whose connection pool to switch
+      #                            (default: ActiveRecord::Base)
+      # @yield  block to run in shard context
+      # @return the return value of +block+
+      def with_shard(shard_name, role: :writing, klass: ActiveRecord::Base, &block)
+        shard_name = shard_name.to_sym
+        assert_ar_version!
+
+        Context.with(shard: shard_name) do
+          klass.prohibit_shard_swapping(true) do
+            klass.connected_to(shard: shard_name, role: role, &block)
+          end
+        end
+      rescue ActiveRecord::ConnectionNotEstablished => e
+        raise ShardNotFound,
+              "Could not connect to shard #{shard_name.inspect}. " \
+              "Make sure it is registered via connects_to. Original error: #{e.message}"
+      end
+
+      # Returns true when the thread is operating inside a +with_shard+ block
+      # (or when the shard was set via thread-local assignment).
+      def active?
+        !Context.shard.nil?
+      end
+
+      # Returns the shard currently set in thread context, or +:default+ when none
+      # is set (mirrors how AR treats the default shard).
+      def current_shard
+        Context.shard || :default
+      end
+
+      # Returns an array of shard names registered on +klass+.
+      # def registered_shards(klass: ActiveRecord::Base)
+      #   klass
+      #     .connection_handler
+      #     .connection_pools
+      #     .map(&:db_config)
+      #     .map(&:name)
+      # rescue StandardError
+      #   []
+      # end
+
+      private
+
+      def assert_ar_version!
+        ar_version = Gem::Version.new(ActiveRecord.version.to_s)
+        return if ar_version >= MinimumARVersion
+
+        raise UnsupportedActiveRecordVersion,
+              "activerecord-multi-tenant shard switching requires ActiveRecord >= 6.1. " \
+              "Current version: #{ar_version}"
+      end
+
+      module_function :assert_ar_version!
+    end
+  end
+end

data/lib/ec/pg/shard_mixin.rb

@@ -0,0 +1,74 @@
+# frozen_string_literal: true
+
+module Ec 
+  module Pg
+    # Mixin for ActiveRecord models that participate in multi-tenancy.
+    #
+    # Include via +acts_as_sharded+ in your model or add it to ApplicationRecord.
+    #
+    # == Usage
+    #
+    #   class ApplicationRecord < ActiveRecord::Base
+    #     include Ec::Pg::ShardMixin
+    #     acts_as_sharded
+    #   end
+    #
+    # == What it does
+    #
+    # * Adds an +around_action+-compatible callback that applies RLS before
+    #   every query on the model when a tenant is active in the thread context.
+    # * Provides a +with_tenant+ class method for explicit scoping.
+    # * Optionally raises TenantNotSet when required: true and no tenant is set.
+    #
+    module ShardMixin
+      extend ActiveSupport::Concern
+
+      included do 
+        class_attribute :is_sharded,   default: false
+      end
+
+      class_methods do
+        # Declares this model as sharded
+        #
+        # @param mode   [Symbol] :solo, or :sharded 
+        # @param writing_database_identifier [Symbol or String] database used for writing
+        # @param reading_database_identifier [Symbol or String] database used for reading (optional)
+        def acts_as_sharded(mode:, writing_database_identifier:, reading_database_identifier: nil)
+          class_attribute :tenant_mode,  default: mode
+          class_attribute :writing_database_identifier,   default: writing_database_identifier
+          class_attribute :reading_database_identifier,   default: reading_database_identifier
+
+          # include InstanceMethods
+          # extend  QueryMethods
+
+          self.is_sharded = true
+          self.abstract_class = true
+          self.connects_to(shards: shards_hash)
+        end
+
+        private
+        def shards_hash
+          {}.tap do |hash|
+            Ec::Pg.configuration.number_of_shards.times do |index|
+              shard_key = "shard_#{index + 1}"
+              hash[shard_key] = {
+                writing: database_identifier_for(writing_database_identifier, index),
+                reading: database_identifier_for(reading_database_identifier, index)
+              }.compact
+            end
+          end.symbolize_keys
+        end
+
+        def database_identifier_for(database_identifier, index)
+          return unless database_identifier.present?
+
+          if tenant_mode.to_sym == :solo
+            database_identifier
+          else
+            [database_identifier, "shard", index + 1].join('_')
+          end.to_sym
+        end
+      end
+    end
+  end
+end

data/lib/ec/pg/tenant_context.rb

@@ -0,0 +1,73 @@
+# frozen_string_literal: true
+
+module Ec
+  module Pg
+    module TenantContext
+      # Raised when require_tenant is true and a query is attempted without a tenant.
+      class TenantNotSet < StandardError; end
+
+      module_function
+
+      # -------------------------------------------------------------------------
+      # Block-scoped (preferred)
+      # -------------------------------------------------------------------------
+
+      # Executes +block+ within the given tenant context.
+      #
+      # @param shard  [Symbol, nil] shard to connect to (skipped when nil)
+      # @param schema [String, nil] Postgres schema to set (skipped when nil)
+      def switch(shard: nil, schema: nil, &block)
+
+        # Layer composition: innermost layer first so that the outermost wrapper
+        # controls the connection that the inner layers receive.
+
+        run = block
+
+        # 1. Schema (innermost)
+        if schema
+          inner_run = run
+          run = -> {SchemaManager.with_schema(schema, &inner_run)}
+        end
+
+        # 2. Shard (outermost — switches the connection first)
+        if shard
+          inner_run = run
+          run = -> {ShardManager.with_shard(shard, &inner_run)}
+        end
+
+        if run.present?
+          run.call
+        else
+          raise TenantNotSet, "No schema or shard specified"
+        end
+      end
+
+      # Convenience alias that shard and only switches schema.
+      def with_schema(schema_name, &block)
+        SchemaManager.with_schema(schema_name, &block)
+      end
+
+      # Convenience alias that skips schema and only switches shard.
+      def with_shard(shard_name, role: :writing, &block)
+        ShardManager.with_shard(shard_name, role: role, &block)
+      end
+
+      # -------------------------------------------------------------------------
+      # Thread-local (stateful)
+      # -------------------------------------------------------------------------
+
+      def current_shard
+        ShardManager.current_shard
+      end
+
+      def current_schema
+        SchemaManager.current_schema
+      end
+
+      # Clears all multi-tenant context for the current thread.
+      def clear_context!
+        Context.clear!
+      end
+    end
+  end
+end

data/lib/ec/pg/version.rb

@@ -0,0 +1,7 @@
+# frozen_string_literal: true
+
+module Ec
+  module Pg
+    VERSION = "0.1.0"
+  end
+end

data/lib/ec/pg.rb

@@ -0,0 +1,53 @@
+# frozen_string_literal: true
+
+require "active_record"
+require "active_support/concern"
+require_relative "pg/version"
+require_relative "pg/context"
+require_relative "pg/tenant_context"
+require_relative "pg/configuration"
+
+if ENV['RAILS_ENV'] == 'test'
+  require 'debug'
+end
+
+if defined?(Rails::Railtie)
+  require_relative "pg/railtie"
+end
+
+I18n.load_path += Dir.glob( File.expand_path("../../../config/locales/*.yml", __FILE__) )
+
+module Ec
+  module Pg
+    class Error < StandardError; end
+    class InvalidType < Error; end
+
+    autoload :SchemaMixin, 'ec/pg/schema_mixin.rb'
+    autoload :SchemaManager, 'ec/pg/schema_manager.rb'
+    autoload :ShardMixin, 'ec/pg/shard_mixin.rb'
+    autoload :ShardManager, 'ec/pg/shard_manager.rb'
+    autoload :RlsMixin, 'ec/pg/rls_mixin.rb'
+    autoload :RlsManager, 'ec/pg/rls_manager.rb'
+    autoload :ContextSwitcher, 'ec/pg/middleware/context_switcher.rb'
+
+    module_function 
+
+    def configure
+      yield configuration
+    end
+
+    def configuration
+      @configuration ||= Ec::Pg::Configuration.new
+    end
+
+    def switch(shard: nil, schema: nil, &block)
+      TenantContext.switch(shard: shard, schema: schema, &block)
+    end
+
+    def current_shard = TenantContext.current_shard
+    def current_schema = TenantContext.current_schema
+
+    # Clears all thread-local context.
+    def clear_context! = TenantContext.clear_context!
+  end
+end

data/sig/ec/pg.rbs

@@ -0,0 +1,6 @@
+module Ec
+  module Pg
+    VERSION: String
+    # See the writing guide of rbs: https://github.com/ruby/rbs#guides
+  end
+end