label_studio/core/utils/db.py - develop - vshmidt/label-studio

label_studio/core/utils/db.py
203 строки8 KB
Sergey Zhuk
ci: ruff format more files
06 мар 2026, 23:53
06 мар 2026, 23:53e69a84d
import itertools
import logging
import time
from typing import Dict, Optional, TypeVar

from django.db import OperationalError, connection, models, transaction
from django.db.models import Model, QuerySet, Subquery
from django.db.models.signals import post_migrate
from django.db.utils import DatabaseError, ProgrammingError
from django.dispatch import receiver

logger = logging.getLogger(__name__)


class SQCount(Subquery):
    template = '(SELECT count(*) FROM (%(subquery)s) _count)'
    output_field = models.IntegerField()


ModelType = TypeVar('ModelType', bound=Model)


def fast_first(queryset: QuerySet[ModelType]) -> Optional[ModelType]:
    """Replacement for queryset.first() when you don't need ordering,
    queryset.first() works slowly in some cases
    """

    if result := queryset[:1]:
        return result[0]
    return None


def fast_first_or_create(model, **model_params) -> Optional[ModelType]:
    """Like get_or_create, but using fast_first instead of first(). Additionally, unlike get_or_create, this method will not raise an exception if more than one model instance matching the given params is returned, making it a safer choice than get_or_create for models that don't have a uniqueness constraint on the fields used."""
    if instance := fast_first(model.objects.filter(**model_params)):
        return instance
    return model.objects.create(**model_params)


def batch_update_with_retry(queryset, batch_size=500, max_retries=3, **update_fields):
    """
    Update objects in batches with retry logic to handle deadlocks.

    Args:
        queryset: QuerySet of objects to update
        batch_size: Number of objects to update in each batch
        max_retries: Maximum number of retry attempts for each batch
        **update_fields: Fields to update (e.g., overlap=1)
    """
    object_ids = list(queryset.values_list('id', flat=True))
    total_objects = len(object_ids)

    for i in range(0, total_objects, batch_size):
        batch_ids = object_ids[i : i + batch_size]
        retry_count = 0
        last_error = None

        while retry_count < max_retries:
            try:
                with transaction.atomic():
                    queryset.model.objects.filter(id__in=batch_ids).update(**update_fields)
                break
            except OperationalError as e:
                last_error = e
                if 'deadlock detected' in str(e):
                    retry_count += 1
                    wait_time = 0.1 * (2**retry_count)  # Exponential backoff
                    logger.warning(
                        f'Deadlock detected, retry {retry_count}/{max_retries} '
                        f'for batch {i}-{i + len(batch_ids)}. Waiting {wait_time}s...'
                    )
                    time.sleep(wait_time)
                else:
                    raise
        else:
            logger.error(f'Failed to update batch after {max_retries} retries. Batch: {i}-{i + len(batch_ids)}')
            raise last_error


def batch_delete(queryset, batch_size=500):
    """
    Delete objects in batches to minimize memory usage and transaction size.

    Args:
        queryset: The queryset to delete
        batch_size: Number of objects to delete in each batch

    Returns:
        int: Total number of deleted objects
    """
    total_deleted = 0

    # Create a database cursor that yields primary keys without loading all into memory
    # The iterator position is maintained between calls to islice
    # Example: if queryset has 1500 records and batch_size=500:
    # - First iteration will get records 1-500
    # - Second iteration will get records 501-1000
    # - Third iteration will get records 1001-1500
    # - Fourth iteration will get empty list (no more records)
    pks_to_delete = queryset.values_list('pk', flat=True).iterator(chunk_size=batch_size)

    # Delete in batches
    while True:
        # Get the next batch of primary keys from where the iterator left off
        # islice advances the iterator's position after taking batch_size items
        batch_iterator = itertools.islice(pks_to_delete, batch_size)

        # Convert the slice iterator to a list we can use
        # This only loads batch_size items into memory at a time
        batch = list(batch_iterator)

        # If no more items to process, we're done
        # This happens when the iterator is exhausted
        if not batch:
            break

        # Delete the batch in a transaction
        with transaction.atomic():
            # Delete all objects whose primary keys are in this batch
            deleted = queryset.model.objects.filter(pk__in=batch).delete()[0]
            total_deleted += deleted

    return total_deleted


# =====================
# Schema helpers
# =====================

_column_presence_cache: Dict[str, Dict[str, Dict[str, bool]]] = {}


def current_db_key() -> str:
    """Return a process-stable identifier for the current DB connection.

    Using vendor + NAME isolates caches between sqlite test DBs and postgres runs,
    avoiding stale lookups across pytest sessions or multi-DB setups.
    """
    try:
        name = str(connection.settings_dict.get('NAME'))
    except Exception as e:
        name = 'unknown'
        logger.error(f'Error getting current DB key: {e}')
    return f'{connection.vendor}:{name}'


def has_column_cached(table_name: str, column_name: str) -> bool:
    """Check if a DB column exists for the given table, with per-process memoization.

    Notes:
    - Uses Django introspection; caches per (table, column) with case-insensitive column keys.
    - Safe during early migrations; returns False on any error.
    - Works in both sync and async contexts by temporarily allowing async-unsafe operations.
    """
    col_key = column_name.lower()
    db_cache = _column_presence_cache.get(current_db_key())
    table_cache = db_cache.get(table_name) if db_cache else None
    if table_cache and col_key in table_cache:
        return table_cache[col_key]

    # Check if we're in an async context and need to allow async-unsafe operations
    import asyncio
    import os

    in_async_context = False
    previous_value = os.environ.get('DJANGO_ALLOW_ASYNC_UNSAFE')

    try:
        asyncio.get_running_loop()
        in_async_context = True
        # Temporarily allow async-unsafe operations for schema introspection
        # This is safe because:
        # 1. We're only reading DB schema metadata, not data
        # 2. This happens once per process during startup and is cached
        # 3. No concurrent access to the same data
        os.environ['DJANGO_ALLOW_ASYNC_UNSAFE'] = 'true'
    except RuntimeError:
        pass  # Not in async context

    try:
        with connection.cursor() as cursor:
            cols = connection.introspection.get_table_description(cursor, table_name)
        present = any(getattr(col, 'name', '').lower() == col_key for col in cols)
    except (DatabaseError, ProgrammingError):
        present = False
    finally:
        # Restore the previous value
        if in_async_context:
            if previous_value is None:
                os.environ.pop('DJANGO_ALLOW_ASYNC_UNSAFE', None)
            else:
                os.environ['DJANGO_ALLOW_ASYNC_UNSAFE'] = previous_value

    _column_presence_cache.setdefault(current_db_key(), {}).setdefault(table_name, {})[col_key] = present
    return present


@receiver(post_migrate)
def signal_clear_column_presence_cache(**_kwargs):
    """If some migration adds a column, we need to clear the column_presence_cache
    so that the next migration can introspect the new column using has_column_cached()."""
    logger.debug('Clearing column presence cache in post_migrate signal')
    _column_presence_cache.clear()