Source code for operators.google_place.place_detail

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 PlaceDetailsBatchOperator(BaseOperator):
    """
    Given a list of place ids, fetch selected details for each:
    name, lat, lon, address, url, phone.
    """

    template_fields = ("place_ids", "batch_size")

    def __init__(
        self,
        *,
        api_key: str,
        place_ids: (
            list[str] | None
        ) = None,  # If None, will pull from upstream XCom (task_id must be set in op_kwargs at DAG use)
        batch_size: int = 20,
        per_call_sleep: float = 0.0,
        timeout: int = 30,
        session: requests.Session | None = None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.api_key = api_key
        self.place_ids = place_ids
        self.batch_size = batch_size
        self.per_call_sleep = per_call_sleep
        self.timeout = timeout
        self._session = session or requests.Session()

    def _fetch_one(self, pid: str) -> dict:
        headers = {
            "X-Goog-Api-Key": self.api_key,
            # Request exactly what we need
            "X-Goog-FieldMask": (
                "id,displayName,formattedAddress,location,googleMapsUri,internationalPhoneNumber"
            ),
        }
        url = f"{PLACES_BASE}/places/{pid}"
        r = self._session.get(url, headers=headers, timeout=self.timeout)
        r.raise_for_status()
        js = r.json()

        name = (js.get("displayName") or {}).get("text")
        loc = js.get("location") or {}
        return {
            "place_id": js.get("id"),
            "name": name,
            "lat": loc.get("latitude"),
            "lon": loc.get("longitude"),
            "address": js.get("formattedAddress"),
            "url": js.get("googleMapsUri"),
            "phone": js.get("internationalPhoneNumber"),
        }



    def execute(self, context: Context) -> list[dict]:
        # Allow pulling IDs from upstream via XCom if not provided at init
        place_ids = self.place_ids
        if place_ids is None:
            # Expect upstream task_id in op_kwargs at DAG usage (see example)
            raise ValueError(
                "place_ids is None. Pass a list at init or pull from XCom in DAG code."
            )

        results: list[dict] = []
        for i, pid in enumerate(place_ids, start=1):
            try:
                results.append(self._fetch_one(pid))
            except requests.HTTPError as e:
                self.log.error("Failed to fetch details for %s: %s", pid, e)
            if self.per_call_sleep:
                time.sleep(self.per_call_sleep)
            if (i % self.batch_size) == 0:
                self.log.info("Fetched %d/%d places...", i, len(place_ids))

        self.log.info(
            "Fetched details for %d places (requested %d).",
            len(results),
            len(place_ids),
        )
        return results