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

bb_integrations_lib/pipelines/steps/get_latest_sync_date.py

@@ -0,0 +1,34 @@
+from datetime import datetime, UTC
+from typing import Dict
+
+import pytz
+from bb_integrations_lib.gravitate.rita_api import GravitateRitaAPI
+from bb_integrations_lib.models.rita.config import MaxSync
+from bb_integrations_lib.protocols.pipelines import Step
+from dateutil.parser import parse
+
+
+class GetLatestSyncDate(Step):
+    def __init__(self, step_configuration: Dict[str, str]):
+        super().__init__(step_configuration)
+        self.tenant_name = step_configuration['tenant_name']
+        self.mode = step_configuration.get('mode', 'production')
+        self.rita_client: GravitateRitaAPI = self.config_manager.environment_from_name(self.tenant_name,
+                                                                                  self.mode).rita.api_client
+        self.config_id = step_configuration.get("config_id")
+        self.test_override = step_configuration.get("test_override", None)
+
+    def describe(self) -> str:
+        return "Get Latest Sync Date"
+
+    async def execute(self, i: str) -> datetime:
+        return await self.get_last_sync_date()
+
+    async def get_last_sync_date(self) -> datetime:
+        if self.test_override:
+            return parse(self.test_override).replace(tzinfo=pytz.UTC)
+        if not self.config_id or self.config_id is None:
+            return datetime.now(UTC)
+        max_sync: MaxSync  = await self.rita_client.get_config_max_sync(config_id=self.config_id)
+        dt = max_sync.max_sync_date.replace(tzinfo=pytz.UTC)
+        return dt

bb_integrations_lib/pipelines/steps/importing/bbd_import_payroll_step.py

