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

bb_integrations_lib/pipelines/steps/create_accessorials_step.py

@@ -0,0 +1,80 @@
+import json
+from time import sleep
+from typing import Dict, List, cast
+
+from loguru import logger
+
+from bb_integrations_lib.gravitate.sd_api import GravitateSDAPI
+from bb_integrations_lib.models.pipeline_structs import BBDUploadResult
+from bb_integrations_lib.protocols.pipelines import Step
+from bb_integrations_lib.util.utils import CustomJSONEncoder
+from bb_integrations_lib.util.config.manager import GlobalConfigManager
+from bb_integrations_lib.util.config.model import GlobalConfig
+
+
+class BBDUploadAccessorialsStep(Step[List, BBDUploadResult, None]):
+
+    def __init__(self, step_configuration: Dict[str, str]):
+        super().__init__(step_configuration)
+        self.config_manager = GlobalConfigManager()
+        self.env_mode = step_configuration.get('mode', "production")
+        self.buffer = step_configuration.get('buffer', 0.5)
+        self.chuk_size = step_configuration.get('chuk_size', 1000)
+        if 'env' in step_configuration:
+            self.bbd_client = GravitateSDAPI.from_config(step_configuration["env"],
+                                                         step_configuration["bbd_username"],
+                                                         step_configuration["bbd_password"])
+            self.bbd_client.username = step_configuration["bbd_username"]
+            self.bbd_client.password = step_configuration["bbd_password"]
+        elif 'tenant_name' in step_configuration:
+            self.secret_data: GlobalConfig = self.config_manager.get_environment(step_configuration["tenant_name"])
+            if self.env_mode == 'production':
+                self.bbd_client = cast(GravitateSDAPI,
+                                       self.config_manager.environment_from_name(step_configuration["tenant_name"],
+                                                                                 "production",
+                                                                                 sd_basic_auth=True).sd.api_client)
+            else:
+                logger.debug("Initializing API in dev mode")
+                self.bbd_client = cast(GravitateSDAPI,
+                                       self.config_manager.environment_from_name(step_configuration["tenant_name"],
+                                                                                 "test",
+                                                                                 sd_basic_auth=True).sd.api_client)
+
+        else:
+            raise Exception("env or tenant is required")
+
+    def describe(self) -> str:
+        return "Upload Accessorials to BBD"
+
+    async def execute(self, accessorials: List[Dict]) -> BBDUploadResult:
+        logs = {"requests": [], "responses": [], "errors": []}
+        try:
+            total_accessorials = len(accessorials)
+            succeeded = []
+            failed_items = []
+
+            for idx, accessorial in enumerate(accessorials):
+                logs["requests"].append(accessorial)
+                try:
+                    resp = await self.bbd_client.call_ep("freight/accessorial/automatic/rate/create", json=accessorial)
+                    resp.raise_for_status()
+                    _json = resp.json()
+                    sleep(self.buffer)
+                    succeeded.append(accessorial)
+                    logger.info(f"Accessorials uploaded successfully: {idx + 1} of {total_accessorials}")
+                    logs["responses"].append({"response": _json, "request": accessorial})
+                except Exception as e:
+                    logs["errors"].append({"record": accessorial, "error": f"Error uploading accessorials: {str(e)} {e.response.content}"})
+                    failed_items.append(accessorial)
+                    continue
+
+            self.pipeline_context.included_files["accessorials data upload"] = json.dumps(logs, cls=CustomJSONEncoder)
+            return BBDUploadResult(
+                succeeded=len(succeeded),
+                failed=len(failed_items),
+                succeeded_items=succeeded
+            )
+
+        except Exception as e:
+            logger.exception(f"Unable to upload | {e}")
+            raise e

bb_integrations_lib/pipelines/steps/distribution_report/distribution_report_datafram_to_raw_data.py

