From a8ea6b87c245632c7fa1be650e5164acb6bf1183 Mon Sep 17 00:00:00 2001
From: Evan Waters <735779+evan-waters@users.noreply.github.com>
Date: Tue, 10 May 2022 12:02:02 -0600
Subject: [PATCH] Support only and except for specific handler methods

---
 lib/twirp/rails/helpers/hooks.rb    | 18 +++++++++++++++---
 lib/twirp/rails/helpers/services.rb |  2 +-
 lib/twirp/rails/routes.rb           |  4 ++--
 lib/twirp/rails/version.rb          |  2 +-
 4 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/lib/twirp/rails/helpers/hooks.rb b/lib/twirp/rails/helpers/hooks.rb
index b1c40f8..7ae9230 100644
--- a/lib/twirp/rails/helpers/hooks.rb
+++ b/lib/twirp/rails/helpers/hooks.rb
@@ -12,14 +12,22 @@ def self.included(klass)
 
         attr_reader :service_wrapper, :options
 
-        def initialize(service_wrapper, **options)
+        def initialize(service_wrapper, only: [], except: [], **options)
           @service_wrapper = service_wrapper
           @options = options
+          @only = only
+          @except = except
+        end
+
+        def bypass?(env)
+          return true if @only.present? && !@only.include?(env[:ruby_method])
+          return true if @except.present? && @except.include?(env[:ruby_method])
+          false
         end
 
         module ClassMethods
-          def attach(service_wrapper, **options)
-            hook_instance = self.new(service_wrapper, **options)
+          def attach(service_wrapper, only: [], except: [], **options)
+            hook_instance = self.new(service_wrapper, only: only, except: except, **options)
 
             unless HOOK_METHODS.any? {|method| hook_instance.respond_to?(method)}
               raise NotImplementedError.new(
@@ -29,24 +37,28 @@ def attach(service_wrapper, **options)
 
             if hook_instance.respond_to?(:before)
               service_wrapper.before do |rack_env, env|
+                next if hook_instance.bypass?(env)
                 hook_instance.before(rack_env, env)
               end
             end
 
             if hook_instance.respond_to?(:on_success)
               service_wrapper.on_success do |env|
+                next if hook_instance.bypass?(env)
                 hook_instance.on_success(env)
               end
             end
 
             if hook_instance.respond_to?(:on_error)
               service_wrapper.on_error do |twerr, env|
+                next if hook_instance.bypass?(env)
                 hook_instance.on_error(twerr, env)
               end
             end
 
             if hook_instance.respond_to?(:exception_raised)
               service_wrapper.exception_raised do |e, env|
+                next if hook_instance.bypass?(env)
                 hook_instance.exception_raised(e, env)
               end
             end
diff --git a/lib/twirp/rails/helpers/services.rb b/lib/twirp/rails/helpers/services.rb
index 3c87a4e..3932648 100644
--- a/lib/twirp/rails/helpers/services.rb
+++ b/lib/twirp/rails/helpers/services.rb
@@ -42,7 +42,7 @@ def bind(service_klass, namespace: nil, context: nil, hooks: [])
               end
             end
 
-            hooks = hooks + all_base_hooks
+            hooks = all_base_hooks + hooks
 
             service_wrapper = Twirp::Rails::ServiceWrapper.new(service_klass.new(new))
 
diff --git a/lib/twirp/rails/routes.rb b/lib/twirp/rails/routes.rb
index e7b14f3..eda60aa 100644
--- a/lib/twirp/rails/routes.rb
+++ b/lib/twirp/rails/routes.rb
@@ -71,8 +71,8 @@ def generate_routes!(mount_namespace, options)
       end
 
       def attach_service_hooks!(service_wrapper, hooks)
-        hooks.each do |hook_klass:, options: {}|
-          hook_klass.attach service_wrapper, options
+        hooks.each do |hook_klass:, only: [], except: [], options: {}|
+          hook_klass.attach(service_wrapper, only: only, except: except, **options)
         end
       end
     end
diff --git a/lib/twirp/rails/version.rb b/lib/twirp/rails/version.rb
index 82baa9e..9962bc0 100644
--- a/lib/twirp/rails/version.rb
+++ b/lib/twirp/rails/version.rb
@@ -1,5 +1,5 @@
 module Twirp
   module Rails
-    VERSION = "0.5.1"
+    VERSION = "0.5.2"
   end
 end