Update the query publish tag to match entity (#221)

databuilder/extractor/neo4j_search_data_extractor.py

@@ -73,14 +73,14 @@ class Neo4jSearchDataExtractor(Extractor):
     #  2. add more fields once we have in the graph; 3. change mode to generic once add more support for dashboard
     DEFAULT_NEO4J_DASHBOARD_CYPHER_QUERY = textwrap.dedent(
         """
-        MATCH (db:Dashboard)
-        OPTIONAL MATCH (db)-[:DASHBOARD_OF]->(dbg:Dashboardgroup)
-        OPTIONAL MATCH (db)-[:DESCRIPTION]->(db_descr:Description)
+        MATCH (dashboard:Dashboard)
+        OPTIONAL MATCH (dashboard)-[:DASHBOARD_OF]->(dbg:Dashboardgroup)
+        OPTIONAL MATCH (dashboard)-[:DESCRIPTION]->(db_descr:Description)
         OPTIONAL MATCH (dbg)-[:DESCRIPTION]->(dbg_descr:Description)
         {publish_tag_filter}
-        with db, dbg, db_descr, dbg_descr
-        where db.name is not null
-        return dbg.name as dashboard_group, db.name as dashboard_name,
+        with dashboard, dbg, db_descr, dbg_descr
+        where dashboard.name is not null
+        return dbg.name as dashboard_group, dashboard.name as dashboard_name,
         coalesce(db_descr.description, '') as description,
         coalesce(dbg.description, '') as dashboard_group_description,
         'mode' as product,
@@ -102,12 +102,12 @@ class Neo4jSearchDataExtractor(Extractor):
         Initialize Neo4jExtractor object from configuration and use that for extraction
         """
         self.conf = conf
+        self.entity = conf.get_string(Neo4jSearchDataExtractor.ENTITY_TYPE, default='table').lower()
         # extract cypher query from conf, if specified, else use default query
         if Neo4jSearchDataExtractor.CYPHER_QUERY_CONFIG_KEY in conf:
             self.cypher_query = conf.get_string(Neo4jSearchDataExtractor.CYPHER_QUERY_CONFIG_KEY)
         else:
-            entity_type = conf.get_string(Neo4jSearchDataExtractor.ENTITY_TYPE, default='table').lower()
-            default_query = Neo4jSearchDataExtractor.DEFAULT_QUERY_BY_ENTITY[entity_type]
+            default_query = Neo4jSearchDataExtractor.DEFAULT_QUERY_BY_ENTITY[self.entity]
             self.cypher_query = self._add_publish_tag_filter(conf.get_string(JOB_PUBLISH_TAG, ''),
                                                              cypher_query=default_query)
 
@@ -148,5 +148,8 @@ class Neo4jSearchDataExtractor(Extractor):
         if not publish_tag:
             publish_tag_filter = ''
         else:
-            publish_tag_filter = """WHERE table.published_tag = '{}'""".format(publish_tag)
+            if not hasattr(self, 'entity'):
+                self.entity = 'table'
+            publish_tag_filter = """WHERE {entity}.published_tag = '{tag}'""".format(entity=self.entity,
+                                                                                     tag=publish_tag)
         return cypher_query.format(publish_tag_filter=publish_tag_filter)

tests/unit/extractor/test_neo4j_search_data_extractor.py

@@ -5,6 +5,7 @@ from pyhocon import ConfigFactory
 from databuilder import Scoped
 from databuilder.extractor.neo4j_extractor import Neo4jExtractor
 from databuilder.extractor.neo4j_search_data_extractor import Neo4jSearchDataExtractor
+from databuilder.publisher.neo4j_csv_publisher import JOB_PUBLISH_TAG
 
 
 class TestNeo4jExtractor(unittest.TestCase):
@@ -42,6 +43,29 @@ class TestNeo4jExtractor(unittest.TestCase):
             self.assertEqual(extractor.cypher_query, Neo4jSearchDataExtractor
                              .DEFAULT_NEO4J_DASHBOARD_CYPHER_QUERY.format(publish_tag_filter=''))
 
+    def test_default_search_query_with_tag(self):
+        # type: (Any) -> None
+        with patch.object(Neo4jExtractor, '_get_driver'):
+            extractor = Neo4jSearchDataExtractor()
+            conf = ConfigFactory.from_dict({
+                'extractor.search_data.extractor.neo4j.{}'.format(Neo4jExtractor.GRAPH_URL_CONFIG_KEY):
+                    'test-endpoint',
+                'extractor.search_data.extractor.neo4j.{}'.format(Neo4jExtractor.NEO4J_AUTH_USER):
+                    'test-user',
+                'extractor.search_data.extractor.neo4j.{}'.format(Neo4jExtractor.NEO4J_AUTH_PW):
+                    'test-passwd',
+                'extractor.search_data.{}'.format(Neo4jSearchDataExtractor.ENTITY_TYPE):
+                    'dashboard',
+                'extractor.search_data.{}'.format(JOB_PUBLISH_TAG):
+                    'test-date',
+            })
+            extractor.init(Scoped.get_scoped_conf(conf=conf,
+                                                  scope=extractor.get_scope()))
+
+            self.assertEqual(extractor.cypher_query,
+                             Neo4jSearchDataExtractor.DEFAULT_NEO4J_DASHBOARD_CYPHER_QUERY.format
+                             (publish_tag_filter="""WHERE dashboard.published_tag = 'test-date'"""))
+
 
 if __name__ == '__main__':
     unittest.main()