Allow loading or dumping a specific tenant in a multitenant setup using Apartment gem

lib/yaml_db/serialization_helper.rb

@@ -59,6 +59,9 @@ def reenable_logger
 
     class Load
       def self.load(io, truncate = true)
+        if defined?(Apartment::Tenant) && ENV['TENANT']
+          Apartment::Tenant.switch! ENV['TENANT']
+        end
         ActiveRecord::Base.connection.transaction do
           load_documents(io, truncate)
         end
@@ -165,6 +168,9 @@ def self.after_table(io, table)
       end
 
       def self.tables
+        if defined?(Apartment::Tenant) && ENV['TENANT']
+          Apartment::Tenant.switch! ENV['TENANT']
+        end
         ActiveRecord::Base.connection.tables.reject { |table| ['schema_info', 'schema_migrations'].include?(table) }.sort
       end
 

lib/yaml_db/version.rb

@@ -1,3 +1,3 @@
 module YamlDb
-  VERSION = "0.6.0"
+  VERSION = "0.6.1"
 end

spec/yaml_db/serialization_helper_dump_spec.rb

@@ -16,6 +16,27 @@ module SerializationHelper
         @io = StringIO.new
       end
 
+      context "with Apartment multi-tenancy" do
+        before(:each) do
+          class_double('Apartment').as_stubbed_const
+          @tenant_klass = class_double('Apartment::Tenant').as_stubbed_const
+          @old_env = ENV['TENANT']
+        end
+        after(:each) do
+          ENV['TENANT'] = @old_env
+        end
+        it "explicitly sets tenant before extracting tables if ENV['TENANT'] set" do
+          tenant_name = 'tenant-33'
+          ENV['TENANT'] = tenant_name
+          expect(@tenant_klass).to receive(:switch!).with(tenant_name)
+          Dump.tables
+        end
+        it "does not explicitly set tenant if ENV['TENANT'] not set" do
+          ENV.delete('TENANT')
+          expect(@tenant_klass).not_to receive(:switch!)
+          Dump.tables
+        end
+      end
       it "returns a list of column names" do
         expect(Dump.table_column_names('mytable')).to eq([ 'a', 'b' ])
       end

spec/yaml_db/serialization_helper_load_spec.rb

@@ -10,6 +10,28 @@ module SerializationHelper
         @io = StringIO.new
       end
 
+      context "with Apartment multi-tenancy" do
+        before(:each) do
+          allow(Load).to receive(:load_documents)
+          class_double('Apartment').as_stubbed_const
+          @tenant_klass = class_double('Apartment::Tenant').as_stubbed_const
+          @old_env = ENV['TENANT']
+        end
+        after(:each) do
+          ENV['TENANT'] = @old_env
+        end
+        it "switches to tenant if ENV[TENANT] set" do
+          tenant_name = 'tenant-23'
+          ENV['TENANT'] = tenant_name
+          expect(@tenant_klass).to receive(:switch!).with(tenant_name)
+          Load.load(@io)
+        end
+        it "does not explicitly switch tenants if ENV[TENANT] is not set" do
+          ENV.delete('TENANT')
+          expect(@tenant_klass).not_to receive(:switch!)
+        end
+      end
+
       it "truncates the table" do
         allow(ActiveRecord::Base.connection).to receive(:execute).with("TRUNCATE mytable").and_return(true)
         expect(ActiveRecord::Base.connection).not_to receive(:execute).with("DELETE FROM mytable")