Kolide's 30 Line Rails Multi-Tenant Strategy

Tue, 09 Nov 2021 00:00:00 +0000

When engineering a new SaaS app, how you plan to handle customer data tenancy is usually one of the first decisions you and your team will need to make. If you are writing a Rails app and decide on a multi-tenant strategy (one instance of the app serves many customers), then this article is for you.

I contend that modern Rails has everything you need to build a multi-tenant strategy that scales, is easy for others to use, and can be written in just a handful of lines of code. Kolide (btw we’re hiring) has been using this simple strategy since the inception of its product, and it’s been one of the most elegant and easiest to understand parts of our code-base.

So before reaching for a gem, consider if the following meets your needs.

The Code

The entire implementation is contained in just two files and requires no additional dependencies.

# app/models/concerns/account_ownable.rb

module AccountOwnable
  extend ActiveSupport::Concern

  included do
    # Account is actually not optional, but we not do want
    # to generate a SELECT query to verify the account is
    # there every time. We get this protection for free
    # because of the `Current.account_or_raise!`
    # and also through FK constraints.
    belongs_to :account, optional: true
    default_scope { where(account: Current.account_or_raise!) }
  end
end

# app/models/current.rb
class Current < ActiveSupport::CurrentAttributes
  attribute :account, :user

  resets { Time.zone = nil }

  class MissingCurrentAccount < StandardError; end

  def account_or_raise!
    raise Current::MissingCurrentAccount, "You must set an account with Current.account=" unless account

    account
  end

  def user=(user)
    super
    self.account      = user.account
    Time.zone         = user.time_zone
  end
end

That’s it. Seriously.

Using the Code In Practice

To use this code, simply mix-in the concern into any standard ActiveRecord model like so…

class ApiKey < ApplicationRecord
  # assumes table has a column named `account_id`
  include AccountOwnable
end

When a user of ours signs in, all we need to do is simply set Current.user in our authentication controller concern which is mixed into our ApplicationController

# app/controllers/concerns/require_authentication.rb
module RequireAuthentication
  extend ActiveSupport::Concern

  included do
    before_action :ensure_authenticated_user
  end

  def ensure_authenticated_user
    if (user = User.find_by_valid_session(session))
      Current.user = user
    else
      redirect_to signin_path
    end
  end
end

For this small amount of effort we now get the following benefits:

1. Because of the default_scope, once a user is signed in, data from sensitive models is automatically scoped to their account. We just don’t need to think about it, no matter how complicated our query chaining gets.

2. Again, because of the default_scope creating new records for these AccountOwnable models will automatically set the account_id for us. One less thing to think about.

3. In situations where we are outside of the standard Rails request/response paradigm (ex: in an ActiveJob) any AccountOwnable models will raise if Current.account is not set. This forces us to constantly think about how we are scoping data for customer needs.

4. The situations where we need to enumerate through more than one tenant’s data at a time are still possible but now require a Model.unscoped which can be easily scanned for in linters requiring engineers to justify their rationale on a per use-case basis.

One thing that became slightly annoying was constantly setting Current.account = in the Rails console. To make that much easier we wrote a simple console command.

# lib/kolide/console.rb
module App
  module Console
    def t(id)
      Current.account = Account.find(id)
      puts "Current account switched to #{Current.account.name} (#{Current.account.id})"
    end
  end
end

# in config/application.rb
console do
  require 'kolide/console'
  Rails::ConsoleMethods.send :include, Kolide::Console
  TOPLEVEL_BINDING.eval('self').extend Kolide::Console # PRY
end

Now we simply run t 1 when we want to switch the tenant with an id of 1. Much better.

In the test suite, you should also reset Current before each spec/test as it’s not done for you automatically. For us that was simply a matter of adding…

# spec/spec_helper.rb
config.before(:all) do
  Current.reset
end

Now we don’t have to worry about our global state being polluted when running our specs serially in the same process.

Concerns we had that didn’t end up being an issue

Kolide has been successfully using this strategy since the inception of our Ruby on Rails SaaS app. While we arrived at this strategy in the first few days of our app’s formation, we definitely were less confident in the approach. Here is a list of concerns we held and how they ended up panning out.

Will this approach be acceptable to our customers?

Kolide is a device security company, and since our buyers are likely to be either security engineers or security minded IT staff, the bar we need to meet is much higher than the normal SaaS company. We were nervous that an app-enforced constraint would feel flimsy, despite how well it works in practice.

In reality, we found the opposite. Customers were mostly ambivalent about our app-enforced constraint approach. Why? It’s because it’s an approach that’s common among their other vendors and matches their pre-conceived expectations about how most SaaS software works. Matched expectations = less concern.

In prior iterations of our app where we did extreme things like spin up separate Kubernetes namespaces and DBs for each customer, we found our efforts were paradoxically met with more concern, not less. This concern manifested as additional process, review, and ultimately unnecessary friction as our buyers grappled to bring more and more technical folks into the procurement process to simply understand the unfamiliar architecture.

With our current approach, our development and deployment story is simpler, and simplicity has significant security advantages.

Current.rb is too magical, will multi-threading in production cause someone’s default_scope to leak to another request?

There is a lot of consternation in the Rails community when DHH introduced the CurrentAttributes paradigm in Rails 5.2. DHH talks about his rationale for adding this in his Youtube video entitled, “Using globals when the price is right”.

Ryan Bigg on the other-hand felt this addition to Rails would cause developers to write a lot of code with unpredictable behavior expressed these views in his blog post entitled, “Rails’ CurrentAttributes considered harmful”

After reading more into the original PR I found a lot of thoughtful consideration for how to make this code truly thread-safe which convinced me to bet big on this approach.

Now over two years later, I can say with confidence that in our app, which serves nearly 1,800 HTTP requests per second across hundreds of different tenants, this code works as described in a real production setting.

Will setting Current.account in asynchronous jobs or in the test suite become tiresome or problematic?

No, the ceremony here is worth it because it forces us to think carefully about how we are acting on our customer’s data. Situations where we need to iterate through more than one customer’s data are also trivial to achieve through an #each block.

ApiKey.unscoped.find_each do |api_key|
  Current.account = api_key.account
  api_key.convert_to_new_format
end

Closing Thoughts

We are now sharing this simplistic approach because it’s worked so well for us at Kolide. I imagine other burgeoning Rails apps considering a multi-tenant strategy will appreciate being able to do this in their own codebase vs offload something so important to an external gem.

We will continue to share our learnings as Kolide grows. In fact prepping for future scale is what I like best about this approach. By marking all of our records with a tenant identifier like account_id we gain the future option of leveraging more sophisticated solutions at the PostgreSQL level like multi-DB sharding or even products like Citus.