Scaling Startup Codebases: Nix those Rails N+1 queries

In a young startup, it is completely normal for MVP code to be written under a tight deadline, for projects to have pivoted, or for corners to be cut due to budget constraints. It’s also likely that the codebase started without a style guide or a set of best practices to adhere to. All of this is expected, and only implies that working in an early stage startup is a wild ride. But, as your startup blossoms into a successful, larger business, you’ll find that the side effects of hastily-written code move from minor annoyances to Big Deals. Fortunately, there are some simple steps that you can take to make things better.

This is part one of a series on helping your codebase handle more — more web traffic, and more engineers. This series will tackle some common problems with startup codebases — problems with readability, slowness, poor abstractions, and unwieldy user interfaces — and help you mitigate them. In this post, we’ll examine how to improve N+1 database queries, queries for each object in a collection that should be combined into one.

We’ll look at three stages of making these queries better:

1. Stop the hemorrhaging with a quick, low risk code change.
2. In a larger refactor, optimize down to a single query.
3. Find a better abstraction to make the app code more readable.

At Beam, two technologies we use are Ruby on Rails and MySQL. Our database and our user base are both getting pretty big. Every time an insurance broker asks our web app for a list of companies and insurance policies under their purview, something like this happens in our logs:

Company Load (1.0ms) SELECT `companies`.* FROM `companies` INNER JOIN `insurance_policies` ON `companies`.`id` = `insurance_policies`.`company_id` WHERE `insurance_policies`.`insurance_plan_id` = ‘plan1’
Company Load (1.2ms) SELECT `companies`.* FROM `companies` INNER JOIN `insurance_policies` ON `companies`.`id` = `insurance_policies`.`company_id` WHERE `insurance_policies`.`insurance_plan_id` = ‘plan2’
Company Load (1.5ms) SELECT `companies`.* FROM `companies` INNER JOIN `insurance_policies` ON `companies`.`id` = `insurance_policies`.`company_id` WHERE `insurance_policies`.`insurance_plan_id` = ‘plan3’
Company Load (3.5ms) SELECT `companies`.* FROM `companies` INNER JOIN `insurance_policies` ON `companies`.`id` = `insurance_policies`.`company_id` WHERE `insurance_policies`.`insurance_plan_id` = ‘plan4’
Company Load (1.3ms) SELECT `companies`.* FROM `companies` INNER JOIN `insurance_policies` ON `companies`.`id` = `insurance_policies`.`company_id` WHERE `insurance_policies`.`insurance_plan_id` = ‘plan5’

This is a classic N+1 query. A prior call to the database returns a list of insurance policies managed by a broker, and an enumerator iterates over the results and queries for the associated company. To guard against N+1 queries like these (and to help find the remaining ones), we use the Bullet gem and see notes like this in our logs:

USE eager loading detected
 InsurancePlan => [:companies]
 Add to your finder: :includes => [:companies]

Rails’ ActiveRecord database ORM has some built-in methods to join tables and to make enumerators over database rows a little smarter. We’ll go over some examples of using these from a Rails console and show how they affect calls to the database. To guard against accidentally spilling the T on user data, these examples will use some made up dental procedure codes and their associated categories. Codes and categories have these database associations:

class ProcedureCode < ApplicationRecord
 belongs_to :procedure_category
end
class ProcedureCategory < ApplicationRecord
 has_many :procedure_codes
end

On the database procedure_codes has a procedure_category_id column to make this relationship work. Here is a naive rails query/map combo to return category names for our first 5 procedure codes:

[1] pry(main)> ProcedureCode.order(:id).limit(5).map { |code| code.procedure_category.name }
 ProcedureCode Load (1.0ms) SELECT `procedure_codes`.* FROM `procedure_codes` ORDER BY `procedure_codes`.`id` ASC LIMIT 5
 ProcedureCategory Load (1.0ms) SELECT `procedure_categories`.`id`, `procedure_categories`.`name`, `procedure_categories`.`created_at`, `procedure_categories`.`updated_at`, `procedure_categories`.`consumer_body` FROM `procedure_categories` WHERE `procedure_categories`.`id` = ‘category1’ LIMIT 1
 ProcedureCategory Load (1.8ms) SELECT `procedure_categories`.`id`, `procedure_categories`.`name`, `procedure_categories`.`created_at`, `procedure_categories`.`updated_at`, `procedure_categories`.`consumer_body` FROM `procedure_categories` WHERE `procedure_categories`.`id` = ‘category2’ LIMIT 1
 ProcedureCategory Load (1.5ms) SELECT `procedure_categories`.`id`, `procedure_categories`.`name`, `procedure_categories`.`created_at`, `procedure_categories`.`updated_at`, `procedure_categories`.`consumer_body` FROM `procedure_categories` WHERE `procedure_categories`.`id` = ‘category3’ LIMIT 1
 ProcedureCategory Load (1.1ms) SELECT `procedure_categories`.`id`, `procedure_categories`.`name`, `procedure_categories`.`created_at`, `procedure_categories`.`updated_at`, `procedure_categories`.`consumer_body` FROM `procedure_categories` WHERE `procedure_categories`.`id` = ‘category4’ LIMIT 1
 ProcedureCategory Load (0.9ms) SELECT `procedure_categories`.`id`, `procedure_categories`.`name`, `procedure_categories`.`created_at`, `procedure_categories`.`updated_at`, `procedure_categories`.`consumer_body` FROM `procedure_categories` WHERE `procedure_categories`.`id` = ‘category5’ LIMIT 1
