Merge pull request #51 from stripe/adunham/rules-loading

Check for and reload rules when we get a log line

doorman/extensions.py

@@ -49,8 +49,8 @@ def handle_result(self, data, **kwargs):
 
 class RuleManager(object):
     def __init__(self, app=None):
-        self.loaded_rules = False
         self.network = None
+        self.last_update = None
 
         if app is not None:
             self.init_app(app)
@@ -70,36 +70,42 @@ def load_alerters(self):
         alerters = self.app.config.get('DOORMAN_ALERTER_PLUGINS', {})
 
         self.alerters = {}
-        with self.app.app_context():
-            for name, (plugin, config) in alerters.items():
-                package, classname = plugin.rsplit('.', 1)
-                module = import_module(package)
-                klass = getattr(module, classname, None)
+        for name, (plugin, config) in alerters.items():
+            package, classname = plugin.rsplit('.', 1)
+            module = import_module(package)
+            klass = getattr(module, classname, None)
 
-                if klass is None:
-                    raise ValueError('Could not find a class named "{0}" in package "{1}"'.format(classname, package))
+            if klass is None:
+                raise ValueError('Could not find a class named "{0}" in package "{1}"'.format(classname, package))
+
+            if not issubclass(klass, AbstractAlerterPlugin):
+                raise ValueError('{0} is not a subclass of AbstractAlerterPlugin'.format(name))
 
-                if not issubclass(klass, AbstractAlerterPlugin):
-                    raise ValueError('{0} is not a subclass of AbstractAlerterPlugin'.format(name))
+            self.alerters[name] = klass(config)
 
-                self.alerters[name] = klass(config)
+    def should_reload_rules(self):
+        """ Checks if we need to reload the set of rules. """
+        from doorman.models import Rule
+
+        if self.last_update is None:
+            return True
+
+        newest_rule = Rule.query.order_by(Rule.updated_at.desc()).limit(1).first()
+        if self.last_update < newest_rule.updated_at:
+            return True
+
+        return False
 
     def load_rules(self):
         """ Load rules from the database. """
         from doorman.rules import Network
         from doorman.models import Rule
         from sqlalchemy.exc import SQLAlchemyError
 
-        self.rules = []
-        with self.app.app_context():
-            try:
-                all_rules = list(Rule.query.all())
-            except SQLAlchemyError:
-                # Ignore DB errors when testing
-                if self.app.config['TESTING']:
-                    all_rules = []
-                else:
-                    raise
+        if not self.should_reload_rules():
+            return
+
+        all_rules = list(Rule.query.all())
 
         self.network = Network()
         for rule in all_rules:
@@ -111,16 +117,19 @@ def load_rules(self):
             # Create the rule.
             self.network.parse_query(rule.conditions, alerters=rule.alerters, rule_id=rule.id)
 
+        # Save the last updated date
+        # Note: we do this here, and not in should_reload_rules, because it's
+        # possible that we've reloaded a rule in between the two functions, and
+        # thus we accidentally don't reload when we should.
+        self.last_update = max(r.updated_at for r in all_rules)
+
     def handle_log_entry(self, entry, node):
         """ The actual entrypoint for handling input log entries. """
         from doorman.models import Rule
         from doorman.rules import RuleMatch
         from doorman.utils import extract_results
 
-        # Need to lazy-load rules
-        if not self.loaded_rules:
-            self.load_rules()
-            self.loaded_rules = True
+        self.load_rules()
 
         to_trigger = []
         for name, action, columns, timestamp in extract_results(entry):

doorman/manage/views.py

@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 import json
+import datetime as dt
 
 from flask import Blueprint, current_app, flash, jsonify, redirect, render_template, request, url_for
 from flask_login import login_required
@@ -23,7 +24,6 @@
     DistributedQuery, DistributedQueryTask,
     FilePath, Node, Pack, Query, Tag, Rule, StatusLog
 )
-from doorman.tasks import reload_rules
 from doorman.utils import (
     create_query_pack_from_upload, flash_errors, get_paginate_options
 )
