Exclude the default tenant table if it exists in the procedure

lib/apartment.rb

@@ -22,7 +22,8 @@ class << self
     extend Forwardable
 
     ACCESSOR_METHODS = %i[use_schemas use_sql seed_after_create prepend_environment default_tenant
-                          append_environment with_multi_server_setup tenant_presence_check active_record_log].freeze
+                          append_environment with_multi_server_setup tenant_presence_check
+                          active_record_log pg_exclude_clone_tables].freeze
 
     WRITER_METHODS = %i[tenant_names database_schema_file excluded_models
                         persistent_schemas connection_class

lib/apartment/adapters/postgresql_adapter.rb

@@ -194,15 +194,13 @@ def copy_schema_migrations
       #   @return {String} raw SQL contaning only postgres schema dump
       #
       def pg_dump_schema
-        # Skip excluded tables? :/
-        # excluded_tables =
-        #   collect_table_names(Apartment.excluded_models)
-        #   .map! {|t| "-T #{t}"}
-        #   .join(' ')
-
-        # `pg_dump -s -x -O -n #{default_tenant} #{excluded_tables} #{dbname}`
-
-        with_pg_env { `pg_dump -s -x -O -n #{default_tenant} #{dbname}` }
+        exclude_table =
+          if Apartment.pg_exclude_clone_tables
+            excluded_tables.map! { |t| "-T #{t}" }.join(' ')
+          else
+            ''
+          end
+        with_pg_env { `pg_dump -s -x -O -n #{default_tenant} #{dbname} #{exclude_table}` }
       end
 
       #   Dump data from schema_migrations table
@@ -254,6 +252,8 @@ def swap_schema_qualifier(sql)
         sql.gsub(/#{default_tenant}\.\w*/) do |match|
           if Apartment.pg_excluded_names.any? { |name| match.include? name }
             match
+          elsif Apartment.pg_exclude_clone_tables && excluded_tables.any?(match)
+            match
           else
             match.gsub("#{default_tenant}.", %("#{current}".))
           end
@@ -266,10 +266,10 @@ def check_input_against_regexps(input, regexps)
         regexps.select { |c| input.match c }
       end
 
-      #   Collect table names from AR Models
+      # Convenience method for excluded table names
       #
-      def collect_table_names(models)
-        models.map do |m|
+      def excluded_tables
+        Apartment.excluded_models.map do |m|
           m.constantize.table_name
         end
       end

spec/adapters/postgresql_adapter_spec.rb

@@ -63,5 +63,61 @@ def tenant_names
       it_behaves_like 'a generic apartment adapter able to handle custom configuration'
       it_behaves_like 'a connection based apartment adapter'
     end
+
+    context 'when using pg_exclude_clone_tables with SQL dump' do
+      before do
+        Apartment.excluded_models = ['Company']
+        Apartment.use_schemas = true
+        Apartment.use_sql = true
+        Apartment.pg_exclude_clone_tables = true
+        ActiveRecord::Base.connection.execute <<-PROCEDURE
+          CREATE OR REPLACE FUNCTION test_function() RETURNS INTEGER AS $function$
+          DECLARE
+            r1 INTEGER;
+            r2 INTEGER;
+          BEGIN
+            SELECT COUNT(*) INTO r1 FROM public.companies;
+            SELECT COUNT(*) INTO r2 FROM public.users;
+            RETURN r1 + r2;
+          END;
+          $function$ LANGUAGE plpgsql;
+        PROCEDURE
+      end
+
+      after do
+        Apartment::Tenant.drop('has-procedure') if Apartment.connection.schema_exists? 'has-procedure'
+        ActiveRecord::Base.connection.execute('DROP FUNCTION IF EXISTS test_function();')
+        # Apartment::Tenant.init creates per model connection.
+        # Remove the connection after testing not to unintentionally keep the connection across tests.
+        Apartment.excluded_models.each do |excluded_model|
+          excluded_model.constantize.remove_connection
+        end
+      end
+
+      # Not sure why, but somehow using let(:tenant_names) memoizes for the whole example group, not just each test
+      def tenant_names
+        ActiveRecord::Base.connection.execute('SELECT nspname FROM pg_namespace;').collect { |row| row['nspname'] }
+      end
+
+      let(:default_tenant) { subject.switch { ActiveRecord::Base.connection.schema_search_path.delete('"') } }
+      let(:c) { rand(5) }
+      let(:u) { rand(5) }
+
+      it_behaves_like 'a generic apartment adapter'
+      it_behaves_like 'a schema based apartment adapter'
+
+      # rubocop:disable RSpec/ExampleLength
+      it 'not change excluded_models in the procedure code' do
+        Apartment::Tenant.init
+        Apartment::Tenant.create('has-procedure')
+        Apartment::Tenant.switch!('has-procedure')
+        c.times { Company.create }
+        u.times { User.create }
+        count = ActiveRecord::Base.connection.execute('SELECT test_function();')[0]['test_function']
+        expect(count).to(eq(Company.count + User.count))
+        Company.delete_all
+      end
+      # rubocop:enable RSpec/ExampleLength
+    end
   end
 end

spec/examples/generic_adapter_examples.rb

@@ -42,7 +42,7 @@
 
     it 'should load schema.rb to new schema' do
       subject.switch(db1) do
-        expect(connection.tables).to include('companies')
+        expect(connection.tables).to include('users')
       end
     end
 

spec/examples/schema_adapter_examples.rb

@@ -63,7 +63,7 @@
   describe '#create' do
     it 'should load schema.rb to new schema' do
       connection.schema_search_path = schema1
-      expect(connection.tables).to include('companies')
+      expect(connection.tables).to include('users')
     end
 
     it 'should yield to block if passed and reset' do