Setting idle connection timeout in neo4j_extractor (#85)

databuilder/extractor/neo4j_extractor.py

@@ -2,7 +2,7 @@ import importlib
 import logging
 from typing import Any, Iterator, Union  # noqa: F401
 
-from pyhocon import ConfigTree  # noqa: F401
+from pyhocon import ConfigTree, ConfigFactory  # noqa: F401
 from neo4j.v1 import GraphDatabase
 
 from databuilder.extractor.base_extractor import Extractor
@@ -19,6 +19,9 @@ class Neo4jExtractor(Extractor):
     MODEL_CLASS_CONFIG_KEY = 'model_class'
     NEO4J_AUTH_USER = 'neo4j_auth_user'
     NEO4J_AUTH_PW = 'neo4j_auth_pw'
+    NEO4J_MAX_CONN_LIFE_TIME_SEC = 'neo4j_max_conn_life_time_sec'
+
+    DEFAULT_CONFIG = ConfigFactory.from_dict({NEO4J_MAX_CONN_LIFE_TIME_SEC: 50, })
 
     def init(self, conf):
         # type: (ConfigTree) -> None
@@ -26,7 +29,7 @@ class Neo4jExtractor(Extractor):
         Establish connections and import data model class if provided
         :param conf:
         """
-        self.conf = conf
+        self.conf = conf.with_fallback(Neo4jExtractor.DEFAULT_CONFIG)
         self.graph_url = conf.get_string(Neo4jExtractor.GRAPH_URL_CONFIG_KEY)
         self.cypher_query = conf.get_string(Neo4jExtractor.CYPHER_QUERY_CONFIG_KEY)
         self.driver = self._get_driver()
@@ -55,6 +58,8 @@ class Neo4jExtractor(Extractor):
         Create a Neo4j connection to Database
         """
         return GraphDatabase.driver(self.graph_url,
+                                    max_connection_life_time=self.conf.get_int(
+                                        Neo4jExtractor.NEO4J_MAX_CONN_LIFE_TIME_SEC),
                                     auth=(self.conf.get_string(Neo4jExtractor.NEO4J_AUTH_USER),
                                           self.conf.get_string(Neo4jExtractor.NEO4J_AUTH_PW)))
 

databuilder/publisher/neo4j_csv_publisher.py

@@ -128,7 +128,7 @@ class Neo4jCsvPublisher(Publisher):
 
         self._driver = \
             GraphDatabase.driver(conf.get_string(NEO4J_END_POINT_KEY),
-                                 max_connection_life_time=50,
+                                 max_connection_life_time=conf.get_int(NEO4J_MAX_CONN_LIFE_TIME_SEC),
                                  auth=(conf.get_string(NEO4J_USER), conf.get_string(NEO4J_PASSWORD)))
         self._transaction_size = conf.get_int(NEO4J_TRANSCATION_SIZE)
         self._session = self._driver.session()
@@ -410,7 +410,6 @@ ON MATCH SET {update_prop_body}""".format(create_prop_body=create_prop_body,
         :param expect_result: By having this True, it will validate if result object is not None.
         :return:
         """
-        LOGGER.info('Executing statement: {} with params {}'.format(stmt, params))
         try:
             if LOGGER.isEnabledFor(logging.DEBUG):
                 LOGGER.debug('Executing statement: {} with params {}'.format(stmt, params))

setup.py

 from setuptools import setup, find_packages
 
 
-__version__ = '1.2.8'
+__version__ = '1.2.9'
 
 
 setup(