bb-integrations-library 3.0.11__py3-none-any.whl

bb_integrations_lib/pipelines/steps/exporting/bbd_export_readings_step.py

@@ -0,0 +1,236 @@
+import asyncio
+from datetime import datetime, UTC, timedelta
+from typing import Dict, Any, Optional, Tuple, Iterable
+from zoneinfo import ZoneInfo
+
+import pandas as pd
+from bson import RawBSONDocument
+from loguru import logger
+from pandas import DataFrame
+from pymongo import MongoClient
+from pymongo.synchronous.database import Database, Collection
+
+from bb_integrations_lib.gravitate.rita_api import GravitateRitaAPI
+from bb_integrations_lib.gravitate.sd_api import GravitateSDAPI
+from bb_integrations_lib.protocols.pipelines import Step
+from bb_integrations_lib.shared.model import ReadingQuery, ExportReadingsWindowMode
+
+
+class BBDExportReadingsStep(Step):
+    def __init__(self, rita_client: GravitateRitaAPI, sd_client: GravitateSDAPI, ims_database: Database,
+                 reading_query: ReadingQuery, timezone: str = "UTC",
+                 window_mode: ExportReadingsWindowMode = ExportReadingsWindowMode.HOURS_BACK, hours_back: float = 6,
+                 batch_hours: float = 2, big_dataset: int = 750, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.rita_client = rita_client
+        self.sd_client = sd_client
+        self.ims_database = ims_database
+
+        self.reading_query = reading_query
+        self.timezone = timezone
+        self.window_mode = window_mode
+        self.hours_back = hours_back
+        self.batch_hours = batch_hours
+        self.big_dataset = big_dataset
+
+    def describe(self) -> str:
+        return "Export tank readings from BBD with a customizable query"
+
+    async def execute(self, _: Any) -> Tuple[Dict, Dict, Iterable]:
+        # Compute start/end windows based on the provided configuration
+        if self.window_mode == ExportReadingsWindowMode.HOURS_BACK or self.window_mode == ExportReadingsWindowMode.LATEST_ONLY:
+            now = datetime.now(UTC)
+            window_start = now - timedelta(hours=self.hours_back)
+            window_end = now
+        elif self.window_mode == ExportReadingsWindowMode.PREVIOUS_DAY:
+            window_end = datetime.now(ZoneInfo(self.timezone)).replace(hour=0, minute=0, second=0, microsecond=0)
+            window_start = window_end - timedelta(days=1)
+        else:
+            raise Exception(f"Unknown window mode {self.window_mode}")
+
+        if self.window_mode == ExportReadingsWindowMode.HOURS_BACK:
+            logger.info(
+                f"Exporting tank readings from BBD for the last {self.hours_back} hours "
+                f"(from {window_start} to {window_end})"
+            )
+        elif self.window_mode == ExportReadingsWindowMode.LATEST_ONLY:
+            logger.info(
+                f"Exporting ONLY the most recent tank reading from BBD for the last {self.hours_back} hours "
+                f"(from {window_start} to {window_end}"
+            )
+        elif self.window_mode == ExportReadingsWindowMode.PREVIOUS_DAY:
+            logger.info(f"Exporting tank readings from BBD for the previous day (from {window_start} to {window_end})")
+
+        store_lkp, tank_lkp, filtered_stores = await self.get_sites_and_tank_ids_to_export()
+
+        if len(filtered_stores) > self.big_dataset and self.reading_query.by_wildcard is None:
+            raise ValueError(
+                f"Large dataset ({len(filtered_stores)} stores) detected. "
+                f"Must use wildcard query strategy (by_wildcard) for datasets > 500 stores."
+            )
+
+        use_or_filter = len(filtered_stores) <= self.big_dataset
+
+        return await self.get_inventory(
+            window_start=window_start,
+            window_end=window_end,
+            latest_only=self.window_mode == ExportReadingsWindowMode.LATEST_ONLY,
+            store_lkp=store_lkp,
+            tank_lkp=tank_lkp,
+            filtered_stores=filtered_stores,
+            use_or_filter=use_or_filter,
+        )
+
+    @staticmethod
+    def collection_as_raw_bson(collection: Collection):
+        return collection.with_options(
+            codec_options=collection.codec_options.with_options(document_class=RawBSONDocument))
+
+    async def get_sites_and_tank_ids_to_export(self, query: Optional[ReadingQuery] = None) -> Tuple[
+        Dict, Dict, DataFrame]:
+        query_to_use = query if query else self.reading_query
+        stores = (await self.sd_client.get_all_stores(include_tanks=True)).json()
+        filtered_stores = self.apply_reading_query_filter(stores=stores, query=query_to_use)
+        filtered_store_numbers = filtered_stores["store_number"].to_list()
+        store_lkp, tank_lkp = BBDExportReadingsStep.lkps(
+            [store for store in stores if store["store_number"] in filtered_store_numbers])
+        return store_lkp, tank_lkp, filtered_stores
+
+    def apply_reading_query_filter(self, stores: list[Dict], query: Optional[ReadingQuery] = None) -> DataFrame:
+        query_to_use = query if query else self.reading_query
+        df = pd.DataFrame(stores)
+        tanks_df = df.explode("tanks")
+        tanks_df = tanks_df[~tanks_df["tanks"].isna()].reset_index()
+        tanks_df["tank_id"] = tanks_df["tanks"].apply(lambda x: x["tank_id"])
+        tanks_df["composite_key"] = tanks_df["store_number"] + ":" + tanks_df["tank_id"].astype(str)
+        if query_to_use.by_wildcard is not None and query_to_use.by_wildcard == "*":
+            return tanks_df[["store_number", "tank_id", "composite_key"]]
+        mask = query_to_use.as_mask(tanks_df)
+        filtered = tanks_df[mask]
+        return filtered[["store_number", "tank_id", "composite_key"]]
+
+    @staticmethod
+    def ims_query_pairs(filtered_stores: DataFrame) -> list:
+        pairs = []
+        for _, row in filtered_stores.iterrows():
+            pairs.append({
+                "store_number": row["store_number"],
+                "tank_id": str(row["tank_id"])
+            })
+        return pairs
+
+    async def get_inventory(self,
+                            window_start: datetime,
+                            window_end: datetime,
+                            latest_only: bool,
+                            store_lkp: dict,
+                            tank_lkp: dict,
+                            filtered_stores: DataFrame | None = None,
+                            use_or_filter: bool = True,  # Add this parameter
+                            ) -> Tuple[Dict, Dict, Iterable]:
+        """
+        Get inventory / tank readings, filtering to the stores and tanks specified by reading_query.
+        :param window_start: Include readings newer than or equal to this datetime.
+        :param window_end: Include readings older than this datetime.
+        :param latest_only: Whether only the latest reading within the window should be provided.
+        :param store_lkp: Optional pre-fetched store lookup dict.
+        :param tank_lkp: Optional pre-fetched tank lookup dict.
+        :param filtered_stores: Optional pre-fetched filtered stores DataFrame.
+        :param use_or_filter: Whether to include $or clause in aggregation for store/tank filtering.
+        :return: A tuple containing a store lookup dict, tank lookup dict, and iterable of tank reading documents.
+        """
+        ims_pairs = BBDExportReadingsStep.ims_query_pairs(filtered_stores)
+        collection = self.ims_database['tank_inventory_log']
+
+        match_conditions: list[dict] = [
+            {
+                "read_time": {
+                    "$gte": window_start,
+                    "$lte": window_end
+                }
+            }
+        ]
+
+        if use_or_filter:
+            match_conditions.append({
+                                        "$or": ims_pairs
+                                        })
+
+        _query: list[dict] = [
+            {
+                "$match": {
+                    "$and": match_conditions
+                }
+            },
+            {
+                "$sort": {
+                    "read_time": 1
+                }
+            },
+            {
+                "$project": {
+                    "tank_agent_name": 1,
+                    "store_number": 1,
+                    "run_time": 1,
+                    "tank_id": 1,
+                    "read_time": 1,
+                    "product": 1,
+                    "monitor_type": 1,
+                    "volume": 1,
+                }
+            }
+        ]
+        # If we're getting only the most recent tank reading, aggregate by store and tank, then grab the latest
+        # reading document from the lookback period.
+        if latest_only:
+            _query.append({
+                "$group": {
+                    "_id": {
+                        "store_number": "$store_number",
+                        "tank_id": "$tank_id"
+                    },
+                    "documents": {
+                        "$last": "$$ROOT"
+                    }
+                }
+            })
+
+        cursor = collection.aggregate(_query)
+        # Have to extract the latest document from the aggregate if we're grabbing the latest, as the query result
+        # set is shaped differently.
+        if latest_only:
+            return store_lkp, tank_lkp, (r["documents"] for r in cursor)
+        return store_lkp, tank_lkp, cursor
+
+    @staticmethod
+    def lkps(stores: Iterable) -> Tuple[Dict, Dict]:
+        store_lkp = {}
+        tank_lkp = {}
+        for store in stores:
+            store_number = store['store_number']
+            tanks = store['tanks']
+            for tank in tanks:
+                tank_id = tank['tank_id']
+                store_lkp[f"{store_number}"] = store
+                tank_lkp[f"{store_number}:{tank_id}"] = tank
+        return store_lkp, tank_lkp
+
+
+if __name__ == "__main__":
+    export = BBDExportReadingsStep(
+        rita_client=GravitateRitaAPI(
+            base_url="",
+            client_id="",
+            client_secret=""
+        ),
+        sd_client=GravitateSDAPI(
+            base_url="",
+            client_id="",
+            client_secret=""
+        ),
+        ims_database=MongoClient("mongodb conn str")["db_name"],
+        reading_query=ReadingQuery(
+            by_store_numbers=["100101"]
+        )
+    )
+    readings = asyncio.run(export.execute("Majors"))

bb_integrations_lib/pipelines/steps/exporting/cargas_wholesale_bundle_upload_step.py

@@ -0,0 +1,33 @@
+from typing import Any
+
+from bb_integrations_lib.protocols.pipelines import Step, Input, StepConfig
+from loguru import logger
+
+from bb_integrations_lib.provider.api.cargas.client import CargasClient
+from bb_integrations_lib.provider.api.cargas.model import CreateWholesaleTicketRequest, \
+    CreateWholesaleTicketRequestBundle
+
+
+class CargasWholesaleBundleUploadStep(Step):
+    def __init__(self, cargas_client: CargasClient, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.cargas_client = cargas_client
+
+    def describe(self):
+        return "Upload a Cargas Wholesale ticket request bundle"
+
+    async def execute(self, i: CreateWholesaleTicketRequestBundle) -> None:
+        cwt_resp = await self.cargas_client.create_wholesale_ticket(i.ticket_request)
+        doc_id = cwt_resp["ResponseValues"]["DocumentID"]
+
+        # Now try to upload each line
+        for index, line in enumerate(i.line_requests):
+            try:
+                line.DocumentID = int(doc_id)
+                await self.cargas_client.create_wholesale_line(line)
+            except Exception as e:
+                logger.error(f"Failed to upload line {index} to Cargas: {e}")
+        print(cwt_resp)
+
+        # TODO: Upload fee line items

bb_integrations_lib/pipelines/steps/exporting/dataframe_flat_file_export.py

@@ -0,0 +1,29 @@
+from datetime import datetime, UTC
+
+from pandas import DataFrame
+
+from bb_integrations_lib.protocols.pipelines import Step
+from bb_integrations_lib.shared.model import FileReference, FileType
+
+
+class DataFrameFlatFileExportStep(Step):
+    def __init__(self, file_path: str, file_type: FileType, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        file_path = file_path
+        if "{date}" in file_path:
+            file_path = file_path.replace("{date}", datetime.now(UTC).strftime("%Y%m%d"))
+
+        self.output = FileReference(file_path, file_type)
+
+    def describe(self) -> str:
+        return f"Exporting DataFrame to flat file {self.output.file_path}"
+
+    async def execute(self, i: DataFrame) -> FileReference:
+        if self.output.file_type == FileType.excel:
+            i.to_excel(self.output.file_path, index=False)
+        elif self.output.file_type == FileType.csv:
+            i.to_csv(self.output.file_path, index=False)
+        else:
+            raise NotImplementedError(f"Unsupported file type: {self.output.file_type}")
+        return self.output

bb_integrations_lib/pipelines/steps/exporting/gcs_bucket_export_file_step.py

@@ -0,0 +1,34 @@
+import os
+from typing import Any, Dict
+
+from bb_integrations_lib.shared.model import FileReference, File
+from bb_integrations_lib.protocols.pipelines import Step
+from bb_integrations_lib.provider.gcp.cloud_storage.client import CloudStorageClient
+
+
+class GCSExportFileStep(Step[Any, Any, None]):
+    def __init__(self, step_configuration: Dict[str, str]) -> None:
+        super().__init__(step_configuration)
+        self.gcs_client = CloudStorageClient()
+        self.bucket = step_configuration['bucket']
+
+    def describe(self) -> str:
+        return "Exporting file to GCS bucket"
+
+    async def execute(self, i: FileReference) -> FileReference:
+        file_name = os.path.basename(i.file_path)
+        with open(i.file_path, "rb") as f:
+            file_data = f.read()
+        file = File(
+            file_name=file_name,
+            file_data=file_data
+        )
+        try:
+            self.gcs_client.upload_file(file, self.bucket)
+            return i
+        except FileExistsError:
+            # If run twice and the file was already archived, we don't need to archive another copy.
+            return i
+        except Exception as e:
+            raise e
+