From c8c7aaca10b5623b26a37606fb20a25cc527018a Mon Sep 17 00:00:00 2001
From: Camilo Lopez <camilo.lopez@shopify.com>
Date: Tue, 19 Dec 2023 19:48:52 +0000
Subject: [PATCH] Return early from allow list when possible

---
 Gemfile                                       |  1 +
 Gemfile.lock                                  |  2 ++
 lib/semian/activerecord_trilogy_adapter.rb    | 31 ++++++++++++++++---
 .../activerecord_trilogy_adapter_test.rb      |  9 +++---
 test/allow_list_bench.rb                      | 20 ++++++++++++
 5 files changed, 54 insertions(+), 9 deletions(-)
 create mode 100644 test/allow_list_bench.rb

diff --git a/Gemfile b/Gemfile
index 4956e424..ce0b9a88 100644
--- a/Gemfile
+++ b/Gemfile
@@ -7,6 +7,7 @@ gem "rake-compiler"
 
 group :test do
   gem "benchmark-memory"
+  gem "benchmark-ips"
   gem "memory_profiler"
   gem "minitest"
   gem "mocha"
diff --git a/Gemfile.lock b/Gemfile.lock
index 09df7555..6c6a8dde 100644
--- a/Gemfile.lock
+++ b/Gemfile.lock
@@ -24,6 +24,7 @@ GEM
   remote: https://rubygems.org/
   specs:
     ast (2.4.2)
+    benchmark-ips (2.13.0)
     benchmark-memory (0.2.0)
       memory_profiler (~> 1)
     byebug (11.1.3)
@@ -112,6 +113,7 @@ PLATFORMS
 
 DEPENDENCIES
   activerecord!
+  benchmark-ips
   benchmark-memory
   grpc (= 1.47.0)
   hiredis (~> 0.6)
diff --git a/lib/semian/activerecord_trilogy_adapter.rb b/lib/semian/activerecord_trilogy_adapter.rb
index 6223d56d..399371ed 100644
--- a/lib/semian/activerecord_trilogy_adapter.rb
+++ b/lib/semian/activerecord_trilogy_adapter.rb
@@ -48,7 +48,7 @@ def initialize(*options)
     end
 
     def raw_execute(sql, *)
-      if query_allowlisted?(sql)
+      if self.query_allowlisted?(sql)
         super
       else
         acquire_semian_resource(adapter: :trilogy_adapter, scope: :query) do
@@ -90,11 +90,34 @@ 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)
+    def self.query_allowlisted?(sql, *)
+      # The common case here NOT to have transcation management statements,
+      # we are exploiting the fact that ActiveRecord will use COMMIT/ROLLBACK and
+      # name savepoints by level of nesting as `active_record_1` ... n.
+      #
+      # Since looking at the last character in a sring is a LOT cheaper than
+      # running a regex we are returning earlier if the last characters of 
+      # the SQL statements are NOT the last characters of the known transaction 
+      # control patterns
+      #
+      # test/alllow_list_bench.rb as of now (Dec 2023 using ruby 3.1) shows 
+      # this method to be 34x faster in the common case
+      #
+      # But what about achoring the regex? 
+      #
+      # This regex is 1.1x faster in benchmark
+      # %r{\A(?:/\*.*?\*/)?\s*(ROLLBACK|COMMIT|RELEASE\s+SAVEPOINT)\Z}i.match?(sql)
+      #
+      # but then we might miss cases and is not nearly as fast as looking at the last 
+      # characters of the sql
+      if !sql.end_with?('T') || !sql.end_with('K') || !sql.end_with?('_1') || !sql.end_with?('_2')
+        false
+      else
+        QUERY_ALLOWLIST.match?(sql)
+      end
+
     rescue ArgumentError
       return false unless sql.valid_encoding?
 
diff --git a/test/adapters/activerecord_trilogy_adapter_test.rb b/test/adapters/activerecord_trilogy_adapter_test.rb
index 44988eee..5e5e7ec5 100644
--- a/test/adapters/activerecord_trilogy_adapter_test.rb
+++ b/test/adapters/activerecord_trilogy_adapter_test.rb
@@ -332,12 +332,11 @@ def test_query_allowlisted_returns_false_for_binary_sql
         refute(@adapter.send(:query_allowlisted?, binary_query))
       end
 
-      def test_semian_allows_rollback_to_safepoint
+      def test_semian_allows_set_savepoint
         @adapter.execute("START TRANSACTION;")
-        @adapter.execute("SAVEPOINT foobar;")
 
         Semian[:mysql_testing].acquire do
-          @adapter.execute("ROLLBACK TO foobar;")
+        @adapter.execute("SAVEPOINT active_record_99;")
         end
 
         @adapter.execute("ROLLBACK;")
@@ -345,10 +344,10 @@ 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_99;")
 
         Semian[:mysql_testing].acquire do
-          @adapter.execute("RELEASE SAVEPOINT foobar;")
+          @adapter.execute("RELEASE SAVEPOINT active_record_99;")
         end
 
         @adapter.execute("ROLLBACK;")
diff --git a/test/allow_list_bench.rb b/test/allow_list_bench.rb
new file mode 100644
index 00000000..ff7be8f4
--- /dev/null
+++ b/test/allow_list_bench.rb
@@ -0,0 +1,20 @@
+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
+