split task to add unit tests

lib/activerecord-multi-tenant/schema_dumper_extension.rb

@@ -1,6 +1,55 @@
 module MultiTenant
   module SchemaDumperExtension
     cattr_accessor :include_distribute_statements, default: true
+
+    def get_distributed_tables(connection)
+      query_distributed = 'SELECT logicalrelid, pg_attribute.attname ' \
+                          'FROM pg_dist_partition ' \
+                          'INNER JOIN pg_attribute ON (logicalrelid=attrelid) ' \
+                          'WHERE partmethod=\'h\' ' \
+                          'AND attnum=substring(partkey from \'%:varattno #"[0-9]+#"%\' for \'#\')::int ' \
+                          'ORDER BY logicalrelid'
+
+      begin
+        return connection.execute(query_distributed).values
+      rescue; end
+    end
+
+    def get_reference_tables(connection)
+      query_reference = "SELECT logicalrelid FROM pg_dist_partition WHERE partmethod = 'n' ORDER BY logicalrelid"
+      begin
+        return connection.execute(query_reference).values
+      rescue; end
+    end
+
+    def get_distribute_statements(connection, reference=false)
+      if reference
+        distributed_tables = get_reference_tables(connection)
+        query = "SELECT create_reference_table('%s');\n"
+      else
+        distributed_tables = get_distributed_tables(connection)
+        query = "SELECT create_distributed_table('%s', '%s');\n"
+      end
+
+      return unless distributed_tables
+
+      schema = ''
+      distributed_tables.each do |distributed_table|
+        attrs = if reference then [distributed_table[0]] else [distributed_table[0], distributed_table[1]] end
+        schema <<  query % attrs
+      end
+
+      schema
+    end
+
+    def get_full_distribute_statements(connection)
+      schema = ActiveRecord::SchemaDumper.get_distribute_statements(connection) || ''
+      schema << (ActiveRecord::SchemaDumper.get_distribute_statements(connection,
+                                                                      reference=true) || '')
+
+      schema
+    end
+
   end
 end
 

lib/activerecord-multi-tenant/tasks/db.rake

@@ -2,26 +2,18 @@ require "active_record"
 
 
 Rake::Task["db:structure:dump"].enhance do
-  return unless ActiveRecord::SchemaDumper.include_distribute_statements
+  puts 'coucou'
+  next unless ActiveRecord::SchemaDumper.include_distribute_statements
 
   if ActiveRecord::VERSION::MAJOR >= 6
     databases = ActiveRecord::Tasks::DatabaseTasks.setup_initial_database_yaml
   else
     databases = [ActiveRecord::Tasks::DatabaseTasks.current_config]
   end
 
-  query_distributed = 'SELECT logicalrelid, pg_attribute.attname ' \
-                      'FROM pg_dist_partition ' \
-                      'INNER JOIN pg_attribute ON (logicalrelid=attrelid) ' \
-                      'WHERE partmethod=\'h\' ' \
-                      'AND attnum=substring(partkey from \'%:varattno #"[0-9]+#"%\' for \'#\')::int;'
-
-
-  query_reference = "SELECT logicalrelid FROM pg_dist_partition WHERE partmethod = 'n';"
-
   databases.each do |db_config|
     # for versions before 6.0, there will only be 1 database in the list
-    connection = ActiveRecord::Base.establish_connection(db_config)
+    connection = ActiveRecord::Base.establish_connection(db_config).connection
     filenames = []
     if ActiveRecord::VERSION::MAJOR >= 6
       Rails.application.config.paths['db'].each do |path|
@@ -35,26 +27,11 @@ Rake::Task["db:structure:dump"].enhance do
       end
     end
 
