switchman 1.5.21 → 2.1.5

data/lib/switchman/active_record/migration.rb

@@ -0,0 +1,80 @@
+# frozen_string_literal: true
+
+module Switchman
+  module ActiveRecord
+    module Migration
+      module Compatibility
+        module V5_0
+          def create_table(*args, **options)
+            unless options.key?(:id)
+              options[:id] = :bigserial
+            end
+            if block_given?
+              super do |td|
+                yield td
+              end
+            else
+              super
+            end
+          end
+        end
+      end
+
+      def connection
+        conn = super
+        if conn.shard != ::ActiveRecord::Base.connection_pool.current_pool.shard
+          ::ActiveRecord::Base.connection_pool.current_pool.switch_database(conn)
+        end
+        conn
+      end
+    end
+
+    module Migrator
+      # significant change: hash shard id, not database name
+      def generate_migrator_advisory_lock_id
+        shard_name_hash = Zlib.crc32(Shard.current.name)
+        ::ActiveRecord::Migrator::MIGRATOR_SALT * shard_name_hash
+      end
+
+      if ::Rails.version >= '6.0'
+        # copy/paste from Rails 6.1
+        def with_advisory_lock
+          lock_id = generate_migrator_advisory_lock_id
+
+          with_advisory_lock_connection do |connection|
+            got_lock = connection.get_advisory_lock(lock_id)
+            raise ::ActiveRecord::ConcurrentMigrationError unless got_lock
+            load_migrated # reload schema_migrations to be sure it wasn't changed by another process before we got the lock
+            yield
+          ensure
+            if got_lock && !connection.release_advisory_lock(lock_id)
+              raise ::ActiveRecord::ConcurrentMigrationError.new(
+                ::ActiveRecord::ConcurrentMigrationError::RELEASE_LOCK_FAILED_MESSAGE
+              )
+            end
+          end
+        end
+
+        # significant change: strip out prefer_secondary from config
+        def with_advisory_lock_connection
+          pool = ::ActiveRecord::ConnectionAdapters::ConnectionHandler.new.establish_connection(
+            ::ActiveRecord::Base.connection_config.except(:prefer_secondary)
+          )
+
+          pool.with_connection { |connection| yield(connection) }
+        ensure
+          pool&.disconnect!
+        end
+      end
+    end
+
+    module MigrationContext
+      def migrations
+        return @migrations if instance_variable_defined?(:@migrations)
+        migrations_cache = Thread.current[:migrations_cache] ||= {}
+        key = Digest::MD5.hexdigest(migration_files.sort.join(','))
+        @migrations = migrations_cache[key] ||= super
+      end
+    end
+  end
+end

data/lib/switchman/active_record/model_schema.rb

@@ -1,10 +1,12 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveRecord
     module ModelSchema
       module ClassMethods
         def quoted_table_name
           @quoted_table_name ||= {}
-          @quoted_table_name[Shard.current.id] ||= connection.quote_table_name(table_name)
+          @quoted_table_name[Shard.current(shard_category).id] ||= connection.quote_table_name(table_name)
         end
       end
     end

data/lib/switchman/active_record/persistence.rb

@@ -1,10 +1,18 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveRecord
     module Persistence
       # touch reads the id attribute directly, so it's not relative to the current shard
-      def touch(*)
+      def touch(*, **)
         shard.activate(self.class.shard_category) { super }
       end
+
+      if ::Rails.version >= '5.2'
+        def update_columns(*)
+          shard.activate(self.class.shard_category) { super }
+        end
+      end
     end
   end
 end

data/lib/switchman/active_record/postgresql_adapter.rb

@@ -1,10 +1,8 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveRecord
     module PostgreSQLAdapter
-      def self.prepended(klass)
-        klass::NATIVE_DATABASE_TYPES[:primary_key] = "bigserial primary key".freeze
-      end
-
       # copy/paste; use quote_local_table_name
       def create_database(name, options = {})
         options = { encoding: 'utf8' }.merge!(options.symbolize_keys)
@@ -42,149 +40,119 @@ module Switchman
         select_values("SELECT * FROM unnest(current_schemas(false))")
       end
 
-      def use_qualified_names?
-        @config[:use_qualified_names]
+      def extract_schema_qualified_name(string)
+        name = ::ActiveRecord::ConnectionAdapters::PostgreSQL::Utils.extract_schema_qualified_name(string.to_s)
+        if string && !name.schema
+          name.instance_variable_set(:@schema, shard.name)
+        end
+        [name.schema, name.identifier]
       end
 
