from apollo_utils.service.exceptions import NotFound
from typing import Optional

from apollo_main_db.apollo.models import ApolloKeyValueStorage
from src.db.base import session


def get_value(key: str) -> Optional[str]:
    """Get value by key from apollo key-value storage.

    Args:
        key: Key.

    Returns:
        Value.

    """
    result = session.query(ApolloKeyValueStorage.value).filter(ApolloKeyValueStorage.key == key).first()
    if not result:
        raise NotFound(f"Key '{key}' is not found.")
    return result[0]


def set_value(key: str, value: str):
    """Set value into apollo key-value storage.

    Args:
        key: Key.
        value: Value.

    """
    result = (
        session.query(ApolloKeyValueStorage)
        .filter(ApolloKeyValueStorage.key == key)
        .update({ApolloKeyValueStorage.value: value}, synchronize_session=False)
    )
    session.commit()
    if not result:
        raise NotFound(f"Key '{key}' is not found.")