activerecord-multi-tenant 1.1.1 → 2.2.0

data/spec/activerecord-multi-tenant/associations_spec.rb

@@ -0,0 +1,21 @@
+require 'spec_helper'
+
+describe MultiTenant, 'Association methods' do
+  let(:account1) { Account.create! name: 'test1' }
+  let(:account2) { Account.create! name: 'test2' }
+  let(:project1) { Project.create! name: 'something1', account: account1 }
+  let(:project2) { Project.create! name: 'something2', account: account2, id: project1.id }
+  let(:task1) { Task.create! name: 'task1', project: project1, account: account1 }
+  let(:task2) { Task.create! name: 'task2', project: project2, account: account2, id: task1.id }
+
+  context 'include the tenant_id in queries and' do
+    it 'creates a task with correct account_id' do
+      expect(project2.tasks.create(name: 'task3').account_id).to eq(account2.id)
+    end
+    it 'return correct account_id' do
+      expect(task1.project.account_id).to_not eq(task2.project.account_id) # belongs_to
+      expect(project2.tasks.count).to eq(1)
+      expect(project2.tasks.first.account_id).to eq(account2.id) # has_many
+    end
+  end
+end

data/spec/activerecord-multi-tenant/model_extensions_spec.rb

@@ -70,6 +70,65 @@ describe MultiTenant do
     it { expect(@partition_key_not_model_task.non_model_id).to be 77 }
   end
 
+
+  describe 'Tenant model with a nonstandard class name' do
+    let(:account_klass) do
+      Class.new(ActiveRecord::Base) do
+        self.table_name = 'account'
+        def self.name
+          'UserAccount'
+        end
+
+        multi_tenant(:account)
+      end
+    end
+    it "does not register the tenant model" do
+      expect(MultiTenant).not_to receive(:register_multi_tenant_model)
+      account_klass
+    end
+  end
+
+  describe 'Changes table_name after multi_tenant called' do
+    before do
+      account_klass.has_many(:posts, anonymous_class: post_klass)
+      post_klass.belongs_to(:account, anonymous_class: account_klass)
+
+      @account1 = account_klass.create! name: 'foo'
+      @account2 = account_klass.create! name: 'bar'
+
+      @post1 = @account1.posts.create! name: 'foobar'
+      @post2 = @account2.posts.create! name: 'baz'
+
+      MultiTenant.current_tenant = @account1
+      @posts = post_klass.all
+    end
+
+    let(:account_klass) do
+      Class.new(Account) do
+        def self.name
+          'Account'
+        end
+      end
+    end
+
+    let(:post_klass) do
+      Class.new(ActiveRecord::Base) do
+        self.table_name = 'unknown'
+
+        multi_tenant(:account)
+
+        self.table_name = 'posts'
+
+        def self.name
+          'Post'
+        end
+      end
+    end
+
+    it { expect(@posts.length).to eq(1) }
+    it { expect(@posts).to eq([@post1]) }
+  end
+
   # Scoping models
   describe 'Project.all should be scoped to the current tenant if set' do
     before do
@@ -144,6 +203,22 @@ describe MultiTenant do
       end
     end
 
+    it 'handles belongs_to with optional: true' do
+      record = OptionalSubTask.create(sub_task_id: sub_task.id)
+      expect(record.reload.sub_task).to eq(sub_task)
+      expect(record.account_id).to eq(nil)
+    end
+
+    it 'handles changing tenant from nil to a value' do
+      record = OptionalSubTask.create(sub_task_id: sub_task.id)
+      expect(record.reload.sub_task).to eq(sub_task)
+      expect(record.account_id).to eq(nil)
+
+      record.account = account
+      record.save!
+      expect(record.reload.account_id).to eq(account.id)
+    end
+
     it 'handles has_many through' do
       MultiTenant.with(account) do
         expect(project.sub_tasks).to eq [sub_task]
@@ -234,6 +309,55 @@ describe MultiTenant do
     end
   end
 