@@ -469,9 +469,9 @@ def add_rule():
         rule = Rule(name=form.name.data,
                     alerters=form.alerters.data,
                     description=form.description.data,
-                    conditions=form.conditions.data)
+                    conditions=form.conditions.data,
+                    updated_at=dt.datetime.utcnow())
         rule.save()
-        reload_rules.delay()
 
         return redirect(url_for('manage.rule', rule_id=rule.id))
 
@@ -489,8 +489,8 @@ def rule(rule_id):
         rule = rule.update(name=form.name.data,
                            alerters=form.alerters.data,
                            description=form.description.data,
-                           conditions=form.conditions.data)
-        reload_rules.delay()
+                           conditions=form.conditions.data,
+                           updated_at=dt.datetime.utcnow())
         return redirect(url_for('manage.rule', rule_id=rule.id))
 
     form = UpdateRuleForm(request.form, obj=rule)

doorman/models.py

@@ -445,12 +445,14 @@ class Rule(SurrogatePK, Model):
     alerters = Column(ARRAY(db.String), nullable=False)
     description = Column(db.String, nullable=True)
     conditions = Column(JSONB)
+    updated_at = Column(db.DateTime, nullable=False, default=dt.datetime.utcnow)
 
-    def __init__(self, name, alerters, description=None, conditions=None):
+    def __init__(self, name, alerters, description=None, conditions=None, updated_at=None):
         self.name = name
         self.description = description
         self.alerters = alerters
         self.conditions = conditions
+        self.updated_at = updated_at
 
 
 class User(UserMixin, SurrogatePK, Model):

doorman/tasks.py

@@ -12,12 +12,6 @@ def analyze_result(result, node):
     return
 
 
-@celery.task()
-def reload_rules():
-    current_app.rule_manager.load_rules()
-    return
-
-
 @celery.task()
 def example_task(one, two):
     print('Adding {0} and {1}'.format(one, two))

migrations/versions/c17f01adbe31_add_updated_at_column_to_rule.py

@@ -0,0 +1,31 @@
+"""Add "updated_at" column to Rule
+
+Revision ID: c17f01adbe31
+Revises: b50c705fea80
+Create Date: 2016-05-27 15:51:58.168840
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = 'c17f01adbe31'
+down_revision = 'b50c705fea80'
+
+from alembic import op
+import sqlalchemy as sa
+import doorman.database
+
+
+def upgrade():
+    ### commands auto generated by Alembic - please adjust! ###
+    op.add_column('rule',
+        sa.Column('updated_at', sa.DateTime(), nullable=False, server_default=sa.func.now()))
+    ### end Alembic commands ###
+
+    op.create_index('idx__rule__updated_at', 'rule', ['updated_at'])
+
+
+def downgrade():
+    op.drop_index('idx__rule__updated_at')
+    ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column('rule', 'updated_at')
+    ### end Alembic commands ###

tests/test_functional.py

@@ -980,20 +980,6 @@ class TestCreateTag:
 
 class TestAddRule:
 
-    def test_will_reload_rules(self, node, app, testapp):
-        from doorman.tasks import reload_rules
-
-        with mock.patch.object(reload_rules, 'delay', return_value=None) as mock_delay:
-            resp = testapp.post(url_for('manage.add_rule'), {
-                'name': 'Test Rule',
-                'type': 'blacklist',
-                'action': 'both',
-                'alerters': 'debug',
-                'config': '{"field_name": "foo", "blacklist": []}',
-            })
-
-        assert mock_delay.called
-
     def test_supports_custom_operators(self, node, app, testapp):
         # Add a rule to the application
         rule = """
@@ -1047,60 +1033,83 @@ def test_supports_custom_operators(self, node, app, testapp):
 
 
 class TestUpdateRule:
+    pass
 
-    def test_will_reload_rules(self, db, node, app, testapp):
-        from doorman.tasks import reload_rules
 