-      def tables(name = nil)
-        schema = shard.name if use_qualified_names?
-
-        query(<<-SQL, 'SCHEMA').map { |row| row[0] }
-          SELECT tablename
-          FROM pg_tables
-          WHERE schemaname = #{schema ? "'#{schema}'" : 'ANY (current_schemas(false))'}
-        SQL
+      # significant change: use the shard name if no explicit schema
+      def quoted_scope(name = nil, type: nil)
+        schema, name = extract_schema_qualified_name(name)
+        type = \
+          case type
+          when "BASE TABLE"
+            "'r','p'"
+          when "VIEW"
+            "'v','m'"
+          when "FOREIGN TABLE"
+            "'f'"
+          end
+        scope = {}
+        scope[:schema] = quote(schema || shard.name)
+        scope[:name] = quote(name) if name
+        scope[:type] = type if type
+        scope
       end
 
-      def table_exists?(name)
-        if ::Rails.version < '4.2'
-          schema, table = ::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter::Utils.extract_schema_and_table(name.to_s)
-          return false unless table
-          schema ||= shard.name if use_qualified_names?
-
-          binds = [[nil, table]]
-          binds << [nil, schema] if schema
-
-          exec_query(<<-SQL, 'SCHEMA').rows.first[0].to_i > 0
-              SELECT COUNT(*)
-              FROM pg_class c
-              LEFT JOIN pg_namespace n ON n.oid = c.relnamespace
-              WHERE c.relkind in ('v','r')
-              AND c.relname = '#{table.gsub(/(^"|"$)/,'')}'
-              AND n.nspname = #{schema ? "'#{schema}'" : 'ANY (current_schemas(false))'}
+      if ::Rails.version < '6.0'
+        def tables(name = nil)
+          query(<<-SQL, 'SCHEMA').map { |row| row[0] }
+            SELECT tablename
+            FROM pg_tables
+            WHERE schemaname = '#{shard.name}'
           SQL
-        else
-          name =  ::ActiveRecord::ConnectionAdapters::PostgreSQL::Utils.extract_schema_qualified_name(name.to_s)
+        end
+
+        def view_exists?(name)
+          name = ::ActiveRecord::ConnectionAdapters::PostgreSQL::Utils.extract_schema_qualified_name(name.to_s)
           return false unless name.identifier
-          if !name.schema && use_qualified_names?
+          if !name.schema
             name.instance_variable_set(:@schema, shard.name)
           end
 
-          exec_query(<<-SQL, 'SCHEMA').rows.first[0].to_i > 0
-              SELECT COUNT(*)
-              FROM pg_class c
-              LEFT JOIN pg_namespace n ON n.oid = c.relnamespace
-              WHERE c.relkind IN ('r','v','m') -- (r)elation/table, (v)iew, (m)aterialized view
-              AND c.relname = '#{name.identifier}'
-              AND n.nspname = #{name.schema ? "'#{name.schema}'" : 'ANY (current_schemas(false))'}
+          select_values(<<-SQL, 'SCHEMA').any?
+            SELECT c.relname
+            FROM pg_class c
+            LEFT JOIN pg_namespace n ON n.oid = c.relnamespace
+            WHERE c.relkind IN ('v','m') -- (v)iew, (m)aterialized view
+            AND c.relname = '#{name.identifier}'
+            AND n.nspname = '#{shard.name}'
           SQL
         end
