db_utils.py

import logging
import os
import re
from typing import Iterator, List, Dict, Optional

from geoalchemy2 import WKTElement
from sqlalchemy import or_
from sqlalchemy import select
from sqlalchemy.orm import joinedload, Session, contains_eager, load_only, selectinload
from sqlalchemy.orm.query import Query
from sqlalchemy.orm.strategy_options import _AbstractLoad
from sqlalchemy import func
from sqlalchemy.sql import and_
from shared.database_gen.sqlacodegen_models import (
    Feed,
    Gtfsdataset,
    Gtfsfeed,
    Location,
    Validationreport,
    Gtfsrealtimefeed,
    Entitytype,
    Redirectingid,
    Feedosmlocationgroup,
    Geopolygon,
    Gbfsfeed,
    Gbfsversion,
    Gbfsvalidationreport,
)
from shared.feed_filters.gtfs_feed_filter import GtfsFeedFilter, LocationFilter
from shared.feed_filters.gtfs_rt_feed_filter import GtfsRtFeedFilter, EntityTypeFilter
from .entity_type_enum import EntityType
from .error_handling import raise_internal_http_validation_error, invalid_bounding_coordinates, invalid_bounding_method
from .iter_utils import batched
from shared.feed_filters.gbfs_feed_filter import GbfsFeedFilter, GbfsVersionFilter


def get_gtfs_feeds_query(
    db_session: Session,
    stable_id: str | None = None,
    limit: int | None = None,
    offset: int | None = None,
    provider: str | None = None,
    producer_url: str | None = None,
    country_code: str | None = None,
    subdivision_name: str | None = None,
    municipality: str | None = None,
    dataset_latitudes: str | None = None,
    dataset_longitudes: str | None = None,
    bounding_filter_method: str | None = None,
    is_official: bool | None = None,
    published_only: bool = True,
    include_options_for_joinedload: bool = True,
) -> Query:
    """Get the DB query to use to retrieve the GTFS feeds.."""
    gtfs_feed_filter = GtfsFeedFilter(
        stable_id=stable_id,
        provider__ilike=provider,
        producer_url__ilike=producer_url,
        location=None,
    )

    subquery = gtfs_feed_filter.filter(select(Gtfsfeed.id))
    subquery = apply_bounding_filtering(
        subquery, dataset_latitudes, dataset_longitudes, bounding_filter_method
    ).subquery()
    feed_query = db_session.query(Gtfsfeed).filter(Gtfsfeed.id.in_(subquery))

    if country_code or subdivision_name or municipality:
        location_filter = LocationFilter(
            country_code=country_code,
            subdivision_name__ilike=subdivision_name,
            municipality__ilike=municipality,
        )
        location_subquery = location_filter.filter(select(Location.id))
        feed_query = feed_query.filter(Gtfsfeed.locations.any(Location.id.in_(location_subquery)))

    if published_only:
        feed_query = feed_query.filter(Gtfsfeed.operational_status == "published")

    feed_query = add_official_filter(feed_query, is_official)

    if include_options_for_joinedload:
        feed_query = feed_query.options(
            # Use selectinload for all collection relationships to avoid a cartesian-product row
            # explosion when multiple one-to-many associations are loaded simultaneously.
            # joinedload on collections multiplies rows (N feeds × M locations × F features …);
            # selectinload issues a separate IN-query per relationship, keeping rows at N per query.
            selectinload(Gtfsfeed.latest_dataset)
            .selectinload(Gtfsdataset.validation_reports)
            .selectinload(Validationreport.features),
            joinedload(Gtfsfeed.visualization_dataset),  # scalar (many-to-one) — joinedload is safe
            *get_selectinload_options(),
        ).order_by(Gtfsfeed.provider, Gtfsfeed.stable_id)

    feed_query = feed_query.limit(limit).offset(offset)
    return feed_query


def apply_most_common_location_filter(query: Query, db_session: Session) -> Query:
    """
    Apply the most common location filter to the query.
    :param query: The query to apply the filter to.
    :param db_session: The database session.

    :return: The query with the most common location filter applied.
    """
    most_common_location_subquery = (
        db_session.query(
            Feedosmlocationgroup.feed_id, func.max(Feedosmlocationgroup.stops_count).label("max_stops_count")
        )
        .group_by(Feedosmlocationgroup.feed_id)
        .subquery()
    )
    return query.outerjoin(Feed.feedosmlocationgroups).filter(
        Feedosmlocationgroup.stops_count == most_common_location_subquery.c.max_stops_count,
        Feedosmlocationgroup.feed_id == most_common_location_subquery.c.feed_id,
    )


