awswrangler/athena/_cache.py

"""Cache Module for Amazon Athena."""

from __future__ import annotations

import datetime
import logging
import re
import threading
from heapq import heappop, heappush
from typing import TYPE_CHECKING, Match, NamedTuple

import boto3

from awswrangler import _utils, typing

if TYPE_CHECKING:
    from mypy_boto3_athena.type_defs import QueryExecutionTypeDef

_logger: logging.Logger = logging.getLogger(__name__)


class _CacheInfo(NamedTuple):
    has_valid_cache: bool
    file_format: str | None = None
    query_execution_id: str | None = None
    query_execution_payload: "QueryExecutionTypeDef" | None = None


class _LocalMetadataCacheManager:
    def __init__(self) -> None:
        self._lock: threading.Lock = threading.Lock()
        self._cache: dict[str, "QueryExecutionTypeDef"] = {}
        self._pqueue: list[tuple[datetime.datetime, str]] = []
        self._max_cache_size = 100

    def update_cache(self, items: list["QueryExecutionTypeDef"]) -> None:
        """
        Update the local metadata cache with new query metadata.

        Parameters
        ----------
        items
            List of query execution metadata which is returned by boto3 `batch_get_query_execution()`.
        """
        with self._lock:
            if self._pqueue:
                oldest_item = self._cache.get(self._pqueue[0][1])
                if oldest_item:
                    items = list(
                        filter(
                            lambda x: x["Status"]["SubmissionDateTime"] > oldest_item["Status"]["SubmissionDateTime"],
                            items,
                        )
                    )

            cache_oversize = len(self._cache) + len(items) - self._max_cache_size
            for _ in range(cache_oversize):
                _, query_execution_id = heappop(self._pqueue)
                self._cache.pop(query_execution_id, None)

            for item in items[: self._max_cache_size]:
                heappush(self._pqueue, (item["Status"]["SubmissionDateTime"], item["QueryExecutionId"]))
                self._cache[item["QueryExecutionId"]] = item

    def sorted_successful_generator(self) -> list["QueryExecutionTypeDef"]:
        """
        Sorts the entries in the local cache based on query Completion DateTime.

        This is useful to guarantee LRU caching rules.

        Returns
        -------
            Returns successful DDL and DML queries sorted by query completion time.
        """
        filtered: list["QueryExecutionTypeDef"] = []
        with self._lock:
            for query in self._cache.values():
                if (query["Status"].get("State") == "SUCCEEDED") and (query.get("StatementType") in ["DDL", "DML"]):
                    filtered.append(query)
        return sorted(filtered, key=lambda e: str(e["Status"]["CompletionDateTime"]), reverse=True)

    def __contains__(self, key: str) -> bool:
        return key in self._cache

    @property
    def max_cache_size(self) -> int:
        """Property max_cache_size."""
        return self._max_cache_size

    @max_cache_size.setter
    def max_cache_size(self, value: int) -> None:
        self._max_cache_size = value


def _parse_select_query_from_possible_ctas(possible_ctas: str) -> str | None:
    """Check if `possible_ctas` is a valid parquet-generating CTAS and returns the full SELECT statement."""
    possible_ctas = possible_ctas.lower()
    parquet_format_regex: str = r"format\s*=\s*\'parquet\'\s*"
    is_parquet_format: Match[str] | None = re.search(pattern=parquet_format_regex, string=possible_ctas)
    if is_parquet_format is not None:
        unstripped_select_statement_regex: str = r"\s+as\s+\(*(select|with).*"
        unstripped_select_statement_match: Match[str] | None = re.search(
            unstripped_select_statement_regex, possible_ctas, re.DOTALL
        )
        if unstripped_select_statement_match is not None:
            stripped_select_statement_match: Match[str] | None = re.search(
                r"(select|with).*", unstripped_select_statement_match.group(0), re.DOTALL
            )
            if stripped_select_statement_match is not None:
                return stripped_select_statement_match.group(0)
    return None


def _compare_query_string(sql: str, other: str) -> bool:
    comparison_query = _prepare_query_string_for_comparison(query_string=other)
    _logger.debug("sql: %s", sql)
    _logger.debug("comparison_query: %s", comparison_query)
    return sql == comparison_query


def _prepare_query_string_for_comparison(query_string: str) -> str:
    """To use cached data, we need to compare queries. Returns a query string in canonical form."""
    # for now this is a simple complete strip, but it could grow into much more sophisticated
    # query comparison data structures
    query_string = "".join(query_string.split()).strip().lower()
    while query_string.startswith("(") and query_string.endswith(")"):
        query_string = query_string[1:-1]
    query_string = query_string[:-1] if query_string.endswith(";") else query_string
    return query_string


