"""Neo4J DB connector."""
from typing import Optional

import config
from lambdacommon.common_config import logger
from neo4j import Driver, GraphDatabase


def _get_driver():
    """Connect to Neo4J and return Neo4J driver."""
    try:
        driver = GraphDatabase.driver(
            config.NEO4J_URL,
            auth=(config.NEO4J_USERNAME, config.NEO4J_PASSWORD),
            max_transaction_retry_time=config.NEO4J_MAX_RETRY_TIME
        )
        logger.info('Successfully connected to Neo4J.')
        return driver
    except Exception:
        logger.exception('Failed to connect to Neo4J.', exc_info=True)
        raise


neo4j_driver: Optional[Driver] = None
if config.ENVIRONMENT != config.TEST_ENVIRONMENT:
    neo4j_driver = _get_driver()