=> [“Diagnostic”, “Preventive”, “Child Orthodontics”, “Minor Restorative”, “Prosthetic Maintenance”]

This approach generated 6 database queries: one to get our list of codes, and five more queries search for a category name for each code. The simplest refactor of this is the method suggested by Bullet: add an includes to the query. Using includes won’t change the initial query, but it will make the subsequent enumerator only touch the database one more time:

[2] pry(main)> ProcedureCode.order(:id).includes(:procedure_category).limit(5).map { |code| code.procedure_category.name }
 ProcedureCode Load (1.4ms) SELECT `procedure_codes`.* FROM `procedure_codes` ORDER BY `procedure_codes`.`id` ASC LIMIT 5
 ProcedureCategory Load (1.1ms) SELECT `procedure_categories`.`id`, `procedure_categories`.`name`, `procedure_categories`.`created_at`, `procedure_categories`.`updated_at`, `procedure_categories`.`consumer_body` FROM `procedure_categories` WHERE `procedure_categories`.`id` IN (‘category1’, ‘category2’, ‘category3’, ‘category4’, ‘category5’)
=> [“Diagnostic”, “Preventive”, “Child Orthodontics”, “Minor Restorative”, “Prosthetic Maintenance”]

Immediately, we’ve saved 4 queries. This is a low-risk, low-hanging fruit that you can sprinkle throughout your codebase every time Bullet complains. If you have a lot of fires to put out, this is a pragmatic place to stop; for the cost of two words, you’ve already saved the database a lot of work. When you have time for a larger refactor, you can usually do everything in a single query with a better, easier to read abstraction. Let’s dig a little deeper…

An alternative to includes is joins, which performs a literal SQL JOIN of two tables. If we look at the SQL returned by our previous query, you’ll see that until we get to the map step; the category table isn’t even referenced. Replacing includes with joins, on the other hand, returns the joined results in one query. Let’s compare the raw SQL returned by both of these:

[3] pry(main)> ProcedureCode.order(:id).includes(:procedure_category).limit(5).to_sql
=> “SELECT `procedure_codes`.* FROM `procedure_codes` ORDER BY `procedure_codes`.`id` ASC LIMIT 5”
[4] pry(main)> ProcedureCode.order(:id).joins(:procedure_category).limit(5).to_sql
=> “SELECT `procedure_codes`.* FROM `procedure_codes` INNER JOIN `procedure_categories` ON `procedure_categories`.`id` = `procedure_codes`.`procedure_category_id` ORDER BY `procedure_codes`.`id` ASC LIMIT 5”

By itself, switching to joins doesn’t take us all the way to our goal, since it doesn’t do anything to the enumerator, and we’re right back where we started:

[5] pry(main)> ProcedureCode.order(:id).joins(:procedure_category).limit(5).map { |code| code.procedure_category.name }
 ProcedureCode Load (1.7ms) SELECT `procedure_codes`.* FROM `procedure_codes` INNER JOIN `procedure_categories` ON `procedure_categories`.`id` = `procedure_codes`.`procedure_category_id` ORDER BY `procedure_codes`.`id` ASC LIMIT 5
 ProcedureCategory Load (1.1ms) SELECT `procedure_categories`.`id`, `procedure_categories`.`name`, `procedure_categories`.`created_at`, `procedure_categories`.`updated_at`, `procedure_categories`.`consumer_body` FROM `procedure_categories` WHERE `procedure_categories`.`id` = ‘category1’ LIMIT 1
 ProcedureCategory Load (1.1ms) SELECT `procedure_categories`.`id`, `procedure_categories`.`name`, `procedure_categories`.`created_at`, `procedure_categories`.`updated_at`, `procedure_categories`.`consumer_body` FROM `procedure_categories` WHERE `procedure_categories`.`id` = ‘category2’ LIMIT 1
 ProcedureCategory Load (0.9ms) SELECT `procedure_categories`.`id`, `procedure_categories`.`name`, `procedure_categories`.`created_at`, `procedure_categories`.`updated_at`, `procedure_categories`.`consumer_body` FROM `procedure_categories` WHERE `procedure_categories`.`id` = ‘category3’ LIMIT 1
 ProcedureCategory Load (1.2ms) SELECT `procedure_categories`.`id`, `procedure_categories`.`name`, `procedure_categories`.`created_at`, `procedure_categories`.`updated_at`, `procedure_categories`.`consumer_body` FROM `procedure_categories` WHERE `procedure_categories`.`id` = ‘category4’ LIMIT 1
 ProcedureCategory Load (1.0ms) SELECT `procedure_categories`.`id`, `procedure_categories`.`name`, `procedure_categories`.`created_at`, `procedure_categories`.`updated_at`, `procedure_categories`.`consumer_body` FROM `procedure_categories` WHERE `procedure_categories`.`id` = ‘category5’ LIMIT 1
