tcell_agent 1.1.3 → 1.1.4

data/spec/lib/tcell_agent/policies/patches_policy_spec.rb

@@ -283,7 +283,7 @@ module TCellAgent
               }
             )
 
-            meta_data = TCellAgent::Patches::MetaData.new(
+            meta_data = TCellAgent::MetaData.new(
               'get',
               '1.3.3.4',
               'route_id',
@@ -294,7 +294,7 @@ module TCellAgent
             )
             expect(@rust_policies.block_request?(meta_data)).to eq(true)
 
-            meta_data = TCellAgent::Patches::MetaData.new(
+            meta_data = TCellAgent::MetaData.new(
               'get',
               '1.1.1.1',
               'route_id',

data/spec/lib/tcell_agent/rails/csrf_exception_spec.rb

@@ -0,0 +1,69 @@
+require 'spec_helper'
+
+module TCellAgent
+  describe 'CsrfExceptionReporter' do
+    class WrapperClass
+      include TCellAgent::CsrfExceptionReporter
+
+      def request; end
+    end
+
+    before(:all) do
+      @csrf_class = WrapperClass.new
+    end
+
+    context 'nil rust policies' do
+      it 'should not set csrf_exception_name' do
+        expect(TCellAgent).to receive(:policy).and_return(nil)
+        expect(@csrf_class).to_not receive(:request)
+
+        @csrf_class.handle_unverified_request
+      end
+    end
+
+    context 'appfirewall_enabled=false' do
+      it 'should not set csrf_exception_name' do
+        rust_policies = double('rust_policies',
+                               :appfirewall_enabled => false)
+
+        expect(TCellAgent).to receive(:policy).and_return(rust_policies)
+        expect(@csrf_class).to_not receive(:request)
+
+        @csrf_class.handle_unverified_request
+      end
+    end
+
+    context 'appfirewall_enabled=true and nil tcell_data' do
+      it 'should not set csrf_exception_name' do
+        rust_policies = double('rust_policies',
+                               :appfirewall_enabled => true)
+        request = double('request',
+                         :env => {})
+
+        expect(TCellAgent).to receive(:policy).and_return(rust_policies)
+        expect(@csrf_class).to receive(:request).and_return(request)
+
+        @csrf_class.handle_unverified_request
+      end
+    end
+
+    context 'appfirewall_enabled=true and tcell_data present' do
+      it 'should set csrf_exception_name' do
+        tcell_data = TCellAgent::Instrumentation::TCellData.new
+        expect(tcell_data.csrf_exception_name).to be(nil)
+
+        rust_policies = double('rust_policies',
+                               :appfirewall_enabled => true)
+        request = double('request',
+                         :env => { TCellAgent::Instrumentation::TCELL_ID => tcell_data })
+
+        expect(TCellAgent).to receive(:policy).and_return(rust_policies)
+        expect(@csrf_class).to receive(:request).and_return(request)
+
+        @csrf_class.handle_unverified_request
+
+        expect(tcell_data.csrf_exception_name).to eq('ActionController::InvalidAuthenticityToken')
+      end
+    end
+  end
+end

data/spec/lib/tcell_agent/rails/dlp_spec.rb

@@ -0,0 +1,1039 @@
+require 'spec_helper'
+
+module TCellAgent
+  module DLP
+    class SomeColumn
+      attr_accessor :name
+      def initialize(name = nil)
+        @name = name
+      end
+    end
+
+    class SomeModel
+      class << self
+        def table_name
+          'some_models'
+        end
+
+        def columns
+          [SomeColumn.new('id'), SomeColumn.new('name'), SomeColumn.new('email')]
+        end
+
+        def connection_config
+          {
+            :database => 'something/test_database'
+          }
+        end
+      end
+    end
+
+    describe '.instrument_pluck' do
+      context 'with empty results' do
+        it 'should skip over instrumentation' do
+          expect(TCellAgent).to_not receive(:configuration)
+
+          TCellAgent::DLP.instrument_pluck([], [:id], SomeModel)
+        end
+      end
+
+      context 'with no dlp policy' do
+        it 'should skip over instrumentation' do
+          configuration = double(
+            'configuration',
+            {
+              :enabled => true,
+              :should_instrument? => true,
+              :should_intercept_requests? => true
+            }
+          )
+          allow(TCellAgent).to receive(:configuration).and_return(configuration)
+          expect(TCellAgent).to receive(:policy).with(TCellAgent::PolicyTypes::DATALOSS).and_return(nil)
+
+          TCellAgent::DLP.instrument_pluck([10], [:id], SomeModel)
+        end
+      end
+
+      context 'with dlp policy disabled' do
+        it 'should skip over instrumentation' do
+          configuration = double(
+            'configuration',
+            {
+              :enabled => true,
+              :should_instrument? => true,
+              :should_intercept_requests? => true
+            }
+          )
+          dataloss_policy = double(
+            'dataloss_policy',
+            {
+              :enabled => false
+            }
+          )
+          allow(TCellAgent).to receive(:configuration).and_return(configuration)
+          expect(TCellAgent).to receive(:policy).with(
+            TCellAgent::PolicyTypes::DATALOSS
+          ).and_return(dataloss_policy)
+
+          TCellAgent::DLP.instrument_pluck([10], [:id], SomeModel)
+        end
+      end
+
+      context 'with dlp policy enabled' do
+        context 'with database_discovery_enabled' do
+          it 'should send model columns to tcell service' do
+            configuration = double(
+              'configuration',
+              {
+                :enabled => true,
+                :should_instrument? => true,
+                :should_intercept_requests? => true,
+                :max_data_ex_db_records_per_request => 1
+              }
+            )
+            dataloss_policy = double(
+              'dataloss_policy',
+              {
+                :enabled => true,
+                :database_discovery_enabled => true
+              }
+            )
+            request_env = double('request_env', {})
+            tcell_context = double(
+              'tcell_context',
+              {
+                :route_id => 'ma_route_id',
+                :database_result_sizes => []
+              }
+            )
+
+            allow(TCellAgent).to receive(:configuration).and_return(configuration)
+            expect(TCellAgent).to receive(:policy).with(TCellAgent::PolicyTypes::DATALOSS).and_return(dataloss_policy)
+            expect(TCellAgent::Instrumentation::Rails::Middleware::ContextMiddleware::THREADS).to receive(
+              :fetch
+            ).and_return(request_env)
+            expect(request_env).to receive(:[]).with(
+              TCellAgent::Instrumentation::TCELL_ID
+            ).and_return(tcell_context)
+            expect(dataloss_policy).to receive(:get_actions_for_table).with(
+              'test_database', '*', 'some_models', 'id', 'ma_route_id'
+            ).and_return(nil)
+
+            expect(TCellAgent).to receive(:discover_database_fields).with(
+              'ma_route_id', 'test_database', '*', 'some_models', ['id']
+            )
+
+            TCellAgent::DLP.instrument_pluck([10], [:id], SomeModel)
+
+            expect(tcell_context.database_result_sizes).to eq([1])
+          end
+        end
+
+        context 'with database_discovery_enabled == false' do
+          context 'with zero columns passed to pluck' do
+            context 'that doesn\'t have any rules set' do
+              it 'should not send anything to tcell' do
+                configuration = double(
+                  'configuration',
+                  {
+                    :enabled => true,
+                    :should_instrument? => true,
+                    :should_intercept_requests? => true,
+                    :max_data_ex_db_records_per_request => 1000
+                  }
+                )
+                dataloss_policy = double(
+                  'dataloss_policy',
+                  {
+                    :enabled => true,
+                    :database_discovery_enabled => false
+                  }
+                )
+                request_env = double('request_env', {})
+                tcell_context = double(
+                  'tcell_context',
+                  {
+                    :route_id => 'ma_route_id',
+                    :database_result_sizes => []
+                  }
+                )
+
+                allow(TCellAgent).to receive(:configuration).and_return(configuration)
+                expect(TCellAgent).to receive(:policy).with(TCellAgent::PolicyTypes::DATALOSS).and_return(dataloss_policy)
+                expect(TCellAgent::Instrumentation::Rails::Middleware::ContextMiddleware::THREADS).to receive(
+                  :fetch
+                ).and_return(request_env)
+                expect(request_env).to receive(:[]).with(
+                  TCellAgent::Instrumentation::TCELL_ID
+                ).and_return(tcell_context)
+                expect(dataloss_policy).to receive(:get_actions_for_table).with(
+                  'test_database', '*', 'some_models', 'id', 'ma_route_id'
+                ).and_return(nil)
+                expect(dataloss_policy).to receive(:get_actions_for_table).with(
+                  'test_database', '*', 'some_models', 'name', 'ma_route_id'
+                ).and_return(nil)
+                expect(dataloss_policy).to receive(:get_actions_for_table).with(
+                  'test_database', '*', 'some_models', 'email', 'ma_route_id'
+                ).and_return(nil)
+                expect(dataloss_policy).to_not receive(:get_actions_for_table)
+
+                expect(TCellAgent).to_not receive(:discover_database_fields)
+
+                TCellAgent::DLP.instrument_pluck([10], [], SomeModel)
+
+                expect(tcell_context.database_result_sizes).to eq([1])
+              end
+            end
+
+            context 'that has a rule set' do
+              it 'should add a response filter' do
+                configuration = double(
+                  'configuration',
+                  {
+                    :enabled => true,
+                    :should_instrument? => true,
+                    :should_intercept_requests? => true,
+                    :max_data_ex_db_records_per_request => 1000
+                  }
+                )
+                dataloss_policy = double(
+                  'dataloss_policy',
+                  {
+                    :enabled => true,
+                    :database_discovery_enabled => false
+                  }
+                )
+                request_env = double('request_env', {})
+                tcell_context = double(
+                  'tcell_context',
+                  {
+                    :route_id => 'ma_route_id',
+                    :database_result_sizes => []
+                  }
+                )
+                id_rule = double('id_rule')
+
+                allow(TCellAgent).to receive(:configuration).and_return(configuration)
+                expect(TCellAgent).to receive(:policy).with(TCellAgent::PolicyTypes::DATALOSS).and_return(dataloss_policy)
+                expect(TCellAgent::Instrumentation::Rails::Middleware::ContextMiddleware::THREADS).to receive(
+                  :fetch
+                ).and_return(request_env)
+                expect(request_env).to receive(:[]).with(
+                  TCellAgent::Instrumentation::TCELL_ID
+                ).and_return(tcell_context)
+                expect(dataloss_policy).to receive(:get_actions_for_table).with(
+                  'test_database', '*', 'some_models', 'id', 'ma_route_id'
+                ).and_return([id_rule])
+                expect(dataloss_policy).to receive(:get_actions_for_table).with(
+                  'test_database', '*', 'some_models', 'name', 'ma_route_id'
+                ).and_return(nil)
+                expect(dataloss_policy).to receive(:get_actions_for_table).with(
+                  'test_database', '*', 'some_models', 'email', 'ma_route_id'
+                ).and_return(nil)
+                expect(dataloss_policy).to_not receive(:get_actions_for_table)
+
+                expect(TCellAgent).to_not receive(:discover_database_fields)
+                expect(tcell_context).to receive(:add_response_db_filter).with(
+                  10, id_rule, 'test_database', '*', 'some_models', 'id'
+                )
+
+                TCellAgent::DLP.instrument_pluck([[10, 'name', 'email']], [], SomeModel)
+
+                expect(tcell_context.database_result_sizes).to eq([1])
+              end
+            end
+
+            context 'that has two rules set' do
+              it 'should add a response filter' do
+                configuration = double(
+                  'configuration',
+                  {
+                    :enabled => true,
+                    :should_instrument? => true,
+                    :should_intercept_requests? => true,
+                    :max_data_ex_db_records_per_request => 1000
+                  }
+                )
+                dataloss_policy = double(
+                  'dataloss_policy',
+                  {
+                    :enabled => true,
+                    :database_discovery_enabled => false
+                  }
+                )
+                request_env = double('request_env', {})
+                tcell_context = double(
+                  'tcell_context',
+                  {
+                    :route_id => 'ma_route_id',
+                    :database_result_sizes => []
+                  }
+                )
+                id_rule = double('id_rule')
+                id_rule_dos = double('id_rule_dos')
+
+                allow(TCellAgent).to receive(:configuration).and_return(configuration)
+                expect(TCellAgent).to receive(:policy).with(TCellAgent::PolicyTypes::DATALOSS).and_return(dataloss_policy)
+                expect(TCellAgent::Instrumentation::Rails::Middleware::ContextMiddleware::THREADS).to receive(
+                  :fetch
+                ).and_return(request_env)
+                expect(request_env).to receive(:[]).with(
+                  TCellAgent::Instrumentation::TCELL_ID
+                ).and_return(tcell_context)
+                expect(dataloss_policy).to receive(:get_actions_for_table).with(
+                  'test_database', '*', 'some_models', 'id', 'ma_route_id'
+                ).and_return([id_rule, id_rule_dos])
+                expect(dataloss_policy).to receive(:get_actions_for_table).with(
+                  'test_database', '*', 'some_models', 'name', 'ma_route_id'
+                ).and_return(nil)
+                expect(dataloss_policy).to receive(:get_actions_for_table).with(
+                  'test_database', '*', 'some_models', 'email', 'ma_route_id'
+                ).and_return(nil)
+                expect(dataloss_policy).to_not receive(:get_actions_for_table)
+
+                expect(TCellAgent).to_not receive(:discover_database_fields)
+                expect(tcell_context).to receive(:add_response_db_filter).with(
+                  10, id_rule, 'test_database', '*', 'some_models', 'id'
+                )
+                expect(tcell_context).to receive(:add_response_db_filter).with(
+                  10, id_rule_dos, 'test_database', '*', 'some_models', 'id'
+                )
+
+                TCellAgent::DLP.instrument_pluck([[10, 'name', 'email']], [], SomeModel)
+
+                expect(tcell_context.database_result_sizes).to eq([1])
+              end
+            end
+          end
+
+          context 'with one column passed to pluck' do
+            context 'that doesn\'t have any rules set' do
+              it 'should not send anything to tcell' do
+                configuration = double(
+                  'configuration',
+                  {
+                    :enabled => true,
+                    :should_instrument? => true,
+                    :should_intercept_requests? => true,
+                    :max_data_ex_db_records_per_request => 1000
+                  }
+                )
+                dataloss_policy = double(
+                  'dataloss_policy',
+                  {
+                    :enabled => true,
+                    :database_discovery_enabled => false
+                  }
+                )
+                request_env = double('request_env', {})
+                tcell_context = double(
+                  'tcell_context',
+                  {
+                    :route_id => 'ma_route_id',
+                    :database_result_sizes => []
+                  }
+                )
+
+                allow(TCellAgent).to receive(:configuration).and_return(configuration)
+                expect(TCellAgent).to receive(:policy).with(TCellAgent::PolicyTypes::DATALOSS).and_return(dataloss_policy)
+                expect(TCellAgent::Instrumentation::Rails::Middleware::ContextMiddleware::THREADS).to receive(
+                  :fetch
+                ).and_return(request_env)
+                expect(request_env).to receive(:[]).with(
+                  TCellAgent::Instrumentation::TCELL_ID
+                ).and_return(tcell_context)
+                expect(dataloss_policy).to receive(:get_actions_for_table).with(
+                  'test_database', '*', 'some_models', 'id', 'ma_route_id'
+                ).and_return(nil)
+                expect(dataloss_policy).to_not receive(:get_actions_for_table)
+
+                expect(TCellAgent).to_not receive(:discover_database_fields)
+
+                TCellAgent::DLP.instrument_pluck([10], [:id], SomeModel)
+
+                expect(tcell_context.database_result_sizes).to eq([1])
+              end
+            end
+
+            context 'that has a rule set' do
+              it 'should add a response filter' do
+                configuration = double(
+                  'configuration',
+                  {
+                    :enabled => true,
+                    :should_instrument? => true,
+                    :should_intercept_requests? => true,
+                    :max_data_ex_db_records_per_request => 1000
+                  }
+                )
+                dataloss_policy = double(
+                  'dataloss_policy',
+                  {
+                    :enabled => true,
+                    :database_discovery_enabled => false
+                  }
+                )
+                request_env = double('request_env', {})
+                tcell_context = double(
+                  'tcell_context',
+                  {
+                    :route_id => 'ma_route_id',
+                    :database_result_sizes => []
+                  }
+                )
+                id_rule = double('id_rule')
+
+                allow(TCellAgent).to receive(:configuration).and_return(configuration)
+                expect(TCellAgent).to receive(:policy).with(TCellAgent::PolicyTypes::DATALOSS).and_return(dataloss_policy)
+                expect(TCellAgent::Instrumentation::Rails::Middleware::ContextMiddleware::THREADS).to receive(
+                  :fetch
+                ).and_return(request_env)
+                expect(request_env).to receive(:[]).with(
+                  TCellAgent::Instrumentation::TCELL_ID
+                ).and_return(tcell_context)
+                expect(dataloss_policy).to receive(:get_actions_for_table).with(
+                  'test_database', '*', 'some_models', 'id', 'ma_route_id'
+                ).and_return([id_rule])
+                expect(dataloss_policy).to_not receive(:get_actions_for_table)
+
+                expect(TCellAgent).to_not receive(:discover_database_fields)
+                expect(tcell_context).to receive(:add_response_db_filter).with(
+                  10, id_rule, 'test_database', '*', 'some_models', 'id'
+                )
+
+                TCellAgent::DLP.instrument_pluck([10], [:id], SomeModel)
+
+                expect(tcell_context.database_result_sizes).to eq([1])
+              end
+            end
+
+            context 'that has two rules set' do
+              it 'should add a response filter' do
+                configuration = double(
+                  'configuration',
+                  {
+                    :enabled => true,
+                    :should_instrument? => true,
+                    :should_intercept_requests? => true,
+                    :max_data_ex_db_records_per_request => 1000
+                  }
+                )
+                dataloss_policy = double(
+                  'dataloss_policy',
+                  {
+                    :enabled => true,
+                    :database_discovery_enabled => false
+                  }
+                )
+                request_env = double('request_env', {})
+                tcell_context = double(
+                  'tcell_context',
+                  {
+                    :route_id => 'ma_route_id',
+                    :database_result_sizes => []
+                  }
+                )
+                id_rule = double('id_rule')
+                id_rule_dos = double('id_rule_dos')
+
+                allow(TCellAgent).to receive(:configuration).and_return(configuration)
+                expect(TCellAgent).to receive(:policy).with(TCellAgent::PolicyTypes::DATALOSS).and_return(dataloss_policy)
+                expect(TCellAgent::Instrumentation::Rails::Middleware::ContextMiddleware::THREADS).to receive(
+                  :fetch
+                ).and_return(request_env)
+                expect(request_env).to receive(:[]).with(
+                  TCellAgent::Instrumentation::TCELL_ID
+                ).and_return(tcell_context)
+                expect(dataloss_policy).to receive(:get_actions_for_table).with(
+                  'test_database', '*', 'some_models', 'id', 'ma_route_id'
+                ).and_return([id_rule, id_rule_dos])
+                expect(dataloss_policy).to_not receive(:get_actions_for_table)
+
+                expect(TCellAgent).to_not receive(:discover_database_fields)
+                expect(tcell_context).to receive(:add_response_db_filter).with(
+                  10, id_rule, 'test_database', '*', 'some_models', 'id'
+                )
+                expect(tcell_context).to receive(:add_response_db_filter).with(
+                  10, id_rule_dos, 'test_database', '*', 'some_models', 'id'
+                )
+
+                TCellAgent::DLP.instrument_pluck([10], [:id], SomeModel)
+
+                expect(tcell_context.database_result_sizes).to eq([1])
+              end
+            end
+
+            context 'that is namespaced' do
+              it 'should add a response filter' do
+                configuration = double(
+                  'configuration',
+                  {
+                    :enabled => true,
+                    :should_instrument? => true,
+                    :should_intercept_requests? => true,
+                    :max_data_ex_db_records_per_request => 1000
+                  }
+                )
+                dataloss_policy = double(
+                  'dataloss_policy',
+                  {
+                    :enabled => true,
+                    :database_discovery_enabled => false
+                  }
+                )
+                request_env = double('request_env', {})
+                tcell_context = double(
+                  'tcell_context',
+                  {
+                    :route_id => 'ma_route_id',
+                    :database_result_sizes => []
+                  }
+                )
+                id_rule = double('id_rule')
+                id_rule_dos = double('id_rule_dos')
+
+                allow(TCellAgent).to receive(:configuration).and_return(configuration)
+                expect(TCellAgent).to receive(:policy).with(TCellAgent::PolicyTypes::DATALOSS).and_return(dataloss_policy)
+                expect(TCellAgent::Instrumentation::Rails::Middleware::ContextMiddleware::THREADS).to receive(
+                  :fetch
+                ).and_return(request_env)
+                expect(request_env).to receive(:[]).with(
+                  TCellAgent::Instrumentation::TCELL_ID
+                ).and_return(tcell_context)
+                expect(dataloss_policy).to receive(:get_actions_for_table).with(
+                  'test_database', '*', 'some_models', 'id', 'ma_route_id'
+                ).and_return([id_rule, id_rule_dos])
+                expect(dataloss_policy).to_not receive(:get_actions_for_table)
+                expect(tcell_context).to receive(:add_response_db_filter).with(
+                  10, id_rule, 'test_database', '*', 'some_models', 'id'
+                )
+                expect(tcell_context).to receive(:add_response_db_filter).with(
+                  10, id_rule_dos, 'test_database', '*', 'some_models', 'id'
+                )
+
+                expect(TCellAgent).to_not receive(:discover_database_fields)
+
+                TCellAgent::DLP.instrument_pluck([10], ['some_models.id'], SomeModel)
+
+                expect(tcell_context.database_result_sizes).to eq([1])
+              end
+            end
+          end
+
+          context 'with two columns passed to pluck' do
+            context 'that doesn\'t have any rules set' do
+              it 'should not send anything to tcell' do
+                configuration = double(
+                  'configuration',
+                  {
+                    :enabled => true,
+                    :should_instrument? => true,
+                    :should_intercept_requests? => true,
+                    :max_data_ex_db_records_per_request => 1000
+                  }
+                )
+                dataloss_policy = double(
+                  'dataloss_policy',
+                  {
+                    :enabled => true,
+                    :database_discovery_enabled => false
+                  }
+                )
+                request_env = double('request_env', {})
+                tcell_context = double(
+                  'tcell_context',
+                  {
+                    :route_id => 'ma_route_id',
+                    :database_result_sizes => []
+                  }
+                )
+
+                allow(TCellAgent).to receive(:configuration).and_return(configuration)
+                expect(TCellAgent).to receive(:policy).with(TCellAgent::PolicyTypes::DATALOSS).and_return(dataloss_policy)
+                expect(TCellAgent::Instrumentation::Rails::Middleware::ContextMiddleware::THREADS).to receive(
+                  :fetch
+                ).and_return(request_env)
+                expect(request_env).to receive(:[]).with(
+                  TCellAgent::Instrumentation::TCELL_ID
+                ).and_return(tcell_context)
+                expect(dataloss_policy).to receive(:get_actions_for_table).with(
+                  'test_database', '*', 'some_models', 'id', 'ma_route_id'
+                ).and_return(nil)
+                expect(dataloss_policy).to receive(:get_actions_for_table).with(
+                  'test_database', '*', 'some_models', 'name', 'ma_route_id'
+                ).and_return(nil)
+                expect(dataloss_policy).to_not receive(:get_actions_for_table)
+
+                expect(TCellAgent).to_not receive(:discover_database_fields)
+
+                TCellAgent::DLP.instrument_pluck([[10, 'dies'], [20, 'veinte']], %i[id name], SomeModel)
+
+                expect(tcell_context.database_result_sizes).to eq([2])
+              end
+            end
+
+            context 'that has a rule set' do
+              it 'should add a response filter' do
+                configuration = double(
+                  'configuration',
+                  {
+                    :enabled => true,
+                    :should_instrument? => true,
+                    :should_intercept_requests? => true,
+                    :max_data_ex_db_records_per_request => 1000
+                  }
+                )
+                dataloss_policy = double(
+                  'dataloss_policy',
+                  {
+                    :enabled => true,
+                    :database_discovery_enabled => false
+                  }
+                )
+                request_env = double('request_env', {})
+                tcell_context = double(
+                  'tcell_context',
+                  {
+                    :route_id => 'ma_route_id',
+                    :database_result_sizes => []
+                  }
+                )
+                id_rule = double('id_rule')
+
+                allow(TCellAgent).to receive(:configuration).and_return(configuration)
+                expect(TCellAgent).to receive(:policy).with(TCellAgent::PolicyTypes::DATALOSS).and_return(dataloss_policy)
+                expect(TCellAgent::Instrumentation::Rails::Middleware::ContextMiddleware::THREADS).to receive(
+                  :fetch
+                ).and_return(request_env)
+                expect(request_env).to receive(:[]).with(
+                  TCellAgent::Instrumentation::TCELL_ID
+                ).and_return(tcell_context)
+                expect(dataloss_policy).to receive(:get_actions_for_table).with(
+                  'test_database', '*', 'some_models', 'id', 'ma_route_id'
+                ).and_return([id_rule])
+                expect(dataloss_policy).to receive(:get_actions_for_table).with(
+                  'test_database', '*', 'some_models', 'name', 'ma_route_id'
+                ).and_return(nil)
+                expect(dataloss_policy).to_not receive(:get_actions_for_table)
+
+                expect(TCellAgent).to_not receive(:discover_database_fields)
+                expect(tcell_context).to receive(:add_response_db_filter).with(
+                  10, id_rule, 'test_database', '*', 'some_models', 'id'
+                )
+                expect(tcell_context).to receive(:add_response_db_filter).with(
+                  20, id_rule, 'test_database', '*', 'some_models', 'id'
+                )
+
+                TCellAgent::DLP.instrument_pluck([[10, 'dies'], [20, 'veinte']], %i[id name], SomeModel)
+
+                expect(tcell_context.database_result_sizes).to eq([2])
+              end
+
+              context 'that is namespaced' do
+                it 'should not send anything to tcell' do
+                  configuration = double(
+                    'configuration',
+                    {
+                      :enabled => true,
+                      :should_instrument? => true,
+                      :should_intercept_requests? => true,
+                      :max_data_ex_db_records_per_request => 1000
+                    }
+                  )
+                  dataloss_policy = double(
+                    'dataloss_policy',
+                    {
+                      :enabled => true,
+                      :database_discovery_enabled => false
+                    }
+                  )
+                  request_env = double('request_env', {})
+                  tcell_context = double(
+                    'tcell_context',
+                    {
+                      :route_id => 'ma_route_id',
+                      :database_result_sizes => []
+                    }
+                  )
+
+                  allow(TCellAgent).to receive(:configuration).and_return(configuration)
+                  expect(TCellAgent).to receive(:policy).with(TCellAgent::PolicyTypes::DATALOSS).and_return(dataloss_policy)
+                  expect(TCellAgent::Instrumentation::Rails::Middleware::ContextMiddleware::THREADS).to receive(
+                    :fetch
+                  ).and_return(request_env)
+                  expect(request_env).to receive(:[]).with(
+                    TCellAgent::Instrumentation::TCELL_ID
+                  ).and_return(tcell_context)
+                  expect(dataloss_policy).to receive(:get_actions_for_table).with(
+                    'test_database', '*', 'some_models', 'id', 'ma_route_id'
+                  ).and_return(nil)
+                  expect(dataloss_policy).to_not receive(:get_actions_for_table)
+
+                  expect(TCellAgent).to_not receive(:discover_database_fields)
+
+                  TCellAgent::DLP.instrument_pluck([10], ['some_other_models.name', 'some_models.id'], SomeModel)
+
+                  expect(tcell_context.database_result_sizes).to eq([1])
+                end
+              end
+            end
+
+            context 'that has two rules set' do
+              it 'should add a response filter' do
+                configuration = double(
+                  'configuration',
+                  {
+                    :enabled => true,
+                    :should_instrument? => true,
+                    :should_intercept_requests? => true,
+                    :max_data_ex_db_records_per_request => 1000
+                  }
+                )
+                dataloss_policy = double(
+                  'dataloss_policy',
+                  {
+                    :enabled => true,
+                    :database_discovery_enabled => false
+                  }
+                )
+                request_env = double('request_env', {})
+                tcell_context = double(
+                  'tcell_context',
+                  {
+                    :route_id => 'ma_route_id',
+                    :database_result_sizes => []
+                  }
+                )
+                id_rule = double('id_rule')
+                id_rule_dos = double('id_rule_dos')
+
+                allow(TCellAgent).to receive(:configuration).and_return(configuration)
+                expect(TCellAgent).to receive(:policy).with(TCellAgent::PolicyTypes::DATALOSS).and_return(dataloss_policy)
+                expect(TCellAgent::Instrumentation::Rails::Middleware::ContextMiddleware::THREADS).to receive(
+                  :fetch
+                ).and_return(request_env)
+                expect(request_env).to receive(:[]).with(
+                  TCellAgent::Instrumentation::TCELL_ID
+                ).and_return(tcell_context)
+                expect(dataloss_policy).to receive(:get_actions_for_table).with(
+                  'test_database', '*', 'some_models', 'id', 'ma_route_id'
+                ).and_return([id_rule, id_rule_dos])
+                expect(dataloss_policy).to receive(:get_actions_for_table).with(
+                  'test_database', '*', 'some_models', 'name', 'ma_route_id'
+                ).and_return(nil)
+                expect(dataloss_policy).to_not receive(:get_actions_for_table)
+
+                expect(TCellAgent).to_not receive(:discover_database_fields)
+                expect(tcell_context).to receive(:add_response_db_filter).with(
+                  10, id_rule, 'test_database', '*', 'some_models', 'id'
+                )
+                expect(tcell_context).to receive(:add_response_db_filter).with(
+                  10, id_rule_dos, 'test_database', '*', 'some_models', 'id'
+                )
+                expect(tcell_context).to receive(:add_response_db_filter).with(
+                  20, id_rule, 'test_database', '*', 'some_models', 'id'
+                )
+                expect(tcell_context).to receive(:add_response_db_filter).with(
+                  20, id_rule_dos, 'test_database', '*', 'some_models', 'id'
+                )
+
+                TCellAgent::DLP.instrument_pluck([[10, 'dies'], [20, 'veinte']], %i[id name], SomeModel)
+
+                expect(tcell_context.database_result_sizes).to eq([2])
+              end
+            end
+          end
+        end
+      end
+    end
+
+    describe '.instrument_find_by_sql' do
+      context 'with empty results' do
+        it 'should skip over instrumentation' do
+          expect(TCellAgent).to_not receive(:configuration)
+
+          TCellAgent::DLP.instrument_find_by_sql([])
+        end
+      end
+
+      context 'with no dlp policy' do
+        it 'should skip over instrumentation' do
+          configuration = double(
+            'configuration',
+            {
+              :enabled => true,
+              :should_instrument? => true,
+              :should_intercept_requests? => true
+            }
+          )
+          allow(TCellAgent).to receive(:configuration).and_return(configuration)
+          expect(TCellAgent).to receive(:policy).with(TCellAgent::PolicyTypes::DATALOSS).and_return(nil)
+
+          TCellAgent::DLP.instrument_find_by_sql([SomeModel.new])
+        end
+      end
+
+      context 'with dlp policy disabled' do
+        it 'should skip over instrumentation' do
+          configuration = double(
+            'configuration',
+            {
+              :enabled => true,
+              :should_instrument? => true,
+              :should_intercept_requests? => true
+            }
+          )
+          dataloss_policy = double(
+            'dataloss_policy',
+            {
+              :enabled => false
+            }
+          )
+          allow(TCellAgent).to receive(:configuration).and_return(configuration)
+          expect(TCellAgent).to receive(:policy).with(
+            TCellAgent::PolicyTypes::DATALOSS
+          ).and_return(dataloss_policy)
+
+          TCellAgent::DLP.instrument_find_by_sql([SomeModel.new])
+        end
+      end
+
+      context 'with dlp policy enabled' do
+        context 'with database_discovery_enabled' do
+          it 'should send model columns to tcell service' do
+            configuration = double(
+              'configuration',
+              {
+                :enabled => true,
+                :should_instrument? => true,
+                :should_intercept_requests? => true,
+                :max_data_ex_db_records_per_request => 1
+              }
+            )
+            dataloss_policy = double(
+              'dataloss_policy',
+              {
+                :enabled => true,
+                :database_discovery_enabled => true
+              }
+            )
+            request_env = double('request_env', {})
+            tcell_context = double(
+              'tcell_context',
+              {
+                :route_id => 'ma_route_id',
+                :database_result_sizes => []
+              }
+            )
+
+            allow(TCellAgent).to receive(:configuration).and_return(configuration)
+            expect(TCellAgent).to receive(:policy).with(TCellAgent::PolicyTypes::DATALOSS).and_return(dataloss_policy)
+            expect(TCellAgent::Instrumentation::Rails::Middleware::ContextMiddleware::THREADS).to receive(
+              :fetch
+            ).and_return(request_env)
+            expect(request_env).to receive(:[]).with(
+              TCellAgent::Instrumentation::TCELL_ID
+            ).and_return(tcell_context)
+            expect(dataloss_policy).to receive(:get_actions_for_table).with(
+              'test_database', '*', 'some_models', 'id', 'ma_route_id'
+            ).and_return(nil)
+            expect(dataloss_policy).to receive(:get_actions_for_table).with(
+              'test_database', '*', 'some_models', 'name', 'ma_route_id'
+            ).and_return(nil)
+            expect(dataloss_policy).to receive(:get_actions_for_table).with(
+              'test_database', '*', 'some_models', 'email', 'ma_route_id'
+            ).and_return(nil)
+
+            expect(TCellAgent).to receive(:discover_database_fields).with(
+              'ma_route_id', 'test_database', '*', 'some_models', %w[id name email]
+            )
+
+            TCellAgent::DLP.instrument_find_by_sql([SomeModel.new])
+
+            expect(tcell_context.database_result_sizes).to eq([1])
+          end
+        end
+
+        context 'with database_discovery_enabled == false' do
+          context 'that doesn\'t have any rules set' do
+            it 'should not send anything to tcell' do
+              configuration = double(
+                'configuration',
+                {
+                  :enabled => true,
+                  :should_instrument? => true,
+                  :should_intercept_requests? => true,
+                  :max_data_ex_db_records_per_request => 1000
+                }
+              )
+              dataloss_policy = double(
+                'dataloss_policy',
+                {
+                  :enabled => true,
+                  :database_discovery_enabled => false
+                }
+              )
+              request_env = double('request_env', {})
+              tcell_context = double(
+                'tcell_context',
+                {
+                  :route_id => 'ma_route_id',
+                  :database_result_sizes => []
+                }
+              )
+
+              allow(TCellAgent).to receive(:configuration).and_return(configuration)
+              expect(TCellAgent).to receive(:policy).with(TCellAgent::PolicyTypes::DATALOSS).and_return(dataloss_policy)
+              expect(TCellAgent::Instrumentation::Rails::Middleware::ContextMiddleware::THREADS).to receive(
+                :fetch
+              ).and_return(request_env)
+              expect(request_env).to receive(:[]).with(
+                TCellAgent::Instrumentation::TCELL_ID
+              ).and_return(tcell_context)
+              expect(dataloss_policy).to receive(:get_actions_for_table).with(
+                'test_database', '*', 'some_models', 'id', 'ma_route_id'
+              ).and_return(nil)
+              expect(dataloss_policy).to receive(:get_actions_for_table).with(
+                'test_database', '*', 'some_models', 'name', 'ma_route_id'
+              ).and_return(nil)
+              expect(dataloss_policy).to receive(:get_actions_for_table).with(
+                'test_database', '*', 'some_models', 'email', 'ma_route_id'
+              ).and_return(nil)
+              expect(dataloss_policy).to_not receive(:get_actions_for_table)
+
+              expect(TCellAgent).to_not receive(:discover_database_fields)
+
+              TCellAgent::DLP.instrument_find_by_sql([SomeModel.new])
+
+              expect(tcell_context.database_result_sizes).to eq([1])
+            end
+          end
+
+          context 'that has a rule set' do
+            it 'should add a response filter' do
+              configuration = double(
+                'configuration',
+                {
+                  :enabled => true,
+                  :should_instrument? => true,
+                  :should_intercept_requests? => true,
+                  :max_data_ex_db_records_per_request => 1000
+                }
+              )
+              dataloss_policy = double(
+                'dataloss_policy',
+                {
+                  :enabled => true,
+                  :database_discovery_enabled => false
+                }
+              )
+              request_env = double('request_env', {})
+              tcell_context = double(
+                'tcell_context',
+                {
+                  :route_id => 'ma_route_id',
+                  :database_result_sizes => []
+                }
+              )
+              id_rule = double('id_rule')
+              some_model = SomeModel.new
+
+              allow(TCellAgent).to receive(:configuration).and_return(configuration)
+              expect(TCellAgent).to receive(:policy).with(TCellAgent::PolicyTypes::DATALOSS).and_return(dataloss_policy)
+              expect(TCellAgent::Instrumentation::Rails::Middleware::ContextMiddleware::THREADS).to receive(
+                :fetch
+              ).and_return(request_env)
+              expect(request_env).to receive(:[]).with(
+                TCellAgent::Instrumentation::TCELL_ID
+              ).and_return(tcell_context)
+              expect(dataloss_policy).to receive(:get_actions_for_table).with(
+                'test_database', '*', 'some_models', 'id', 'ma_route_id'
+              ).and_return([id_rule])
+              expect(dataloss_policy).to receive(:get_actions_for_table).with(
+                'test_database', '*', 'some_models', 'name', 'ma_route_id'
+              ).and_return(nil)
+              expect(dataloss_policy).to receive(:get_actions_for_table).with(
+                'test_database', '*', 'some_models', 'email', 'ma_route_id'
+              ).and_return(nil)
+              expect(dataloss_policy).to_not receive(:get_actions_for_table)
+
+              expect(TCellAgent).to_not receive(:discover_database_fields)
+              expect(some_model).to receive(:[]).with(:id).and_return(10)
+              expect(tcell_context).to receive(:add_response_db_filter).with(
+                10, id_rule, 'test_database', '*', 'some_models', 'id'
+              )
+
+              TCellAgent::DLP.instrument_find_by_sql([some_model])
+
+              expect(tcell_context.database_result_sizes).to eq([1])
+            end
+          end
+
+          context 'that has two rules set' do
+            it 'should add a response filter' do
+              configuration = double(
+                'configuration',
+                {
+                  :enabled => true,
+                  :should_instrument? => true,
+                  :should_intercept_requests? => true,
+                  :max_data_ex_db_records_per_request => 1000
+                }
+              )
+              dataloss_policy = double(
+                'dataloss_policy',
+                {
+                  :enabled => true,
+                  :database_discovery_enabled => false
+                }
+              )
+              request_env = double('request_env', {})
+              tcell_context = double(
+                'tcell_context',
+                {
+                  :route_id => 'ma_route_id',
+                  :database_result_sizes => []
+                }
+              )
+              id_rule = double('id_rule')
+              id_rule_dos = double('id_rule_dos')
+              some_model = SomeModel.new
+
+              allow(TCellAgent).to receive(:configuration).and_return(configuration)
+              expect(TCellAgent).to receive(:policy).with(TCellAgent::PolicyTypes::DATALOSS).and_return(dataloss_policy)
+              expect(TCellAgent::Instrumentation::Rails::Middleware::ContextMiddleware::THREADS).to receive(
+                :fetch
+              ).and_return(request_env)
+              expect(request_env).to receive(:[]).with(
+                TCellAgent::Instrumentation::TCELL_ID
+              ).and_return(tcell_context)
+              expect(dataloss_policy).to receive(:get_actions_for_table).with(
+                'test_database', '*', 'some_models', 'id', 'ma_route_id'
+              ).and_return([id_rule, id_rule_dos])
+              expect(dataloss_policy).to receive(:get_actions_for_table).with(
+                'test_database', '*', 'some_models', 'name', 'ma_route_id'
+              ).and_return(nil)
+              expect(dataloss_policy).to receive(:get_actions_for_table).with(
+                'test_database', '*', 'some_models', 'email', 'ma_route_id'
+              ).and_return(nil)
+              expect(dataloss_policy).to_not receive(:get_actions_for_table)
+
+              expect(TCellAgent).to_not receive(:discover_database_fields)
+              expect(some_model).to receive(:[]).with(:id).and_return(10)
+              expect(some_model).to receive(:[]).with(:id).and_return(10)
+              expect(tcell_context).to receive(:add_response_db_filter).with(
+                10, id_rule, 'test_database', '*', 'some_models', 'id'
+              )
+              expect(tcell_context).to receive(:add_response_db_filter).with(
+                10, id_rule_dos, 'test_database', '*', 'some_models', 'id'
+              )
+
+              TCellAgent::DLP.instrument_find_by_sql([some_model])
+
+              expect(tcell_context.database_result_sizes).to eq([1])
+            end
+          end
+        end
+      end
+    end
+  end
+end