def get_geopolygons(db_session: Session, feeds: List[Feed], include_geometry: bool = False) -> Dict[str, Geopolygon]:
    """
    Get the geolocations for the given feeds.
    :param db_session: The database session.
    :param feeds: The feeds to get the geolocations for.
    :param include_geometry: Whether to include the geometry in the result.

    :return: The geolocations for the given location groups.
    """
    location_groups = [feed.feedosmlocationgroups for feed in feeds]
    location_groups = [item for sublist in location_groups for item in sublist]

    if not location_groups:
        return dict()
    geo_polygons_osm_ids = []
    for location_group in location_groups:
        split_ids = location_group.group_id.split(".")
        if not split_ids:
            continue
        geo_polygons_osm_ids += [int(split_id) for split_id in split_ids if split_id.isdigit()]
    if not geo_polygons_osm_ids:
        return dict()
    geo_polygons_osm_ids = list(set(geo_polygons_osm_ids))
    query = db_session.query(Geopolygon).filter(Geopolygon.osm_id.in_(geo_polygons_osm_ids))
    if not include_geometry:
        query = query.options(
            load_only(Geopolygon.osm_id, Geopolygon.name, Geopolygon.iso_3166_2_code, Geopolygon.iso_3166_1_code)
        )
    query = query.order_by(Geopolygon.admin_level)
    geopolygons = query.all()
    geopolygon_map = {str(geopolygon.osm_id): geopolygon for geopolygon in geopolygons}
    return geopolygon_map


def get_all_gtfs_feeds(
    db_session: Session,
    published_only: bool = True,
    w_extracted_locations_only: bool = False,
) -> Iterator[Gtfsfeed]:
    """
    Fetch all GTFS feeds.

    :param db_session: The database session.
    :param published_only: Include only the published feeds.
    :param w_extracted_locations_only: Whether to include only feeds with extracted locations.

    :return: The GTFS feeds in an iterator.
    """
    batch_size = int(os.getenv("BATCH_SIZE", "50"))

    # We fetch in small batches and stream results to avoid loading the whole table in memory.
    # stream_results=True lets SQLAlchemy iterate rows without buffering them all at once.
    # We also clear the session cache between batches (see expunge_all() below) to prevent
    # memory from growing indefinitely when many ORM objects are loaded.
    batch_query = db_session.query(Gtfsfeed).order_by(Gtfsfeed.stable_id).execution_options(stream_results=True)
    if published_only:
        batch_query = batch_query.filter(Gtfsfeed.operational_status == "published")

    processed = 0

    for batch_num, batch in enumerate(batched(batch_query, batch_size), start=1):
        start_index = processed + 1
        end_index = processed + len(batch)
        logging.info("Processing feeds %d - %d", start_index, end_index)

        # Convert to a list intentionally: we want to "materialize" IDs now to make any cost
        # visible here (and keep the logic simple). This also avoids subtle lazy-evaluation
        # effects that can hide where time/memory is really spent.
        stable_ids = [f.stable_id for f in batch]
        if not stable_ids:
            processed += len(batch)
            continue

        if w_extracted_locations_only:
            feed_query = apply_most_common_location_filter(db_session.query(Gtfsfeed), db_session)
            inner_q = feed_query.filter(Gtfsfeed.stable_id.in_(stable_ids)).options(
                # See note above: selectinload is chosen for collections to keep memory and row
                # counts under control when streaming.
                selectinload(Gtfsfeed.latest_dataset)
                .selectinload(Gtfsdataset.validation_reports)
                .selectinload(Validationreport.features),
                selectinload(Gtfsfeed.bounding_box_dataset),
                *get_selectinload_options(include_extracted_location_entities=True),
            )
        else:
            inner_q = (
                db_session.query(Gtfsfeed)
                .outerjoin(Gtfsfeed.gtfsdatasets)
                .filter(Gtfsfeed.stable_id.in_(stable_ids))
                .options(
                    selectinload(Gtfsfeed.latest_dataset)
                    .selectinload(Gtfsdataset.validation_reports)
                    .selectinload(Validationreport.features),
                    selectinload(Gtfsfeed.bounding_box_dataset),
                    *get_selectinload_options(include_extracted_location_entities=False),
                )
            )

        # Iterate and stream rows out; the options above ensure related data is preloaded in
        # a few small queries per batch, rather than one giant join.
        for item in inner_q.execution_options(stream_results=True):
            yield item

        # Clear the Session identity map so objects from this batch can be GC'd. Without this,
        # the Session will keep references and memory usage will grow with each batch.
        try:
            db_session.expunge_all()
        except Exception:
            logging.getLogger("get_all_gtfs_feeds").exception("Failed to expunge session after batch %d", batch_num)

        processed += len(batch)