def _get_last_query_infos(
    max_remote_cache_entries: int,
    boto3_session: boto3.Session | None = None,
    workgroup: str | None = None,
) -> list["QueryExecutionTypeDef"]:
    """Return an iterator of `query_execution_info`s run by the workgroup in Athena."""
    client_athena = _utils.client(service_name="athena", session=boto3_session)
    page_size = 50
    args: dict[str, str | dict[str, int]] = {
        "PaginationConfig": {"MaxItems": max_remote_cache_entries, "PageSize": page_size}
    }
    if workgroup is not None:
        args["WorkGroup"] = workgroup
    paginator = client_athena.get_paginator("list_query_executions")
    uncached_ids: list[str] = []
    for page in paginator.paginate(**args):  # type: ignore[arg-type]
        _logger.debug("paginating Athena's queries history...")
        query_execution_id_list: list[str] = page["QueryExecutionIds"]
        for query_execution_id in query_execution_id_list:
            if query_execution_id not in _cache_manager:
                uncached_ids.append(query_execution_id)
    if uncached_ids:
        new_execution_data: list[QueryExecutionTypeDef] = []
        for i in range(0, len(uncached_ids), page_size):
            new_execution_data.extend(
                client_athena.batch_get_query_execution(  # type: ignore[arg-type]
                    QueryExecutionIds=uncached_ids[i : i + page_size],
                ).get("QueryExecutions")
            )
        _cache_manager.update_cache(new_execution_data)
    return _cache_manager.sorted_successful_generator()


def _check_for_cached_results(
    sql: str,
    boto3_session: boto3.Session | None,
    workgroup: str | None,
    athena_cache_settings: typing.AthenaCacheSettings | None = None,
) -> _CacheInfo:
    """
    Check whether `sql` has been run before, within the `max_cache_seconds` window, by the `workgroup`.

    If so, returns a dict with Athena's `query_execution_info` and the data format.
    """
    athena_cache_settings = athena_cache_settings or {}

    max_cache_seconds = athena_cache_settings.get("max_cache_seconds", 0)
    max_cache_query_inspections = athena_cache_settings.get("max_cache_query_inspections", 50)
    max_remote_cache_entries = athena_cache_settings.get("max_remote_cache_entries", 50)
    max_local_cache_entries = athena_cache_settings.get("max_local_cache_entries", 100)
    max_remote_cache_entries = min(max_remote_cache_entries, max_local_cache_entries)

    _cache_manager.max_cache_size = max_local_cache_entries

    if max_cache_seconds <= 0:
        return _CacheInfo(has_valid_cache=False)
    num_executions_inspected: int = 0
    comparable_sql: str = _prepare_query_string_for_comparison(sql)
    current_timestamp: datetime.datetime = datetime.datetime.now(datetime.timezone.utc)
    _logger.debug("current_timestamp: %s", current_timestamp)
    for query_info in _get_last_query_infos(
        max_remote_cache_entries=max_remote_cache_entries,
        boto3_session=boto3_session,
        workgroup=workgroup,
    ):
        query_execution_id: str = query_info["QueryExecutionId"]
        query_timestamp: datetime.datetime = query_info["Status"]["CompletionDateTime"]
        _logger.debug("query_timestamp: %s", query_timestamp)
        if (current_timestamp - query_timestamp).total_seconds() > max_cache_seconds:
            return _CacheInfo(
                has_valid_cache=False, query_execution_id=query_execution_id, query_execution_payload=query_info
            )
        statement_type: str | None = query_info.get("StatementType")
        if statement_type == "DDL" and query_info["Query"].startswith("CREATE TABLE"):
            parsed_query: str | None = _parse_select_query_from_possible_ctas(possible_ctas=query_info["Query"])
            if parsed_query is not None:
                if _compare_query_string(sql=comparable_sql, other=parsed_query):
                    return _CacheInfo(
                        has_valid_cache=True,
                        file_format="parquet",
                        query_execution_id=query_execution_id,
                        query_execution_payload=query_info,
                    )
        elif statement_type == "DML" and not query_info["Query"].startswith("INSERT"):
            if _compare_query_string(sql=comparable_sql, other=query_info["Query"]):
                return _CacheInfo(
                    has_valid_cache=True,
                    file_format="csv",
                    query_execution_id=query_execution_id,
                    query_execution_payload=query_info,
                )
        num_executions_inspected += 1
        _logger.debug("num_executions_inspected: %s", num_executions_inspected)
        if num_executions_inspected >= max_cache_query_inspections:
            return _CacheInfo(has_valid_cache=False)
    return _CacheInfo(has_valid_cache=False)


_cache_manager = _LocalMetadataCacheManager()