-        rule_conds = {
-          "condition": "AND",
-          "rules": [
-            {
-              "id": "query_name",
-              "field": "query_name",
-              "type": "string",
-              "input": "text",
-              "operator": "equal",
-              "value": "foo",
-            },
-          ],
+class TestRuleManager:
+    def test_will_load_rules_on_each_call(self, app, db):
+        """
+        Verifies that each call to handle_log_entry will result in a call to load_rules.
+        """
+        from doorman.rules import Network
+
+        mgr = app.rule_manager
+        now = dt.datetime.utcnow()
+
+        with mock.patch.object(mgr, 'load_rules', wraps=lambda: []) as mock_load_rules:
+            with mock.patch.object(mgr, 'network', wraps=Network()) as mock_network:
+                for i in range(0, 2):
+                    mgr.handle_log_entry({
+                        'data': [
+                            {
+                                "diffResults": {
+                                    "added": [{'op': 'added'}],
+                                    "removed": "",
+                                },
+                                "name": "fake",
+                                "hostIdentifier": "hostname.local",
+                                "calendarTime": "%s %s" % (now.ctime(), "UTC"),
+                                "unixTime": now.strftime('%s')
+                            }
+                        ]
+                    }, {'host_identifier': 'yes'})
+
+                assert mock_load_rules.call_count == 2
+
+    def test_will_reload_when_changed(self, app, db):
+        from doorman.models import Rule
+
+        mgr = app.rule_manager
+        dummy_rule = {
+            "id": "query_name",
+            "field": "query_name",
+            "type": "string",
+            "input": "text",
+            "operator": "equal",
+            "value": "dummy-query",
         }
 
-        r = Rule(
-            name='Test-Rule',
-            alerters=['debug'],
-            description='A test rule',
-            conditions=rule_conds
+        now = dt.datetime.utcnow()
+        next = now + dt.timedelta(minutes=5)
+
+        # Insert a first rule.
+        rule = Rule(
+            name='foo',
+            alerters=[],
+            conditions={'condition': 'AND', 'rules': [dummy_rule]},
+            updated_at=now
         )
-        db.session.add(r)
+        db.session.add(rule)
+        db.session.commit()
+
+        # Verify that we will reload these rules
+        assert mgr.should_reload_rules() is True
+
+        # Actually load them
+        mgr.load_rules()
+
+        # Verify that (with no changes made), we should NOT reload.
+        assert mgr.should_reload_rules() is False
+
+        # Make a change to a rule.
+        rule.update(
+            conditions={'condition': 'OR', 'rules': [dummy_rule]},
+            updated_at=next)
+        db.session.add(rule)
         db.session.commit()
 
-        # Manually reload the rules here, and verify that we have the right
-        # rule in our list
-        app.rule_manager.load_rules()
-        assert len(app.rule_manager.network.conditions) == 2
-
-        condition_classes = [x.__class__.__name__ for x in app.rule_manager.network.conditions.values()]
-        assert sorted(condition_classes) == ['AndCondition', 'EqualCondition']
-
-        # Fake wrapper that just calls reload
-        def real_reload(*args, **kwargs):
-            app.rule_manager.load_rules()
-
-        # Update the rule
-        rule_conds['condition'] = 'OR'
-        with mock.patch.object(reload_rules, 'delay', wraps=real_reload) as mock_delay:
-            resp = testapp.post(url_for('manage.rule', rule_id=r.id), {
-                'name': 'Test-Rule',
-                'alerters': 'debug',
-                'conditions': json.dumps(rule_conds),
-            })
-
-        assert mock_delay.called
-
-        # Trigger a manual reload again, and verify that it's been updated
-        app.rule_manager.load_rules()
-        condition_classes = [x.__class__.__name__ for x in app.rule_manager.network.conditions.values()]
-        assert sorted(condition_classes) == ['EqualCondition', 'OrCondition']
+        # Verify that we will now reload
+        assert mgr.should_reload_rules() is True
 
 
 class TestRuleEndToEnd: