__package__ = "archivebox.search" import asyncio import hashlib import json import threading import time from copy import copy from queue import Full, Queue from urllib.parse import urlsplit from uuid import UUID from django.core.cache import cache from django.db import close_old_connections, connections from django.db.models import Q from django.http import HttpResponseForbidden, QueryDict, StreamingHttpResponse from archivebox.search.config import get_search_mode, get_search_mode_base from archivebox.search.query import crawl_config_values_search_wave, iter_query_search_ids SEARCH_RESULT_CACHE_TTL = 60 URL_PREFIX_SEARCH_LIMIT = 500 def get_admin_search_cache_key(request, url: str | None = None) -> str: """Build the cache key for one user and changelist URL.""" # Search streams publish IDs for one exact changelist URL. Keeping the URL # whole makes sidebar filters, ordering, and user scope part of the key. payload = json.dumps( { "user": str(request.user.pk or "anon"), "url": url or request.get_full_path(), }, sort_keys=True, ) return f"abx:admin-search:{hashlib.sha256(payload.encode()).hexdigest()}" def get_public_search_cache_key(request, url: str | None = None) -> str: """Build the cache key for one public search URL.""" payload = json.dumps( { "url": url or request.get_full_path(), }, sort_keys=True, ) return f"abx:public-search:{hashlib.sha256(payload.encode()).hexdigest()}" def get_cached_admin_search_ids(request) -> list[str] | None: """Return streamed admin search IDs from Django cache.""" cached = cache.get(get_admin_search_cache_key(request)) if isinstance(cached, dict): return cached.get("ids") or [] return None def get_cached_public_search_ids(request) -> list[str] | None: """Return streamed public search IDs from Django cache.""" cached = get_cached_public_search_state(request) if isinstance(cached, dict): return cached.get("ids") or [] return None def get_cached_public_search_state(request) -> dict | None: """Return streamed public search state from Django cache.""" cached = cache.get(get_public_search_cache_key(request)) return cached if isinstance(cached, dict) else None def iter_url_search_prefixes(query: str): """Yield URL prefixes that can use indexed startswith scans for common search input.""" query = query.strip().lower() if not query or any(char.isspace() for char in query): return prefixes = [] def add(prefix: str): if prefix and prefix not in prefixes: prefixes.append(prefix) add(query) if "://" in query: parsed = urlsplit(query) if parsed.scheme and parsed.netloc: host = parsed.netloc path = parsed.path or "" if parsed.query: path = f"{path}?{parsed.query}" if host.startswith("www."): add(f"{parsed.scheme}://{host[4:]}{path}") else: add(f"{parsed.scheme}://www.{host}{path}") else: trimmed = query.lstrip("/") for scheme in ("https://", "http://"): add(f"{scheme}{trimmed}") if trimmed.startswith("www."): add(f"{scheme}{trimmed[4:]}") else: add(f"{scheme}www.{trimmed}") yield from prefixes def url_prefix_upper_bound(prefix: str) -> str: """Return the exclusive upper bound for an indexed URL prefix range.""" if not prefix: return prefix return f"{prefix[:-1]}{chr(ord(prefix[-1]) + 1)}" def iter_url_prefix_search_ids(prefix: str, queryset): """Yield IDs for one URL prefix using the URL index, then apply caller filters.""" if not prefix: return model = queryset.model db_alias = queryset.db connection = connections[db_alias] table = connection.ops.quote_name(model._meta.db_table) pk_column = connection.ops.quote_name(model._meta.pk.column) url_column = connection.ops.quote_name(model._meta.get_field("url").column) raw_ids = [] if connection.vendor == "sqlite": # Bytewise range comparison uses the plain url btree index directly. where_clause = f"{url_column} >= %s AND {url_column} < %s" where_params = [prefix, url_prefix_upper_bound(prefix)] else: # Range comparisons are collation-dependent on postgres (linguistic # collations don't compare bytewise), so use LIKE with escaped # wildcards instead — correct under any collation and able to use the # url pattern-ops index. from archivebox.search.query import escape_like_query where_clause = f"{url_column} LIKE %s ESCAPE '\\'" where_params = [f"{escape_like_query(prefix)}%"] with connection.cursor() as cursor: cursor.execute( f""" SELECT {pk_column} FROM {table} WHERE {where_clause} ORDER BY {url_column} LIMIT %s """, [*where_params, URL_PREFIX_SEARCH_LIMIT], ) raw_ids = [str(row[0]).replace("-", "") for row in cursor.fetchall()] if not raw_ids: return valid_ids = {str(pk).replace("-", "") for pk in queryset.filter(pk__in=raw_ids).values_list("pk", flat=True)} for snapshot_id in raw_ids: if snapshot_id in valid_ids: yield snapshot_id def iter_meta_search_ids(query, queryset): """Yield metadata search matches from a filtered Snapshot queryset.""" seen = set() try: snapshot_id = UUID(query) except ValueError: snapshot_id = None if snapshot_id: for pk in queryset.filter(pk=snapshot_id).values_list("pk", flat=True): seen.add(pk) yield pk for prefix in iter_url_search_prefixes(query): for pk in iter_url_prefix_search_ids(prefix, queryset): if pk in seen: continue seen.add(pk) yield pk waves = [ Q(timestamp__startswith=query) | Q(title__istartswith=query), Q(url__icontains=query), Q(title__icontains=query), Q(tags__name__icontains=query), Q(notes__icontains=query), ] for wave in waves: for pk in queryset.filter(wave).values_list("pk", flat=True).distinct().iterator(chunk_size=500): if pk in seen: continue seen.add(pk) yield pk crawl_metadata_wave = Q(crawl__notes__icontains=query) | Q(crawl__label__icontains=query) | Q(crawl__created_by__username=query) if not seen: for pk in queryset.filter(crawl_metadata_wave).values_list("pk", flat=True).distinct().iterator(chunk_size=500): seen.add(pk) yield pk config_wave = crawl_config_values_search_wave(query) if config_wave is not None and not seen: for pk in queryset.filter(config_wave).values_list("pk", flat=True).distinct().iterator(chunk_size=500): seen.add(pk) yield pk def normalize_search_result_id(snapshot_id) -> str | None: """Return a compact Snapshot ID string from a search provider result.""" snapshot_id = str(snapshot_id).strip().lower().replace("-", "") if len(snapshot_id) != 32: return None return snapshot_id def iter_filtered_search_result_ids(iterator, queryset, *, flush_max_delay=0.05): """Yield provider IDs that still match the filtered queryset. This is the single intersection/dedupe path used for metadata and every search backend. It flushes by elapsed time so sparse providers stream rows as soon as IDs are found instead of waiting for a fixed batch size. """ batch = [] seen = set() queued = set() last_flush_at = 0.0 def flush_batch(): nonlocal batch, queued, last_flush_at if not batch: return batch_ids = batch batch = [] queued = set() last_flush_at = time.monotonic() valid = {str(pk).replace("-", "") for pk in queryset.filter(pk__in=batch_ids).values_list("pk", flat=True)} for snapshot_id in batch_ids: if snapshot_id in valid and snapshot_id not in seen: seen.add(snapshot_id) yield snapshot_id for snapshot_id in iterator: snapshot_id = normalize_search_result_id(snapshot_id) if not snapshot_id or snapshot_id in seen or snapshot_id in queued: continue batch.append(snapshot_id) queued.add(snapshot_id) if not seen or time.monotonic() - last_flush_at >= flush_max_delay: yield from flush_batch() if batch: yield from flush_batch() def iter_search_result_ids(query, base_queryset, *, search_mode, config): """Yield filtered Snapshot IDs from the selected search provider.""" search_mode_base = get_search_mode_base(search_mode, config=config) provider = ( iter_meta_search_ids(query, base_queryset) if search_mode_base == "meta" else iter_query_search_ids(query, search_mode=search_mode, config=config) ) yield from iter_filtered_search_result_ids(provider, base_queryset) def snapshot_search_stream_response(query, base_queryset, *, search_mode, config, cache_key, thread_name): """Stream Snapshot search progress and cache matching IDs for a list view.""" if not query: return StreamingHttpResponse((), content_type="text/plain") async def snapshot_ids(): seen = set() ids = [] last_sent = 0 last_sent_at = time.monotonic() stream_max_delay = 0.05 stream_padding = " " * 4096 cache.set(cache_key, {"ids": [], "done": False}, SEARCH_RESULT_CACHE_TTL) queue = Queue(maxsize=8) stop_event = threading.Event() def emit(item): while not stop_event.is_set(): try: queue.put(item, timeout=0.1) return except Full: continue def publish_count(done=False): nonlocal last_sent, last_sent_at cache.set(cache_key, {"ids": list(ids), "done": done}, SEARCH_RESULT_CACHE_TTL) last_sent = len(ids) last_sent_at = time.monotonic() emit(f"{last_sent}{stream_padding}\n") def run_search(): iterator = None try: close_old_connections() iterator = iter_search_result_ids(query, base_queryset, search_mode=search_mode, config=config) for snapshot_id in iterator: if stop_event.is_set(): break snapshot_id = normalize_search_result_id(snapshot_id) if not snapshot_id or snapshot_id in seen: continue seen.add(snapshot_id) ids.append(snapshot_id) if len(ids) == 1 or time.monotonic() - last_sent_at >= stream_max_delay: publish_count() if not stop_event.is_set() and len(ids) != last_sent: publish_count(done=True) except BaseException as err: emit(err) finally: if iterator is not None: try: iterator.close() except AttributeError: pass cache.set(cache_key, {"ids": list(ids), "done": True}, SEARCH_RESULT_CACHE_TTL) close_old_connections() emit(None) threading.Thread(target=run_search, name=thread_name, daemon=True).start() yield f"0{stream_padding}\n" try: while True: item = await asyncio.to_thread(queue.get) if item is None: break if isinstance(item, BaseException): raise item yield item finally: stop_event.set() response = StreamingHttpResponse(snapshot_ids(), content_type="text/plain") response["X-Accel-Buffering"] = "no" return response def admin_snapshot_search_stream_view(model_admin, request): """Stream admin Snapshot search progress and cache matching IDs.""" query = (request.GET.get("q") or "").strip() config = request.archivebox_config search_mode = get_search_mode(request.GET.get("search_mode"), config=config) search_url = request.GET.get("search_url") or request.get_full_path() target_url = urlsplit(search_url) target_get = QueryDict(target_url.query, mutable=True) for key in ("q", "search_mode", "p", "search_url"): target_get.pop(key, None) filter_request = copy(request) filter_request.path = target_url.path or request.path filter_request.path_info = target_url.path or request.path_info filter_request.GET = target_get filter_request.archivebox_config = config # Build the same filtered base queryset the changelist uses, but with the # search params stripped. The stream intersects each wave with this queryset # before writing IDs into the short-lived cache consumed by the changelist. current_request = model_admin.__dict__.get("request") try: base_queryset = model_admin.get_changelist_instance(filter_request).queryset finally: model_admin.request = current_request return snapshot_search_stream_response( query, base_queryset, search_mode=search_mode, config=config, cache_key=get_admin_search_cache_key(request, search_url), thread_name="admin-snapshot-search-stream", ) def public_snapshot_search_stream_view(request): """Stream public Snapshot search progress and cache matching IDs.""" from archivebox.config.common import get_request_config from archivebox.core.models import Snapshot from archivebox.core.permissions import public_snapshots_queryset config = getattr(request, "archivebox_config", None) or get_request_config(request, resolve_plugins=False) if not request.user.is_authenticated and not config.PUBLIC_INDEX: return HttpResponseForbidden("Public index is disabled") query = (request.GET.get("q") or "").strip() search_mode = get_search_mode(request.GET.get("search_mode"), config=config) search_url = request.GET.get("search_url") or request.get_full_path() base_queryset = public_snapshots_queryset(Snapshot.objects.all()) return snapshot_search_stream_response( query, base_queryset, search_mode=search_mode, config=config, cache_key=get_public_search_cache_key(request, search_url), thread_name="public-snapshot-search-stream", )