Replace Alias Method Chain With Prepend for SchemaDumper

lib/activerecord-multi-tenant/migrations.rb

@@ -81,52 +81,3 @@ def create_table(table_name, options = {}, &block)
   end
 end
 ActiveRecord::ConnectionAdapters::SchemaStatements.prepend(MultiTenant::SchemaStatementsExtensions)
-
-module ActiveRecord
-  class SchemaDumper
-    private
-
-    alias initialize_without_citus initialize
-
-    def initialize(connection, options = {})
-      initialize_without_citus(connection, options)
-
-      citus_version =
-        begin
-          ActiveRecord::Migration.citus_version
-        rescue StandardError
-          # Handle the case where this gem is used with MySQL https://github.com/citusdata/activerecord-multi-tenant/issues/166
-          nil
-        end
-      @distribution_columns =
-        if citus_version.present?
-          query_to_execute = <<-SQL.strip
-            SELECT logicalrelid::regclass AS table_name,
-                   column_to_column_name(logicalrelid, partkey) AS dist_col_name
-            FROM pg_dist_partition
-          SQL
-          @connection.execute(query_to_execute).to_h do |v|
-            [v['table_name'], v['dist_col_name']]
-          end
-        else
-          {}
-        end
-    end
-
-    # Support for create_distributed_table & create_reference_table
-    alias table_without_citus table
-
-    def table(table, stream)
-      table_without_citus(table, stream)
-      table_name = remove_prefix_and_suffix(table)
-      distribution_column = @distribution_columns[table_name]
-      if distribution_column
-        stream.puts "  create_distributed_table(#{table_name.inspect}, #{distribution_column.inspect})"
-        stream.puts
-      elsif @distribution_columns.key?(table_name)
-        stream.puts "  create_reference_table(#{table_name.inspect})"
-        stream.puts
-      end
-    end
-  end
-end

lib/activerecord-multi-tenant/schema_dumper.rb

@@ -0,0 +1,47 @@
+module MultiTenant
+  module SchemaDumper
+    private
+
+    def initialize(connection, options = {})
+      super
+
+      citus_version =
+        begin
+          ActiveRecord::Migration.citus_version
+        rescue StandardError
+          # Handle the case where this gem is used with MySQL https://github.com/citusdata/activerecord-multi-tenant/issues/166
+          nil
+        end
+      @distribution_columns =
+        if citus_version.present?
+          query_to_execute = <<-SQL.strip
+            SELECT logicalrelid::regclass AS table_name,
+                   column_to_column_name(logicalrelid, partkey) AS dist_col_name
+            FROM pg_dist_partition
+          SQL
+          @connection.execute(query_to_execute).to_h do |v|
+            [v['table_name'], v['dist_col_name']]
+          end
+        else
+          {}
+        end
+    end
+
+    def table(table, stream)
+      super
+      table_name = remove_prefix_and_suffix(table)
+      distribution_column = @distribution_columns[table_name]
+      if distribution_column
+        stream.puts "  create_distributed_table(#{table_name.inspect}, #{distribution_column.inspect})"
+        stream.puts
+      elsif @distribution_columns.key?(table_name)
+        stream.puts "  create_reference_table(#{table_name.inspect})"
+        stream.puts
+      end
+    end
+  end
+end
+
+ActiveSupport.on_load(:active_record) do
+  ActiveRecord::SchemaDumper.prepend(MultiTenant::SchemaDumper)
+end

lib/activerecord-multi-tenant/version.rb

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

lib/activerecord_multi_tenant.rb

@@ -12,3 +12,4 @@
 require_relative 'activerecord-multi-tenant/version'
 require_relative 'activerecord-multi-tenant/habtm'
 require_relative 'activerecord-multi-tenant/relation_extension'
+require_relative 'activerecord-multi-tenant/schema_dumper'