Source code for neosqlite.collection.jsonb_support

"""
JSONB support detection utilities for NeoSQLite.

Provides efficient detection of JSONB capabilities with automatic caching
to avoid redundant database queries.
"""

import logging

from .._sqlite import sqlite3

logger = logging.getLogger(__name__)

# Module-level cache for JSONB support detection results
# Key: connection id (int), Value: dict with support flags
_jsonb_support_cache: dict[int, dict[str, bool]] = {}




def clear_jsonb_cache(db_connection=None) -> None:
    """
    Clear the JSONB support cache.

    Args:
        db_connection: Optional specific connection to clear from cache.
                      If None, clears the entire cache.
    """
    if db_connection is None:
        _jsonb_support_cache.clear()
    else:
        conn_id = id(db_connection)
        _jsonb_support_cache.pop(conn_id, None)





def supports_jsonb(db_connection) -> bool:
    """
    Check if the SQLite connection supports JSONB functions.

    This function tests whether the SQLite installation has JSONB support
    by attempting to call the jsonb() function. Results are cached to avoid
    redundant queries.

    Args:
        db_connection: SQLite database connection to test

    Returns:
        bool: True if JSONB is supported, False otherwise
    """
    conn_id = id(db_connection)

    if conn_id not in _jsonb_support_cache:
        _jsonb_support_cache[conn_id] = {}

    cache = _jsonb_support_cache[conn_id]

    if "jsonb_supported" not in cache:
        try:
            db_connection.execute('SELECT jsonb(\'{"test": "value"}\')')
            cache["jsonb_supported"] = True
        except sqlite3.OperationalError as e:
            logger.debug(f"JSONB not supported: {e}")
            cache["jsonb_supported"] = False

    return cache["jsonb_supported"]





def supports_jsonb_each(db_connection) -> bool:
    """
    Check if the SQLite connection supports jsonb_each() table-valued function.

    The jsonb_each() and jsonb_tree() functions were added in SQLite 3.51.0.
    This function tests whether they're available by attempting to use them.
    Results are cached to avoid redundant queries.

    Args:
        db_connection: SQLite database connection to test

    Returns:
        bool: True if jsonb_each is supported, False otherwise
    """
    conn_id = id(db_connection)

    if conn_id not in _jsonb_support_cache:
        _jsonb_support_cache[conn_id] = {}

    cache = _jsonb_support_cache[conn_id]

    if "jsonb_each_supported" not in cache:
        try:
            db_connection.execute(
                "SELECT key, value FROM jsonb_each('[1,2,3]')"
            )
            cache["jsonb_each_supported"] = True
        except sqlite3.OperationalError as e:
            logger.debug(f"jsonb_each not supported: {e}")
            cache["jsonb_each_supported"] = False

    return cache["jsonb_each_supported"]





def _get_json_function_prefix(jsonb_supported: bool) -> str:
    """
    Get the appropriate JSON function prefix based on JSONB support.

    Args:
        jsonb_supported: Whether JSONB functions are supported

    Returns:
        str: "jsonb" if JSONB is supported, "json" otherwise
    """
    return "jsonb" if jsonb_supported else "json"





def _get_json_each_function(
    jsonb_supported: bool, jsonb_each_supported: bool
) -> str:
    """
    Get the appropriate json_each function name based on support.

    Args:
        jsonb_supported: Whether JSONB functions are supported
        jsonb_each_supported: Whether jsonb_each is supported (SQLite 3.51.0+)

    Returns:
        str: "jsonb_each" if supported, "json_each" otherwise
    """
    return (
        "jsonb_each"
        if (jsonb_supported and jsonb_each_supported)
        else "json_each"
    )





def _get_json_tree_function(
    jsonb_supported: bool, jsonb_each_supported: bool
) -> str:
    """
    Get the appropriate json_tree function name based on support.

    Args:
        jsonb_supported: Whether JSONB functions are supported
        jsonb_each_supported: Whether jsonb_each/jsonb_tree is supported (SQLite 3.51.0+)

    Returns:
        str: "jsonb_tree" if supported, "json_tree" otherwise
    """
    return (
        "jsonb_tree"
        if (jsonb_supported and jsonb_each_supported)
        else "json_tree"
    )





def _get_json_group_array_function(jsonb_supported: bool) -> str:
    """
    Get the appropriate json_group_array function name based on support.

    Args:
        jsonb_supported: Whether JSONB functions are supported

    Returns:
        str: "jsonb_group_array" if supported, "json_group_array" otherwise
    """
    return "jsonb_group_array" if jsonb_supported else "json_group_array"





def should_use_json_functions(
    query: dict | None = None, jsonb_supported: bool = False
) -> bool:
    """
    Determine if we should use json_* functions instead of jsonb_* functions.

    This function determines whether to use json_* or jsonb_* functions based on:
    1. JSONB support availability
    2. Query content (specifically text search queries which require FTS compatibility)

    Args:
        query: MongoDB query dictionary to check for text search operations
        jsonb_supported: Whether JSONB functions are supported by the database

    Returns:
        bool: True if json_* functions should be used, False if jsonb_* functions should be used
    """
    # If JSONB is not supported, we must use json_* functions
    if not jsonb_supported:
        return True

    # If no query provided, default to using jsonb_* functions for better performance
    if query is None:
        return False

    # Check if query contains text search operations which require FTS compatibility
    return _contains_text_operator(query)





def _contains_text_operator(query: dict) -> bool:
    """
    Check if a query contains any $text operators, including nested in logical operators.

    This method recursively traverses a MongoDB query specification to detect the presence
    of $text operators, which require special handling and fallback to Python implementation.
    It checks both top-level $text operators and those nested within logical operators
    ($and, $or, $nor, $not).

    Args:
        query: The query to check

    Returns:
        True if the query contains $text operators, False otherwise
    """
    if not isinstance(query, dict):
        return False

    for field, value in query.items():
        if field in ("$and", "$or", "$nor"):
            # Check each condition in logical operators
            if isinstance(value, list):
                for condition in value:
                    if isinstance(condition, dict) and _contains_text_operator(
                        condition
                    ):
                        return True
        elif field == "$not":
            # Check the condition in $not operator
            if isinstance(value, dict) and _contains_text_operator(value):
                return True
        elif field == "$text":
            # Found a $text operator
            return True
    return False





def json_data_column(jsonb_supported: bool, source: str = "data") -> str:
    """Return the data column expression for SELECT queries.

    When JSONB is supported, wraps the source with json() to convert
    JSONB to text.  Otherwise returns the source unchanged.

    When JSONB is supported, also handles malformed JSON gracefully by
    checking json_valid() first for text values. This prevents the entire
    query from failing when a single row has corrupted data. JSONB blobs
    are always passed through json() since they're stored in SQLite's
    binary JSON format.
    """
    if jsonb_supported:
        # Handle three cases:
        # 1. JSONB blobs (typeof='blob'): always valid, convert with json()
        # 2. Valid JSON text: convert with json() for consistency
        # 3. Malformed text or other types: return as-is (Python handles errors)
        return (
            f"CASE WHEN typeof({source})='blob' THEN json({source}) "
            f"WHEN typeof({source})='text' AND json_valid({source}) "
            f"THEN json({source}) ELSE {source} END"
        )
    return source