def get_gtfs_rt_feeds_query(
    limit: int | None,
    offset: int | None,
    provider: str | None,
    producer_url: str | None,
    entity_types: str | None,
    country_code: str | None,
    subdivision_name: str | None,
    municipality: str | None,
    is_official: bool | None,
    published_only: bool = True,
    db_session: Session = None,
) -> Query:
    """Get some (or all) GTFS Realtime feeds from the Mobility Database."""
    entity_types_list = entity_types.split(",") if entity_types else None

    # Validate entity types using the EntityType enum
    if entity_types_list:
        try:
            entity_types_list = [EntityType(et.strip()).value for et in entity_types_list]
        except ValueError:
            raise_internal_http_validation_error(
                "Entity types must be the value 'vp', 'sa', or 'tu'. "
                "When provided a list values must be separated by commas."
            )

    gtfs_rt_feed_filter = GtfsRtFeedFilter(
        stable_id=None,
        provider__ilike=provider,
        producer_url__ilike=producer_url,
        entity_types=EntityTypeFilter(name__in=entity_types_list),
        location=LocationFilter(
            country_code=country_code,
            subdivision_name__ilike=subdivision_name,
            municipality__ilike=municipality,
        ),
    )
    subquery = gtfs_rt_feed_filter.filter(
        select(Gtfsrealtimefeed.id)
        .join(Location, Gtfsrealtimefeed.locations, isouter=True)
        .join(Entitytype, Gtfsrealtimefeed.entitytypes)
    ).subquery()
    feed_query = db_session.query(Gtfsrealtimefeed).filter(Gtfsrealtimefeed.id.in_(subquery))

    if published_only:
        feed_query = feed_query.filter(Gtfsrealtimefeed.operational_status == "published")

    feed_query = feed_query.options(
        selectinload(Gtfsrealtimefeed.entitytypes),
        selectinload(Gtfsrealtimefeed.gtfs_feeds),
        *get_selectinload_options(),
    )
    feed_query = add_official_filter(feed_query, is_official)

    feed_query = feed_query.limit(limit).offset(offset)
    return feed_query


def add_official_filter(query: Query, is_official: bool | None) -> Query:
    """
    Add the is_official filter to the query if necessary
    """
    if is_official is not None:
        if is_official:
            query = query.filter(Feed.official.is_(True))
        else:
            query = query.filter(or_(Feed.official.is_(False), Feed.official.is_(None)))
    return query


def get_all_gtfs_rt_feeds(
    db_session: Session,
    published_only: bool = True,
    batch_size: int = 250,
    w_extracted_locations_only: bool = False,
) -> Iterator[Gtfsrealtimefeed]:
    """
    Fetch all GTFS realtime feeds.

    :param db_session: The database session.
    :param published_only: Include only the published feeds.
    :param batch_size: The number of feeds to fetch from the database at a time.
        A lower value means less memory but more queries.
    :param w_extracted_locations_only: Whether to include only feeds with extracted locations.

    :return: The GTFS realtime feeds in an iterator.
    """
    batched_query = (
        db_session.query(Gtfsrealtimefeed.stable_id)
        .order_by(Gtfsrealtimefeed.stable_id)
        .yield_per(batch_size)
        .execution_options(stream_results=True)
    )
    if published_only:
        batched_query = batched_query.filter(Gtfsrealtimefeed.operational_status == "published")

    for batch in batched(batched_query, batch_size):
        stable_ids = (f.stable_id for f in batch)
        if w_extracted_locations_only:
            feed_query = apply_most_common_location_filter(db_session.query(Gtfsrealtimefeed), db_session)
            yield from (
                feed_query.filter(Gtfsrealtimefeed.stable_id.in_(stable_ids))
                .options(
                    selectinload(Gtfsrealtimefeed.entitytypes),
                    selectinload(Gtfsrealtimefeed.gtfs_feeds),
                    *get_joinedload_options(include_extracted_location_entities=True),
                )
                .order_by(Gtfsfeed.stable_id)
            )
        else:
            yield from (
                db_session.query(Gtfsrealtimefeed)
                .filter(Gtfsrealtimefeed.stable_id.in_(stable_ids))
                .options(
                    selectinload(Gtfsrealtimefeed.entitytypes),
                    selectinload(Gtfsrealtimefeed.gtfs_feeds),
                    *get_selectinload_options(include_extracted_location_entities=False),
                )
            )


