Source code for operators.cost_quest.cqa_get_location_id

from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from config import Config
from airflow.models import BaseOperator

MAX_LOCATIONS = 1000  # Maximum number of locations to process in a single request



class CQAGetLocationId(BaseOperator):
    ui_color = "#A6E6A6"

    @apply_defaults
    def __init__(
        self,
        query: str,
        longitude_col_name: str,
        latitude_col_name: str,
        verbose: bool = True,
        *args,
        **kwargs,
    ):
        """
        An operator to get the location id from CostQuest API from longitude and latitude

        :param query: the SQL query to get the longitude and latitude
        :type query: str
        :param longitude_col_name: the name of the column containing the longitude
        :type longitude_col_name: str
        :param latitude_col_name: the name of the column containing the latitude
        :type latitude_col_name: str
        :param verbose: whether to print debug information, defaults to True
        :type verbose: bool, optional
        """

        self.query = query
        self.longitude_col_name = longitude_col_name
        self.latitude_col_name = latitude_col_name
        self.verbose = verbose

        super().__init__(*args, **kwargs)



    def get_api_token(self) -> str:
        from airflow.hooks.base import BaseHook

        return BaseHook.get_connection(Config.CQA_API_KEY).password

    


    def p(self, message: str) -> None:
        """
        Print a message if verbose is True
        """
        if self.verbose:
            print(message)

    


    def store_longitudes_latitudes(self) -> None:
        import awswrangler as wr

        self.p(f"Running SQL to get longitudes and latitudes:\n{self.query}")

        df = wr.athena.read_sql_query(sql=self.query, database=Config.DATA_LAKE_DB)

        if self.longitude_col_name not in df.columns or self.latitude_col_name not in df.columns:
            raise ValueError(
                f"Longitude and latitude columns must be in the query result. "
                f"Found columns: {df.columns.tolist()}"
            )
    
        if df[self.longitude_col_name].isnull().any() or df[self.latitude_col_name].isnull().any():
            raise ValueError(
                "Longitude and latitude columns must not contain null or NaN values. "
                "Please check your query and ensure it returns valid data."
            )
        
        self.longitudes = df[self.longitude_col_name].astype(float).tolist()
        self.latitudes = df[self.latitude_col_name].astype(float).tolist()

        if not self.longitudes or not self.latitudes:
            raise ValueError(
                "No longitudes or latitudes found in the query result. "
                "Please check your query and ensure it returns valid data."
            )
        
        if len(self.longitudes) != len(self.latitudes):
            raise ValueError(
                "The number of longitudes and latitudes must match. "
                f"Found {len(self.longitudes)} longitudes and {len(self.latitudes)} latitudes."
            )

    


    def fetch_location_id(self, longitudes: list[float], latitudes: list[float], token: str, retries=3) -> int:
        import requests
        import time

        url = 'https://api.costquest.com/fabricext/202406/locate'
        headers = {
            "apikey": token,
            'Content-Type': "application/json",
            'User-Agent': "insomnia/11.1.0"
        }

        payload = [
            {"longitude": lon, "latitude": lat, "sourcekey": i}
            for i, (lon, lat) in enumerate(zip(longitudes, latitudes))
        ]

        for attempt in range(retries):
            response = requests.post(url, headers=headers, json=payload)
            if response.status_code == 200:
                data = response.json()
                return [item['location_id'] for item in data]
            elif response.status_code == 429:
                self.p(f"Rate limit exceeded for {url}. Waiting for retry...")
                time.sleep(60)
            else:
                self.p(f"Error: {response.status_code} - {response.text}")
                break

        self.p(f"Max retries exceeded for fetching location ids from CQA.")
        return []




    def execute(self, context) -> list[int]:
        self.store_longitudes_latitudes()
        token = self.get_api_token()

        # Split longitudes and latitudes into chunks of MAX_LOCATIONS
        chunks = [
            (self.longitudes[i:i + MAX_LOCATIONS], self.latitudes[i:i + MAX_LOCATIONS])
            for i in range(0, len(self.longitudes), MAX_LOCATIONS)
        ]

        location_ids = []
        for i, (longitudes_chunk, latitudes_chunk) in enumerate(chunks):
            self.p(f"Fetching location ids for chunk {i + 1}/{len(chunks)} with {len(longitudes_chunk)} locations.")
            location_ids.extend(self.fetch_location_id(longitudes_chunk, latitudes_chunk, token))
        
        self.p(f"Fetched {len(location_ids)} location ids from CostQuest API.")
        if len(location_ids) != len(self.longitudes):
            raise ValueError(
                f"Expected {len(self.longitudes)} location ids, but got {len(location_ids)}. "
                "Please check the API response and ensure all locations were processed."
            )
        
        return location_ids