Merge pull request #1209 from airbnb/dw--csq

Make scheduled_queries/ directory configurable from global.json

conf/global.json

@@ -10,6 +10,9 @@
     ],
     "rule_locations": [
       "rules"
+    ],
+    "scheduled_query_locations": [
+      "scheduled_queries"
     ]
   },
   "infrastructure": {

docs/source/config-global.rst

@@ -66,19 +66,23 @@ Configuration
       ],
       "rule_locations": [
         "rules"
+      ],
+      "scheduled_query_locations": [
+        "scheduled_queries"
       ]
     }
   }
 
 
 Options
 -------
-======================  ============  =================  ===============
-**Key**                 **Required**  **Default**        **Description**
-----------------------  ------------  -----------------  ---------------
-``matcher_locations``   Yes           ``["matchers"]``   List of local paths where ``matchers`` are defined
-``rule_locations``      Yes           ``["rules"]``      List of local paths where ``rules`` are defined
-======================  ============  =================  ===============
+=============================  =============  =========================  ===============
+**Key**                        **Required**   **Default**                **Description**
+-----------------------------  -------------  -------------------------  ---------------
+``matcher_locations``          Yes            ``["matchers"]``           List of local paths where ``matchers`` are defined
+``rule_locations``             Yes            ``["rules"]``              List of local paths where ``rules`` are defined
+``scheduled_query_locations``  Yes            ``["scheduled_queries"]``  List of local paths where ``scheduled_queries`` are defined
+=============================  =============  =========================  ===============
 
 
 **************

streamalert/scheduled_queries/config/services.py

@@ -32,6 +32,7 @@
 from streamalert.scheduled_queries.state.state_manager import StateManager, StepFunctionStateManager
 from streamalert.scheduled_queries.streamalert.kinesis import KinesisClient
 from streamalert.scheduled_queries.support.clock import Clock
+from streamalert.shared.config import load_config
 
 
 # FIXME (Ryxias)
@@ -86,6 +87,11 @@ def configure_container(container):
     container.register(ServiceDefinition('clock', _make_clock))
     container.register(ServiceDefinition('boto3_athena_client', _make_boto3_athena_client))
     container.register(ServiceDefinition('boto3_kinesis_client', _make_boto3_kinesis_client))
+    container.register(ServiceDefinition('config', _load_config))
+
+
+def _load_config(_):
+    return load_config()
 
 
 def _make_command_processor(container):
@@ -135,9 +141,16 @@ def _make_param_generator(container):
     return QueryParameterGenerator(container.get('logger'), container.get('clock'))
 
 
-def _make_query_pack_repo(_):
+def _make_query_pack_repo(container):
     repo = QueryPackRepository
-    repo.load_packs()
+
+    config = container.get('config')
+    query_directories = [
+        item
+        for item in config['global']['general'].get('scheduled_query_locations', [])
+    ]
+
+    repo.load_packs(query_directories)
     return repo
 
 

streamalert/scheduled_queries/query_packs/configuration.py

@@ -15,8 +15,6 @@
 """
 from streamalert.shared.importer import import_folders
 
-PACKS_DIRECTORY = 'scheduled_queries/'
-
 
 class QueryPackConfiguration:
 
@@ -112,5 +110,5 @@ def register(cls, config):
         cls.QUERY_PACKS[name] = config
 
     @classmethod
-    def load_packs(cls):
-        import_folders(PACKS_DIRECTORY)
+    def load_packs(cls, directories):
+        import_folders(*directories)

tests/unit/conf/global.json

@@ -7,7 +7,8 @@
   },
   "general": {
     "matcher_locations": [],
-    "rule_locations": []
+    "rule_locations": [],
+    "scheduled_query_locations": []
   },
   "infrastructure": {
     "alerts_table": {

tests/unit/streamalert/scheduled_queries/query_packs/test_configuration.py

@@ -59,6 +59,6 @@ class TestQueryPackRepository:
     @staticmethod
     def test_load_and_get_packs():
         """StreamQuery - QueryPackRepository - get_packs"""
-        QueryPackRepository.load_packs()
+        QueryPackRepository.load_packs(['scheduled_queries/'])
 
         assert_true(len(QueryPackRepository.get_packs()) >= 1)