def apply_bounding_filtering(
    query: Query,
    bounding_latitudes: str,
    bounding_longitudes: str,
    bounding_filter_method: str,
) -> Query:
    """Create a new query based on the bounding parameters."""

    if not bounding_latitudes or not bounding_longitudes or not bounding_filter_method:
        return query

    # Parse tokens explicitly to satisfy static analyzers and keep error messages clear.
    bounding_latitudes_tokens = bounding_latitudes.split(",")
    bounding_longitudes_tokens = bounding_longitudes.split(",")
    if len(bounding_latitudes_tokens) != 2 or len(bounding_longitudes_tokens) != 2:
        raise_internal_http_validation_error(
            invalid_bounding_coordinates.format(bounding_latitudes, bounding_longitudes)
        )
    min_latitude, max_latitude = bounding_latitudes_tokens
    min_longitude, max_longitude = bounding_longitudes_tokens
    try:
        min_latitude = float(min_latitude)
        max_latitude = float(max_latitude)
        min_longitude = float(min_longitude)
        max_longitude = float(max_longitude)
    except ValueError:
        raise_internal_http_validation_error(
            invalid_bounding_coordinates.format(bounding_latitudes, bounding_longitudes)
        )
    points = [
        (min_longitude, min_latitude),
        (min_longitude, max_latitude),
        (max_longitude, max_latitude),
        (max_longitude, min_latitude),
        (min_longitude, min_latitude),
    ]
    wkt_polygon = f"POLYGON(({', '.join(f'{lon} {lat}' for lon, lat in points)}))"
    bounding_box = WKTElement(
        wkt_polygon,
        srid=Gtfsdataset.bounding_box.type.srid,
    )
    query = query.join(Gtfsdataset, Gtfsdataset.feed_id == Gtfsfeed.id)

    if bounding_filter_method == "partially_enclosed":
        return query.filter(
            or_(
                Gtfsdataset.bounding_box.ST_Overlaps(bounding_box),
                Gtfsdataset.bounding_box.ST_Contains(bounding_box),
            )
        )
    elif bounding_filter_method == "completely_enclosed":
        return query.filter(bounding_box.ST_Covers(Gtfsdataset.bounding_box))
    elif bounding_filter_method == "disjoint":
        return query.filter(Gtfsdataset.bounding_box.ST_Disjoint(bounding_box))
    else:
        raise_internal_http_validation_error(invalid_bounding_method.format(bounding_filter_method))


def get_joinedload_options(include_extracted_location_entities: bool = False) -> [_AbstractLoad]:
    """
    Returns common joinedload options for feeds queries.
    :param include_extracted_location_entities: Whether to include extracted location entities.

    :return: A list of joinedload options.
    """
    joinedload_options = []
    if include_extracted_location_entities:
        joinedload_options = [contains_eager(Feed.feedosmlocationgroups).joinedload(Feedosmlocationgroup.group)]
    return joinedload_options + [
        joinedload(Feed.locations),
        joinedload(Feed.externalids),
        joinedload(Feed.feedrelatedlinks),
        joinedload(Feed.redirectingids).joinedload(Redirectingid.target),
        joinedload(Feed.officialstatushistories),
    ]


def get_selectinload_options(include_extracted_location_entities: bool = False) -> List[_AbstractLoad]:
    """
    Returns common joinedload options for feeds queries.
    :param include_extracted_location_entities: Whether to include extracted location entities.

    :return: A list of joinedload options.
    """
    # NOTE: For collections we prefer selectinload to avoid row explosion and high memory usage
    # during streaming. When callers explicitly join some paths (e.g., most common locations),
    # we use contains_eager on that specific path to tell SQLAlchemy the data came from a JOIN.
    loaders = []
    if include_extracted_location_entities:
        loaders.append(contains_eager(Feed.feedosmlocationgroups).joinedload(Feedosmlocationgroup.group))

    # collections -> selectinload; scalar relationships can remain joinedload
    loaders.extend(
        [
            selectinload(Feed.locations),
            selectinload(Feed.externalids),
            selectinload(Feed.feedrelatedlinks),
            selectinload(Feed.redirectingids).selectinload(Redirectingid.target),
            selectinload(Feed.officialstatushistories),
        ]
    )
    return loaders


