Merge pull request #515 from Shopify/micro-opt-allowlist

Return early from allow list when possible

.rubocop.yml

@@ -12,6 +12,8 @@ inherit_gem:
 AllCops:
   TargetRubyVersion: 2.7
   NewCops: enable
+  Exclude:
+    - test/allow_list_bench.rb
 
 Layout/LineLength:
   Exclude:

Gemfile

@@ -7,6 +7,7 @@ gem "rake-compiler"
 
 group :test do
   gem "benchmark-memory"
+  gem "benchmark-ips"
   gem "memory_profiler"
   gem "minitest"
   gem "mocha"

lib/semian/activerecord_trilogy_adapter.rb

@@ -29,6 +29,39 @@ module ActiveRecordTrilogyAdapter
     ResourceBusyError = ::ActiveRecord::ConnectionAdapters::TrilogyAdapter::ResourceBusyError
     CircuitOpenError = ::ActiveRecord::ConnectionAdapters::TrilogyAdapter::CircuitOpenError
 
+    QUERY_ALLOWLIST = %r{\A(?:/\*.*?\*/)?\s*(ROLLBACK|COMMIT|RELEASE\s+SAVEPOINT)}i
+
+    # The common case here is NOT to have transaction management statements, therefore
+    # we are exploiting the fact that Active Record will use COMMIT/ROLLBACK as
+    # the suffix of the command string and
+    # name savepoints by level of nesting as `active_record_1` ... n.
+    #
+    # Since looking at the last characters in a string using `end_with?` is a LOT cheaper than
+    # running a regex, we are returning early if the last characters of
+    # the SQL statements are NOT the last characters of the known transaction
+    # control statements.
+    class << self
+      def query_allowlisted?(sql, *)
+        # COMMIT, ROLLBACK
+        tx_command_statement = sql.end_with?("T") || sql.end_with?("K")
+
+        # RELEASE SAVEPOINT. Nesting past _3 levels won't get bypassed.
+        # Active Record does not send trailing spaces or `;`, so we are in the realm of hand crafted queries here.
+        savepoint_statement = sql.end_with?("_1") || sql.end_with?("_2")
+        unclear = sql.end_with?(" ") || sql.end_with?(";")
+
+        if !tx_command_statement && !savepoint_statement && !unclear
+          false
+        else
+          QUERY_ALLOWLIST.match?(sql)
+        end
+      rescue ArgumentError
+        return false unless sql.valid_encoding?
+
+        raise
+      end
+    end
+
     attr_reader :raw_semian_options, :semian_identifier
 
     def initialize(*options)
@@ -48,7 +81,7 @@ def initialize(*options)
     end
 
     def raw_execute(sql, *)
-      if query_allowlisted?(sql)
+      if Semian::ActiveRecordTrilogyAdapter.query_allowlisted?(sql)
         super
       else
         acquire_semian_resource(adapter: :trilogy_adapter, scope: :query) do
@@ -90,17 +123,6 @@ def resource_exceptions
       ]
     end
 
-    # TODO: share this with Mysql2
-    QUERY_ALLOWLIST = %r{\A(?:/\*.*?\*/)?\s*(ROLLBACK|COMMIT|RELEASE\s+SAVEPOINT)}i
-
-    def query_allowlisted?(sql, *)
-      QUERY_ALLOWLIST.match?(sql)
-    rescue ArgumentError
-      return false unless sql.valid_encoding?
-
-      raise
-    end
-
     def connect(*args)
       acquire_semian_resource(adapter: :trilogy_adapter, scope: :connection) do
         super

test/adapters/activerecord_trilogy_adapter_test.rb

@@ -306,52 +306,52 @@ def test_semian_allows_rollback
         @adapter.execute("START TRANSACTION;")
 
         Semian[:mysql_testing].acquire do
-          @adapter.execute("ROLLBACK;")
+          @adapter.execute("ROLLBACK")
         end
       end
 
       def test_semian_allows_rollback_with_marginalia
         @adapter.execute("START TRANSACTION;")
 
         Semian[:mysql_testing].acquire do
-          @adapter.execute("/*foo:bar*/ ROLLBACK;")
+          @adapter.execute("/*foo:bar*/ ROLLBACK")
         end
       end
 
       def test_semian_allows_commit
         @adapter.execute("START TRANSACTION;")
 
         Semian[:mysql_testing].acquire do
-          @adapter.execute("COMMIT;")
+          @adapter.execute("COMMIT")
         end
       end
 
       def test_query_allowlisted_returns_false_for_binary_sql
         binary_query = File.read(File.expand_path("../../fixtures/binary.sql", __FILE__))
 
-        refute(@adapter.send(:query_allowlisted?, binary_query))
+        refute(Semian::ActiveRecordTrilogyAdapter.query_allowlisted?(binary_query))
       end
 
-      def test_semian_allows_rollback_to_safepoint
+      def test_semian_allows_release_savepoint
         @adapter.execute("START TRANSACTION;")
-        @adapter.execute("SAVEPOINT foobar;")
+        @adapter.execute("SAVEPOINT active_record_2;")
 
         Semian[:mysql_testing].acquire do
-          @adapter.execute("ROLLBACK TO foobar;")
+          @adapter.execute("RELEASE SAVEPOINT active_record_2")
         end
 
         @adapter.execute("ROLLBACK;")
       end
 
-      def test_semian_allows_release_savepoint
+      def test_semian_allows_rollback_to_savepoint
         @adapter.execute("START TRANSACTION;")
-        @adapter.execute("SAVEPOINT foobar;")
+        @adapter.execute("SAVEPOINT active_record_1;")
 
         Semian[:mysql_testing].acquire do
-          @adapter.execute("RELEASE SAVEPOINT foobar;")
+          @adapter.execute("ROLLBACK TO SAVEPOINT active_record_1")
         end
 
-        @adapter.execute("ROLLBACK;")
+        @adapter.execute("ROLLBACK")
       end
 
       def test_changes_timeout_when_half_open_and_configured

test/allow_list_bench.rb

@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+require "semian"
+require "active_record"
+require "semian/activerecord_trilogy_adapter"
+require "benchmark/ips"
+
+# sql_commit = "/*key:value,key:morevalue,morekey:morevalue,id:IDIDIDIDI,anotherkey:anothervalue,k:v,k:v*/ COMMIT"
+
+# Common case
+sql_not_commit = "/*key:value,key:morevalue,morekey:morevalue,id:IDIDIDIDID,anotherkey:anothervalue,k:v,k:v*/ SELECT /*+ 1111111111111111111111111*/ `line_items`.`id`, `line_items`.`shop_id`, `line_items`.`title`, `line_items`.`sku`, `line_items`.`vendor`,`line_items`.`variant_id`, `line_items`.`variant_title`, `line_items`.`order_id`, `line_items`.`currency`, `line_items`.`presentment_price`, `line_items`.`price`, `line_items`.`gift_card` FROM `line_items` WHERE `line_items`.`id` = ASKJAKJSDASDKJ"
+
+Benchmark.ips do |x|
+  x.report("end-with-case?") do
+    Semian::ActiveRecordTrilogyAdapter.query_allowlisted?(sql_not_commit)
+  end
+
+  x.report("regex") { Semian::ActiveRecordTrilogyAdapter::QUERY_ALLOWLIST.match?(sql_not_commit) }
+  x.compare!
+end