Source code for operators.google_geocode_batch

"""This operator performs batch geocoding using the Google Place API's findplacefromtext endpoint."""

import requests
import pandas as pd
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults




class BatchGooglePlaceFindPlaceOperator(BaseOperator):
    """
    Airflow Operator to perform batch geocoding using the Google Place API's findplacefromtext endpoint.

    This operator expects a pandas DataFrame with the columns "place_name" and "address".
    For each row, it will query the API by combining these two columns and append the following
    columns to the DataFrame:
        - formatted_address
        - latitude
        - longitude

    :param input_df: The input pandas DataFrame with columns "place_name" and "address".
    :type input_df: pandas.DataFrame
    :param api_key: Your Google API key.
    :type api_key: str
    :param language: (optional) The language for the request. Default is 'en'.
    :type language: str
    :param fields: (optional) A comma-separated list of fields to return (e.g., "formatted_address,name,geometry").
    :type fields: str
    :return: The input DataFrame with additional columns for formatted_address, latitude, and longitude.
    """

    @apply_defaults
    def __init__(
        self,
        input_df: pd.DataFrame,
        api_key: str,
        language: str = "en",
        fields: str = "formatted_address,name,geometry",
        *args,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.input_df = input_df
        self.api_key = api_key
        self.language = language
        self.fields = fields



    def execute(self, context):
        endpoint = "https://maps.googleapis.com/maps/api/place/findplacefromtext/json"

        # Create a copy of the dataframe to avoid modifying the original
        df = self.input_df.copy()

        # Initialize new columns with default values
        df["formatted_address"] = None
        df["latitude"] = None
        df["longitude"] = None

        for idx, row in df.iterrows():
            # Combine the place_name and address for a more robust query
            input_text = f"{row['place_name']}, {row['address']}"
            params = {
                "input": input_text,
                "inputtype": "textquery",
                "language": self.language,
                "key": self.api_key,
            }
            if self.fields:
                params["fields"] = self.fields

            self.log.info("Requesting place search for input: %s", input_text)
            try:
                response = requests.get(endpoint, params=params)
                response.raise_for_status()
            except Exception as e:
                self.log.error("Request failed for input '%s': %s", input_text, str(e))
                continue

            data = response.json()
            if data.get("status") != "OK":
                error_msg = data.get("error_message", "Unknown error")
                self.log.error(
                    "Place API error for input '%s': %s", input_text, error_msg
                )
                continue

            candidates = data.get("candidates", [])
            if not candidates:
                self.log.warning("No candidates found for input: %s", input_text)
                continue

            candidate = candidates[0]  # Use the first candidate returned
            self.log.info("Found candidate for input '%s': %s", input_text, candidate)

            # Extract required information
            formatted_address = candidate.get("formatted_address")
            geometry = candidate.get("geometry", {})
            location = geometry.get("location", {})
            latitude = location.get("lat")
            longitude = location.get("lng")

            # Update the DataFrame
            df.at[idx, "formatted_address"] = formatted_address
            df.at[idx, "latitude"] = latitude
            df.at[idx, "longitude"] = longitude

        # Optionally, the result can be pushed via XCom
        return df