fix: cypher statement param issue in Neo4jStalenessRemovalTask (#307)

* fix: cypher statement param issue * test: update existing test cases according to change in marker param
amundsen-io · Jul 30, 2020 · 0078761 · 0078761
1 parent d6714b7
commit 0078761
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 11 deletions.
diff --git a/databuilder/task/neo4j_staleness_removal_task.py b/databuilder/task/neo4j_staleness_removal_task.py
@@ -91,7 +91,7 @@ def init(self, conf):
             self.ms_to_expire = conf.get_int(MS_TO_EXPIRE)
             if self.ms_to_expire < conf.get_int(MIN_MS_TO_EXPIRE):
                 raise Exception('{} is too small'.format(MS_TO_EXPIRE))
-            self.marker = '(timestamp() - {})'.format(conf.get_int(MS_TO_EXPIRE))
+            self.marker = self.ms_to_expire
         else:
             self.marker = conf.get_string(JOB_PUBLISH_TAG)
 
@@ -144,7 +144,7 @@ def _decorate_staleness(self, statement):
         """
         if self.ms_to_expire:
             return statement.format(textwrap.dedent("""
-            n.publisher_last_updated_epoch_ms < ${marker}
+            n.publisher_last_updated_epoch_ms < (timestamp() - ${marker})
             OR NOT EXISTS(n.publisher_last_updated_epoch_ms)""".format(marker=MARKER_VAR_NAME)))
 
         return statement.format(textwrap.dedent("""

diff --git a/tests/unit/task/test_neo4j_staleness_removal_task.py b/tests/unit/task/test_neo4j_staleness_removal_task.py
@@ -136,7 +136,7 @@ def test_marker(self):
 
             task.init(job_config)
             self.assertIsNotNone(task.ms_to_expire)
-            self.assertEqual(task.marker, '(timestamp() - 86400000)')
+            self.assertEqual(task.marker, 86400000)
 
     def test_validation_statement_publish_tag(self):
         with patch.object(GraphDatabase, 'driver'), patch.object(Neo4jStalenessRemovalTask, '_execute_cypher_query') \
@@ -215,22 +215,22 @@ def test_validation_statement_ms_to_expire(self):
             RETURN head(node) as type, count
             """))
 
-            mock_execute.assert_any_call(param_dict={'marker': '(timestamp() - 9876543210)'},
+            mock_execute.assert_any_call(param_dict={'marker': 9876543210},
                                          statement=textwrap.dedent("""
             MATCH (n)
             WHERE 
-            n.publisher_last_updated_epoch_ms < $marker
+            n.publisher_last_updated_epoch_ms < (timestamp() - $marker)
             OR NOT EXISTS(n.publisher_last_updated_epoch_ms)
             WITH DISTINCT labels(n) as node, count(*) as count
             RETURN head(node) as type, count
             """))
 
             task._validate_relation_staleness_pct()
-            mock_execute.assert_any_call(param_dict={'marker': '(timestamp() - 9876543210)'},
+            mock_execute.assert_any_call(param_dict={'marker': 9876543210},
                                          statement=textwrap.dedent("""
             MATCH ()-[n]-()
             WHERE 
-            n.publisher_last_updated_epoch_ms < $marker
+            n.publisher_last_updated_epoch_ms < (timestamp() - $marker)
             OR NOT EXISTS(n.publisher_last_updated_epoch_ms)
             RETURN type(n) as type, count(*) as count
             """))
@@ -313,23 +313,23 @@ def test_delete_statement_ms_to_expire(self):
             task._delete_stale_relations()
 
             mock_execute.assert_any_call(dry_run=False,
-                                         param_dict={'marker': '(timestamp() - 9876543210)', 'batch_size': 100},
+                                         param_dict={'marker': 9876543210, 'batch_size': 100},
                                          statement=textwrap.dedent("""
             MATCH (n:Foo)
             WHERE 
-            n.publisher_last_updated_epoch_ms < $marker
+            n.publisher_last_updated_epoch_ms < (timestamp() - $marker)
             OR NOT EXISTS(n.publisher_last_updated_epoch_ms)
             WITH n LIMIT $batch_size
             DETACH DELETE (n)
             RETURN COUNT(*) as count;
             """))
 
             mock_execute.assert_any_call(dry_run=False,
-                                         param_dict={'marker': '(timestamp() - 9876543210)', 'batch_size': 100},
+                                         param_dict={'marker': 9876543210, 'batch_size': 100},
                                          statement=textwrap.dedent("""
             MATCH ()-[n:BAR]-()
             WHERE 
-            n.publisher_last_updated_epoch_ms < $marker
+            n.publisher_last_updated_epoch_ms < (timestamp() - $marker)
             OR NOT EXISTS(n.publisher_last_updated_epoch_ms)
             WITH n LIMIT $batch_size
             DELETE n