Dynamically Retrieve Credentials

Author: zprobst

nodestream_plugin_neo4j/neo4j_database.py

@@ -1,24 +1,52 @@
 from logging import getLogger
-from typing import Iterable
+from typing import Iterable, Union, Tuple, Awaitable
+
+from nodestream.file_io import LazyLoadedArgument
 
 from neo4j import AsyncDriver, AsyncGraphDatabase, AsyncSession, Record, RoutingControl
+from neo4j.auth_management import AsyncAuthManagers
 
 from .query import Query
 
 
+def auth_provider_factory(
+    username: Union[str, LazyLoadedArgument],
+    password: Union[str, LazyLoadedArgument],
+) -> Awaitable[Tuple[str, str]]:
+    logger = getLogger(__name__)
+
+    async def auth_provider():
+        logger.info("Fetching new neo4j credentials")
+
+        if isinstance(username, LazyLoadedArgument):
+            logger.debug("Fetching username since value is lazy loaded")
+            current_username = username.get_value()
+        else:
+            current_username = username
+
+        if isinstance(password, LazyLoadedArgument):
+            logger.debug("Fetching password since value is lazy loaded")
+            current_password = password.get_value()
+        else:
+            current_password = password
+
+        return current_username, current_password
+
+    return auth_provider
+
+
 class Neo4jDatabaseConnection:
     @classmethod
     def from_configuration(
         cls,
         uri: str,
-        username: str,
-        password: str,
+        username: Union[str, LazyLoadedArgument],
+        password: Union[str, LazyLoadedArgument],
         database_name: str = "neo4j",
         **driver_kwargs
     ):
-        driver = AsyncGraphDatabase.driver(
-            uri, auth=(username, password), **driver_kwargs
-        )
+        auth = AsyncAuthManagers.basic(auth_provider_factory(username, password))
+        driver = AsyncGraphDatabase.driver(uri, auth=auth, **driver_kwargs)
         return cls(driver, database_name)
 
     def __init__(self, driver: AsyncDriver, database_name: str) -> None:

tests/unit/test_neo4j_database.py

@@ -2,7 +2,12 @@
 from hamcrest import assert_that, equal_to
 from neo4j import AsyncDriver, RoutingControl
 
-from nodestream_plugin_neo4j.neo4j_database import Neo4jDatabaseConnection
+from nodestream.file_io import LazyLoadedArgument
+
+from nodestream_plugin_neo4j.neo4j_database import (
+    Neo4jDatabaseConnection,
+    auth_provider_factory,
+)
 from nodestream_plugin_neo4j.query import Query
 
 
@@ -34,3 +39,23 @@ async def test_session(database_connection):
     session = database_connection.session()
     assert_that(session, equal_to(database_connection.driver.session.return_value))
     database_connection.driver.session.assert_called_once_with(database="neo4j")
+
+
+@pytest.mark.asyncio
+async def test_auth_provider_factory_with_dynamic_values(mocker):
+    username = mocker.Mock(LazyLoadedArgument)
+    password = mocker.Mock(LazyLoadedArgument)
+    provider = auth_provider_factory(username, password)
+    retrieved_username, retrieved_password = await provider()
+    assert_that(retrieved_username, equal_to(username.get_value.return_value))
+    assert_that(retrieved_password, equal_to(password.get_value.return_value))
+
+
+@pytest.mark.asyncio
+async def test_auth_provider_factory_with_static_values():
+    username = "neo4j"
+    password = "password"
+    provider = auth_provider_factory(username, password)
+    retrieved_username, retrieved_password = await provider()
+    assert_that(retrieved_username, equal_to(username))
+    assert_that(retrieved_password, equal_to(password))