diff --git a/Gemfile.lock b/Gemfile.lock
index 659c063..df0cc3d 100644
--- a/Gemfile.lock
+++ b/Gemfile.lock
@@ -25,7 +25,7 @@ GEM
     pp (0.6.2)
       prettyprint
     prettyprint (0.2.0)
-    prism (1.4.0)
+    prism (1.9.0)
     psych (5.2.3)
       date
       stringio
diff --git a/lib/taql/cli.rb b/lib/taql/cli.rb
index 1831fdd..be82ff9 100644
--- a/lib/taql/cli.rb
+++ b/lib/taql/cli.rb
@@ -5,8 +5,10 @@ class Cli
     attr_reader :options, :query
 
     def initialize(argv)
+      @argv = argv
       @options = {markdown: false}
       @query = parse!
+      raise ArgumentError, "Usage: taql [--markdown] QUERY" if @query.empty?
     end
 
     def run
@@ -24,7 +26,7 @@ def environment_path
     def parse!
       OptionParser.new do |parser|
         parser.on("-m", "--markdown", TrueClass, "Output table in Markdown")
-      end.parse!(into: options)
+      end.parse!(@argv, into: options)
     end
 
     def silence
diff --git a/lib/taql/railtie.rb b/lib/taql/railtie.rb
index 92b345b..e4a8d5d 100644
--- a/lib/taql/railtie.rb
+++ b/lib/taql/railtie.rb
@@ -6,8 +6,16 @@ class Railtie < Rails::Railtie
 
     initializer "taql.initialize" do
       ActiveSupport.on_load(:active_record) do
-        Taql.instance_variable_set(:@default_connection, -> { ActiveRecord::Base.connection_pool.connection })
+        Taql.instance_variable_set(:@default_connection, method(:connection))
       end
     end
+
+    def self.connection
+      pool.respond_to?(:lease_connection) ? pool.lease_connection : pool.connection
+    end
+
+    def self.pool
+      ActiveRecord::Base.connection_pool
+    end
   end
 end
diff --git a/lib/taql/table.rb b/lib/taql/table.rb
index 84b42be..d2d9783 100644
--- a/lib/taql/table.rb
+++ b/lib/taql/table.rb
@@ -17,13 +17,13 @@ def body
     end
 
     def columns
-      headers.map do |header|
+      @columns ||= headers.map do |header|
         [header, *entries.map { |entry| entry[header] }]
       end
     end
 
     def headers
-      entries.map(&:keys).uniq.flatten
+      @headers ||= entries.flat_map(&:keys).uniq
     end
 
     def print
@@ -43,7 +43,7 @@ def border
     end
 
     def column_widths
-      columns.map { |column| column.map(&:length).max }
+      @column_widths ||= columns.map { |column| column.map(&:length).max }
     end
 
     def edge
diff --git a/test/test_cli.rb b/test/test_cli.rb
index f7858cb..2cdbce9 100644
--- a/test/test_cli.rb
+++ b/test/test_cli.rb
@@ -2,20 +2,20 @@
 
 class TestCli < Minitest::Test
   def test_that_it_parses_queries
-    ARGV.replace(["SELECT * FROM users"])
-
-    cli = Taql::Cli.new(ARGV)
+    cli = Taql::Cli.new(["SELECT * FROM users"])
 
     assert_equal "SELECT * FROM users", *cli.query
     assert_equal false, cli.options[:markdown]
   end
 
   def test_that_it_parses_options
-    ARGV.replace(["--markdown", "SELECT * FROM users"])
-
-    cli = Taql::Cli.new(ARGV)
+    cli = Taql::Cli.new(["--markdown", "SELECT * FROM users"])
 
     assert_equal "SELECT * FROM users", *cli.query
     assert_equal true, cli.options[:markdown]
   end
+
+  def test_that_it_raises_without_query
+    assert_raises(ArgumentError) { Taql::Cli.new([]) }
+  end
 end
diff --git a/test/test_helper.rb b/test/test_helper.rb
index 82615ea..d3e11d7 100644
--- a/test/test_helper.rb
+++ b/test/test_helper.rb
@@ -2,3 +2,4 @@
 
 require "taql"
 require "minitest/autorun"
+require "minitest/mock"
diff --git a/test/test_table.rb b/test/test_table.rb
index 472bd8a..127eaf1 100644
--- a/test/test_table.rb
+++ b/test/test_table.rb
@@ -68,4 +68,13 @@ def test_that_it_renders_nothing
 
     assert_output("\n") { puts empty_table.print }
   end
+
+  def test_headers_with_heterogeneous_entries
+    table = Taql::Table.new([
+      {"name" => "Alice", "age" => 30},
+      {"name" => "Bob", "email" => "bob@example.com"}
+    ])
+
+    assert_equal ["name", "age", "email"], table.headers
+  end
 end