@@ -0,0 +1,33 @@
+from datetime import datetime, UTC
+from io import BytesIO
+from typing import Tuple
+
+import pandas as pd
+
+from bb_integrations_lib.protocols.pipelines import Step
+from bb_integrations_lib.shared.model import RawData
+
+
+class DistributionReportDfToRawData(Step):
+    def __init__(self, file_base_name: str, file_name_date_format: str, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.file_base_name = file_base_name
+        self.file_name_date_format = file_name_date_format
+
+    def describe(self) -> str:
+        return "Distribution Report Dataframe to File"
+
+    @property
+    def file_name(self) -> str:
+        return f"{self.file_base_name}_{datetime.now(UTC).strftime(self.file_name_date_format)}.xlsx"
+
+    async def execute(self, data: Tuple[pd.DataFrame, pd.DataFrame]) -> RawData:
+        buff = BytesIO()
+        df_summary, df_detailed = data
+        with pd.ExcelWriter(buff, engine='openpyxl') as writer:
+            df_summary.to_excel(writer, sheet_name='Summary', index=False)
+            df_detailed.to_excel(writer, sheet_name='Details', index=False)
+        buff.seek(0)
+        return RawData(data=buff.read(), file_name=self.file_name)
+
+

bb_integrations_lib/pipelines/steps/distribution_report/get_model_history_step.py

@@ -0,0 +1,50 @@
+import json
+from datetime import datetime, timedelta, UTC
+
+from pymongo.synchronous.database import Database
+
+from bb_integrations_lib.protocols.pipelines import Step
+from bb_integrations_lib.util.utils import CustomJSONEncoder
+
+
+class GetModelHistoryStep(Step):
+    def __init__(self, mongo_database: Database, hours_back: float | None = None, include_model_mode: str = "latest_only", *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.mongo_database = mongo_database
+        self.hours_back = hours_back
+        self.include_model_mode = include_model_mode
+
+    def describe(self) -> str:
+        return f"Get model history for {self.tenant_name}"
+
+    async def execute(self, i: None) -> str | list:
+        if self.include_model_mode == "latest_only":
+            return self.get_latest_model_id()
+        else:
+            if self.hours_back is not None:
+                return self.get_last_n_models(int(self.hours_back))
+            else:
+                raise NotImplementedError("Please specify n_hours_back in the configuration")
+
+    def get_latest_model_id(self) -> str:
+        collection = self.mongo_database["model_history"]
+        latest_model = collection.find_one(
+            {'status': 'Success'},
+            sort=[('time_ran', -1)]
+        )
+        latest_model_id = str(latest_model['_id'])
+        self.pipeline_context.included_files[f'{self.__class__.__name__} result'] = json.dumps(latest_model, cls=CustomJSONEncoder)
+        self.pipeline_context.extra_data["latest_model_id"] = latest_model_id
+        return latest_model
+
+
+    def get_last_n_models(self, n: int) -> list:
+        n_hours_ago = datetime.now(UTC) - timedelta(hours=n)
+        collection = self.mongo_database["model_history"]
+        models = collection.find(
+            {'status': 'Success', 'time_ran': {'$gte': n_hours_ago}},
+            sort=[('time_ran', -1)]
+        )
+        models = list(models)
+        self.pipeline_context.included_files[f'{self.__class__.__name__} result'] = json.dumps(models, cls=CustomJSONEncoder)
+        return models

bb_integrations_lib/pipelines/steps/distribution_report/get_order_by_site_product_step.py

@@ -0,0 +1,62 @@
+from datetime import datetime
+
+from loguru import logger
+from pandas.core.interchange.dataframe_protocol import DataFrame
+
+from bb_integrations_lib.gravitate.sd_api import GravitateSDAPI
+from bb_integrations_lib.protocols.pipelines import Step
+
+
+class GetOrderBySiteProductStep(Step):
+    def __init__(self, sd_client: GravitateSDAPI, include_model_mode: str = "latest_only", state: str = "accepted",
+                 start_date: datetime | None = None, end_date: datetime | None = None, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.sd_client = sd_client
+        self.include_model_mode = include_model_mode
+        self.state = state
+        self.start_date = start_date
+        self.end_date = end_date
+
+    def describe(self) -> str:
+        return f"Get orders by site product"
+
+    async def execute(self, latest_model: dict | list) -> str | DataFrame:
+        if "orders_by_site_product" not in self.pipeline_context.extra_data:
+            self.pipeline_context.extra_data["orders_by_site_product"] = []
+        if isinstance(latest_model, dict):
+            return await self.get_orders_in_model_id(str(latest_model["_id"]))
+        elif isinstance(latest_model, list):
+            return [await self.get_orders_in_model_id(str(model["_id"])) for model in latest_model]
+        else:
+            raise ValueError("latest_model must be a dict or a list of dicts")
+
+    async def get_orders_in_model_id(self, model_id: str, market: str | None = None) -> str | DataFrame:
+        filter = self.build_filter()
+        json_data = {
+            'filter': {"lp_relationship.solver_id": str(model_id)
+                       },
+            'market': market if market else "",
+        }
+
+        response = await self.sd_client.call_ep(url="order/export_by_site_product",
+                                                json=json_data
+                                                )
+        orders = response.content.decode("utf-8")
+        if not hasattr(self, "custom_parser"):
+            response = orders
+        else:
+            logger.info(f"Using custom parser for {self.__class__.__name__}")
+            parser = self.custom_parser()
+            response = await parser.parse(orders)
+        self.pipeline_context.extra_data["orders_by_site_product"].append(response)
+        return response
+
+    def build_filter(self):
+        filter = {
+            "state": self.state,
+        }
+        if self.start_date:
+            filter["From"] = self.start_date.isoformat()
+        if self.end_date:
+            filter["To"] = self.end_date.isoformat()
+        return filter

bb_integrations_lib/pipelines/steps/distribution_report/get_tank_configs_step.py

@@ -0,0 +1,40 @@
+import json
+from typing import Dict, List
+
+from pymongo.synchronous.database import Database
+
+from bb_integrations_lib.protocols.pipelines import Step
+from bb_integrations_lib.util.utils import init_db, CustomJSONEncoder
+from loguru import logger
+from pandas import DataFrame
+
+
+class GetTankConfigsStep(Step):
+    def __init__(self, mongo_database: Database, include_model_mode: str = "latest_only", *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.database = mongo_database
+        self.include_model_mode = include_model_mode
+
+    def describe(self) -> str:
+        return f"Get tank configs for {self.tenant_name}"
+
+    async def execute(self, i: None) -> str | DataFrame:
+        return await self.get_tank_configs()
+
+    async def get_tank_configs(self) -> List[Dict] | DataFrame:
+        dos_columns = ['store_number', 'product', 'daily_lifting_estimate', 'measured_inventory']
+        collection = self.database["tank_config"]
+        tank_configs = collection.find(
+            {},
+            {col: 1 for col in dos_columns}
+        )
+        tank_configs = list(tank_configs)
+        self.pipeline_context.included_files[f'{self.__class__.__name__} result'] = json.dumps(tank_configs, cls=CustomJSONEncoder)
+        if not hasattr(self, "custom_parser"):
+            tc = tank_configs
+        else:
+            logger.info(f"Using custom parser for {self.__class__.__name__}")
+            parser = self.custom_parser()
+            tc =  await parser.parse(tank_configs)
+        self.pipeline_context.extra_data["tank_configs"] = tc
+        return tc

bb_integrations_lib/pipelines/steps/distribution_report/join_distribution_order_dos_step.py

@@ -0,0 +1,85 @@
+from datetime import datetime, UTC
+from typing import Dict, Tuple
+import pandas as pd
+from bb_integrations_lib.protocols.pipelines import Step
+
+
+
+class JoinDistributionOrderDosStep(Step):
+    def __init__(self, client_name: str, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.client_name = client_name
+
+    def describe(self) -> str:
+        return "Join Distribution Order with DOS"
+
+    async def execute(self, latest_model: dict | list) -> Tuple[pd.DataFrame, pd.DataFrame]:
+        tc_df = self.pipeline_context.extra_data["tank_configs"]
+        orders_data = self.pipeline_context.extra_data["orders_by_site_product"]
+
+        if isinstance(latest_model, dict):
+            pivot_df = orders_data
+            df_detailed = pd.merge(
+                pivot_df,
+                tc_df,
+                how='left',
+                left_on=['site', 'component_product'],
+                right_on=['store_number', 'product'],
+            )
+
+            df_detailed['dos_bucket'] = df_detailed['dos_bucket'].fillna('N/A')
+            df_summary = df_detailed.groupby(['dos_bucket'])[
+                ['component_volume_contract', 'component_volume_rack']].sum().reset_index()
+            df_detailed = df_detailed.groupby(['dos_bucket', 'component_product'])[
+                ['component_volume_contract', 'component_volume_rack']].sum().reset_index()
+
+            df_summary = self.contract_rack_split(df_summary, latest_model)
+            df_detailed = self.contract_rack_split(df_detailed, latest_model)
+
+            return df_summary, df_detailed
+        else:
+            all_summaries = []
+            all_details = []
+
+            for i, model in enumerate(latest_model):
+                pivot_df = orders_data[i]
+
+                df_detailed = pd.merge(
+                    pivot_df,
+                    tc_df,
+                    how='left',
+                    left_on=['site', 'component_product'],
+                    right_on=['store_number', 'product'],
+                )
+
+                df_detailed['dos_bucket'] = df_detailed['dos_bucket'].fillna('N/A')
+                df_summary = df_detailed.groupby(['dos_bucket'])[
+                    ['component_volume_contract', 'component_volume_rack']].sum().reset_index()
+                df_detailed_agg = df_detailed.groupby(['dos_bucket', 'component_product'])[
+                    ['component_volume_contract', 'component_volume_rack']].sum().reset_index()
+
+                df_summary = self.contract_rack_split(df_summary, model)
+                df_detailed_agg = self.contract_rack_split(df_detailed_agg, model)
+
+                all_summaries.append(df_summary)
+                all_details.append(df_detailed_agg)
+
+            combined_summary = pd.concat(all_summaries, ignore_index=True)
+            combined_detailed = pd.concat(all_details, ignore_index=True)
+
+            return combined_summary, combined_detailed
+
+    def contract_rack_split(self, df, latest_model):
+        markets = latest_model['markets']
+        _id = str(latest_model['_id'])
+        time_ran = latest_model['time_ran']
+
+        df['total'] = df['component_volume_contract'] + df['component_volume_rack']
+        df['pct_contract'] = df['component_volume_contract'] / df['total']
+        df['pct_rack'] = df['component_volume_rack'] / df['total']
+        df['ingested_at'] = datetime.now(UTC).replace(tzinfo=None)
+        df['markets'] = markets
+        df['model_id'] = _id
+        df['run_time'] = time_ran
+        df['client_name'] = self.client_name
+        return df

bb_integrations_lib/pipelines/steps/distribution_report/upload_distribution_report_datafram_to_big_query.py

@@ -0,0 +1,47 @@
+from typing import Dict, Tuple
+import pandas as pd
+from google.oauth2 import service_account
+from loguru import logger
+from bb_integrations_lib.protocols.pipelines import Step
+from bb_integrations_lib.provider.gcp.model import GoogleCredential
+from bb_integrations_lib.util.utils import load_credentials
+import pandas_gbq
+
+
+class UploadDistributionReportToBigQuery(Step):
+    def __init__(self, gbq_table_details: str, gbq_table_summary: str, google_project_id: str, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.gbq_table_details = gbq_table_details
+        self.gbq_table_summary = gbq_table_summary
+        self.google_project_id = google_project_id
+
+    def describe(self) -> str:
+        return "Upload Distribution Report to GBQ"
+
+    @property
+    def credentials(self) -> GoogleCredential:
+        return load_credentials(credential_type="google.credentials")
+
+    async def execute(self, data: Tuple[pd.DataFrame, pd.DataFrame]) -> None:
+        credentials = service_account.Credentials.from_service_account_info(self.credentials.model_dump())
+        df_summary, df_detailed = data
+        try:
+            pandas_gbq.to_gbq(
+                df_summary,
+                destination_table=self.gbq_table_summary,
+                project_id=self.google_project_id,
+                if_exists='append',
+                credentials=credentials,
+            )
+            pandas_gbq.to_gbq(
+                df_detailed,
+                destination_table='bb_reporting.contract_rack_util_product_detail',
+                project_id=self.google_project_id,
+                if_exists='append',
+                credentials=credentials,
+            )
+        except Exception as e:
+            logger.error(f"Failed to upload distribution report to BigQuery: {e}")
+            raise e
+
+

bb_integrations_lib/pipelines/steps/echo_step.py

@@ -0,0 +1,14 @@
+from bb_integrations_lib.protocols.pipelines import Step, Input
+from loguru import logger
+
+
+class EchoStep(Step[Input, Input]):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def describe(self):
+        return "Echo step input at debug priority"
+
+    async def execute(self, i: Input) -> Input:
+        logger.debug(i)
+        return i

bb_integrations_lib/pipelines/steps/export_dataframe_to_rawdata_step.py

@@ -0,0 +1,28 @@
+from typing import Callable
+
+import pandas as pd
+import polars as pl
+from bb_integrations_lib.protocols.pipelines import Step
+from bb_integrations_lib.shared.model import RawData
+
+
+class ExportDataFrameToRawDataStep(Step):
+    def __init__(self, pandas_export_function: str, pandas_export_kwargs: dict,
+                 file_name: str | Callable[[], str], *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.pandas_export_function = pandas_export_function
+        self.pandas_export_kwargs = pandas_export_kwargs
+        self.file_name = file_name
+
+    def _get_file_name(self) -> str:
+        if callable(self.file_name):
+            return self.file_name()
+        else:
+            return self.file_name
+
+    def describe(self) -> str:
+        return "Export a DataFrame to a file wrapped in a RawData object"
+
+    async def execute(self, i: pd.DataFrame | pl.DataFrame) -> RawData:
+        func = getattr(i, self.pandas_export_function)
+        return RawData(data=func(**self.pandas_export_kwargs).encode("utf-8"), file_name=self._get_file_name())

bb_integrations_lib/pipelines/steps/exporting/bbd_export_payroll_file_step.py

@@ -0,0 +1,107 @@
+from io import BytesIO
+from loguru import logger
+
+import pandas as pd
+
+from bb_integrations_lib.gravitate.sd_api import GravitateSDAPI
+from bb_integrations_lib.protocols.pipelines import Step, ParserBase, Input, Output
+from bb_integrations_lib.shared.model import RawData
+from pandas import DataFrame
+from datetime import datetime, date, UTC
+
+
+class BBDExportPayrollStep(Step):
+    def __init__(self, sd_client: GravitateSDAPI, file_name: str, parser: type[ParserBase] | None = None,
+                 parser_kwargs: dict | None = None, export_file: bool = True, target_date: str | None = None,
+                 additional_ep_params: dict | None = None, *args, **kwargs):
+        """
+        Export a payroll file from S&D.
+
+        :param export_file: If True, call the endpoint to export the file as an XLSX. If False, use the endpoint that
+          provides JSON.
+        """
+        super().__init__(*args, **kwargs)
+        self.sd_client = sd_client
+        self.file_name = file_name
+        self.export_file_only = export_file
+        self.additional_endpoint_parameters = additional_ep_params or {
+            "status": None,
+            "driver_id": None,
+            "updated_after": None,
+        }
+        if parser:
+            self.custom_parser = parser
+            self.custom_parser_kwargs = parser_kwargs or {}
+        self.target_date = target_date
+
+    def describe(self):
+        return "Export Payroll from Supply and Dispatch"
+
+    async def execute(self, last_sync_date: datetime | None = None) -> RawData:
+        last_sync_date = self.last_sync_date(last_sync_date)
+        dt = datetime.now(UTC)
+        if self.export_file_only is True:
+            logger.info(f"Exporting payroll file only for {last_sync_date}")
+            df = await self.export_payroll_file_only(last_sync_date)
+        else:
+            logger.info(f"Exporting payroll json for {last_sync_date}")
+            df =  await self.export_json(last_sync_date)
+        return RawData(data=df.to_csv().encode("utf-8"), file_name=f"{self.file_name}_{dt.strftime("%Y%m%d%H%M%S")}")
+
+    def last_sync_date(self, last_sync_date: datetime | None = None) -> str:
+        if last_sync_date is None:
+            if self.target_date is not None:
+                last_sync_date = self.target_date
+            elif self.pipeline_context.max_sync is not None:
+                last_sync_date = self.pipeline_context.max_sync.max_sync_date.isoformat()
+            else:
+                last_sync_date = datetime.combine(date.today(), datetime.min.time()).isoformat()
+        else:
+            last_sync_date = last_sync_date.isoformat()
+        return last_sync_date
+
+    async def export_payroll_file_only(self, dt: datetime) -> DataFrame:
+        resp = await self.sd_client.payroll_export_file(date=dt)
+        df = pd.read_excel(BytesIO(resp.content),
+                           dtype={"driver_source_id": str},
+                           engine="openpyxl",
+                           keep_default_na=False)
+        if df.empty:
+            raise Exception("No payroll data found")
+        if hasattr(self, "custom_parser"):
+            parser = self.custom_parser(tenant_name=self.tenant_name, **self.custom_parser_kwargs)
+            df = await parser.parse(df)
+            parser_logs = parser.get_logs()
+            self.pipeline_context.included_files["parser_logs"] = parser_logs
+        return df
+
+    async def export_json(self, dt: datetime) -> DataFrame:
+        resp = await self.sd_client.payroll_export(date=dt, **self.additional_endpoint_parameters)
+        json_resp = resp.json()
+        df = self.pre_parse_json(json_resp)
+        if df.empty:
+            raise Exception("No payroll data found")
+        return df
+
+    def pre_parse_json(self, json_resp: dict) -> DataFrame:
+        all_rows = []
+        for payroll_record in json_resp:
+            header_info = {k: v for k, v in payroll_record.items() if k != 'detail'}
+            details = payroll_record.get('detail', [])
+            if not details:
+                all_rows.append(header_info)
+            else:
+                for detail in details:
+                    row = header_info.copy()
+                    row.update(detail)
+                    all_rows.append(row)
+        df = pd.DataFrame(all_rows)
+        datetime_columns = [
+            'start_date', 'end_date', 'updated', 'shift_start',
+            'shift_actual_start', 'shift_actual_end', 'overridden_datetime'
+        ]
+        for col in datetime_columns:
+            if col in df.columns:
+                df[col] = pd.to_datetime(df[col], errors='coerce')
+
+        return df