def get_gbfs_feeds_query(
    db_session: Session,
    stable_id: Optional[str] = None,
    provider: Optional[str] = None,
    producer_url: Optional[str] = None,
    country_code: Optional[str] = None,
    subdivision_name: Optional[str] = None,
    municipality: Optional[str] = None,
    system_id: Optional[str] = None,
    version: Optional[str] = None,
) -> Query:
    gbfs_feed_filter = GbfsFeedFilter(
        stable_id=stable_id,
        provider__ilike=provider,
        producer_url__ilike=producer_url,
        system_id=system_id,
        location=LocationFilter(
            country_code=country_code,
            subdivision_name__ilike=subdivision_name,
            municipality__ilike=municipality,
        )
        if country_code or subdivision_name or municipality
        else None,
        version=GbfsVersionFilter(
            version=version,
        )
        if version
        else None,
    )
    # Subquery: latest report per version
    latest_report_subq = (
        db_session.query(
            Gbfsvalidationreport.gbfs_version_id.label("gbfs_version_id"),
            func.max(Gbfsvalidationreport.validated_at).label("latest_validated_at"),
        )
        .group_by(Gbfsvalidationreport.gbfs_version_id)
        .subquery()
    )

    # Join validation reports filtered by latest `validated_at`
    query = gbfs_feed_filter.filter(
        db_session.query(Gbfsfeed)
        .outerjoin(Location, Gbfsfeed.locations)
        .outerjoin(Gbfsfeed.gbfsversions)
        .outerjoin(latest_report_subq, Gbfsversion.id == latest_report_subq.c.gbfs_version_id)
        .outerjoin(
            Gbfsvalidationreport,
            and_(
                Gbfsversion.id == Gbfsvalidationreport.gbfs_version_id,
                Gbfsvalidationreport.validated_at == latest_report_subq.c.latest_validated_at,
            ),
        )
        .options(
            contains_eager(Gbfsfeed.gbfsversions).contains_eager(Gbfsversion.gbfsvalidationreports),
            contains_eager(Gbfsfeed.gbfsversions).joinedload(Gbfsversion.gbfsendpoints),
            joinedload(Feed.locations),
            joinedload(Feed.externalids),
            joinedload(Feed.redirectingids).joinedload(Redirectingid.target),
        )
    )
    return query


def normalize_url(url_column) -> str:
    """
    Normalize a URL by removing the protocol (http:// or https://), 'www.' prefix, and trailing slash.
    This function generates a SQLAlchemy expression that can be used in queries.
    Args:
        url_column: The SQLAlchemy column representing the URL.
    Returns:
        A SQLAlchemy expression that normalizes the URL.
    """
    return func.regexp_replace(
        func.regexp_replace(
            func.regexp_replace(url_column, r"^https?://", "", "gi"),
            r"^www\.",
            "",
            "gi",
        ),
        r"/$",
        "",
        "g",
    )


def normalize_url_str(url: str | None) -> str:
    """Normalize a license URL for matching.
    Steps:
    - Trim whitespace and quotes
    - Remove BOM characters
    - Strip fragments and query parameters
    - Remove scheme (http/https) and www prefix
    - Lowercase the host
    """
    u = (url or "").strip().strip("'\"").replace("\ufeff", "")
    u = re.sub(r"#.*$", "", u)
    u = re.sub(r"\?.*$", "", u)
    u = re.sub(r"^https?://", "", u, flags=re.I)
    u = re.sub(r"^www\.", "", u, flags=re.I)
    # remove trailing slashes
    u = re.sub(r"/+$", "", u)
    return u.lower()


def get_feed_query_by_normalized_url(url: str, db_session: Session) -> Query:
    """
    Get a query to find the feed by normalized URL and exclude deprecated feeds.
    Args:
        url: The URL to normalize and search for.
        db_session: SQLAlchemy session.
    """
    return db_session.query(Feed).filter(
        normalize_url_str(url) == func.lower(func.trim(normalize_url(Feed.producer_url))),
        Feed.status != "deprecated",
    )


def get_feed_by_normalized_url(url: str, db_session: Session) -> Feed | None:
    """
    Query the feed by normalized URL and exclude deprecated feeds.
    Args:
        url: The URL to normalize and search for.
        db_session: SQLAlchemy session.
    """
    return get_feed_query_by_normalized_url(url, db_session).first()