Source code for operators.google_geocode

"""This is an operator to get correct latitude and longitude for each address in the input file."""

import requests
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults




class GooglePlaceFindPlaceOperator(BaseOperator):
    """
    Airflow Operator to search for a place using the Google Place API's findplacefromtext endpoint.

    :param input_text: The text query for the place search.
    :type input_text: str
    :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: A dictionary with the input text and the candidate place.
    """

    @apply_defaults
    def __init__(
        self,
        input_text: str = None,
        api_key: str = None,
        language: str = "en",
        fields: str = "formatted_address,name,geometry",
        *args,
        **kwargs
    ):
        super().__init__(*args, **kwargs)
        self.input_text = input_text
        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"
        input_text = self.input_text
        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)
        response = requests.get(endpoint, params=params)
        response.raise_for_status()

        data = response.json()
        if data.get("status") != "OK":
            error_msg = data.get("error_message", "Unknown error")
            self.log.error("Place API returned an error: %s", error_msg)
            raise Exception("Error from Place API: " + error_msg)

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

        candidate = candidates[0]  # only return the first candidate
        self.log.info("Found candidate: %s", candidate)

        # Optionally, push the candidate info to XCom for downstream tasks.
        return {"input_text": input_text, "candidate": candidate}