-      end
-
-      def indexes(table_name)
-        schema = shard.name if use_qualified_names?
-
-        result = query(<<-SQL, 'SCHEMA')
-           SELECT distinct i.relname, d.indisunique, d.indkey, pg_get_indexdef(d.indexrelid), t.oid
-           FROM pg_class t
-           INNER JOIN pg_index d ON t.oid = d.indrelid
-           INNER JOIN pg_class i ON d.indexrelid = i.oid
-           WHERE i.relkind = 'i'
-             AND d.indisprimary = 'f'
-             AND t.relname = '#{table_name}'
-             AND i.relnamespace IN (SELECT oid FROM pg_namespace WHERE nspname = #{schema ? "'#{schema}'" : 'ANY (current_schemas(false))'} )
-          ORDER BY i.relname
-        SQL
-
-
-        result.map do |row|
-          index_name = row[0]
-          unique = row[1] == 't'
-          indkey = row[2].split(" ")
-          inddef = row[3]
-          oid = row[4]
-
-          columns = Hash[query(<<-SQL, "SCHEMA")]
-          SELECT a.attnum, a.attname
-          FROM pg_attribute a
-          WHERE a.attrelid = #{oid}
-          AND a.attnum IN (#{indkey.join(",")})
-          SQL
 
-          column_names = columns.stringify_keys.values_at(*indkey).compact
-
-          unless column_names.empty?
-            # add info on sort order for columns (only desc order is explicitly specified, asc is the default)
-            desc_order_columns = inddef.scan(/(\w+) DESC/).flatten
-            orders = desc_order_columns.any? ? Hash[desc_order_columns.map {|order_column| [order_column, :desc]}] : {}
-            where = inddef.scan(/WHERE (.+)$/).flatten[0]
-            using = inddef.scan(/USING (.+?) /).flatten[0].to_sym
-
-            ::ActiveRecord::ConnectionAdapters::IndexDefinition.new(table_name, index_name, unique, column_names, [], orders, where, nil, using)
-          end
-        end.compact
-      end
-
-      def index_name_exists?(table_name, index_name, default)
-        schema = shard.name if use_qualified_names?
-
-        exec_query(<<-SQL, 'SCHEMA').rows.first[0].to_i > 0
-            SELECT COUNT(*)
+        def indexes(table_name)
+          result = query(<<-SQL, 'SCHEMA')
+            SELECT distinct i.relname, d.indisunique, d.indkey, pg_get_indexdef(d.indexrelid), t.oid
             FROM pg_class t
             INNER JOIN pg_index d ON t.oid = d.indrelid
             INNER JOIN pg_class i ON d.indexrelid = i.oid
             WHERE i.relkind = 'i'
-              AND i.relname = '#{index_name}'
+              AND d.indisprimary = 'f'
               AND t.relname = '#{table_name}'
-              AND i.relnamespace IN (SELECT oid FROM pg_namespace WHERE nspname = #{schema ? "'#{schema}'" : 'ANY (current_schemas(false))'} )
-        SQL
-      end
-
-      def quote_local_table_name(name)
-        # postgres quotes tables and columns the same; just pass through
-        # (differs from quote_table_name below by no logic to explicitly
-        # qualify the table)
-        quote_column_name(name)
-      end
-
-      def quote_table_name name
-        if ::Rails.version < '4.2'.freeze
-          schema, name_part = extract_pg_identifier_from_name(name.to_s)
+              AND i.relnamespace IN (SELECT oid FROM pg_namespace WHERE nspname = '#{shard.name}' )
+            ORDER BY i.relname
+          SQL
 
-          if !name_part && use_qualified_names? && shard.name
-            schema, name_part = shard.name, schema
-          end
 
-          unless name_part
-            quote_column_name(schema)
-          else
-            table_name, name_part = extract_pg_identifier_from_name(name_part)
-            "#{quote_column_name(schema)}.#{quote_column_name(table_name)}"
-          end
-        else
+          result.map do |row|
+            index_name = row[0]
+            unique = row[1] == true || row[1] == 't'
+            indkey = row[2].split(" ")
+            inddef = row[3]
+            oid = row[4]
+
+            columns = Hash[query(<<-SQL, "SCHEMA")]
+            SELECT a.attnum, a.attname
+            FROM pg_attribute a
+            WHERE a.attrelid = #{oid}
+            AND a.attnum IN (#{indkey.join(",")})
+            SQL
+
+            column_names = columns.stringify_keys.values_at(*indkey).compact
+
+            unless column_names.empty?
+              # add info on sort order for columns (only desc order is explicitly specified, asc is the default)
+              desc_order_columns = inddef.scan(/(\w+) DESC/).flatten
+              orders = desc_order_columns.any? ? Hash[desc_order_columns.map {|order_column| [order_column, :desc]}] : {}
+              where = inddef.scan(/WHERE (.+)$/).flatten[0]
+              using = inddef.scan(/USING (.+?) /).flatten[0].to_sym
+
+              if ::Rails.version >= "5.2"
+                ::ActiveRecord::ConnectionAdapters::IndexDefinition.new(table_name, index_name, unique, column_names, orders: orders, where: where, using: using)
+              else
+                ::ActiveRecord::ConnectionAdapters::IndexDefinition.new(table_name, index_name, unique, column_names, [], orders, where, nil, using)
+              end
+            end
+          end.compact
+        end
 
-          name = ::ActiveRecord::ConnectionAdapters::PostgreSQL::Utils.extract_schema_qualified_name(name.to_s)
-          if !name.schema && use_qualified_names?
-            name.instance_variable_set(:@schema, shard.name)
-          end
-          name.quoted
+        def index_name_exists?(table_name, index_name, _default = nil)
+          exec_query(<<-SQL, 'SCHEMA').rows.first[0].to_i > 0
+              SELECT COUNT(*)
+              FROM pg_class t
+              INNER JOIN pg_index d ON t.oid = d.indrelid
+              INNER JOIN pg_class i ON d.indexrelid = i.oid
+              WHERE i.relkind = 'i'
+                AND i.relname = '#{index_name}'
+                AND t.relname = '#{table_name}'
+                AND i.relnamespace IN (SELECT oid FROM pg_namespace WHERE nspname = '#{shard.name}' )
+          SQL
         end
-      end
 
-      if ::Rails.version >= '4.2'
         def foreign_keys(table_name)
-          schema = shard.name if use_qualified_names?
-
           # mostly copy-pasted from AR - only change is to the nspname condition for qualified names support
           fk_info = select_all <<-SQL.strip_heredoc
             SELECT t2.oid::regclass::text AS to_table, a1.attname AS column, a2.attname AS primary_key, c.conname AS name, c.confupdtype AS on_update, c.confdeltype AS on_delete
@@ -196,7 +164,7 @@ module Switchman
             JOIN pg_namespace t3 ON c.connamespace = t3.oid
             WHERE c.contype = 'f'
               AND t1.relname = #{quote(table_name)}
-              AND t3.nspname = #{schema ? "'#{schema}'" : 'ANY (current_schemas(false))'}
+              AND t3.nspname = '#{shard.name}'
             ORDER BY c.conname
           SQL
 
@@ -210,9 +178,85 @@ module Switchman
             options[:on_delete] = extract_foreign_key_action(row['on_delete'])
             options[:on_update] = extract_foreign_key_action(row['on_update'])
 
-            ::ActiveRecord::ConnectionAdapters::ForeignKeyDefinition.new(table_name, row['to_table'], options)
+            # strip the schema name from to_table if it matches
+            to_table = row['to_table']
+            to_table_qualified_name = ::ActiveRecord::ConnectionAdapters::PostgreSQL::Utils.extract_schema_qualified_name(to_table)
+            if to_table_qualified_name.schema == shard.name
+              to_table = to_table_qualified_name.identifier
+            end
+
+            ::ActiveRecord::ConnectionAdapters::ForeignKeyDefinition.new(table_name, to_table, options)
           end
         end
+      else
+        def foreign_keys(table_name)
+          super.each do |fk|
+            to_table_qualified_name = ::ActiveRecord::ConnectionAdapters::PostgreSQL::Utils.extract_schema_qualified_name(fk.to_table)
+            if to_table_qualified_name.schema == shard.name
+              fk.to_table = to_table_qualified_name.identifier
+            end
+          end
+        end
+      end
+
+      def quote_local_table_name(name)
+        # postgres quotes tables and columns the same; just pass through
+        # (differs from quote_table_name_with_shard below by no logic to
+        # explicitly qualify the table)
+        quote_column_name(name)
+      end
+
+      def quote_table_name(name)
+        return quote_local_table_name(name) if @use_local_table_name
+        name = ::ActiveRecord::ConnectionAdapters::PostgreSQL::Utils.extract_schema_qualified_name(name.to_s)
+        if !name.schema
+          name.instance_variable_set(:@schema, shard.name)
+        end
+        name.quoted
+      end
+
+      def with_global_table_name(&block)
+        with_local_table_name(false, &block)
+      end
+
+      def with_local_table_name(enable = true)
+        old_value = @use_local_table_name
+        @use_local_table_name = enable
+        yield
+      ensure
+        @use_local_table_name = old_value
+      end
+
+      def add_index_options(_table_name, _column_name, **)
+        index_name, index_type, index_columns, index_options, algorithm, using = super
+        algorithm = nil if DatabaseServer.creating_new_shard && algorithm == "CONCURRENTLY"
+        [index_name, index_type, index_columns, index_options, algorithm, using]
+      end
+
+      def rename_table(table_name, new_name)
+        clear_cache!
+        execute "ALTER TABLE #{quote_table_name(table_name)} RENAME TO #{quote_local_table_name(new_name)}"
+        pk, seq = pk_and_sequence_for(new_name)
+        if pk
+          idx = "#{table_name}_pkey"
+          new_idx = "#{new_name}_pkey"
+          execute "ALTER INDEX #{quote_table_name(idx)} RENAME TO #{quote_local_table_name(new_idx)}"
+          if seq && seq.identifier == "#{table_name}_#{pk}_seq"
+            new_seq = "#{new_name}_#{pk}_seq"
+            execute "ALTER TABLE #{seq.quoted} RENAME TO #{quote_local_table_name(new_seq)}"
+          end
+        end
+        rename_table_indexes(table_name, new_name)
+      end
+
+      def rename_index(table_name, old_name, new_name)
+        validate_index_length!(table_name, new_name)
+
+        execute "ALTER INDEX #{quote_table_name(old_name)} RENAME TO #{quote_local_table_name(new_name)}"
+      end
+
+      def columns(*)
+        with_local_table_name(false) { super }
       end
     end
   end

data/lib/switchman/active_record/predicate_builder.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveRecord
     module PredicateBuilder
@@ -13,7 +15,7 @@ module Switchman
         def convert_to_id(value)
           case value
           when ::ActiveRecord::Base
-            value.send(primary_key)
+            value.id
           else
             super
           end

data/lib/switchman/active_record/query_cache.rb

@@ -1,99 +1,34 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveRecord
     module QueryCache
-      # thread local accessors to replace @query_cache_enabled
-      def query_cache
-        thread_cache = Thread.current[:query_cache] ||= {}
-        thread_cache[self.object_id] ||= Hash.new { |h,sql| h[sql] = {} }
-      end
-
-      def query_cache_enabled
-        Thread.current[:query_cache_enabled]
-      end
-
-      def query_cache_enabled=(value)
-        Thread.current[:query_cache_enabled] = value
-      end
-
-      # basically wholesale repeat of the methods from the original (see
-      # https://github.com/rails/rails/blob/master/activerecord/lib/active_record/connection_adapters/abstract/query_cache.rb),
-      # but with self.query_cache_enabled and self.query_cache_enabled= instead
-      # of @query_cache_enabled.
-
-      def enable_query_cache!
-        self.query_cache_enabled = true
-      end
-
-      def disable_query_cache!
-        self.query_cache_enabled = false
-      end
-
-      def cache
-        old, self.query_cache_enabled = query_cache_enabled, true
-        yield
-      ensure
-        self.query_cache_enabled = old
-        clear_query_cache unless self.query_cache_enabled
-      end
-
-      def uncached
-        old, self.query_cache_enabled = query_cache_enabled, false
-        yield
-      ensure
-        self.query_cache_enabled = old
-      end
-
-      def clear_query_cache
-        Thread.current[:query_cache].try(:clear)
-      end
-
-      def select_all(arel, name = nil, binds = [], preparable: nil)
-        if self.query_cache_enabled && !locked?(arel)
-          arel, binds = binds_from_relation(arel, binds)
-          sql = to_sql(arel, binds)
-          if ::Rails.version >= '5'
-            cache_sql(sql, binds) { super(sql, name, binds, preparable: preparable) }
-          else
-            cache_sql(sql, binds) { super(sql, name, binds) }
-          end
-        else
-          if ::Rails.version >= '5'
-            super
-          else
-            super(arel, name, binds)
-          end
-        end
-      end
-
-      # no reason to define these on the including class directly. the super
-      # works just as well from a method on the included module
-      [:insert, :update, :delete].each do |method_name|
-        class_eval <<-end_code, __FILE__, __LINE__ + 1
-          def #{method_name}(*args)
-            clear_query_cache if self.query_cache_enabled
-            super
-          end
-        end_code
-      end
 
       private
 
-      def cache_sql(sql, binds)
+      def cache_sql(sql, name, binds)
         # have to include the shard id in the cache key because of switching dbs on the same connection
         sql = "#{self.shard.id}::#{sql}"
-        result =
-            if query_cache[sql].key?(binds)
-              ::ActiveSupport::Notifications.instrument("sql.active_record",
-                                                      :sql => sql, :binds => binds, :name => "CACHE", :connection_id => object_id)
-              query_cache[sql][binds]
-            else
-              query_cache[sql][binds] = yield
-            end
-
-        if ::ActiveRecord::Result === result
+        @lock.synchronize do
+          result =
+              if query_cache[sql].key?(binds)
+                args = {
+                  sql: sql,
+                  binds: binds,
+                  name: name,
+                  connection_id: object_id,
+                  cached: true
+                }
+                args[:type_casted_binds] = -> { type_casted_binds(binds) } if ::Rails.version >= '5.1.5'
+                ::ActiveSupport::Notifications.instrument(
+                    "sql.active_record",
+                    args
+                )
+                query_cache[sql][binds]
+              else
+                query_cache[sql][binds] = yield
+              end
           result.dup
-        else
-          result.collect { |row| row.dup }
         end
       end
     end