22 lines
607 B
Python
22 lines
607 B
Python
from airflow.hooks.base import BaseHook
|
|
from neo4j import Driver, GraphDatabase
|
|
from task_01.dto import ConnectionInfo
|
|
|
|
|
|
def get_neo4j_connection(conn_id: str) -> ConnectionInfo:
|
|
conn = BaseHook.get_connection(conn_id)
|
|
uri = f'bolt://{conn.host}:{conn.port}'
|
|
return ConnectionInfo(
|
|
uri,
|
|
username=conn.login,
|
|
password=conn.password,
|
|
)
|
|
|
|
|
|
def get_driver(conn_id: str) -> Driver:
|
|
connection_info = get_neo4j_connection(conn_id)
|
|
return GraphDatabase.driver(
|
|
uri=connection_info.uri,
|
|
auth=(connection_info.username, connection_info.password),
|
|
)
|