-    schema = ''
-    begin
-      distributed_tables = connection.connection.execute(query_distributed)
-      reference_tables = connection.connection.execute(query_reference)
-    rescue
-      # citus is not installed, failing because pg_dist_partition not found
-    else
-      distributed_tables.values.each do |distributed_table|
-        schema <<  "SELECT create_distributed_table('%s', '%s');\n" % [distributed_table[0], distributed_table[1]]
-      end
-
-      reference_tables.values.each do |reference_table|
-        schema <<  "SELECT create_reference_table('%s');\n" % [reference_table[0]]
-      end
-    end
+    schema = ActiveRecord::SchemaDumper.get_full_distribute_statements(connection)
 
     filenames.each do |filename|
       File.open(filename, "a") { |f| f.puts schema }
     end
     puts "Added distribute statements to #{filenames}"
   end
-
 end

spec/activerecord-multi-tenant/schema_dumper_spec.rb

@@ -0,0 +1,96 @@
+require 'spec_helper'
+require 'rake'
+
+
+describe 'Schema Dumper enhancement' do
+  let(:file_like_object) { double("file like object") }
+
+  it 'should list distributed tables' do
+    distributed_tables  = ActiveRecord::SchemaDumper.get_distributed_tables(ActiveRecord::Base.connection)
+
+    distributed_result = [["accounts", "id"],
+                          ["projects", "account_id"],
+                          ["managers", "account_id"],
+                          ["tasks", "account_id"],
+                          ["sub_tasks", "account_id"],
+                          ["aliased_tasks", "account_id"],
+                          ["custom_partition_key_tasks", "accountID"],
+                          ["comments", "account_id"],
+                          ["partition_key_not_model_tasks", "non_model_id"],
+                          ["subclass_tasks", "non_model_id"],
+                          ["uuid_records", "organization_id"],
+                          ["allowed_places", "account_id"],
+                          ["project_categories", "account_id"]]
+
+    expect(distributed_tables.to_set).to eq(distributed_result.to_set)
+  end
+
+  it 'should list reference tables' do
+    reference_tables = ActiveRecord::SchemaDumper.get_reference_tables(ActiveRecord::Base.connection)
+    reference_result = [["categories"]]
+    expect(reference_tables.to_set).to eq(reference_result.to_set)
+  end
+
+  it 'distribute statements' do
+    distributed_statements = ActiveRecord::SchemaDumper.get_distribute_statements(ActiveRecord::Base.connection)
+    expect(distributed_statements).to eq("SELECT create_distributed_table('accounts', 'id');\nSELECT create_distributed_table('projects', 'account_id');\nSELECT create_distributed_table('managers', 'account_id');\nSELECT create_distributed_table('tasks', 'account_id');\nSELECT create_distributed_table('sub_tasks', 'account_id');\nSELECT create_distributed_table('aliased_tasks', 'account_id');\nSELECT create_distributed_table('custom_partition_key_tasks', 'accountID');\nSELECT create_distributed_table('comments', 'account_id');\nSELECT create_distributed_table('partition_key_not_model_tasks', 'non_model_id');\nSELECT create_distributed_table('subclass_tasks', 'non_model_id');\nSELECT create_distributed_table('uuid_records', 'organization_id');\nSELECT create_distributed_table('project_categories', 'account_id');\nSELECT create_distributed_table('allowed_places', 'account_id');\n")
+  end
+
+  it 'reference tables statements' do
+    distributed_statements = ActiveRecord::SchemaDumper.get_distribute_statements(ActiveRecord::Base.connection, reference=true)
+    expect(distributed_statements).to eq("SELECT create_reference_table('categories');\n")
+  end
+
+  it 'no distributed tables' do
+    ActiveRecord::SchemaDumper.stub(:get_distributed_tables).with(anything()) {[]}
+    distributed_statements = ActiveRecord::SchemaDumper.get_distribute_statements(ActiveRecord::Base.connection)
+    expect(distributed_statements).to eq("")
+  end
+
+  it 'no citus metadata tables' do
+    ActiveRecord::SchemaDumper.stub(:get_distributed_tables).with(anything()) {nil}
+    distributed_statements = ActiveRecord::SchemaDumper.get_distribute_statements(ActiveRecord::Base.connection)
+    expect(distributed_statements).to eq(nil)
+  end
+
+  it 'no reference tables' do
+    ActiveRecord::SchemaDumper.stub(:get_reference_tables).with(anything()) {[]}
+    distributed_statements = ActiveRecord::SchemaDumper.get_distribute_statements(ActiveRecord::Base.connection, reference=true)
+    expect(distributed_statements).to eq("")
+
+  end
+
+  it 'no citus metadata tables for reference' do
+    ActiveRecord::SchemaDumper.stub(:get_reference_tables).with(anything()) {nil}
+    distributed_statements = ActiveRecord::SchemaDumper.get_distribute_statements(ActiveRecord::Base.connection, reference=true)
+    expect(distributed_statements).to eq(nil)
+  end
+
+
+  it 'full statements' do
+    distributed_statements = ActiveRecord::SchemaDumper.get_full_distribute_statements(ActiveRecord::Base.connection)
+    expect(distributed_statements).to eq("SELECT create_distributed_table('accounts', 'id');\nSELECT create_distributed_table('projects', 'account_id');\nSELECT create_distributed_table('managers', 'account_id');\nSELECT create_distributed_table('tasks', 'account_id');\nSELECT create_distributed_table('sub_tasks', 'account_id');\nSELECT create_distributed_table('aliased_tasks', 'account_id');\nSELECT create_distributed_table('custom_partition_key_tasks', 'accountID');\nSELECT create_distributed_table('comments', 'account_id');\nSELECT create_distributed_table('partition_key_not_model_tasks', 'non_model_id');\nSELECT create_distributed_table('subclass_tasks', 'non_model_id');\nSELECT create_distributed_table('uuid_records', 'organization_id');\nSELECT create_distributed_table('project_categories', 'account_id');\nSELECT create_distributed_table('allowed_places', 'account_id');\nSELECT create_reference_table('categories');\n")
+
+  end
+
+  it 'full statements no reference' do
+    ActiveRecord::SchemaDumper.stub(:get_reference_tables).with(anything()) {[]}
+    distributed_statements = ActiveRecord::SchemaDumper.get_full_distribute_statements(ActiveRecord::Base.connection)
+    expect(distributed_statements).to eq("SELECT create_distributed_table('accounts', 'id');\nSELECT create_distributed_table('projects', 'account_id');\nSELECT create_distributed_table('managers', 'account_id');\nSELECT create_distributed_table('tasks', 'account_id');\nSELECT create_distributed_table('sub_tasks', 'account_id');\nSELECT create_distributed_table('aliased_tasks', 'account_id');\nSELECT create_distributed_table('custom_partition_key_tasks', 'accountID');\nSELECT create_distributed_table('comments', 'account_id');\nSELECT create_distributed_table('partition_key_not_model_tasks', 'non_model_id');\nSELECT create_distributed_table('subclass_tasks', 'non_model_id');\nSELECT create_distributed_table('uuid_records', 'organization_id');\nSELECT create_distributed_table('project_categories', 'account_id');\nSELECT create_distributed_table('allowed_places', 'account_id');\n")
+  end
+
+  it 'full statements no distributed' do
+    ActiveRecord::SchemaDumper.stub(:get_distributed_tables).with(anything()) {nil}
+    distributed_statements = ActiveRecord::SchemaDumper.get_full_distribute_statements(ActiveRecord::Base.connection)
+    expect(distributed_statements).to eq("SELECT create_reference_table('categories');\n")
+  end
+
+  it 'full statements no citus' do
+    ActiveRecord::SchemaDumper.stub(:get_distributed_tables).with(anything()) {nil}
+    ActiveRecord::SchemaDumper.stub(:get_reference_tables).with(anything()) {nil}
+
+    distributed_statements = ActiveRecord::SchemaDumper.get_full_distribute_statements(ActiveRecord::Base.connection)
+    expect(distributed_statements).to eq("")
+
+  end
+end