+  # Joins
+  describe 'joins for models' do
+    context 'for models with where condition in associations' do
+      let(:account) { Account.create!(name: 'Account 1') }
+
+      it 'should add tenant condition to the queries when tenant is set' do
+        expected_join_sql = <<-SQL.strip
+          SELECT "comments".* FROM "comments" INNER JOIN "tasks" ON "tasks"."id" = "comments"."commentable_id" AND "comments"."commentable_type" = 'Task' AND "tasks"."account_id" = 1 WHERE "comments"."account_id" = 1
+        SQL
+
+        MultiTenant.with(account) do
+          expect(Comment.joins(:task).to_sql).to eq(expected_join_sql)
+        end
+      end
+
+      it 'should add tenant condition to the queries when tenant is not set' do
+        MultiTenant.without do
+          expected_join_sql = <<-SQL.strip
+            SELECT "comments".* FROM "comments" INNER JOIN "tasks" ON "tasks"."id" = "comments"."commentable_id" AND "comments"."commentable_type" = 'Task' AND "comments"."account_id" = "tasks"."account_id"
+          SQL
+          expect(Comment.joins(:task).to_sql).to eq(expected_join_sql)
+        end
+      end
+    end
+
+    context 'for models with default associations' do
+      let(:account) { Account.create!(name: 'Account 1') }
+
+      it 'should add tenant condition to the queries when tenant is set' do
+        expected_join_sql = <<-SQL.strip
+          SELECT "projects".* FROM "projects" INNER JOIN "tasks" ON "tasks"."project_id" = "projects"."id" AND "tasks"."account_id" = 1 WHERE "projects"."account_id" = 1
+        SQL
+
+        MultiTenant.with(account) do
+          expect(Project.joins(:tasks).to_sql).to eq(expected_join_sql)
+        end
+      end
+
+      it 'should add tenant condition to the queries when tenant is not set' do
+        MultiTenant.without do
+          expected_join_sql = <<-SQL.strip
+            SELECT "projects".* FROM "projects" INNER JOIN "tasks" ON "tasks"."project_id" = "projects"."id" AND "projects"."account_id" = "tasks"."account_id"
+          SQL
+          expect(Project.joins(:tasks).to_sql).to eq(expected_join_sql)
+        end
+      end
+    end
+  end
+
   # ::with
   describe "::with" do
     it "should set current_tenant to the specified tenant inside the block" do
@@ -348,9 +472,12 @@ describe MultiTenant do
   end
 
   it "applies the team_id conditions in the where clause" do
-    expected_sql = <<-sql
-                     SELECT "sub_tasks".* FROM "sub_tasks" INNER JOIN "tasks" ON "sub_tasks"."task_id" = "tasks"."id" AND "sub_tasks"."account_id" = "tasks"."account_id" WHERE "tasks"."account_id" = 1 AND "sub_tasks"."account_id" = 1 AND "tasks"."project_id" = 1
-                     sql
+    option1 = <<-sql.strip
+      SELECT "sub_tasks".* FROM "sub_tasks" INNER JOIN "tasks" ON "sub_tasks"."task_id" = "tasks"."id" AND "tasks"."account_id" = "sub_tasks"."account_id" WHERE "tasks"."project_id" = 1 AND "sub_tasks"."account_id" = 1 AND "tasks"."account_id" = 1
+    sql
+    option2 = <<-sql.strip
+      SELECT "sub_tasks".* FROM "sub_tasks" INNER JOIN "tasks" ON "sub_tasks"."task_id" = "tasks"."id" AND "tasks"."account_id" = "sub_tasks"."account_id" WHERE "sub_tasks"."account_id" = 1 AND "tasks"."project_id" = 1 AND "tasks"."account_id" = 1
+    sql
 
     account1 = Account.create! name: 'Account 1'
 
@@ -358,13 +485,13 @@ describe MultiTenant do
       project1 = Project.create! name: 'Project 1'
       task1 = Task.create! name: 'Task 1', project: project1
       subtask1 = SubTask.create! task: task1
-      expect(project1.sub_tasks.to_sql).to eq(expected_sql.strip)
+      expect(project1.sub_tasks.to_sql).to eq(option1).or(eq(option2))
       expect(project1.sub_tasks).to include(subtask1)
     end
 
     MultiTenant.without do
       expected_sql = <<-sql