@@ -0,0 +1,30 @@
+from typing import Dict, Any
+
+import loguru
+from httpx import HTTPStatusError
+from bb_integrations_lib.protocols.pipelines import Step
+from bb_integrations_lib.shared.model import FileReference, FileType
+from bb_integrations_lib.gravitate.sd_api import GravitateSDAPI
+
+
+class BBDImportPayrollStep(Step):
+    def __init__(self, sd_client: GravitateSDAPI, output_file_path: str, bbd_date_argument: str, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.sd_client = sd_client
+        self.temp_file_path = output_file_path
+        self.bbd_date_argument = bbd_date_argument
+
+    def describe(self) -> str:
+        return "Import Payroll from BBD"
+
+    async def execute(self, _: Any) -> FileReference:
+        try:
+            payroll_resp = await self.sd_client.payroll_export(self.bbd_date_argument)
+            # Download the file from the response
+            if payroll_resp.status_code == 200:
+                with open(self.temp_file_path, "wb") as f:
+                    f.write(payroll_resp.content)
+
+            return FileReference(self.temp_file_path, FileType.excel)
+        except HTTPStatusError as e:
+            loguru.logger.error(e.response.content)

bb_integrations_lib/pipelines/steps/importing/get_order_numbers_to_export_step.py

@@ -0,0 +1,138 @@
+import enum
+from datetime import datetime, UTC, timedelta
+from typing import Any, List
+
+import pandas_gbq
+from google.oauth2 import service_account
+from loguru import logger
+import pandas as pd
+from pymongo import MongoClient
+from pymongo.synchronous.database import Database
+
+from bb_integrations_lib.models.pipeline_structs import StopPipeline, StopBranch
+from bb_integrations_lib.protocols.pipelines import Step
+from bb_integrations_lib.util.utils import init_db
+
+class OrderType(enum.Enum):
+    ANY = enum.auto()
+    NOT_BACKHAUL = enum.auto()
+    BACKHAUL_ONLY = enum.auto()
+
+    def to_mongo_query(self):
+        match self:
+            case OrderType.ANY:
+                return {}
+            case OrderType.NOT_BACKHAUL:
+                return {"type": {"$ne": "backhaul"}}
+            case OrderType.BACKHAUL_ONLY:
+                return {"type": "backhaul"}
+
+    def __str__(self) -> str:
+        return {
+            OrderType.ANY: "any",
+            OrderType.NOT_BACKHAUL: "not backhaul",
+            OrderType.BACKHAUL_ONLY: "backhaul only",
+        }[self]
+
+class GetOrderNumbersToExportStep(Step):
+    def __init__(self, mongo_database: Database, order_type_filter: OrderType, exported_order_table_name: str,
+                 project_id: str, gcp_credentials_file: str, testing_date_min: datetime | None = None,
+                 testing_date_max: datetime | None = None, lookback_days: int = 60, use_old_change_query: bool = False,
+                 check_for_updated_orders: bool = True, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.mongo_database = mongo_database
+        self.order_type_filter = order_type_filter
+        self.exported_order_table_name = exported_order_table_name
+        self.project_id = project_id
+        self.gcp_credentials_file = gcp_credentials_file
+        self.testing_date_min = testing_date_min
+        self.testing_date_max = testing_date_max
+        self.lookback_days = lookback_days
+        self.use_old_change_query = use_old_change_query
+        self.check_for_updated_orders = check_for_updated_orders
+        if not self.gcp_credentials_file.endswith(".json"):
+            self.gcp_credentials_file += ".json"
+
+    def describe(self) -> str:
+        return "Determine which order numbers to export."
+
+    async def execute(self, i: Any) -> List[int]:
+        collection = self.mongo_database["order_v2"]
+        if not self.testing_date_min and not self.testing_date_max:
+            lookback_date = datetime.now(UTC) - timedelta(days=self.lookback_days)
+            date_query = {'$or': [
+                {'updated_on': {'$gte': lookback_date}},
+                {'movement_updated': {'$gte': lookback_date}}
+            ]}
+            logger.debug(f"Looking for orders in the previous {self.lookback_days} days.")
+        elif self.testing_date_min and self.testing_date_max:
+            date_query = {'$or': [
+                {'updated_on': {'$gte': self.testing_date_min, '$lte': self.testing_date_max}},
+                {'movement_updated': {'$gte': self.testing_date_min, '$lte': self.testing_date_max}}
+            ]}
+            logger.debug(f"Looking for orders between {self.testing_date_min} and {self.testing_date_max}")
+        else:
+            raise RuntimeError("testing_date_max and testing_date_min must provided together, or both not provided.")
+
+        if self.use_old_change_query:
+            change_query = {"change_date": "$updated_on"}
+        else:
+            change_query = {"change_date": { "$max": ["$updated_on", "$movement_updated"] }}
+
+        # Over time the number of orders will grow. To keep step performance constant we're only going to look at
+        # orders updated within the last 60 days
+        assert isinstance(self.order_type_filter, OrderType)
+        logger.debug(f"Searching for orders of type: {str(self.order_type_filter)}")
+        orders = list(collection.find({**date_query,
+                                       **self.order_type_filter.to_mongo_query(),
+                                       "state": "complete"},
+                                      {
+                                          "_id": 0,
+                                          "number": 1,
+                                          **change_query,
+                                      }))
+        logger.debug(f"There are {len(orders)} orders updated in the last {self.lookback_days} days.")
+        if len(orders) == 0:
+            raise StopBranch()
+        # We need to look up the list of exported orders from GBQ
+        credentials = service_account.Credentials.from_service_account_file(self.gcp_credentials_file)
+        sql = (f"select order_number as number, max(date) as exported_date"
+               f" from `{self.exported_order_table_name}`"
+               f" group by order_number;")
+        gbq_df = pandas_gbq.read_gbq(sql, project_id=self.project_id, credentials=credentials, progress_bar_type=None)
+        logger.debug(f"Received {gbq_df.shape[0]} GBQ records.")
+
+        # Create a dataframe from the list of orders and do a left join on the orders exported from GBQ
+        orders_df = pd.DataFrame.from_records(orders)
+        joined = pd.merge(orders_df, gbq_df, on="number", how="left")
+
+        never_exported = joined[joined["exported_date"].isna()]
+        logger.debug(f"There are {never_exported.shape[0]} new orders to export.")
+
+        if self.check_for_updated_orders:
+            updated_after_export = joined[joined["exported_date"] < joined["change_date"]]
+            logger.debug(f"There are {updated_after_export.shape[0]} orders that have been updated and need to be exported again.")
+            to_export = set(never_exported["number"]).union(set(updated_after_export["number"]))
+        else:
+            to_export = set(never_exported["number"])
+
+        output = list(to_export)
+        logger.debug(output)
+        if len(output) == 0:
+            # There are no orders to export. Cancel the pipeline this step is part of.
+            raise StopBranch()
+        return output
+
+
+if __name__ == "__main__":
+    import asyncio
+    async def main():
+        s = GetOrderNumbersToExportStep(
+            mongo_database=MongoClient("mongodb conn str")["db_name"],
+            order_type_filter=OrderType.ANY,
+            exported_order_table_name="gravitate-harms-prod.harms_order_exports.exported_backhaul_orders",
+            project_id="gravitate-harms-prod",
+            gcp_credentials_file="google.credentials.json"
+        )
+        await s.execute(None)
+    asyncio.run(main())

bb_integrations_lib/pipelines/steps/importing/load_file_to_dataframe_step.py

@@ -0,0 +1,46 @@
+import asyncio
+from typing import Dict
+
+import pandas as pd
+
+from bb_integrations_lib.protocols.pipelines import Step
+from bb_integrations_lib.shared.model import FileReference, FileType, RawData
+from pandas import DataFrame
+
+
+class LoadFileToDataFrameStep(Step):
+    def __init__(self, sheet_name: str | int = 0, file_type: FileType | None = None, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.sheet_name = sheet_name
+        self.file_type = file_type
+
+    def describe(self) -> str:
+        return "Load file into dataframe"
+
+    async def execute(self, i: FileReference | RawData) -> DataFrame:
+        if isinstance(i, FileReference):
+            if i.file_type == FileType.excel:
+                if i.sheet_name is None:
+                    return pd.read_excel(i.file_path, sheet_name=0)
+                else:
+                    return pd.read_excel(i.file_path, sheet_name=i.sheet_name)
+            elif i.file_type == FileType.csv:
+                return pd.read_csv(i.file_path)
+            else:
+                raise NotImplementedError()
+        elif isinstance(i, RawData):
+            if self.file_type is not None:
+                if self.file_type == "csv":
+                    return pd.read_csv(i.data)
+            if i.file_name.endswith("csv"):
+                return pd.read_csv(i.data)
+            elif i.file_name.endswith("xlsx") or i.file_name.endswith("xls"):
+                return pd.read_excel(i.data, sheet_name=0)
+
+if __name__ == "__main__":
+    async def main():
+        input = FileReference("/home/ben-allen/Downloads/herdrich-payroll-test.xlsx", FileType.excel, "payroll")
+        output = await LoadFileToDataFrameStep().execute(input)
+        print(output.head(5))
+
+    asyncio.run(main())

bb_integrations_lib/pipelines/steps/importing/load_imap_attachment_step.py

@@ -0,0 +1,172 @@
+from time import sleep
+from typing import Any, AsyncIterator
+
+from loguru import logger
+
+from bb_integrations_lib.gravitate.rita_api import GravitateRitaAPI
+from bb_integrations_lib.models.pipeline_structs import NoPipelineData
+from bb_integrations_lib.protocols.pipelines import GeneratorStep
+from bb_integrations_lib.provider.imap.client import IMAPClient
+from bb_integrations_lib.secrets import IMAPCredential
+from bb_integrations_lib.secrets.credential_models import IMAPAuthSimple
+from bb_integrations_lib.shared.model import RawData, FileConfigRawData
+
+
+class LoadIMAPAttachmentStep(GeneratorStep):
+    def __init__(
+            self,
+            rita_client: GravitateRitaAPI,
+            imap_client: IMAPClient,
+            attachment_extension: str,
+            to_email_address: str | None = None,
+            from_email_address: str | None = None,
+            delivered_to_email_address: str | None = None,
+            retries: int = 3,
+            email_subject: str | None = None,
+            bucket_name: str | None = None,
+            config_names: list[str] | None = None,
+            raise_on_no_data: bool = False,
+            *args, **kwargs
+
+    ):
+        """
+        Load attachments from an IMAP folder / mailbox. The search criteriae are reductive; more filters are ANDed and
+        will result in fewer results. If no filters are provided, the default is to search for all emails in the inbox.
+
+        :param rita_client: The RITA client to retrieve fileconfigs with.
+        :param imap_client: The IMAP account which owns the mailbox.
+        :param attachment_extension: Sets the extension on the returned RawData.
+        :param to_email_address: Matches emails sent TO this address, if provided.
+        :param from_email_address: Matches emails sent FROM this address, if provided.
+        :param delivered_to_email_address: Matches emails "Delivered-To" this address, if provided. Differs from the
+          ``to_email_address`` argument - this matches on the Delivered-To mail header, which some email services set if
+          they are automatically forwarding from their mailbox to ours.
+        :param retries: How many times to retry a mail query that fails.
+        :param email_subject: The subject of the email to search for.
+        :param bucket_name: The config bucket that holds the config_names configs.
+        :param config_names: The names of FileConfigs to use to search for.
+        :param raise_on_no_data: Whether to raise an error when there are no matching emails.
+        """
+        super().__init__(*args, **kwargs)
+        self.rita_client = rita_client
+        self.imap_client = imap_client
+
+        self.to_email_address = to_email_address
+        self.from_email_address = from_email_address
+        self.delivered_to_email_address = delivered_to_email_address
+        self.email_subject = email_subject
+        self.attachment_extension = attachment_extension
+        self.retries = retries
+        criteria_parts = ['UNSEEN']
+
+        if self.from_email_address:
+            criteria_parts.append(f'FROM {self.from_email_address}')
+
+        if self.delivered_to_email_address:
+            criteria_parts.append(f'HEADER Delivered-To {self.delivered_to_email_address}')
+        if self.to_email_address:
+            criteria_parts.append(f'TO {self.to_email_address}')
+
+        if self.email_subject:
+            if ' ' in self.email_subject:
+                criteria_parts.append(f'SUBJECT "{self.email_subject}"')
+            else:
+                criteria_parts.append(f'SUBJECT {self.email_subject}')
+        self.criteria = '(' + ' '.join(criteria_parts) + ')'
+
+        self.bucket_name = bucket_name
+        self.config_names = config_names
+
+        self.raise_on_no_data = raise_on_no_data
+
+    def describe(self) -> str:
+        return "Get attachments from IMAP folder"
+
+    async def generator(self, i: Any) -> AsyncIterator[RawData]:
+        for config_name in self.config_names:
+            async for result in self.load_using_config(config_name):
+                yield result
+
+    async def load_using_config(self, config_name: str) -> AsyncIterator[RawData]:
+        file_config = (
+            await self.rita_client.get_fileconfig_by_name(self.bucket_name, config_name)
+        )[config_name]
+        logger.info(f"Searching with criteria {self.criteria}")
+        message_indexes = self.imap_client.search(self.criteria)
+        logger.info(f"Found {len(message_indexes)} new emails in inbox meeting search criteria")
+        if self.raise_on_no_data and not message_indexes:
+            raise NoPipelineData("No new emails found in inbox meeting search criteria")
+        for idx in message_indexes:
+            logger.info(f"Fetching mail {idx}")
+            for retry in range(self.retries):
+                try:
+                    message = self.imap_client.fetch(idx)
+                    attachment_rd = self.imap_client.get_attachment_from_message(
+                        message,
+                        extension=self.attachment_extension,
+                        return_rawdata=True
+                    )
+                    if attachment_rd is not None:
+                        logger.info(f"Fetched attachment from email id: {idx}")
+                        if file_config:
+                            attachment_rd = FileConfigRawData(
+                                file_name=attachment_rd.file_name,
+                                data=attachment_rd.data,
+                                file_config=file_config
+                            )
+                            self.pipeline_context.file_config = file_config
+                        yield attachment_rd
+                        break
+                    else:
+                        raise Exception("Returned attachment was None")
+                except Exception as e:
+                    logger.error(f"Error reading data from mail: {e}")
+                    sleep(3)
+                    self.imap_client.mark_unseen(idx)
+
+    @staticmethod
+    def _unplus_email(email_str: str) -> str:
+        """
+        Remove plus-addressing from an email to get the 'base' email account. If there is no plus, the original email
+        string is returned. Note that in either case, the email must have an @ symbol in it.
+        """
+        split_at = email_str.split("@")
+        if len(split_at) != 2:
+            raise ValueError(f"Email address '{email_str}' is not in the expected format")
+        first_half = split_at[0]
+        has_plus = "+" in first_half
+        if has_plus:
+            unplussed = first_half[:first_half.rfind("+")]
+            return f"{unplussed}@{split_at[1]}"
+        return email_str
+
+
+if __name__ == "__main__":
+    async def main():
+        step = LoadIMAPAttachmentStep(
+            rita_client=GravitateRitaAPI(
+                base_url="",
+                client_id="",
+                client_secret=""
+            ),
+            imap_client=IMAPClient(
+                credentials=IMAPCredential(
+                    host="",
+                    port=993,
+                    email_address="",
+                    auth=IMAPAuthSimple(
+                        password="",
+                    )
+                )
+            ),
+            attachment_extension=".csv",
+            bucket_name="/Inventory",
+            config_names=["wawa"]
+        )
+        async for result in step.generator(None):
+            print(result)
+
+
+    import asyncio
+
+    asyncio.run(main())

bb_integrations_lib/pipelines/steps/importing/pe_bulk_sync_price_structure_step.py

@@ -0,0 +1,68 @@
+import json
+from typing import List
+
+from loguru import logger
+
+from bb_integrations_lib.gravitate.pe_api import GravitatePEAPI
+from bb_integrations_lib.models.pipeline_structs import UploadResult
+from bb_integrations_lib.models.rita.issue import IssueBase, IssueCategory
+from bb_integrations_lib.protocols.flat_file import PeBulkSyncIntegration
+from bb_integrations_lib.protocols.pipelines import Step
+from bb_integrations_lib.util.utils import CustomJSONEncoder
+
+
+class PEBulksSyncPriceStructure(Step):
+    def __init__(self, pe_client: GravitatePEAPI, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.pe_client = pe_client
+
+    def describe(self) -> str:
+        return "Bulk Sync Price Structure in Pricing Engine"
+
+    async def execute(self, i: List[PeBulkSyncIntegration]) -> UploadResult:
+        failed_rows: List = []
+        success_rows: List = []
+        responses: List = []
+        try:
+            for row in i:
+                row_dump = row.model_dump(exclude_none=True)
+                row_dump = self.gen_unique_price_dtos(row_dump)
+                try:
+                    response = await self.pe_client.bulk_sync_price_structure(row_dump)
+                    responses.append(response)
+                    success_rows.append({**row_dump, "response": response})
+                except Exception as e:
+                    logger.error(f"Failed to bulk sync row: {e}")
+                    failed_rows.append({**row_dump, "response": str(e)})
+                    continue
+        except Exception as e:
+            if irc := self.pipeline_context.issue_report_config:
+                fc = self.pipeline_context.file_config
+                key = f"{irc.key_base}_{fc.config_id}_failed_to_upload"
+                self.pipeline_context.issues.append(IssueBase(
+                    key=key,
+                    config_id=fc.config_id,
+                    name="Failed to bulk sync data",
+                    category=IssueCategory.TANK_READING,
+                    problem_short=f"{len(failed_rows)} rows failed to bulk sync",
+                    problem_long=json.dumps(failed_rows)
+                ))
+        logs = {
+            "request": [l.model_dump() for l in i],
+            "response": responses
+        }
+        self.pipeline_context.included_files["price merge data"] = json.dumps(logs, cls=CustomJSONEncoder)
+        return UploadResult(succeeded=len(success_rows), failed=len(failed_rows),
+                               succeeded_items=list(success_rows))
+
+    def gen_unique_price_dtos(self, row_dump: dict):
+        deduped = []
+        keys = set()
+        for elem in row_dump['IntegrationDtos'][0]['PriceInstrumentDTOs']:
+            key = f"{elem['ProductLookup']['SourceIdString']} @ {elem['LocationLookup']['SourceIdString']} - {elem['CounterPartyLookup']['SourceIdString']}"
+            if key in keys:
+                continue
+            keys.add(key)
+            deduped.append(elem)
+        row_dump['IntegrationDtos'][0]['PriceInstrumentDTOs'] = deduped
+        return row_dump

bb_integrations_lib/pipelines/steps/importing/pe_price_merge_step.py

@@ -0,0 +1,86 @@
+import json
+from typing import Dict, List
+
+from bb_integrations_lib.gravitate.pe_api import GravitatePEAPI
+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
+from loguru import logger
+from bb_integrations_lib.models.pipeline_structs import BBDUploadResult, UploadResult
+from bb_integrations_lib.models.rita.issue import IssueBase, IssueCategory
+from bb_integrations_lib.protocols.pipelines import Step
+from bb_integrations_lib.protocols.flat_file import PePriceMergeIntegration
+
+
+class PEPriceMerge(Step[List[PePriceMergeIntegration], 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.tenant_name = step_configuration.get('tenant_name', None)
+        if 'pe_client_base_url' in step_configuration:
+            self.client_base_url = step_configuration['pe_client_base_url']
+            self.username = step_configuration['pe_username']
+            self.password = step_configuration['pe_password']
+            self.client_name = step_configuration['client_name']
+            self.pe_integration_client = GravitatePEAPI(
+                base_url=self.client_base_url,
+                username=self.username,
+                password=self.password)
+
+        else:
+            if self.tenant_name is None:
+                raise ValueError("'tenant_name' or pe client base url are required")
+            self.secret_data: GlobalConfig = self.config_manager.get_environment(self.tenant_name)
+            if self.env_mode == 'production':
+                self.pe_integration_client = GravitatePEAPI(
+                    base_url=self.secret_data.prod.pe.base_url,
+                    username=self.secret_data.prod.pe.username,
+                    password=self.secret_data.prod.pe.password,
+                )
+            else:
+                logger.debug("Initializing API in dev mode")
+                self.pe_integration_client = GravitatePEAPI(
+                    base_url=self.secret_data.test.pe.base_url,
+                    username=self.secret_data.test.pe.username,
+                    password=self.secret_data.test.pe.password,
+                )
+
+    def describe(self) -> str:
+        return "Merge Prices in Pricing Engine"
+
+    async def execute(self, i: List[PePriceMergeIntegration]) -> BBDUploadResult:
+        failed_rows: List = []
+        success_rows: List = []
+        responses: List = []
+        try:
+            for row in i:
+                row_dump = row.model_dump(exclude_none=True)
+                try:
+                    response = await self.pe_integration_client.merge_prices(row_dump)
+                    success_rows.append({**row_dump, "response": response})
+                    responses.append(response)
+                except Exception as e:
+                    logger.error(f"Failed to merge row: {e}")
+                    failed_rows.append(row_dump)
+                    continue
+        except Exception as e:
+            if irc := self.pipeline_context.issue_report_config:
+                fc = self.pipeline_context.file_config
+                key = f"{irc.key_base}_{fc.config_id}_failed_to_upload"
+                self.pipeline_context.issues.append(IssueBase(
+                    key=key,
+                    config_id=fc.config_id,
+                    name="Failed to merge price row",
+                    category=IssueCategory.PRICE,
+                    problem_short=f"{len(failed_rows)} rows failed to price merge",
+                    problem_long=json.dumps(failed_rows)
+                ))
+        logs = {
+            "request": [l.model_dump() for l in i],
+            "response": responses
+        }
+        self.pipeline_context.included_files["price merge data"] = json.dumps(logs, cls=CustomJSONEncoder)
+        return UploadResult(succeeded=len(success_rows), failed=len(failed_rows),
+                            succeeded_items=list(success_rows))
+