SL-66188: Melody migrations failing locally

lib/activerecord-multi-tenant.rb

@@ -1,13 +1,19 @@
-if Object.const_defined?(:ActionController)
-  require_relative 'activerecord-multi-tenant/controller_extensions'
-end
+require_relative 'activerecord-multi-tenant/controller_extensions' if Object.const_defined?(:ActionController)
 require_relative 'activerecord-multi-tenant/copy_from_client'
 require_relative 'activerecord-multi-tenant/fast_truncate'
-require_relative 'activerecord-multi-tenant/migrations'
 require_relative 'activerecord-multi-tenant/model_extensions'
 require_relative 'activerecord-multi-tenant/multi_tenant'
 require_relative 'activerecord-multi-tenant/query_rewriter'
 require_relative 'activerecord-multi-tenant/query_monitor'
 require_relative 'activerecord-multi-tenant/version'
 require_relative 'activerecord-multi-tenant/with_lock'
 require_relative 'activerecord-multi-tenant/delete_operations'
+
+ActiveSupport.on_load(:active_record) do
+  require_relative 'activerecord-multi-tenant/migrations_extension'
+  require_relative 'activerecord-multi-tenant/schema_dumper'
+  require_relative 'activerecord-multi-tenant/schema_statements'
+
+  ActiveRecord::Migration.include MultiTenant::MigrationExtensions
+  ActiveRecord::SchemaDumper.prepend(MultiTenant::SchemaDumper)
+end

lib/activerecord-multi-tenant/migrations_extension.rb

@@ -0,0 +1,64 @@
+module MultiTenant
+  module MigrationExtensions
+    def create_distributed_table(table_name, partition_key)
+      return unless citus_version.present?
+
+      reversible do |dir|
+        dir.up do
+          execute "SELECT create_distributed_table($$#{table_name}$$, $$#{partition_key}$$)"
+        end
+        dir.down do
+          undistribute_table(table_name)
+        end
+      end
+    end
+
+    def create_reference_table(table_name)
+      return unless citus_version.present?
+
+      reversible do |dir|
+        dir.up do
+          execute "SELECT create_reference_table($$#{table_name}$$)"
+        end
+        dir.down do
+          undistribute_table(table_name)
+        end
+      end
+    end
+
+    def undistribute_table(table_name)
+      return unless citus_version.present?
+
+      execute "SELECT undistribute_table($$#{table_name}$$))"
+    end
+
+    def rebalance_table_shards
+      return unless citus_version.present?
+
+      execute 'SELECT rebalance_table_shards()'
+    end
+
+    def execute_on_all_nodes(sql)
+      execute sql
+
+      case citus_version
+      when '6.0'
+        execute "SELECT citus_run_on_all_workers($$#{sql}$$)" # initial citus_tools.sql with different names
+      when nil
+        # Do nothing, this is regular Postgres
+      else # 6.1 and newer
+        execute "SELECT run_command_on_workers($$#{sql}$$)"
+      end
+    end
+
+    def enable_extension_on_all_nodes(extension)
+      execute_on_all_nodes "CREATE EXTENSION IF NOT EXISTS \"#{extension}\""
+    end
+
+    def citus_version
+      execute("SELECT extversion FROM pg_extension WHERE extname = 'citus'").getvalue(0, 0).try(:split, '-').try(:first)
+    rescue ArgumentError => e
+      raise unless e.message == 'invalid tuple number 0'
+    end
+  end
+end

lib/activerecord-multi-tenant/schema_dumper.rb

@@ -0,0 +1,35 @@
+module MultiTenant
+  module SchemaDumper
+    def initialize(connection, options = {})
+      super(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?
+          @connection.execute('SELECT logicalrelid::regclass AS table_name, column_to_column_name(logicalrelid, partkey) AS dist_col_name FROM pg_dist_partition').to_h do |v|
+            [v['table_name'], v['dist_col_name']]
+          end
+        else
+          {}
+        end
+    end
+
+    def table(table, stream)
+      super(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_statements.rb

@@ -0,0 +1,15 @@
+module ActiveRecord
+  module ConnectionAdapters # :nodoc:
+    module SchemaStatements
+      alias orig_create_table create_table
+      def create_table(table_name, options = {}, &block)
+        ret = orig_create_table(table_name, **options.except(:partition_key), &block)
+        if options[:id] != false && options[:partition_key] && options[:partition_key].to_s != 'id'
+          execute "ALTER TABLE #{table_name} DROP CONSTRAINT #{table_name}_pkey"
+          execute "ALTER TABLE #{table_name} ADD PRIMARY KEY(\"#{options[:partition_key]}\", id)"
+        end
+        ret
+      end
+    end
+  end
+end