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),
    )