Source code for operators.google_place.text_search

# places_operators.py
import time
import typing as T
import requests

from airflow.models import BaseOperator
from airflow.utils.context import Context

PLACES_BASE = "https://places.googleapis.com/v1"




class PlacesTextSearchOperator(BaseOperator):
    """
    Uses Places API v1 Text Search (new) to return a list of place ids.
    """

    template_fields = ("text_query", "max_pages", "page_pause_seconds")

    def __init__(
        self,
        *,
        api_key: str,
        text_query: str,
        max_pages: int = 5,
        page_pause_seconds: float = 3.0,
        timeout: int = 30,
        session: requests.Session | None = None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.api_key = api_key
        self.text_query = text_query
        self.max_pages = max_pages
        self.page_pause_seconds = page_pause_seconds
        self.timeout = timeout
        self._session = session or requests.Session()



    def execute(self, context: Context) -> list[str]:
        headers = {
            "Content-Type": "application/json",
            "X-Goog-Api-Key": self.api_key,
            "X-Goog-FieldMask": "places.id,nextPageToken",
        }
        url = f"{PLACES_BASE}/places:searchText"

        # build the base payload ONCE (if you add filters later, put them here too)
        base_payload: dict[str, T.Any] = {"textQuery": self.text_query}

        def post_json(payload):
            # small retry with backoff; re-raise the last HTTPError cleanly
            delay = max(self.page_pause_seconds, 2.0)
            last_err = None
            for attempt in range(5):
                try:
                    r = self._session.post(
                        url, json=payload, headers=headers, timeout=self.timeout
                    )
                    r.raise_for_status()
                    return r.json()
                except requests.HTTPError as e:
                    last_err = e
                    body = getattr(e.response, "text", "")
                    self.log.warning(
                        "POST failed (attempt %d). Status=%s Body=%s",
                        attempt + 1,
                        getattr(e.response, "status_code", "?"),
                        body[:500],
                    )
                    time.sleep(delay)
                    delay = min(delay * 1.8, 8.0)
            # after retries, raise the last seen error
            raise last_err

        place_ids: list[str] = []
        seen_ids: set[str] = set()
        seen_tokens: set[str] = set()

        page = 0
        payload = dict(base_payload)  # page 1: textQuery only
        while True:
            self.log.info("Now processing page: %d", page + 1)
            body = post_json(payload)

            for p in body.get("places", []):
                pid = p.get("id")
                if pid and pid not in seen_ids:
                    seen_ids.add(pid)
                    place_ids.append(pid)

            tok = body.get("nextPageToken")
            page += 1
            if not tok or page >= self.max_pages:
                break

            if tok in seen_tokens:
                self.log.warning("Duplicate page token seen; stopping pagination.")
                break
            seen_tokens.add(tok)

            # let token become valid, then request NEXT page
            time.sleep(max(self.page_pause_seconds, 3.0))

            # IMPORTANT: subsequent pages must include the SAME params as page 1 + pageToken
            payload = {**base_payload, "pageToken": tok}

        self.log.info(
            "Collected %d unique place ids across %d page(s).", len(place_ids), page
        )
        return place_ids