/cognee/infrastructure/databases/relational/create_relational_engine.py

from functools import lru_cache
from cognee.shared.lru_cache import DATABASE_MAX_LRU_CACHE_SIZE

from sqlalchemy import URL

from .sqlalchemy.SqlAlchemyAdapter import SQLAlchemyAdapter


@lru_cache(maxsize=DATABASE_MAX_LRU_CACHE_SIZE)
def create_relational_engine(
    db_path: str,
    db_name: str,
    db_host: str,
    db_port: str,
    db_username: str,
    db_password: str,
    db_provider: str,
    database_connect_args: tuple = None,
    pool_args: tuple = None,
) -> SQLAlchemyAdapter:
    """
    Create a relational database engine based on the specified parameters.

    Parameters:
    -----------

        - db_path (str): The file path to the database directory, applicable for SQLite.
        - db_name (str): The name of the database to be accessed or created.
        - db_host (str): The hostname or IP address of the database server, required for
          PostgreSQL.
        - db_port (str): The port number on which the database server is listening, required
          for PostgreSQL.
        - db_username (str): The username for database authentication, required for
          PostgreSQL.
        - db_password (str): The password for database authentication, required for
          PostgreSQL.
        - db_provider (str): The type of database provider (e.g., 'sqlite' or 'postgres').
        - database_connect_args (dict, optional): Database driver connection arguments.

    Returns:
    --------

        Returns a SQLAlchemyAdapter instance for the specified database connection.
    """
    # Transform pool_args and database_connect_args from tuple of key-value pairs back to dictionary format if they are provided
    database_connect_args = dict(database_connect_args) if database_connect_args else {}
    pool_args = dict(pool_args) if pool_args else {}

    if db_provider == "sqlite":
        connection_string = f"sqlite+aiosqlite:///{db_path}/{db_name}"

    elif db_provider == "postgres":
        try:
            # Test if asyncpg is available
            import asyncpg

            # Handle special characters in username and password like # or @
            connection_string = URL.create(
                "postgresql+asyncpg",
                username=db_username,
                password=db_password,
                host=db_host,
                port=int(db_port),
                database=db_name,
            )

        except ImportError:
            raise ImportError(
                "PostgreSQL dependencies are not installed. Please install with 'pip install cognee\"[postgres]\"' or 'pip install cognee\"[postgres-binary]\"' to use PostgreSQL functionality."
            )

    else:
        raise ConnectionError("unsupported DB type: " + db_provider)

    return SQLAlchemyAdapter(
        connection_string,
        connect_args=database_connect_args,
        pool_args=pool_args,
    )