From f3f98643ed1609873f52a7a79f3e6f8acb06237e Mon Sep 17 00:00:00 2001 From: Nurit Date: Mon, 18 Mar 2019 16:49:04 -0400 Subject: [PATCH 1/4] adding filter function to sources to filter message before saving them to the db --- comet_core/app.py | 30 +++++++++++++++++++++++++++++- 1 file changed, 29 insertions(+), 1 deletion(-) diff --git a/comet_core/app.py b/comet_core/app.py index 206a93b..2426e30 100644 --- a/comet_core/app.py +++ b/comet_core/app.py @@ -144,6 +144,7 @@ def __init__(self, database_uri='sqlite://'): self.inputs = list() self.instantiated_inputs = list() self.hydrators = dict() + self.filters = dict() self.parsers = dict() self.routers = SourceTypeFunction() self.escalators = SourceTypeFunction() @@ -189,8 +190,13 @@ def message_callback(self, source_type, message): if hydrate: hydrate(event) + # Filter event + filter_event = self.filters.get(source_type) + event = filter_event(event) + # Add to datastore - self.data_store.add_record(event.get_record()) + if event: + self.data_store.add_record(event.get_record()) return True def set_config(self, source_type, config): @@ -278,6 +284,28 @@ def decorator(func): else: self.hydrators[source_type] = func + def register_filter(self, source_type, func=None): + """Register a filter function to filter events before saving them to the db. + + This method can be used either as a decorator or with a filter function passed in. + + Args: + source_type (str): the source type to register the filter for + func (Optional[function]): a function that filter a message of type source_type, or None if used as a + decorator + Return: + function or None: if no func is given returns a decorator function, otherwise None + """ + if not func: + # pylint: disable=missing-docstring, missing-return-doc, missing-return-type-doc + def decorator(func): + self.filters[source_type] = func + return func + + return decorator + else: + self.filters[source_type] = func + def register_router(self, source_types=None, func=None): """Register a hydrator. From d4c107d9cb4d96e74831e4e168578530203d4b0f Mon Sep 17 00:00:00 2001 From: Nurit Date: Tue, 19 Mar 2019 11:10:10 -0400 Subject: [PATCH 2/4] check if there is a filter function --- comet_core/app.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/comet_core/app.py b/comet_core/app.py index 2426e30..ded837d 100644 --- a/comet_core/app.py +++ b/comet_core/app.py @@ -192,7 +192,8 @@ def message_callback(self, source_type, message): # Filter event filter_event = self.filters.get(source_type) - event = filter_event(event) + if filter_event: + event = filter_event(event) # Add to datastore if event: From 8e906600b487a93141864b0c96417ffd62c2d516 Mon Sep 17 00:00:00 2001 From: Nurit Date: Mon, 6 May 2019 13:16:40 -0400 Subject: [PATCH 3/4] adding tests for filter decorator --- tests/test_app.py | 39 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/tests/test_app.py b/tests/test_app.py index 173e072..54c1264 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -135,11 +135,34 @@ def loads(self, msg): hydrator_mock = mock.Mock() app.register_hydrator('test', hydrator_mock) + filter_return_value = EventContainer('test', {"a": "b"}) + filter_mock = mock.Mock(return_value=filter_return_value) + app.register_filter('test', filter_mock) + assert not app.message_callback('test1', '{}') assert not app.message_callback('test', '{ "c": "d" }') app.message_callback('test', '{ "a": "b" }') assert hydrator_mock.called + assert filter_mock.called + assert filter_mock.return_value, filter_return_value + + +def test_message_callback_filter(app): + @app.register_parser('test') + class TestParser: + def loads(self, msg): + ev = json.loads(msg) + if 'a' in ev: + return ev, None + return None, 'fail' + + filter_mock = mock.Mock(return_value=None) + app.register_filter('test', filter_mock) + + app.message_callback('test', '{ "a": "b" }') + assert filter_mock.called + assert filter_mock.return_value is None def test_register_input(app): @@ -184,6 +207,22 @@ def test_hydrator(*args): assert len(app.hydrators) == 2, app.hydrators +def test_register_filter(app): + assert not app.filters + + @app.register_filter('test1') + def test_filter(*args): + pass + + # Override existing + app.register_filter('test1', test_filter) + assert len(app.filters) == 1, app.filters + + # Add another + app.register_filter('test2', test_filter) + assert len(app.filters) == 2, app.filters + + def test_set_config(app): assert not app.specific_configs From f4f42a443c4e905bdec9c7723facb637bcd54081 Mon Sep 17 00:00:00 2001 From: Fredrik Gustafsson Date: Tue, 7 May 2019 09:47:57 +0200 Subject: [PATCH 4/4] bump version --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index d69f3f9..63dc556 100644 --- a/setup.py +++ b/setup.py @@ -25,7 +25,7 @@ setuptools.setup( name="comet-core", - version="1.0.10", + version="1.0.11", url="https://github.com/spotify/comet-core", author="Spotify Platform Security",