-                        SELECT "sub_tasks".* FROM "sub_tasks" INNER JOIN "tasks" ON "sub_tasks"."task_id" = "tasks"."id" AND "sub_tasks"."account_id" = "tasks"."account_id" WHERE "tasks"."project_id" = 1
+                        SELECT "sub_tasks".* FROM "sub_tasks" INNER JOIN "tasks" ON "sub_tasks"."task_id" = "tasks"."id" AND "tasks"."account_id" = "sub_tasks"."account_id" WHERE "tasks"."project_id" = 1
                      sql
 
       project = Project.first
@@ -373,9 +500,13 @@ describe MultiTenant do
   end
 
   it "tests joins between distributed and reference table" do
-    expected_sql = <<-sql
-                     SELECT "categories".* FROM "categories" INNER JOIN "project_categories" ON "categories"."id" = "project_categories"."category_id" WHERE "project_categories"."account_id" = 1 AND "project_categories"."project_id" = 1
-                   sql
+    option1 = <<-sql.strip
+      SELECT "categories".* FROM "categories" INNER JOIN "project_categories" ON "categories"."id" = "project_categories"."category_id" WHERE "project_categories"."project_id" = 1 AND "project_categories"."account_id" = 1
+    sql
+    option2 = <<-sql.strip
+      SELECT "categories".* FROM "categories" INNER JOIN "project_categories" ON "categories"."id" = "project_categories"."category_id" WHERE "project_categories"."account_id" = 1 AND "project_categories"."project_id" = 1
+    sql
+
     account1 = Account.create! name: 'Account 1'
     category1 = Category.create! name: 'Category 1'
 
@@ -383,7 +514,7 @@ describe MultiTenant do
       project1 = Project.create! name: 'Project 1'
       projectcategory = ProjectCategory.create! name: 'project cat 1', project: project1, category: category1
 
-      expect(project1.categories.to_sql).to eq(expected_sql.strip)
+      expect(project1.categories.to_sql).to eq(option1).or(eq(option2))
       expect(project1.categories).to include(category1)
       expect(project1.project_categories).to include(projectcategory)
     end
@@ -398,7 +529,7 @@ describe MultiTenant do
       expect(project.categories).to include(category1)
 
       expected_sql = <<-sql
-                        SELECT "projects".* FROM "projects" INNER JOIN "project_categories" ON "project_categories"."project_id" = "projects"."id" AND "project_categories"."account_id" = "projects"."account_id" INNER JOIN "categories" ON "categories"."id" = "project_categories"."category_id" WHERE "projects"."account_id" = 1
+                        SELECT "projects".* FROM "projects" INNER JOIN "project_categories" ON "project_categories"."project_id" = "projects"."id" AND "projects"."account_id" = "project_categories"."account_id" INNER JOIN "categories" ON "categories"."id" = "project_categories"."category_id" WHERE "projects"."account_id" = 1
                     sql
 
       expect(Project.where(account_id: 1).joins(:categories).to_sql).to eq(expected_sql.strip)
@@ -412,21 +543,18 @@ describe MultiTenant do
     account1 = Account.create! name: 'Account 1'
     category1 = Category.create! name: 'Category 1'
 
-    expected_sql = if uses_prepared_statements? && (ActiveRecord::VERSION::MAJOR == 5 || (ActiveRecord::VERSION::MAJOR == 6 && ActiveRecord::VERSION::MINOR >= 1))
-                     <<-sql
-                     SELECT "projects"."id" AS t0_r0, "projects"."account_id" AS t0_r1, "projects"."name" AS t0_r2, "categories"."id" AS t1_r0, "categories"."name" AS t1_r1 FROM "projects" LEFT OUTER JOIN "project_categories" ON "project_categories"."project_id" = "projects"."id" AND "project_categories"."account_id" = 1 AND "projects"."account_id" = 1 LEFT OUTER JOIN "categories" ON "categories"."id" = "project_categories"."category_id" AND "project_categories"."account_id" = 1 WHERE "projects"."account_id" = 1
-                     sql
-                   else
-                     <<-sql
-                     SELECT "projects"."id" AS t0_r0, "projects"."account_id" AS t0_r1, "projects"."name" AS t0_r2, "categories"."id" AS t1_r0, "categories"."name" AS t1_r1 FROM "projects" LEFT OUTER JOIN "project_categories" ON "project_categories"."account_id" = 1 AND "project_categories"."project_id" = "projects"."id" AND "projects"."account_id" = 1 LEFT OUTER JOIN "categories" ON "categories"."id" = "project_categories"."category_id" AND "project_categories"."account_id" = 1 WHERE "projects"."account_id" = 1
-                     sql
-                   end
+    option1 = <<-sql.strip
+      SELECT "projects"."id" AS t0_r0, "projects"."account_id" AS t0_r1, "projects"."name" AS t0_r2, "categories"."id" AS t1_r0, "categories"."name" AS t1_r1 FROM "projects" LEFT OUTER JOIN "project_categories" ON "project_categories"."project_id" = "projects"."id" AND "project_categories"."account_id" = 1 AND "projects"."account_id" = 1 LEFT OUTER JOIN "categories" ON "categories"."id" = "project_categories"."category_id" AND "project_categories"."account_id" = 1 WHERE "projects"."account_id" = 1
+    sql
+    option2 = <<-sql.strip
+      SELECT "projects"."id" AS t0_r0, "projects"."account_id" AS t0_r1, "projects"."name" AS t0_r2, "categories"."id" AS t1_r0, "categories"."name" AS t1_r1 FROM "projects" LEFT OUTER JOIN "project_categories" ON "project_categories"."account_id" = 1 AND "project_categories"."project_id" = "projects"."id" AND "projects"."account_id" = 1 LEFT OUTER JOIN "categories" ON "categories"."id" = "project_categories"."category_id" AND "project_categories"."account_id" = 1 WHERE "projects"."account_id" = 1
+    sql
 
     MultiTenant.with(account1) do
       project1 = Project.create! name: 'Project 1'
       projectcategory = ProjectCategory.create! name: 'project cat 1', project: project1, category: category1
 
-      expect(Project.eager_load(:categories).to_sql).to eq(expected_sql.strip)
+      expect(Project.eager_load(:categories).to_sql).to eq(option1).or(eq(option2))
 
       project = Project.eager_load(:categories).first
       expect(project.categories).to include(category1)
@@ -435,7 +563,7 @@ describe MultiTenant do
 
     MultiTenant.without do
       expected_sql = <<-sql
-                     SELECT "projects"."id" AS t0_r0, "projects"."account_id" AS t0_r1, "projects"."name" AS t0_r2, "categories"."id" AS t1_r0, "categories"."name" AS t1_r1 FROM "projects" LEFT OUTER JOIN "project_categories" ON "project_categories"."project_id" = "projects"."id" AND "project_categories"."account_id" = "projects"."account_id" LEFT OUTER JOIN "categories" ON "categories"."id" = "project_categories"."category_id" WHERE "projects"."account_id" = 1
+                     SELECT "projects"."id" AS t0_r0, "projects"."account_id" AS t0_r1, "projects"."name" AS t0_r2, "categories"."id" AS t1_r0, "categories"."name" AS t1_r1 FROM "projects" LEFT OUTER JOIN "project_categories" ON "project_categories"."project_id" = "projects"."id" AND "projects"."account_id" = "project_categories"."account_id" LEFT OUTER JOIN "categories" ON "categories"."id" = "project_categories"."category_id" WHERE "projects"."account_id" = 1
                      sql
 
       expect(Project.where(account_id: 1).eager_load(:categories).to_sql).to eq(expected_sql.strip)
@@ -451,26 +579,23 @@ describe MultiTenant do
     category1 = Category.create! name: 'Category 1'
 
     MultiTenant.with(account1) do
-      expected_sql =  if uses_prepared_statements? && (ActiveRecord::VERSION::MAJOR == 5 || (ActiveRecord::VERSION::MAJOR == 6 && ActiveRecord::VERSION::MINOR >= 1))
-                        <<-sql
-                        SELECT "tasks".* FROM "tasks" INNER JOIN "projects" ON "projects"."id" = "tasks"."project_id" AND "projects"."account_id" = 1 LEFT JOIN project_categories pc ON project.category_id = pc.id WHERE "tasks"."account_id" = 1
-                        sql
-                      else
-                        <<-sql
-                        SELECT "tasks".* FROM "tasks" INNER JOIN "projects" ON "projects"."account_id" = 1 AND "projects"."id" = "tasks"."project_id" LEFT JOIN project_categories pc ON project.category_id = pc.id WHERE "tasks"."account_id" = 1
-                        sql
-                      end
+      option1 = <<-sql.strip
+        SELECT "tasks".* FROM "tasks" INNER JOIN "projects" ON "projects"."id" = "tasks"."project_id" AND "projects"."account_id" = 1 LEFT JOIN project_categories pc ON project.category_id = pc.id WHERE "tasks"."account_id" = 1
+      sql
+      option2 = <<-sql.strip
+        SELECT "tasks".* FROM "tasks" INNER JOIN "projects" ON "projects"."account_id" = 1 AND "projects"."id" = "tasks"."project_id" LEFT JOIN project_categories pc ON project.category_id = pc.id WHERE "tasks"."account_id" = 1
+      sql
 
       project1 = Project.create! name: 'Project 1'
       projectcategory = ProjectCategory.create! name: 'project cat 1', project: project1, category: category1
 
       project1.tasks.create! name: 'baz'
-      expect(Task.joins(:project).joins('LEFT JOIN project_categories pc ON project.category_id = pc.id').to_sql).to eq(expected_sql.strip)
+      expect(Task.joins(:project).joins('LEFT JOIN project_categories pc ON project.category_id = pc.id').to_sql).to eq(option1).or(eq(option2))
     end
 
     MultiTenant.without do
       expected_sql = <<-sql
-                     SELECT "tasks".* FROM "tasks" INNER JOIN "projects" ON "projects"."id" = "tasks"."project_id" AND "projects"."account_id" = "tasks"."account_id" LEFT JOIN project_categories pc ON project.category_id = pc.id WHERE "tasks"."account_id" = 1
+                     SELECT "tasks".* FROM "tasks" INNER JOIN "projects" ON "projects"."id" = "tasks"."project_id" AND "tasks"."account_id" = "projects"."account_id" LEFT JOIN project_categories pc ON project.category_id = pc.id WHERE "tasks"."account_id" = 1
                      sql
 
       expect(Task.where(account_id: 1).joins(:project).joins('LEFT JOIN project_categories pc ON project.category_id = pc.id').to_sql).to eq(expected_sql.strip)
@@ -485,32 +610,39 @@ describe MultiTenant do
     project2 = Project.create! name: 'Project 2', account: Account.create!(name: 'Account2')
 
     MultiTenant.with(account) do
-      expected_sql = if uses_prepared_statements? && ActiveRecord::VERSION::MAJOR > 5
-                       <<-sql.strip
-                                SELECT "projects".* FROM "projects" WHERE "projects"."account_id" = #{account.id} AND "projects"."id" = $1 LIMIT $2
-                         sql
-                     else
-                       <<-sql.strip
-                         SELECT  "projects".* FROM "projects" WHERE "projects"."account_id" = #{account.id} AND "projects"."id" = $1 LIMIT $2
-                         sql
-                     end
-
-      expect(Project).to receive(:find_by_sql).with(expected_sql, any_args).and_call_original
+      option1 = <<-sql.strip
+        SELECT "projects".* FROM "projects" WHERE "projects"."account_id" = #{account.id} AND "projects"."id" = $1 LIMIT $2
+      sql
+      option2 = <<-sql.strip
+        SELECT "projects".* FROM "projects" WHERE "projects"."id" = $1 AND "projects"."account_id" = #{account.id} LIMIT $2
+      sql
+      option3 = <<-sql.strip
+        SELECT  "projects".* FROM "projects" WHERE "projects"."id" = $1 AND "projects"."account_id" = #{account.id} LIMIT $2
+      sql
+
+      # Couldn't make the following line pass for some reason, so came up with an uglier alternative
+      # expect(Project).to receive(:find_by_sql).with(eq(option1).or(eq(option2)).or(eq(option3)), any_args).and_call_original
+      expect(Project).to receive(:find_by_sql).and_wrap_original do |m, *args|
+        expect(args[0]).to(eq(option1).or(eq(option2)).or(eq(option3)))
+        m.call(args[0], args[1], preparable:args[2][:preparable])
+      end
       expect(Project.find(project.id)).to eq(project)
     end
 
     MultiTenant.without do