=> [“Diagnostic”, “Preventive”, “Child Orthodontics”, “Minor Restorative”, “Prosthetic Maintenance”]

The select method can be added to our query, which will inject data from the joined table into the result set:

[6] pry(main)> ProcedureCode.order(:id).joins(:procedure_category).select(‘procedure_categories.name’).limit(5)
 ProcedureCode Load (3.0ms) SELECT procedure_categories.name FROM `procedure_codes` INNER JOIN `procedure_categories` ON `procedure_categories`.`id` = `procedure_codes`.`procedure_category_id` ORDER BY `procedure_codes`.`id` ASC LIMIT 5
 => [#<ProcedureCode:0x00007f84c36b8290 id: nil, name: “Diagnostic”>,
 #<ProcedureCode:0x00007f84c36ebf50 id: nil, name: “Preventive”>,
 #<ProcedureCode:0x00007f84c36ebd70 id: nil, name: “Child Orthodontics”>,
 #<ProcedureCode:0x00007f84c36ebc30 id: nil, name: “Minor Restorative”>,
 #<ProcedureCode:0x00007f84c36ebaf0 id: nil, name: “Prosthetic Maintenance”>]

The SQL is much nicer with this approach, and this is the first time we’re querying for just the name column in procedure_categories. It looks like we have everything we need from the query, and all we need to do is pluck the name field from the results:

[7] pry(main)> ProcedureCode.order(:id).joins(:procedure_category).select(‘procedure_categories.name’).limit(5).pluck :name
 (1.6ms) SELECT `procedure_codes`.`name` FROM `procedure_codes` INNER JOIN `procedure_categories` ON `procedure_categories`.`id` = `procedure_codes`.`procedure_category_id` ORDER BY `procedure_codes`.`id` ASC LIMIT 5
 => [“code1”, “code2”, “code3”, “code4”, “code5”]

Surprise! After adding pluck, notice that the generated SQL is now asking for the procedure code’s name, not for the category name. The earlier results are an ActiveRecord relation, not an array, and pluck is still referencing the ProcedureCode class, thanks to some magic going on in ActiveRecord::Calculations#pluck. Still, mapping over the results works like you’d expect:

[8] pry(main)> ProcedureCode.order(:id).joins(:procedure_category).select(‘procedure_categories.name’).limit(5).map(&:name)
 ProcedureCode Load (2.4ms) SELECT procedure_categories.name FROM `procedure_codes` INNER JOIN `procedure_categories` ON `procedure_categories`.`id` = `procedure_codes`.`procedure_category_id` ORDER BY `procedure_codes`.`id` ASC LIMIT 5
 => [“Diagnostic”, “Preventive”, “Child Orthodontics”, “Minor Restorative”, “Prosthetic Maintenance”]

This is an improvement, and the compiled SQL is pretty clean, but we can still do a little better. I’d like for this to be a little easier for the next engineer to parse, and to start from the category rather than the procedure code. Starting from the table you care about can lead to a more natural, understandable query. For instance:

[9] pry(main)> ProcedureCategory.joins(:procedure_codes).order(‘procedure_codes.id’).limit(5).pluck :name
 (3.9ms) SELECT `procedure_categories`.`name` FROM `procedure_categories` INNER JOIN `procedure_codes` ON `procedure_codes`.`procedure_category_id` = `procedure_categories`.`id` ORDER BY procedure_codes.id LIMIT 5
 => [“Diagnostic”, “Preventive”, “Child Orthodontics”, “Minor Restorative”, “Prosthetic Maintenance”]

The next engineer will see immediately that we’re interested in the category, and that the codes are only useful for ordering the results. Pluck works the way we want in this case, so we don’t need to resort to enumerator magic. The generated SQL asks for the join in a more natural order, and it still only asks for the column we need. In my opinion, this pattern is the clear winner.

Not all N+1 refactors will be this straightforward, but they’re worth the effort to get right. Any database query you skip now pushes back the date where your company needs to pivot its data management strategy. Keeping your database queries natural and readable will keep your fellow engineers happy and push back the date where everyone starts asking for a ground-up rewrite. Those days might still come, but the easier it is for your company’s leadership to grow the business without breaking the tech stack or causing Legacy Code Fatigue, the better for everyone.