-      expected_sql = if uses_prepared_statements? && ActiveRecord::VERSION::MAJOR > 5
-                       <<-sql.strip
-                         SELECT "projects".* FROM "projects" WHERE "projects"."id" = $1 LIMIT $2
-                         sql
-                     else
-                       <<-sql.strip
-                         SELECT  "projects".* FROM "projects" WHERE "projects"."id" = $1 LIMIT $2
-                         sql
-                     end
-
-      expect(Project).to receive(:find_by_sql).with(expected_sql, any_args).and_call_original
+      option1 = <<-sql.strip
+        SELECT "projects".* FROM "projects" WHERE "projects"."id" = $1 LIMIT $2
+      sql
+      option2 = <<-sql.strip
+        SELECT  "projects".* FROM "projects" WHERE "projects"."id" = $1 LIMIT $2
+      sql
+
+      # Couldn't make the following line pass for some reason, so came up with an uglier alternative
+      # expect(Project).to receive(:find_by_sql).with(eq(option1).or(eq(option2)), any_args).and_call_original
+      expect(Project).to receive(:find_by_sql).and_wrap_original do |m, *args|
+        expect(args[0]).to(eq(option1).or(eq(option2)))
+        m.call(args[0], args[1], preparable:args[2][:preparable])
+      end
       expect(Project.find(project2.id)).to eq(project2)
     end
   end

data/spec/activerecord-multi-tenant/query_rewriter_spec.rb

@@ -98,4 +98,21 @@ describe "Query Rewriter" do
       }.not_to raise_error
     end
   end
+
+  context "when joining with a model with a default scope" do
+    let!(:account) { Account.create!(name: "Test Account") }
+
+    it "fetches only records within the default scope" do
+      alive = Domain.create(name: "alive", account: account)
+      deleted = Domain.create(name: "deleted", deleted: true, account: account)
+      page_in_alive_domain = Page.create(name: "alive", account: account, domain: alive)
+      page_in_deleted_domain = Page.create(name: "deleted", account: account, domain: deleted)
+
+      expect(
+        MultiTenant.with(account) do
+          Page.joins(:domain).pluck(:id)
+        end
+      ).to eq([page_in_alive_domain.id])
+    end
+  end
 end

data/spec/activerecord-multi-tenant/record_modifications_spec.rb

@@ -54,6 +54,25 @@ describe MultiTenant, 'Record modifications' do
     end
   end
 
+  it 'should not update other objects with same id when calling object.update_columns' do
+    # When two records with same id but different account_id are updated, it should only update the current one
+    expect(project.account).to eq(account)
+    expect(project2.account).to eq(account2)
+    expect(project.id).to eq(project2.id)
+
+    MultiTenant.without do
+      project2.update_columns(name: 'newthing2')
+      expect(project.reload.name).to eq('something')
+      expect(project2.reload.name).to eq('newthing2')
+    end
+  end
+
+  it 'should return the same object when calling object.reload' do
+    # When two records with same id but different account_id are updated, it should not return the other object
+    expect(project.reload.account_id).to eq(account.id)
+    expect(project2.reload.account_id).to eq(account2.id)
+  end
+
   it 'test delete for reference tables' do
     category1 = Category.create! name: 'Category 1'
     expect(Category.count).to eq(1)

data/spec/activerecord-multi-tenant/sidekiq_spec.rb

@@ -5,6 +5,9 @@ require 'activerecord-multi-tenant/sidekiq'
 describe MultiTenant, 'Sidekiq' do
   let(:server) { Sidekiq::Middleware::MultiTenant::Server.new }
   let(:account) { Account.create(name: 'test') }
+  let(:deleted_acount) { Account.create(name: 'deleted') }
+
+  before { deleted_acount.destroy! }
 
   describe 'server middleware' do
     it 'sets the multitenant context when provided in message' do
@@ -15,6 +18,14 @@ describe MultiTenant, 'Sidekiq' do
       end
     end
 
+    it 'sets the multitenant context (id) even if tenant not found' do
+      server.call(double,{'bogus' => 'message',
+        'multi_tenant' => { 'class' => deleted_acount.class.name, 'id' => deleted_acount.id}},
+        'bogus_queue') do
+        expect(MultiTenant.current_tenant).to eq(deleted_acount.id)
+      end
+    end
+
     it 'does not set the multitenant context when no tenant provided' do
       server.call(double, {'bogus' => 'message'}, 'bogus_queue') do
         expect(MultiTenant.current_tenant).to be_nil

data/spec/schema.rb

@@ -34,6 +34,13 @@ ARGV.grep(/\w+_spec\.rb/).empty? && ActiveRecord::Schema.define(version: 1) do
     t.column :type, :string
   end
 
+  create_table :optional_sub_tasks, force: true do |t|
+    t.references :account, :integer
+    t.column :sub_task_id, :integer
+    t.column :name, :string
+    t.column :type, :string
+  end
+
   create_table :countries, force: true do |t|
     t.column :name, :string
   end
@@ -89,10 +96,26 @@ ARGV.grep(/\w+_spec\.rb/).empty? && ActiveRecord::Schema.define(version: 1) do
     t.column :category_id, :integer
   end
 
-
   create_table :allowed_places, force: true, id: false do |t|
-  t.string :account_id, :integer
-  t.string :name, :string
+    t.string :account_id, :integer
+    t.string :name, :string
+  end
+
+  create_table :domains, force: true, partition_key: :account_id do |t|
+    t.column :account_id, :integer
+    t.column :name, :string
+    t.column :deleted, :boolean, default: false
+  end
+
+  create_table :pages, force: true, partition_key: :account_id do |t|
+    t.column :account_id, :integer
+    t.column :name, :string
+    t.column :domain_id, :integer
+  end
+
+  create_table :posts, force: true, partition_key: :account_id do |t|
+    t.column :account_id, :integer
+    t.column :name, :string
   end
 
   create_distributed_table :accounts, :id
@@ -108,6 +131,9 @@ ARGV.grep(/\w+_spec\.rb/).empty? && ActiveRecord::Schema.define(version: 1) do
   create_distributed_table :uuid_records, :organization_id
   create_distributed_table :project_categories, :account_id
   create_distributed_table :allowed_places, :account_id
+  create_distributed_table :domains, :account_id
+  create_distributed_table :pages, :account_id
+  create_distributed_table :posts, :account_id
   create_reference_table :categories
 end
 
@@ -115,6 +141,7 @@ class Account < ActiveRecord::Base
   multi_tenant :account
   has_many :projects
   has_one :manager, inverse_of: :account
+  has_many :optional_sub_tasks
 end
 
 class Project < ActiveRecord::Base
@@ -146,6 +173,14 @@ class SubTask < ActiveRecord::Base
   multi_tenant :account
   belongs_to :task
   has_one :project, through: :task
+  has_many :optional_sub_tasks
+end
+
+with_belongs_to_required_by_default do
+  class OptionalSubTask < ActiveRecord::Base
+    multi_tenant :account, optional: true
+    belongs_to :sub_task
+  end
 end
 
 class StiSubTask < SubTask
@@ -185,6 +220,7 @@ class Comment < ActiveRecord::Base
 end
 
 class Organization < ActiveRecord::Base
+  multi_tenant :organization
   has_many :uuid_records
 end
 
@@ -204,7 +240,17 @@ class ProjectCategory < ActiveRecord::Base
   belongs_to :account
 end
 
-
 class AllowedPlace < ActiveRecord::Base
   multi_tenant :account
 end
+
+class Domain < ActiveRecord::Base
+  multi_tenant :account
+  has_many :pages
+  default_scope { where(deleted: false) }
+end
+
+class Page < ActiveRecord::Base
+  multi_tenant :account
+  belongs_to :domain
+end

data/spec/spec_helper.rb

@@ -46,4 +46,11 @@ def uses_prepared_statements?
   ActiveRecord::Base.connection.prepared_statements
 end
 
+def with_belongs_to_required_by_default(&block)
+  default_value = ActiveRecord::Base.belongs_to_required_by_default
+  ActiveRecord::Base.belongs_to_required_by_default = true
+  yield
+ensure
+  ActiveRecord::Base.belongs_to_required_by_default = default_value
+end
 require 'schema'